Skip to content

Commit

Permalink
Refactor Orchestrator class in core.py to include before and after ex…
Browse files Browse the repository at this point in the history
…periment plugins
  • Loading branch information
cyanxiao committed May 2, 2024
1 parent 906ac83 commit f956add
Showing 1 changed file with 29 additions and 0 deletions.
29 changes: 29 additions & 0 deletions midori/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,20 +16,26 @@ def __init__(
hostname: str,
username: str,
password: str,
before_experiment_plugins: List[Type[PluginHelper]],
before_trial_cooling_time: int,
setup_plugins: List[Type[PluginHelper]],
trial_timespan: int,
end_trial_plugins: List[Type[PluginHelper]],
after_trial_cooling_time: int,
end_experiment_plugins: List[Type[PluginHelper]],
variables: Dict[str, List[str]],
subject_path: str,
) -> None:
self.__before_experiment_plugins: List[Type[PluginHelper]] = (
before_experiment_plugins
)
self.__subject_path: str = subject_path
self.__before_trial_cooling_time: int = before_trial_cooling_time
self.__trial_timespan: int = trial_timespan
self.__setup_plugins: List[Type[PluginHelper]] = setup_plugins
self.__after_trial_cooling_time: int = after_trial_cooling_time
self.__end_trial_plugins: List[Type[PluginHelper]] = end_trial_plugins
self.__end_experiment_plugins: List[Type[PluginHelper]] = end_experiment_plugins
self.__ssh: paramiko.SSHClient = paramiko.SSHClient()
self.__ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())

Expand All @@ -43,6 +49,18 @@ def __init__(

def run(self) -> None:
print("Starting the experiment...")

# Before Experiment plugins
output = None
for Plugin in self.__before_experiment_plugins:
plugin = Plugin(
ssh=self.__ssh,
subject_path=self.__subject_path,
treatment="",
previous_output=output,
)
output = plugin.execute()

print(f"Treatments: {self.treatments}")
for treatment in self.treatments:
if not is_a_branch_exist(
Expand Down Expand Up @@ -85,4 +103,15 @@ def run(self) -> None:
# After Trial Cooling time
pause(interval=self.__after_trial_cooling_time)

# End Experiment plugins
output = None
for Plugin in self.__end_experiment_plugins:
plugin = Plugin(
ssh=self.__ssh,
subject_path=self.__subject_path,
treatment=treatment,
previous_output=output,
)
output = plugin.execute()

close(ssh=self.__ssh)

0 comments on commit f956add

Please sign in to comment.