toil 6.1.0a1__py3-none-any.whl → 8.0.0__py3-none-any.whl

toil/server/wes/tasks.py

@@ -20,7 +20,7 @@ import subprocess
 import sys
 import tempfile
 import zipfile
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any, Optional, Union
 from urllib.parse import urldefrag
 
 from celery.exceptions import SoftTimeLimitExceeded  # type: ignore
@@ -29,13 +29,15 @@ import toil.server.wes.amazon_wes_utils as amazon_wes_utils
 from toil.common import Toil
 from toil.jobStores.utils import generate_locator
 from toil.server.celery_app import celery
-from toil.server.utils import (WorkflowStateMachine,
-                               connect_to_workflow_state_store,
-                               download_file_from_internet,
-                               download_file_from_s3,
-                               get_file_class,
-                               get_iso_time,
-                               link_file)
+from toil.server.utils import (
+    WorkflowStateMachine,
+    connect_to_workflow_state_store,
+    download_file_from_internet,
+    download_file_from_s3,
+    get_file_class,
+    get_iso_time,
+    link_file,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -45,6 +47,7 @@ logger = logging.getLogger(__name__)
 # our patience for CANCELING WES workflows to time out to CANCELED.
 WAIT_FOR_DEATH_TIMEOUT = 20
 
+
 class ToilWorkflowRunner:
     """
     A class to represent a workflow runner to run the requested workflow.
@@ -53,7 +56,14 @@ class ToilWorkflowRunner:
     that command, and collecting the outputs of the resulting workflow run.
     """
 
-    def __init__(self, base_scratch_dir: str, state_store_url: str, workflow_id: str, request: Dict[str, Any], engine_options: List[str]):
+    def __init__(
+        self,
+        base_scratch_dir: str,
+        state_store_url: str,
+        workflow_id: str,
+        request: dict[str, Any],
+        engine_options: list[str],
+    ):
         """
         Make a new ToilWorkflowRunner to actually run a workflow leader based
         on a WES request.
@@ -85,8 +95,11 @@ class ToilWorkflowRunner:
         self.out_dir = os.path.join(self.scratch_dir, "outputs")
 
         # Compose the right kind of job store to use it the user doesn't specify one.
-        default_type = os.getenv('TOIL_WES_JOB_STORE_TYPE', 'file')
-        self.default_job_store = generate_locator(default_type, local_suggestion=os.path.join(self.scratch_dir, "toil_job_store"))
+        default_type = os.getenv("TOIL_WES_JOB_STORE_TYPE", "file")
+        self.default_job_store = generate_locator(
+            default_type,
+            local_suggestion=os.path.join(self.scratch_dir, "toil_job_store"),
+        )
 
         self.job_store = self.default_job_store
 
@@ -112,7 +125,9 @@ class ToilWorkflowRunner:
             logger.info("Linking workflow from filesystem.")
             link_file(src=src_url[7:], dest=dest)
         elif src_url.startswith(("http://", "https://")):
-            logger.info("Downloading workflow_url from the Internet.")
+            logger.info(
+                f"Downloading workflow_url from the Internet.  src: {src_url} dst: {dest}"
+            )
             download_file_from_internet(src=src_url, dest=dest, content_type="text/")
         elif src_url.startswith("s3://"):
             logger.info("Downloading workflow_url from Amazon S3.")
@@ -127,7 +142,9 @@ class ToilWorkflowRunner:
 
         return dest
 
-    def sort_options(self, workflow_engine_parameters: Optional[Dict[str, Optional[str]]] = None) -> List[str]:
+    def sort_options(
+        self, workflow_engine_parameters: Optional[dict[str, Optional[str]]] = None
+    ) -> list[str]:
         """
         Sort the command line arguments in the order that can be recognized by
         the workflow execution engine.
@@ -188,7 +205,7 @@ class ToilWorkflowRunner:
 
         return options
 
-    def initialize_run(self) -> List[str]:
+    def initialize_run(self) -> list[str]:
         """
         Write workflow and input files and construct a list of shell commands
         to be executed. Return that list of shell commands that should be
@@ -204,56 +221,69 @@ class ToilWorkflowRunner:
         # Obtain main workflow file to a path (no-scheme file URL).
         workflow_url = self.write_workflow(src_url=self.request["workflow_url"])
 
-        if os.path.basename(workflow_url) == "workflow.zip" and zipfile.is_zipfile(workflow_url):
+        if os.path.basename(workflow_url) == "workflow.zip" and zipfile.is_zipfile(
+            workflow_url
+        ):
             # We've been sent a zip file. We should interpret this as an Amazon Genomics CLI-style zip file.
             # Extract everything next to the zip and find and open relvant files.
             logger.info("Extracting and parsing Amazon-style workflow bundle...")
-            zip_info = amazon_wes_utils.parse_workflow_zip_file(workflow_url, self.wf_type)
+            zip_info = amazon_wes_utils.parse_workflow_zip_file(
+                workflow_url, self.wf_type
+            )
 
             # Now parse Amazon's internal format into our own.
 
             # Find the real workflow source for the entrypoint file
-            if 'workflowSource' in zip_info['files']:
+            if "workflowSource" in zip_info["files"]:
                 # They sent a file, which has been opened, so grab its path
-                workflow_url = zip_info['files']['workflowSource'].name
+                workflow_url = zip_info["files"]["workflowSource"].name
                 logger.info("Workflow source file: '%s'", workflow_url)
-            elif 'workflowUrl' in zip_info['data']:
+            elif "workflowUrl" in zip_info["data"]:
                 # They are pointing to another URL.
                 # TODO: What does Amazon expect this to mean? Are we supposed to recurse?
                 # For now just forward it.
-                workflow_url = zip_info['data']['workflowUrl']
+                workflow_url = zip_info["data"]["workflowUrl"]
                 logger.info("Workflow reference URL: '%s'", workflow_url)
             else:
                 # The parser is supposed to throw so we can't get here
-                raise RuntimeError("Parser could not find workflow source or URL in zip")
-
-            if 'workflowInputFiles' in zip_info['files'] and len(zip_info['files']['workflowInputFiles']) > 0:
+                raise RuntimeError(
+                    "Parser could not find workflow source or URL in zip"
+                )
+
+            if (
+                "workflowInputFiles" in zip_info["files"]
+                and len(zip_info["files"]["workflowInputFiles"]) > 0
+            ):
                 # The bundle contains a list of input files.
                 # We interpret these as JSON, and layer them on top of each
                 # other, and then apply workflow_params from the request.
-                logger.info("Workflow came with %d bundled inputs files; coalescing final parameters",
-                            len(zip_info['files']['workflowInputFiles']))
+                logger.info(
+                    "Workflow came with %d bundled inputs files; coalescing final parameters",
+                    len(zip_info["files"]["workflowInputFiles"]),
+                )
 
                 coalesced_parameters = {}
-                for binary_file in zip_info['files']['workflowInputFiles']:
+                for binary_file in zip_info["files"]["workflowInputFiles"]:
                     try:
                         # Read each input file as a JSON
                         loaded_parameters = json.load(binary_file)
                     except json.JSONDecodeError as e:
-                        raise RuntimeError(f"Could not parse inputs JSON {os.path.basename(binary_file.name)}: {e}")
+                        raise RuntimeError(
+                            f"Could not parse inputs JSON {os.path.basename(binary_file.name)}: {e}"
+                        )
                     # And merge them together in order
                     coalesced_parameters.update(loaded_parameters)
                 # Then apply and replace the parameters that came with the request
                 coalesced_parameters.update(workflow_params)
                 workflow_params = coalesced_parameters
 
-            if 'workflowOptions' in zip_info['files']:
+            if "workflowOptions" in zip_info["files"]:
                 # The bundle contains an options JSON. We interpret these as
                 # defaults for workflow_engine_parameters.
                 logger.info(f"Workflow came with bundled options JSON")
                 try:
                     # Read as a JSON
-                    loaded_options = json.load(zip_info['files']['workflowOptions'])
+                    loaded_options = json.load(zip_info["files"]["workflowOptions"])
                 except json.JSONDecodeError as e:
                     raise RuntimeError(f"Could not parse options JSON: {e}")
                 # Apply and replace the engine parameters that came with the
@@ -275,13 +305,9 @@ class ToilWorkflowRunner:
 
         # construct and return the command to run
         if self.wf_type == "cwl":
-            command_args = (
-                ["toil-cwl-runner"] + options + [workflow_url, input_json]
-            )
+            command_args = ["toil-cwl-runner"] + options + [workflow_url, input_json]
         elif self.wf_type == "wdl":
-            command_args = (
-                ["toil-wdl-runner"] + options + [workflow_url, input_json]
-            )
+            command_args = ["toil-wdl-runner"] + options + [workflow_url, input_json]
         elif self.wf_type == "py":
             command_args = ["python", workflow_url] + options
         else:
@@ -290,7 +316,9 @@ class ToilWorkflowRunner:
 
         return command_args
 
-    def call_cmd(self, cmd: Union[List[str], str], cwd: str) -> "subprocess.Popen[bytes]":
+    def call_cmd(
+        self, cmd: Union[list[str], str], cwd: str
+    ) -> "subprocess.Popen[bytes]":
         """
         Calls a command with Popen. Writes stdout, stderr, and the command to
         separate files.
@@ -300,7 +328,9 @@ class ToilWorkflowRunner:
 
         with open(stdout_f, "w") as stdout, open(stderr_f, "w") as stderr:
             logger.info(f"Calling: '{' '.join(cmd)}'")
-            process = subprocess.Popen(cmd, stdout=stdout, stderr=stderr, close_fds=True, cwd=cwd)
+            process = subprocess.Popen(
+                cmd, stdout=stdout, stderr=stderr, close_fds=True, cwd=cwd
+            )
 
         return process
 
@@ -348,7 +378,7 @@ class ToilWorkflowRunner:
         self.store.set("end_time", get_iso_time())
         self.store.set("exit_code", str(exit_code))
 
-        logger.info('Toil child finished with code %s', exit_code)
+        logger.info("Toil child finished with code %s", exit_code)
 
         if exit_code == 0:
             self.state_machine.send_complete()
@@ -387,7 +417,14 @@ class ToilWorkflowRunner:
 
         self.write_scratch_file("outputs.json", json.dumps(output_obj))
 
-def run_wes_task(base_scratch_dir: str, state_store_url: str, workflow_id: str, request: Dict[str, Any], engine_options: List[str]) -> str:
+
+def run_wes_task(
+    base_scratch_dir: str,
+    state_store_url: str,
+    workflow_id: str,
+    request: dict[str, Any],
+    engine_options: list[str],
+) -> str:
     """
     Run a requested workflow.
 
@@ -402,8 +439,13 @@ def run_wes_task(base_scratch_dir: str, state_store_url: str, workflow_id: str,
 
     logger.info("Starting WES workflow")
 
-    runner = ToilWorkflowRunner(base_scratch_dir, state_store_url, workflow_id,
-                                request=request, engine_options=engine_options)
+    runner = ToilWorkflowRunner(
+        base_scratch_dir,
+        state_store_url,
+        workflow_id,
+        request=request,
+        engine_options=engine_options,
+    )
 
     try:
         runner.run()
@@ -416,24 +458,27 @@ def run_wes_task(base_scratch_dir: str, state_store_url: str, workflow_id: str,
             runner.write_output_files()
     except (KeyboardInterrupt, SystemExit, SoftTimeLimitExceeded):
         # We canceled the workflow run
-        logger.info('Canceling the workflow')
+        logger.info("Canceling the workflow")
         runner.state_machine.send_canceled()
     except Exception:
         # The workflow run broke. We still count as the executor here.
-        logger.exception('Running Toil produced an exception.')
+        logger.exception("Running Toil produced an exception.")
         runner.state_machine.send_executor_error()
         raise
 
     return runner.get_state()
 
+
 # Wrap the task function as a Celery task
 run_wes = celery.task(name="run_wes")(run_wes_task)
 
+
 def cancel_run(task_id: str) -> None:
     """
     Send a SIGTERM signal to the process that is running task_id.
     """
-    celery.control.terminate(task_id, signal='SIGUSR1')
+    celery.control.terminate(task_id, signal="SIGUSR1")
+
 
 class TaskRunner:
     """
@@ -443,13 +488,13 @@ class TaskRunner:
     """
 
     @staticmethod
-    def run(args: Tuple[str, str, str, Dict[str, Any], List[str]], task_id: str) -> None:
+    def run(
+        args: tuple[str, str, str, dict[str, Any], list[str]], task_id: str
+    ) -> None:
         """
         Run the given task args with the given ID on Celery.
         """
-        run_wes.apply_async(args=args,
-                            task_id=task_id,
-                            ignore_result=True)
+        run_wes.apply_async(args=args, task_id=task_id, ignore_result=True)
 
     @staticmethod
     def cancel(task_id: str) -> None:
@@ -467,8 +512,10 @@ class TaskRunner:
         # Nothing to do for Celery
         return True
 
+
 # If Celery can't be set up, we can just use this fake version instead.
 
+
 class MultiprocessingTaskRunner(TaskRunner):
     """
     Version of TaskRunner that just runs tasks with Multiprocessing.
@@ -478,11 +525,13 @@ class MultiprocessingTaskRunner(TaskRunner):
     ToilWorkflowRunner) don't poll for it.
     """
 
-    _id_to_process: Dict[str, multiprocessing.Process] = {}
-    _id_to_log: Dict[str, str] = {}
+    _id_to_process: dict[str, multiprocessing.Process] = {}
+    _id_to_log: dict[str, str] = {}
 
     @staticmethod
-    def set_up_and_run_task(output_path: str, args: Tuple[str, str, str, Dict[str, Any], List[str]]) -> None:
+    def set_up_and_run_task(
+        output_path: str, args: tuple[str, str, str, dict[str, Any], list[str]]
+    ) -> None:
         """
         Set up logging for the process into the given file and then call
         run_wes_task with the given arguments.
@@ -499,8 +548,8 @@ class MultiprocessingTaskRunner(TaskRunner):
         # that we were told about, so the server can come get the log if we
         # unexpectedly die.
 
-        output_file = open(output_path, 'w')
-        output_file.write('Initializing task log\n')
+        output_file = open(output_path, "w")
+        output_file.write("Initializing task log\n")
         output_file.flush()
 
         # Take over logging.
@@ -523,22 +572,24 @@ class MultiprocessingTaskRunner(TaskRunner):
         signal.signal(signal.SIGTERM, handle_sigterm)
 
         try:
-            logger.info('Running task')
+            logger.info("Running task")
             output_file.flush()
             run_wes_task(*args)
         except Exception:
-            logger.exception('Exception in task!')
+            logger.exception("Exception in task!")
             raise
         else:
             # If the task does not crash, clean up the log
             os.unlink(output_path)
         finally:
-            logger.debug('Finishing task log')
+            logger.debug("Finishing task log")
             output_file.flush()
             output_file.close()
 
     @classmethod
-    def run(cls, args: Tuple[str, str, str, Dict[str, Any], List[str]], task_id: str) -> None:
+    def run(
+        cls, args: tuple[str, str, str, dict[str, Any], list[str]], task_id: str
+    ) -> None:
         """
         Run the given task args with the given ID.
         """
@@ -549,9 +600,13 @@ class MultiprocessingTaskRunner(TaskRunner):
         # Store the log filename before the process, like is_ok() expects.
         cls._id_to_log[task_id] = path
 
-        logger.info("Starting task %s in a process that should log to %s", task_id, path)
+        logger.info(
+            "Starting task %s in a process that should log to %s", task_id, path
+        )
 
-        cls._id_to_process[task_id] = multiprocessing.Process(target=cls.set_up_and_run_task, args=(path, args))
+        cls._id_to_process[task_id] = multiprocessing.Process(
+            target=cls.set_up_and_run_task, args=(path, args)
+        )
         cls._id_to_process[task_id].start()
 
     @classmethod
@@ -581,9 +636,14 @@ class MultiprocessingTaskRunner(TaskRunner):
         # being canceled by cancel(), then it is OK.
         ACCEPTABLE_EXIT_CODES = [0, -signal.SIGTERM]
 
-        if process.exitcode is not None and process.exitcode not in ACCEPTABLE_EXIT_CODES:
+        if (
+            process.exitcode is not None
+            and process.exitcode not in ACCEPTABLE_EXIT_CODES
+        ):
             # Something went wring in the task and it couldn't handle it.
-            logger.error("Process for running %s failed with code %s", task_id, process.exitcode)
+            logger.error(
+                "Process for running %s failed with code %s", task_id, process.exitcode
+            )
             try:
                 for line in open(cls._id_to_log[task_id]):
                     # Dump the task log
@@ -596,5 +656,3 @@ class MultiprocessingTaskRunner(TaskRunner):
             return False
 
         return True
-
-