pybiolib 1.2.339__tar.gz → 1.2.340.dev1__tar.gz

{pybiolib-1.2.339 → pybiolib-1.2.340.dev1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: pybiolib
-Version: 1.2.339
+Version: 1.2.340.dev1
 Summary: BioLib Python Client
 License: MIT
 Keywords: biolib

{pybiolib-1.2.339 → pybiolib-1.2.340.dev1}/biolib/_internal/utils/multinode.py

@@ -84,6 +84,8 @@ def fasta_batch_records(records, work_per_batch_min, work_per_residue=1, verbose
             n_seqs, n_res = batch_dict['records'], batch_dict['residues']
             print(f'Batch {i+1}: {n_res} residues from {n_seqs} sequences')
 
+    # records = SeqUtil.parse_fasta(fasta_file)
+
     batches = []
     batch = []
     current_work_units = 0

{pybiolib-1.2.339 → pybiolib-1.2.340.dev1}/biolib/compute_node/job_worker/executors/docker_executor.py

@@ -27,7 +27,7 @@ from biolib.compute_node import utils as compute_node_utils
 from biolib.compute_node.cloud_utils import CloudUtils
 from biolib.compute_node.job_worker.docker_image_cache import DockerImageCache
 from biolib.compute_node.job_worker.executors.docker_types import DockerDiffKind
-from biolib.compute_node.job_worker.executors.types import LocalExecutorOptions, StatusUpdate
+from biolib.compute_node.job_worker.executors.types import LocalExecutorOptions, StatusUpdate, MetadataToSaveOutput
 from biolib.compute_node.job_worker.mappings import Mappings, path_without_first_folder
 from biolib.compute_node.job_worker.utilization_reporter_thread import UtilizationReporterThread
 from biolib.compute_node.job_worker.utils import ComputeProcessException
@@ -61,9 +61,12 @@ class DockerExecutor:
         os.makedirs(user_data_tar_dir, exist_ok=True)
 
         self._docker_container: Optional[Container] = None
+        self._docker_api_client = BiolibDockerClient.get_docker_client().api
         self._runtime_tar_path = f'{user_data_tar_dir}/runtime_{self._random_docker_id}.tar'
         self._input_tar_path = f'{user_data_tar_dir}/input_{self._random_docker_id}.tar'
 
+        self._metadata_for_save_output_on_cancel: Optional[MetadataToSaveOutput] = None
+
         if utils.IS_RUNNING_IN_CLOUD and not utils.BIOLIB_SECRETS_TMPFS_PATH:
             error_message = 'Running in cloud but no TMPFS path has been set for secrets'
             logger_no_user_data.error(error_message)
@@ -81,12 +84,12 @@ class DockerExecutor:
             raise Exception('Docker container was None')
         return self._docker_container
 
-    def execute_module(self, module_input_path: str, module_output_path: str) -> None:
+    def execute_module(self) -> None:
         try:
             job_uuid = self._options['job']['public_id']
             send_status_update = self._options['send_status_update']
             logger_no_user_data.debug(f'Reading module input of {job_uuid}.')
-            with open(module_input_path, 'rb') as fp:
+            with open(self._options['module_input_path'], 'rb') as fp:
                 module_input_tmp = ModuleInput(fp.read())
                 logger_no_user_data.debug(f'Deserialing module input of {job_uuid}...')
                 module_input = module_input_tmp.deserialize()
@@ -100,7 +103,7 @@ class DockerExecutor:
 
             logger_no_user_data.debug(f'Starting execution of {job_uuid}.')
             try:
-                self._execute_helper(module_input, module_output_path)
+                self._execute_helper(module_input)
             except docker.errors.NotFound as docker_error:
                 raise DockerContainerNotFoundDuringExecutionException from docker_error
             except Exception as exception:
@@ -161,7 +164,8 @@ class DockerExecutor:
             may_contain_user_data=False,
         )
 
-    def _execute_helper(self, module_input, module_output_path: str) -> None:
+
+    def _execute_helper(self, module_input) -> None:
         job_uuid = self._options['job']['public_id']
         logger_no_user_data.debug(f'Initializing container for {job_uuid}.')
         self._initialize_docker_container(module_input)
@@ -187,10 +191,9 @@ class DockerExecutor:
         logger_no_user_data.debug(f'_map_and_copy_input_files_to_container for {job_uuid}.')
         self._map_and_copy_input_files_to_container(module_input['files'], module_input['arguments'])
 
-        docker_api_client = BiolibDockerClient.get_docker_client().api
         logger_no_user_data.debug(f'Attaching Docker container for {job_uuid}')
 
-        stdout_and_stderr_stream = docker_api_client.attach(
+        stdout_and_stderr_stream = self._docker_api_client.attach(
             container=self._container.id,
             stderr=True,
             stdout=True,
@@ -211,6 +214,12 @@ class DockerExecutor:
             logger.debug(f'Warning: Job "{job_uuid}" failed to start container. Hit error: {startup_error_string}')
             # even though the container start failed we should still be able to call logs() and wait() on it, so we pass
 
+        self._metadata_for_save_output_on_cancel = MetadataToSaveOutput(
+            arguments=module_input['arguments'],
+            startup_error_string=startup_error_string,
+            pre_start_diff=pre_start_diff,
+        )
+
         if self._options['job']['app_version'].get('stdout_render_type') != 'markdown':
             logger_no_user_data.debug(f'Streaming stdout for {job_uuid}')
             for stdout_and_stderr in stdout_and_stderr_stream:
@@ -220,7 +229,14 @@ class DockerExecutor:
                 self._send_stdout_and_stderr(stdout_and_stderr)
 
         logger_no_user_data.debug(f'Waiting on docker for {job_uuid}')
-        docker_result = docker_api_client.wait(self._container.id)
+        try:
+            docker_result = self._docker_api_client.wait(self._container.id)
+        except docker.errors.NotFound as error:
+            if self._is_cleaning_up:
+                return
+            else:
+                raise error
+
         logger_no_user_data.debug(f'Got result from docker for {job_uuid}')
         exit_code = docker_result['StatusCode']
         # 137 is the error code from linux OOM killer (Should catch 90% of OOM errors)
@@ -232,22 +248,26 @@ class DockerExecutor:
             )
 
         logger_no_user_data.debug(f'Docker container exited with code {exit_code} for {job_uuid}')
+        self._save_module_output_from_container(exit_code, self._metadata_for_save_output_on_cancel)
+
 
-        full_stdout = docker_api_client.logs(self._container.id, stdout=True, stderr=False)
-        full_stderr = docker_api_client.logs(self._container.id, stdout=False, stderr=True)
+    def _save_module_output_from_container(self, exit_code: int, metadata: MetadataToSaveOutput) -> None:
+        full_stdout = self._docker_api_client.logs(self._container.id, stdout=True, stderr=False)
+        full_stderr = self._docker_api_client.logs(self._container.id, stdout=False, stderr=True)
 
-        if startup_error_string:
-            full_stderr = full_stderr + startup_error_string.encode()
+        if metadata['startup_error_string']:
+            full_stderr = full_stderr + metadata['startup_error_string'].encode()
 
         self._write_module_output_to_file(
-            arguments=module_input['arguments'],
+            arguments=metadata['arguments'],
             exit_code=exit_code,
-            module_output_path=module_output_path,
+            module_output_path=self._options['module_output_path'],
             stderr=full_stderr,
             stdout=full_stdout,
-            pre_start_diff=pre_start_diff,
+            pre_start_diff=metadata['pre_start_diff'],
         )
 
+
     def cleanup(self):
         # Don't clean up if already in the process of doing so, or done doing so
         if self._is_cleaning_up:
@@ -255,6 +275,22 @@ class DockerExecutor:
         else:
             self._is_cleaning_up = True
 
+        if self._metadata_for_save_output_on_cancel is not None:
+            try:
+                logger_no_user_data.debug('Attempting to save results')
+                self._docker_container.stop()
+                self._docker_container.reload()
+                logger_no_user_data.debug(f'Container state {self._docker_container.status}')
+                self._save_module_output_from_container(
+                    exit_code=self._docker_container.attrs['State']['ExitCode'],
+                    metadata=self._metadata_for_save_output_on_cancel
+                )
+                logger_no_user_data.debug('Saved results')
+            except BaseException as error:
+                logger_no_user_data.error(f'Failed to save results on cancel with error: {error}')
+        else:
+            logger_no_user_data.debug('Missing metadata, cannot save results')
+
         tar_time = time.time()
         for path_to_delete in [self._input_tar_path, self._runtime_tar_path]:
             if os.path.exists(path_to_delete):

{pybiolib-1.2.339 → pybiolib-1.2.340.dev1}/biolib/compute_node/job_worker/executors/types.py

@@ -34,9 +34,16 @@ class LocalExecutorOptions(TypedDict):
     cloud_job: Optional[CloudJob]
     large_file_systems: Dict[str, LargeFileSystem]
     module: Module
+    module_input_path: str
+    module_output_path: str
     remote_host_proxies: List[RemoteHostProxy]
     root_job_id: str
     runtime_zip_bytes: Optional[bytes]  # TODO: replace this with a module_source_serialized
     send_status_update: SendStatusUpdateType
     send_system_exception: SendSystemExceptionType
     send_stdout_and_stderr: SendStdoutAndStderrType
+
+class MetadataToSaveOutput(TypedDict):
+    arguments: List[str]
+    startup_error_string: Optional[str]
+    pre_start_diff: List[Dict]

{pybiolib-1.2.339 → pybiolib-1.2.340.dev1}/biolib/compute_node/job_worker/job_worker.py

@@ -400,15 +400,15 @@ class JobWorker:
                 cloud_job=self._root_job_wrapper['cloud_job'],
                 large_file_systems=lfs_dict,
                 module=main_module,
+                module_input_path=module_input_path,
+                module_output_path=main_module_output_path,
                 remote_host_proxies=self._remote_host_proxies,
                 root_job_id=root_job_id,
                 runtime_zip_bytes=runtime_zip_bytes,
                 send_status_update=self._send_status_update,
                 send_system_exception=self.send_system_exception,
                 send_stdout_and_stderr=self.send_stdout_and_stderr,
-            ),
-            module_input_path,
-            main_module_output_path,
+            )
         )
 
         if utils.IS_RUNNING_IN_CLOUD:
@@ -418,11 +418,11 @@ class JobWorker:
     def _run_module(
             self,
             options: LocalExecutorOptions,
-            module_input_path: str,
-            module_output_path: str,
     ) -> None:
         module = options['module']
         job_id = options['job']['public_id']
+        module_output_path = options['module_output_path']
+        module_input_path = options['module_input_path']
         logger_no_user_data.debug(f'Job "{job_id}" running module "{module["name"]}"...')
 
         executor_instance: DockerExecutor
@@ -469,7 +469,7 @@ class JobWorker:
             # Log memory and disk before pulling and executing module
             log_disk_and_memory_usage_info()
 
-        executor_instance.execute_module(module_input_path, module_output_path)
+        executor_instance.execute_module()
 
     def _connect_to_parent(self):
         try:

{pybiolib-1.2.339 → pybiolib-1.2.340.dev1}/biolib/compute_node/webserver/worker_thread.py

@@ -7,6 +7,7 @@ import sys
 import threading
 import time
 from queue import Queue
+from typing import Optional
 
 from biolib import api, utils
 from biolib.biolib_binary_format import ModuleOutputV2, SystemException, SystemStatusUpdate
@@ -55,6 +56,26 @@ class WorkerThread(threading.Thread):
     def _job_temporary_dir(self):
         return self.compute_state['job_temporary_dir']
 
+    def _upload_module_output_and_get_exit_code(self) -> Optional[int]:
+        exit_code = None
+        try:
+            module_output_path = os.path.join(
+                self._job_temporary_dir,
+                JobStorage.module_output_file_name,
+            )
+            if os.path.exists(module_output_path):
+                module_output = ModuleOutputV2(buffer=LocalFileIndexableBuffer(filename=module_output_path))
+                exit_code = module_output.get_exit_code()
+                logger_no_user_data.debug(f'Got exit code: {exit_code}')
+                if utils.IS_RUNNING_IN_CLOUD:
+                    JobStorage.upload_module_output(
+                        job_temporary_dir=self._job_temporary_dir,
+                        job_uuid=self._job_uuid,
+                    )
+        except Exception as error:
+            logger_no_user_data.error(f'Could not upload module output or get exit code: {error}')
+        return exit_code
+
     def run(self):
         try:
             while True:
@@ -77,25 +98,7 @@ class WorkerThread(threading.Thread):
 
                     # If 'Computation Finished'
                     if progress == 94:
-                        # Get Job exit code
-                        try:
-                            module_output_path = os.path.join(
-                                self._job_temporary_dir,
-                                JobStorage.module_output_file_name,
-                            )
-                            module_output = ModuleOutputV2(buffer=LocalFileIndexableBuffer(filename=module_output_path))
-                            self.compute_state['exit_code'] = module_output.get_exit_code()
-                            logger_no_user_data.debug(f"Got exit code: {self.compute_state['exit_code']}")
-
-                        except Exception as error:  # pylint: disable=broad-except
-                            logger_no_user_data.error(f'Could not get exit_code from module output due to: {error}')
-
-                        if utils.IS_RUNNING_IN_CLOUD:
-                            JobStorage.upload_module_output(
-                                job_temporary_dir=self._job_temporary_dir,
-                                job_uuid=self._job_uuid,
-                            )
-
+                        self.compute_state['exit_code'] = self._upload_module_output_and_get_exit_code()
                         self._set_status_update(progress=95, log_message='Result Ready')
                         self.compute_state['is_completed'] = True
                         self.terminate()
@@ -174,13 +177,13 @@ class WorkerThread(threading.Thread):
 
     def terminate(self) -> None:
         cloud_job_uuid = self.compute_state['cloud_job_id']
-        exit_code = self.compute_state.get('exit_code')
         system_exception_code = self.compute_state['status'].get('error_code')
-        if utils.IS_RUNNING_IN_CLOUD:
+
+        if utils.IS_RUNNING_IN_CLOUD and system_exception_code != SystemExceptionCodes.CANCELLED_BY_USER.value:
             CloudUtils.finish_cloud_job(
                 cloud_job_id=cloud_job_uuid,
                 system_exception_code=system_exception_code,
-                exit_code=exit_code,
+                exit_code=self.compute_state.get('exit_code', None),
             )
 
         deregistered_due_to_error = False
@@ -211,6 +214,15 @@ class WorkerThread(threading.Thread):
                     CloudUtils.deregister(error='job_cleanup_timed_out')
                     deregistered_due_to_error = True
 
+            elif system_exception_code == SystemExceptionCodes.CANCELLED_BY_USER.value:
+                self.compute_state['exit_code'] = self._upload_module_output_and_get_exit_code()
+                CloudUtils.finish_cloud_job(
+                    cloud_job_id=cloud_job_uuid,
+                    system_exception_code=system_exception_code,
+                    exit_code=self.compute_state.get('exit_code', None),
+                )
+
+
         # Delete result as error occurred
         if system_exception_code and os.path.exists(self._job_temporary_dir):
             shutil.rmtree(self._job_temporary_dir)

{pybiolib-1.2.339 → pybiolib-1.2.340.dev1}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "pybiolib"
-version = "1.2.339"
+version = "1.2.340.dev1"
 description = "BioLib Python Client"
 readme = "PYPI_README.md"
 license = "MIT"