From 154175acdda47ff4be1729a7cfaad22bc2724c5a Mon Sep 17 00:00:00 2001 From: Leonardo Schwarz Date: Fri, 28 Feb 2025 12:06:39 +0100 Subject: [PATCH] improve state logic --- .../bfabric_app/submitter.py | 3 +- .../submitter/slurm_submitter.py | 31 +++++++++++++------ .../bfabric_app/test_submitter.py | 3 +- 3 files changed, 24 insertions(+), 13 deletions(-) diff --git a/bfabric_app_runner/src/bfabric_app_runner/bfabric_app/submitter.py b/bfabric_app_runner/src/bfabric_app_runner/bfabric_app/submitter.py index f58dedf7..9d21dd1b 100644 --- a/bfabric_app_runner/src/bfabric_app_runner/bfabric_app/submitter.py +++ b/bfabric_app_runner/src/bfabric_app_runner/bfabric_app/submitter.py @@ -89,8 +89,7 @@ def run(self) -> None: workunit_config=workunit_config, ) submitter = SlurmSubmitter(slurm_config_template) - submitter.submit(workunit_wrapper_data=workunit_wrapper_data) - submitter.create_log_resource(workunit_wrapper_data=workunit_wrapper_data, client=self._client) + submitter.submit(workunit_wrapper_data=workunit_wrapper_data, client=self._client) @use_client diff --git a/bfabric_app_runner/src/bfabric_app_runner/submitter/slurm_submitter.py b/bfabric_app_runner/src/bfabric_app_runner/submitter/slurm_submitter.py index b1bf7724..649e3122 100644 --- a/bfabric_app_runner/src/bfabric_app_runner/submitter/slurm_submitter.py +++ b/bfabric_app_runner/src/bfabric_app_runner/submitter/slurm_submitter.py @@ -23,6 +23,7 @@ id mkdir -p "{working_directory}" cd "{working_directory}" +{activate_log_resource} set +x tee app_version.yml < str: return f"{script_header}\n\n{main_command}" def _get_main_command( - self, workunit_wrapper_data: WorkunitWrapperData, working_directory: str, submitter_config: SubmitterSlurmSpec + self, + workunit_wrapper_data: WorkunitWrapperData, + working_directory: str, + submitter_config: SubmitterSlurmSpec, + log_resource_id: int | None, ) -> str: app_version_yml = yaml.safe_dump(workunit_wrapper_data.app_version.model_dump(mode="json")) workunit_definition_yml = yaml.safe_dump(workunit_wrapper_data.workunit_definition.model_dump(mode="json")) @@ -65,6 +70,10 @@ def _get_main_command( if not submitter_config.config.force_storage else f"--force-storage {str(submitter_config.config.force_storage)}" ) + if log_resource_id is not None: + activate_log_resource = f"bfabric-cli api save resource {log_resource_id} status available" + else: + activate_log_resource = "" render_args = dict( app_version_yml=app_version_yml, workunit_definition_yml=workunit_definition_yml, @@ -72,6 +81,7 @@ def _get_main_command( working_directory=working_directory, force_storage_flags=force_storage_flags, workunit_id=workunit_wrapper_data.workunit_definition.registration.workunit_id, + activate_log_resource=activate_log_resource, ) logger.info("Render args: {}", render_args) return _MAIN_BASH_TEMPLATE.format(**render_args) @@ -90,7 +100,7 @@ def evaluate_config(self, workunit_wrapper_data: WorkunitWrapperData) -> SlurmCo workunit = VariablesWorkunit(id=workunit_wrapper_data.workunit_definition.registration.workunit_id) return self._config_template.evaluate(app=app, workunit=workunit) - def submit(self, workunit_wrapper_data: WorkunitWrapperData) -> None: + def submit(self, workunit_wrapper_data: WorkunitWrapperData, client: Bfabric) -> None: # Evaluate the config slurm_config = self.evaluate_config(workunit_wrapper_data=workunit_wrapper_data) logger.info("Slurm Config: {}", slurm_config) @@ -99,6 +109,9 @@ def submit(self, workunit_wrapper_data: WorkunitWrapperData) -> None: workunit_id = workunit_wrapper_data.workunit_definition.registration.workunit_id script_path = slurm_config.submitter_config.config.local_script_dir / f"workunitid-{workunit_id}_run.bash" + # Create the log resource + log_resource_id = self._create_log_resource(config=slurm_config, workunit_id=workunit_id, client=client) + # Determine the working directory. working_directory = slurm_config.get_scratch_dir() @@ -107,6 +120,7 @@ def submit(self, workunit_wrapper_data: WorkunitWrapperData) -> None: workunit_wrapper_data=workunit_wrapper_data, working_directory=working_directory, submitter_config=slurm_config.submitter_config, + log_resource_id=log_resource_id, ) script = self._compose_script(main_command=main_command, slurm_config=slurm_config) script_path.write_text(script) @@ -120,20 +134,19 @@ def submit(self, workunit_wrapper_data: WorkunitWrapperData) -> None: logger.info("Running {}", shlex.join(cmd)) subprocess.run(cmd, env=env, check=True) - def create_log_resource(self, workunit_wrapper_data: WorkunitWrapperData, client: Bfabric) -> None: - config = self.evaluate_config(workunit_wrapper_data=workunit_wrapper_data) + def _create_log_resource(self, config: SlurmConfig, workunit_id: int, client: Bfabric) -> int | None: + """Creates the log resource and returns its id, or if no log storage ID is provided, returns None.""" if config.submitter_config.config.log_storage_id is None: logger.info("No log storage ID provided, skipping log resource creation") - return + return None - workunit_id = workunit_wrapper_data.workunit_definition.registration.workunit_id - client.save( + return client.save( "resource", { "name": f"Developer Log WU{workunit_id}", "relativepath": config.submitter_config.config.log_storage_filename, "storageid": config.submitter_config.config.log_storage_id, "workunitid": workunit_id, - "status": "available", + "status": "pending", }, - ) + )[0]["id"] diff --git a/tests/bfabric_app_runner/bfabric_app/test_submitter.py b/tests/bfabric_app_runner/bfabric_app/test_submitter.py index a538179e..5f2e27a4 100644 --- a/tests/bfabric_app_runner/bfabric_app/test_submitter.py +++ b/tests/bfabric_app_runner/bfabric_app/test_submitter.py @@ -143,7 +143,6 @@ def test_run( mock_submitter.run() # Verify - mock_slurm_submitter.submit.assert_called_once_with(workunit_wrapper_data=mock_workunit_wrapper_data) - mock_slurm_submitter.create_log_resource.assert_called_once_with( + mock_slurm_submitter.submit.assert_called_once_with( workunit_wrapper_data=mock_workunit_wrapper_data, client=mock_client )