virtool-workflow 6.0.0a8__py3-none-any.whl → 6.0.0a9__py3-none-any.whl

virtool_workflow/data/ml.py

@@ -1,4 +1,6 @@
 """A fixture and dataclass for working with machine learning models in workflows."""
+import asyncio
+import shutil
 from dataclasses import dataclass
 from pathlib import Path
 
@@ -8,15 +10,14 @@ from virtool_core.models.ml import MLModelRelease
 
 from virtool_workflow.api.client import APIClient
 from virtool_workflow.data.analyses import WFAnalysis
-from virtool_workflow.utils import make_directory
+from virtool_workflow.utils import make_directory, move_all_model_files, untar
 
 logger = get_logger("api")
 
 
 @dataclass
 class WFMLModelRelease:
-    """
-    A machine learning model.
+    """A machine learning model.
 
     This class represents a machine learning model and the selected release of that
     model in the workflow.
@@ -39,7 +40,9 @@ class WFMLModelRelease:
 
 @fixture
 async def ml(
-    _api: APIClient, analysis: WFAnalysis, work_path: Path
+    _api: APIClient,
+    analysis: WFAnalysis,
+    work_path: Path,
 ) -> WFMLModelRelease | None:
     if analysis.ml is None:
         return None
@@ -50,11 +53,11 @@ async def ml(
     log = logger.bind(model_id=analysis.ml.id, model_release_id=model_release_id)
 
     model_release_json = await _api.get_json(
-        f"/ml/{model_id}/releases/{model_release_id}"
+        f"/ml/{model_id}/releases/{model_release_id}",
     )
     model_release = MLModelRelease(**model_release_json)
 
-    log.info("Fetched ML model release json")
+    log.info("fetched ml model release json")
 
     release = WFMLModelRelease(
         id=model_release.id,
@@ -65,9 +68,14 @@ async def ml(
     await make_directory(release.path)
 
     await _api.get_file(
-        f"/ml/{model_id}/releases/{model_release_id}/model.tar.gz", release.file_path
+        f"/ml/{model_id}/releases/{model_release_id}/model.tar.gz",
+        release.file_path,
     )
 
-    log.info("Downloaded ML model data file")
+    await asyncio.to_thread(untar, release.file_path, release.path)
+    await asyncio.to_thread(move_all_model_files, release.path / "model", release.path)
+    await asyncio.to_thread(shutil.rmtree, release.path / "model")
+
+    log.info("downloaded ml model release file")
 
     return release

virtool_workflow/data/samples.py

@@ -7,13 +7,13 @@ from pyfixtures import fixture
 from structlog import get_logger
 from virtool_core.models.enums import LibraryType
 from virtool_core.models.job import Job
-from virtool_core.models.samples import Sample, Quality
+from virtool_core.models.samples import Quality, Sample
 
 from virtool_workflow.analysis.utils import ReadPaths
 from virtool_workflow.api.client import APIClient
-from virtool_workflow.files import VirtoolFileFormat
 from virtool_workflow.data.uploads import WFUploads
 from virtool_workflow.errors import JobsAPINotFound
+from virtool_workflow.files import VirtoolFileFormat
 
 logger = get_logger("api")
 
@@ -42,8 +42,7 @@ class WFSample:
 
     @property
     def min_length(self) -> int | None:
-        """
-        The minimum observed read length in the sample sequencing data.
+        """The minimum observed read length in the sample sequencing data.
 
         Returns ``None`` if the sample is still being created and no quality data is available.
 
@@ -52,8 +51,7 @@ class WFSample:
 
     @property
     def max_length(self) -> int | None:
-        """
-        The maximum observed read length in the sample sequencing data.
+        """The maximum observed read length in the sample sequencing data.
 
         Returns ``None`` if the sample is still being created and no quality data is available.
 
@@ -85,10 +83,9 @@ class WFNewSample:
 
 @fixture
 async def sample(
-    _api: APIClient, job: Job, uploads: WFUploads, work_path: Path
+    _api: APIClient, job: Job, uploads: WFUploads, work_path: Path,
 ) -> WFSample:
     """The sample associated with the current job."""
-
     id_ = job.args["sample_id"]
 
     base_url_path = f"/samples/{id_}"
@@ -104,7 +101,7 @@ async def sample(
     await asyncio.to_thread(reads_path.mkdir, exist_ok=True, parents=True)
 
     await _api.get_file(
-        f"{base_url_path}/reads/reads_1.fq.gz", reads_path / "reads_1.fq.gz"
+        f"{base_url_path}/reads/reads_1.fq.gz", reads_path / "reads_1.fq.gz",
     )
 
     if sample.paired:
@@ -113,7 +110,7 @@ async def sample(
             reads_path / "reads_2.fq.gz",
         )
         await _api.get_file(
-            f"{base_url_path}/reads/reads_2.fq.gz", reads_path / "reads_2.fq.gz"
+            f"{base_url_path}/reads/reads_2.fq.gz", reads_path / "reads_2.fq.gz",
         )
     else:
         read_paths = (reads_path / "reads_1.fq.gz",)
@@ -130,10 +127,9 @@ async def sample(
 
 @fixture
 async def new_sample(
-    _api: APIClient, job: Job, uploads: WFUploads, work_path: Path
+    _api: APIClient, job: Job, uploads: WFUploads, work_path: Path,
 ) -> WFNewSample:
     """The sample associated with the current job."""
-
     id_ = job.args["sample_id"]
 
     log = logger.bind(resource="sample", id=id_)

virtool_workflow/data/subtractions.py

@@ -7,8 +7,8 @@ from pyfixtures import fixture
 from structlog import get_logger
 from virtool_core.models.job import Job
 from virtool_core.models.subtraction import (
-    Subtraction,
     NucleotideComposition,
+    Subtraction,
     SubtractionFile,
 )
 
@@ -22,8 +22,7 @@ logger = get_logger("api")
 
 @dataclass
 class WFSubtraction:
-    """
-    A Virtool subtraction that has been loaded into the workflow environment.
+    """A Virtool subtraction that has been loaded into the workflow environment.
 
     The subtraction files are downloaded to the workflow's local work path so they can
     be used for analysis.
@@ -47,7 +46,7 @@ class WFSubtraction:
     path: Path
     """
     The path to the subtraction directory.
- 
+
     The subtraction directory contains the FASTA and Bowtie2 files for the subtraction.
     """
 
@@ -58,8 +57,7 @@ class WFSubtraction:
 
     @property
     def bowtie2_index_path(self) -> Path:
-        """
-        The path to Bowtie2 prefix in the running workflow's work_path
+        """The path to Bowtie2 prefix in the running workflow's work_path
 
         For example, ``/work/subtractions/<id>/subtraction`` refers to the Bowtie2
         index that comprises the files:
@@ -85,7 +83,7 @@ class WFNewSubtraction:
     A callable that deletes the subtraction from Virtool.
 
     This should be called if the subtraction creation fails before the subtraction is
-    finalized.    
+    finalized.
     """
 
     finalize: Callable[[dict[str, int | float], int], Coroutine[None, None, None]]
@@ -120,10 +118,9 @@ class WFNewSubtraction:
 
 @fixture
 async def subtractions(
-    _api: APIClient, analysis: WFAnalysis, work_path: Path
+    _api: APIClient, analysis: WFAnalysis, work_path: Path,
 ) -> list[WFSubtraction]:
     """The subtractions to be used for the current analysis job."""
-
     subtraction_work_path = work_path / "subtractions"
     await asyncio.to_thread(subtraction_work_path.mkdir)
 
@@ -161,10 +158,9 @@ async def subtractions(
 
 @fixture
 async def new_subtraction(
-    _api: APIClient, job: Job, uploads: WFUploads, work_path: Path
+    _api: APIClient, job: Job, uploads: WFUploads, work_path: Path,
 ) -> WFNewSubtraction:
-    """
-    A new subtraction that will be created during the current job.
+    """A new subtraction that will be created during the current job.
 
     Currently only used for the `create-subtraction` workflow.
     """
@@ -193,8 +189,7 @@ async def new_subtraction(
         await _api.delete(f"/subtractions/{subtraction_.id}")
 
     async def finalize(gc: dict[str, int | float], count: int):
-        """
-        Finalize the subtraction by setting the gc.
+        """Finalize the subtraction by setting the gc.
 
         :param gc: the nucleotide composition of the subtraction
         :param count: the number of sequences in the FASTA file
@@ -205,8 +200,7 @@ async def new_subtraction(
         await _api.patch_json(url_path, {"gc": gc_.dict(), "count": count})
 
     async def upload(path: Path):
-        """
-        Upload a file relating to this subtraction.
+        """Upload a file relating to this subtraction.
 
         Filenames must be one of:
             - subtraction.fa.gz
@@ -227,7 +221,7 @@ async def new_subtraction(
         log.info("Uploading subtraction file")
 
         await _api.put_file(
-            f"/subtractions/{subtraction_.id}/files/{filename}", path, "unknown"
+            f"/subtractions/{subtraction_.id}/files/{filename}", path, "unknown",
         )
 
         log.info("Finished uploading subtraction file")

virtool_workflow/data/uploads.py

@@ -18,14 +18,13 @@ class WFUploads:
 
 @fixture
 async def uploads(_api: APIClient) -> WFUploads:
-    """
-    Provides access to files that have been uploaded to the Virtool instance.
+    """Provides access to files that have been uploaded to the Virtool instance.
 
     Files can be downloaded into the workflow environment be calling
     :meth:`.WFUploads.download`.
 
     Example:
-
+    -------
     .. code-block:: python
 
         @step

virtool_workflow/decorators.py

@@ -6,8 +6,7 @@ from virtool_workflow.workflow import Workflow
 
 
 def step(func: Callable = None, *, name: str | None = None) -> Callable:
-    """
-    Mark a function as a workflow step function.
+    """Mark a function as a workflow step function.
 
     :param func: the workflow step function
     :param name: the display name of the workflow step. A name
@@ -23,13 +22,11 @@ def step(func: Callable = None, *, name: str | None = None) -> Callable:
 
 
 def collect(module: ModuleType) -> Workflow:
-    """
-    Build a :class:`.Workflow` object from a workflow module.
+    """Build a :class:`.Workflow` object from a workflow module.
 
     :param module: A workflow module
     :return: A workflow object
     """
-
     workflow = Workflow()
 
     markers = [

virtool_workflow/errors.py

@@ -4,7 +4,6 @@ from subprocess import SubprocessError
 class IllegalJobArguments(ValueError):
     """The `job.args` dict is in an illegal state."""
 
-    ...
 
 
 class InsufficientJobRights(Exception):
@@ -14,7 +13,7 @@ class InsufficientJobRights(Exception):
 class JobAlreadyAcquired(Exception):
     def __init__(self, job_id: str):
         super(JobAlreadyAcquired, self).__init__(
-            f"Job {job_id} is has already been acquired."
+            f"Job {job_id} is has already been acquired.",
         )
 
 
@@ -25,40 +24,33 @@ class JobAlreadyFinalized(Exception):
 class JobsAPIError(Exception):
     """A base exception for errors due to HTTP errors from the jobs API."""
 
-    ...
 
 
 class JobsAPIBadRequest(JobsAPIError):
     """A ``400 Bad Request`` response from the jobs API."""
 
     status = 400
-    ...
 
 
 class JobsAPIForbidden(JobsAPIError):
     status = 403
-    ...
 
 
 class JobsAPINotFound(JobsAPIError):
     status = 404
-    ...
 
 
 class JobsAPIConflict(JobsAPIError):
     status = 409
-    ...
 
 
 class JobsAPIServerError(JobsAPIError):
     status = 500
-    ...
 
 
 class MissingJobArgument(ValueError):
     """The `job.args` dict is missing a required key for some funcionality."""
 
-    ...
 
 
 class NotFound(KeyError):

virtool_workflow/hooks.py

@@ -1,5 +1,4 @@
-"""
-Hooks provide a way to do things when events happen during the workflow lifecycle.
+"""Hooks provide a way to do things when events happen during the workflow lifecycle.
 """
 
 from virtool_workflow.runtime.hook import Hook
@@ -126,8 +125,7 @@ __all__ = [
 
 
 def cleanup_builtin_status_hooks():
-    """
-    Clear callbacks for built-in status hooks.
+    """Clear callbacks for built-in status hooks.
 
     This prevents carryover of hooks between tests. Carryover won't be encountered in
     production because workflow processes exit after one run.

virtool_workflow/pytest_plugin/data.py

@@ -15,8 +15,8 @@ from virtool_core.models.samples import (
 )
 from virtool_core.models.subtraction import (
     Subtraction,
-    SubtractionNested,
     SubtractionFile,
+    SubtractionNested,
 )
 
 from virtool_workflow.pytest_plugin.utils import SUBTRACTION_FILENAMES
@@ -55,9 +55,9 @@ class Data:
     """An un-finalized subtraction for testing subtraction creation workflows."""
 
 
-@pytest.fixture
+@pytest.fixture()
 def data(
-    virtool_workflow_example_path: Path, static_datetime: datetime.datetime
+    virtool_workflow_example_path: Path, static_datetime: datetime.datetime,
 ) -> Data:
     class AnalysisFactory(ModelFactory):
         __model__ = Analysis

virtool_workflow/pytest_plugin/subprocess.py

@@ -3,6 +3,6 @@ import pytest
 import virtool_workflow.runtime.run_subprocess
 
 
-@pytest.fixture
+@pytest.fixture()
 def run_subprocess() -> virtool_workflow.runtime.run_subprocess.RunSubprocess:
     return virtool_workflow.runtime.run_subprocess.run_subprocess()

virtool_workflow/runtime/discover.py

@@ -1,6 +1,6 @@
 import sys
 from importlib import import_module
-from importlib.util import spec_from_file_location, module_from_spec
+from importlib.util import module_from_spec, spec_from_file_location
 from pathlib import Path
 from types import ModuleType
 
@@ -13,8 +13,7 @@ logger = get_logger("runtime")
 
 
 def discover_workflow(path: Path) -> Workflow:
-    """
-    Find an instance of :class:`.Workflow` in the python module located at the given
+    """Find an instance of :class:`.Workflow` in the python module located at the given
     path.
 
     :param path: The path to a Python module.
@@ -32,8 +31,7 @@ def discover_workflow(path: Path) -> Workflow:
 
 
 def load_builtin_fixtures():
-    """
-    Load built-in fixtures.
+    """Load built-in fixtures.
 
     This function is called before any fixtures defined in a workflow's
     ``fixtures.py`` file. It is used to provide built-in fixtures that are
@@ -59,8 +57,7 @@ def load_custom_fixtures():
 
 
 def load_workflow_from_file() -> Workflow:
-    """
-    Load a workflow from a Python file at ``./workflow.py`` and return a :class:`.Workflow` object.
+    """Load a workflow from a Python file at ``./workflow.py`` and return a :class:`.Workflow` object.
 
     :raises FileNotFoundError: If no workflow.py file is found.
     :return: The workflow.
@@ -75,8 +72,7 @@ def load_workflow_from_file() -> Workflow:
 
 
 def import_module_from_file(module_name: str, path: Path) -> ModuleType:
-    """
-    Import a module from a file.
+    """Import a module from a file.
 
     The parent directory of `path` will also be added to `sys.path` prior to importing.
     This ensures that modules and packages defined in that directory can be properly

virtool_workflow/runtime/hook.py

@@ -16,8 +16,7 @@ class Hook:
     """Used to hook into the workflow lifecycle."""
 
     def __init__(self, hook_name: str):
-        """
-        A set of functions to be called as a group upon a particular event.
+        """A set of functions to be called as a group upon a particular event.
 
         The signature of any functions added via :func:`Hook.callback` or
         :func:`Hook.__call__` are validated to match the types provided.
@@ -31,8 +30,7 @@ class Hook:
         self.clear = self.callbacks.clear
 
     def __call__(self, callback_: Callable = None, until=None, once=False):
-        """
-        Add a callback function to this Hook that will be called when the hook is
+        """Add a callback function to this Hook that will be called when the hook is
         triggered.
 
         :param callback_: The callback function to register.
@@ -75,8 +73,7 @@ class Hook:
         return _temporary_callback
 
     async def trigger(self, scope: FixtureScope, suppress=False, **kwargs) -> list[Any]:
-        """
-        Trigger the hook.
+        """Trigger the hook.
 
         Bind fixtures from `scope` to each callback function and invoke them.
 
@@ -122,7 +119,7 @@ class Hook:
                 return await callback(*args, **kwargs)
 
         results = await gather(
-            *[call_callback(callback) for callback in callbacks], return_exceptions=True
+            *[call_callback(callback) for callback in callbacks], return_exceptions=True,
         )
 
         for error in results:

virtool_workflow/runtime/path.py

@@ -1,7 +1,7 @@
 import asyncio
 from contextlib import asynccontextmanager
-from shutil import rmtree
 from pathlib import Path
+from shutil import rmtree
 
 from virtool_workflow.runtime.config import RunConfig
 

virtool_workflow/runtime/ping.py

@@ -35,8 +35,7 @@ async def _ping_periodically(api: APIClient, job_id: str):
 
 @asynccontextmanager
 async def ping_periodically(api: APIClient, job_id: str):
-    """
-    Ping the API to keep the job alive.
+    """Ping the API to keep the job alive.
 
     While the context manager is open, a task runs that pings the API every 5 seconds.
     When the context manager is closed, the task is cleanly cancelled.

virtool_workflow/runtime/run.py

@@ -18,20 +18,20 @@ from virtool_workflow.api.client import api_client
 from virtool_workflow.hooks import (
     cleanup_builtin_status_hooks,
     on_cancelled,
-    on_success,
-    on_step_start,
-    on_terminated,
     on_error,
     on_failure,
-    on_result,
     on_finish,
-    on_workflow_start,
+    on_result,
     on_step_finish,
+    on_step_start,
+    on_success,
+    on_terminated,
+    on_workflow_start,
 )
 from virtool_workflow.runtime.config import RunConfig
 from virtool_workflow.runtime.discover import (
-    load_custom_fixtures,
     load_builtin_fixtures,
+    load_custom_fixtures,
     load_workflow_from_file,
 )
 from virtool_workflow.runtime.events import Events
@@ -48,8 +48,7 @@ logger = get_logger("runtime")
 
 
 def configure_status_hooks():
-    """
-    Configure built-in job status hooks.
+    """Configure built-in job status hooks.
 
     Push status updates to API when various lifecycle hooks are triggered.
 
@@ -77,8 +76,7 @@ def configure_status_hooks():
 
 
 async def execute(workflow: Workflow, scope: FixtureScope, events: Events):
-    """
-    Execute a workflow.
+    """Execute a workflow.
 
     :param workflow: The workflow to execute
     :param scope: The :class:`FixtureScope` to use for fixture injection
@@ -118,7 +116,7 @@ async def execute(workflow: Workflow, scope: FixtureScope, events: Events):
 
             if not events.terminated.is_set():
                 logger.warning(
-                    "workflow terminated without sigterm. this should not happen."
+                    "workflow terminated without sigterm. this should not happen.",
                 )
 
             await asyncio.gather(
@@ -162,7 +160,7 @@ async def run_workflow(
     job = await acquire_job_by_id(config.jobs_api_connection_string, job_id)
 
     async with api_client(
-        config.jobs_api_connection_string, job.id, job.key
+        config.jobs_api_connection_string, job.id, job.key,
     ) as api, FixtureScope() as scope:
         # These fixtures should not be used directly by the workflow. They are used
         # by other built-in fixtures.
@@ -200,8 +198,7 @@ async def start_runtime(
     work_path: Path,
     workflow_loader: Callable[[], Workflow] = load_workflow_from_file,
 ):
-    """
-    Start the workflow runtime.
+    """Start the workflow runtime.
 
     The runtime loads the workflow and fixtures. It then waits for a job ID to be pushed
     to the configured Redis list.
@@ -261,7 +258,7 @@ async def start_runtime(
             job_id,
             workflow,
             events,
-        )
+        ),
     )
 
     def terminate_workflow(*_):
@@ -278,7 +275,7 @@ async def start_runtime(
 
     async with configure_redis(redis_connection_string) as redis:
         cancellation_task = asyncio.create_task(
-            wait_for_cancellation(redis, job_id, cancel_workflow)
+            wait_for_cancellation(redis, job_id, cancel_workflow),
         )
 
         await run_workflow_task

virtool_workflow/runtime/run_subprocess.py

@@ -15,12 +15,11 @@ logger = get_logger("subprocess")
 
 class LineOutputHandler(Protocol):
     async def __call__(self, line: str):
-        """
-        Handle input from stdin, or stderr, line by line.
+        """Handle input from stdin, or stderr, line by line.
 
         :param line: A line of output from the stream.
         """
-        raise NotImplementedError()
+        raise NotImplementedError
 
 
 class RunSubprocess(Protocol):
@@ -32,8 +31,7 @@ class RunSubprocess(Protocol):
         stderr_handler: LineOutputHandler | None = None,
         stdout_handler: LineOutputHandler | None = None,
     ) -> Process:
-        """
-        Run a shell command in a subprocess.
+        """Run a shell command in a subprocess.
 
         :param command: A shell command
         :param stdout_handler: A function to handle stdout output line by line
@@ -43,14 +41,13 @@ class RunSubprocess(Protocol):
         :raise SubprocessFailed: The subprocess has exited with a non-zero exit code
         :return: An :class:`.Process` instance
         """
-        raise NotImplementedError()
+        raise NotImplementedError
 
 
 async def watch_pipe(
-    stream: asyncio.StreamReader, handler: Callable[[bytes], Awaitable[None]]
+    stream: asyncio.StreamReader, handler: Callable[[bytes], Awaitable[None]],
 ):
-    """
-    Watch the stdout or stderr stream and pass lines to the `handler` callback function.
+    """Watch the stdout or stderr stream and pass lines to the `handler` callback function.
 
     :param stream: a stdout or stderr file object
     :param handler: a handler coroutine for output lines
@@ -127,7 +124,7 @@ async def _run_subprocess(
     if process.returncode not in [0, 15, -15]:
         raise SubprocessFailed(
             f"{command[0]} failed with exit code {process.returncode}\n"
-            f"arguments: {command}\n"
+            f"arguments: {command}\n",
         )
 
     log.info(

virtool_workflow/runtime/sentry.py

@@ -9,8 +9,7 @@ logger = get_logger("runtime")
 
 
 def configure_sentry(dsn: str):
-    """
-    Initialize Sentry for log aggregation.
+    """Initialize Sentry for log aggregation.
     """
     if dsn:
         logger.info("initializing sentry", dsn=f"{dsn[:15]}...")
@@ -20,7 +19,7 @@ def configure_sentry(dsn: str):
             integrations=[
                 LoggingIntegration(
                     event_level=logging.WARNING,
-                )
+                ),
             ],
             release=pkg_resources.get_distribution("virtool-workflow").version,
             traces_sample_rate=0.2,