Coverage for src/shephex/executor/executor.py: 98%

48 statements  

« prev     ^ index     » next       coverage.py v7.6.1, created at 2025-03-30 12:30 +0200

1""" 

2Executor base class. 

3""" 

4from abc import ABC, abstractmethod 

5from pathlib import Path 

6from typing import List, Optional, Sequence, Union 

7 

8from shephex.decorators import disable_decorators 

9from shephex.experiment import DryResult, Experiment, ExperimentResult 

10from shephex.experiment.status import Pending, Status 

11 

12from shephex.experiment.chain_iterator import ChainableExperimentIterator 

13 

14 

15class Executor(ABC): 

16 """ 

17 Executor base class. 

18 """ 

19 def __init__(self) -> None: 

20 pass 

21 

22 def execute( 

23 self, 

24 experiments: Union[Experiment, Sequence[Experiment]], 

25 dry: bool = False, 

26 execution_directory: Optional[Union[Path, str]] = None, 

27 valid_statuses: Optional[Sequence[Status]] = None, 

28 ) -> Union[ExperimentResult, List[ExperimentResult]]: 

29 """ 

30 Execute a set of experiments. 

31 

32 Parameters 

33 ----------- 

34 experiments: Experiment or Sequence[Experiment] 

35 The experiments to be executed. 

36 dry: bool 

37 If True, the experiments will not be executed, only information about 

38 them will be printed. 

39 """ 

40 if isinstance(experiments, Experiment): 

41 experiments = [experiments] 

42 elif isinstance(experiments, ChainableExperimentIterator): 

43 experiments = list(experiments) 

44 

45 if valid_statuses is None: 

46 valid_statuses = [Pending()] 

47 

48 valid_experiments = [] 

49 for experiment in experiments: 

50 if experiment.status in valid_statuses: 

51 valid_experiments.append(experiment) 

52 else: 

53 print(f"Experiment {experiment.identifier} has status {experiment.status}, skipping.") 

54 

55 for experiment in valid_experiments: 

56 if not experiment.shephex_directory.exists(): 

57 experiment.dump() 

58 

59 results = self._sequence_execute( 

60 valid_experiments, dry=dry, execution_directory=execution_directory 

61 ) 

62 

63 return results 

64 

65 def _sequence_execute( 

66 self, 

67 experiments: Sequence[Experiment], 

68 dry: bool = False, 

69 execution_directory: Optional[Union[Path, str]] = None, 

70 ) -> Sequence[ExperimentResult]: 

71 results = [] 

72 for experiment in experiments: 

73 result = self._execute( 

74 experiment, dry=dry, execution_directory=execution_directory 

75 ) 

76 results.append(result) 

77 return results 

78 

79 @abstractmethod 

80 def _single_execute( 

81 self, 

82 experiment: Experiment, 

83 dry: bool = False, 

84 execution_directory: Optional[Union[Path, str]] = None, 

85 ) -> ExperimentResult: 

86 raise NotImplementedError # pragma: no cover 

87 

88 def _execute( 

89 self, 

90 experiment: Experiment, 

91 dry: bool = False, 

92 execution_directory: Optional[Union[Path, str]] = None, 

93 ) -> ExperimentResult: 

94 result = self._single_execute( 

95 experiment, dry=dry, execution_directory=execution_directory 

96 ) 

97 return result 

98 

99 

100class LocalExecutor(Executor): 

101 """ 

102 Executor that runs the experiment locally, in the current process (or subprocess), 

103 without any parallelization. 

104 """ 

105 

106 def __init__(self) -> None: 

107 super().__init__() 

108 

109 def _single_execute( 

110 self, 

111 experiment: Experiment, 

112 dry: bool = False, 

113 execution_directory: Optional[Union[Path, str]] = None, 

114 ) -> ExperimentResult: 

115 

116 if dry: 

117 print(f'Experiment {experiment.identifier} to be executed locally.') 

118 return DryResult() 

119 

120 with disable_decorators(): 

121 result = experiment._execute(execution_directory=execution_directory) 

122 

123 return result