virtool-workflow 7.1.5__py3-none-any.whl → 7.2.1__py3-none-any.whl

virtool_workflow/api/client.py

@@ -3,15 +3,18 @@ from pathlib import Path
 
 import aiofiles
 from aiohttp import BasicAuth, ClientSession
+from structlog import get_logger
 
 from virtool_workflow.api.utils import (
+    API_CHUNK_SIZE,
     decode_json_response,
     raise_exception_by_status_code,
+    retry,
 )
 from virtool_workflow.errors import JobsAPIError
 from virtool_workflow.files import VirtoolFileFormat
 
-CHUNK_SIZE = 1024 * 1024 * 2
+logger = get_logger("http")
 
 
 class APIClient:
@@ -19,26 +22,29 @@ class APIClient:
         self.http = http
         self.jobs_api_connection_string = jobs_api_connection_string
 
+    @retry
     async def get_json(self, path: str) -> dict:
         """Get the JSON response from the provided API ``path``."""
         async with self.http.get(f"{self.jobs_api_connection_string}{path}") as resp:
             await raise_exception_by_status_code(resp)
             return await decode_json_response(resp)
 
+    @retry
     async def get_file(self, path: str, target_path: Path):
-        """Download the file at URL ``path`` to the local filesystem path ``target_path``.
-        """
+        """Download the file at URL ``path`` to the local ``target_path``."""
         async with self.http.get(f"{self.jobs_api_connection_string}{path}") as resp:
             if resp.status != 200:
                 raise JobsAPIError(
                     f"Encountered {resp.status} while downloading '{path}'",
                 )
+
             async with aiofiles.open(target_path, "wb") as f:
-                async for chunk in resp.content.iter_chunked(CHUNK_SIZE):
+                async for chunk in resp.content.iter_chunked(API_CHUNK_SIZE):
                     await f.write(chunk)
 
             return target_path
 
+    @retry
     async def patch_json(self, path: str, data: dict) -> dict:
         """Make a patch request against the provided API ``path`` and return the response
         as a dictionary of decoded JSON.
@@ -48,11 +54,13 @@ class APIClient:
         :return: the response as a dictionary of decoded JSON
         """
         async with self.http.patch(
-            f"{self.jobs_api_connection_string}{path}", json=data,
+            f"{self.jobs_api_connection_string}{path}",
+            json=data,
         ) as resp:
             await raise_exception_by_status_code(resp)
             return await decode_json_response(resp)
 
+    @retry
     async def post_file(
         self,
         path: str,
@@ -73,13 +81,16 @@ class APIClient:
         ) as response:
             await raise_exception_by_status_code(response)
 
+    @retry
     async def post_json(self, path: str, data: dict) -> dict:
         async with self.http.post(
-            f"{self.jobs_api_connection_string}{path}", json=data,
+            f"{self.jobs_api_connection_string}{path}",
+            json=data,
         ) as resp:
             await raise_exception_by_status_code(resp)
             return await decode_json_response(resp)
 
+    @retry
     async def put_file(
         self,
         path: str,
@@ -100,13 +111,16 @@ class APIClient:
         ) as response:
             await raise_exception_by_status_code(response)
 
+    @retry
     async def put_json(self, path: str, data: dict) -> dict:
         async with self.http.put(
-            f"{self.jobs_api_connection_string}{path}", json=data,
+            f"{self.jobs_api_connection_string}{path}",
+            json=data,
         ) as resp:
             await raise_exception_by_status_code(resp)
             return await decode_json_response(resp)
 
+    @retry
     async def delete(self, path: str) -> dict | None:
         """Make a delete request against the provided API ``path``."""
         async with self.http.delete(f"{self.jobs_api_connection_string}{path}") as resp:
@@ -124,8 +138,7 @@ async def api_client(
     job_id: str,
     key: str,
 ):
-    """An authenticated :class:``APIClient`` to make requests against the jobs API.
-    """
+    """An authenticated :class:``APIClient`` to make requests against the jobs API."""
     async with ClientSession(
         auth=BasicAuth(login=f"job-{job_id}", password=key),
     ) as http:

virtool_workflow/api/utils.py

@@ -1,11 +1,10 @@
 import asyncio
-import functools
+from functools import wraps
 
 from aiohttp import (
-    ClientConnectorError,
+    ClientError,
     ClientResponse,
     ContentTypeError,
-    ServerDisconnectedError,
 )
 from structlog import get_logger
 
@@ -19,40 +18,76 @@ from virtool_workflow.errors import (
 
 logger = get_logger("api")
 
+API_CHUNK_SIZE = 1024 * 1024 * 2
+"""The size of chunks to use when downloading files from the API in bytes."""
 
-def retry(func):
-    """Retry an API call five times when encountering the following exceptions:
-      * ``ConnectionRefusedError``.
-      * ``ClientConnectorError``.
-      * ``ServerDisconnectedError``.
+API_MAX_RETRIES = 5
+"""The maximum number of retries for API requests."""
 
-    These are probably due to transient issues in the cluster network.
+API_RETRY_BASE_DELAY = 5.0
+"""The base delay in seconds between retries for API requests."""
 
+
+def retry(
+    func=None,
+    *,
+    max_retries: int = API_MAX_RETRIES,
+    base_delay: float = API_RETRY_BASE_DELAY,
+):
+    """Retry the decorated function on connection errors.
+
+    :param func: The function to decorate (when used without parentheses)
+    :param max_retries: Maximum number of retry attempts before giving up (default: 5)
+    :param base_delay: Base delay in seconds between retries (default: 5.0)
     """
 
-    @functools.wraps(func)
-    async def wrapped(*args, **kwargs):
-        attempts = 0
+    def decorator(f):
+        @wraps(f)
+        async def wrapper(*args, **kwargs):
+            last_exception = None
+
+            log = logger.bind(func_name=f.__name__, max_retries=max_retries)
+
+            for attempt in range(max_retries + 1):
+                try:
+                    return await f(*args, **kwargs)
+                except (
+                    ClientError,
+                    ConnectionError,
+                ) as e:
+                    last_exception = e
+
+                    if attempt == max_retries:
+                        log.warning(
+                            "max retries reached for function",
+                            exception=str(e),
+                        )
+                        raise
 
-        try:
-            return await func(*args, **kwargs)
-        except (
-            ConnectionRefusedError,
-            ClientConnectorError,
-            ServerDisconnectedError,
-        ) as err:
-            if attempts == 5:
-                raise
+                    # Use exponential backoff if base_delay != 5.0, otherwise use
+                    # fixed delay.
+                    if base_delay == API_RETRY_BASE_DELAY:
+                        delay = base_delay
+                    else:
+                        delay = base_delay * (2**attempt)
 
-            attempts += 1
-            get_logger("runtime").info(
-                f"Encountered {type(err).__name__}. Retrying in 5 seconds.",
-            )
-            await asyncio.sleep(5)
+                    log.info(
+                        "retrying after connection error",
+                        exception=str(e),
+                        retries=attempt,
+                        retrying_in=delay,
+                    )
 
-            return await func(*args, **kwargs)
+                    await asyncio.sleep(delay)
 
-    return wrapped
+            raise last_exception
+
+        return wrapper
+
+    if func is None:
+        return decorator
+
+    return decorator(func)
 
 
 async def decode_json_response(resp: ClientResponse) -> dict | list | None:
@@ -103,7 +138,6 @@ async def raise_exception_by_status_code(resp: ClientResponse):
 
         if resp.status in status_exception_map:
             raise status_exception_map[resp.status](message)
-        else:
-            raise ValueError(
-                f"Status code {resp.status} not handled for response\n {resp}",
-            )
+        raise ValueError(
+            f"Status code {resp.status} not handled for response\n {resp}",
+        )

virtool_workflow/data/hmms.py

@@ -7,7 +7,6 @@ from functools import cached_property
 from pathlib import Path
 from shutil import which
 
-import aiofiles
 from pyfixtures import fixture
 from virtool.hmm.models import HMM
 from virtool.utils import decompress_file
@@ -84,25 +83,22 @@ async def hmms(
     hmms_path = work_path / "hmms"
     await asyncio.to_thread(hmms_path.mkdir, parents=True, exist_ok=True)
 
+    annotations_path = hmms_path / "annotations.json"
     compressed_annotations_path = hmms_path / "annotations.json.gz"
     await _api.get_file("/hmms/files/annotations.json.gz", compressed_annotations_path)
-
-    annotations_path = hmms_path / "annotations.json"
     await asyncio.to_thread(
         decompress_file,
         compressed_annotations_path,
         annotations_path,
         proc,
     )
+    annotations = await asyncio.to_thread(
+        lambda: [HMM(**hmm) for hmm in json.loads(annotations_path.read_text())],
+    )
 
     profiles_path = hmms_path / "profiles.hmm"
     await _api.get_file("/hmms/files/profiles.hmm", profiles_path)
-
-    async with aiofiles.open(annotations_path) as f:
-        annotations = [HMM(**hmm) for hmm in json.loads(await f.read())]
-
     p = await run_subprocess(["hmmpress", str(profiles_path)])
-
     if p.returncode != 0:
         raise RuntimeError("hmmpress command failed")
 

virtool_workflow/data/indexes.py

@@ -3,7 +3,6 @@ import json
 from dataclasses import dataclass
 from pathlib import Path
 
-import aiofiles
 from pyfixtures import fixture
 from structlog import get_logger
 from virtool.analyses.models import Analysis
@@ -245,8 +244,7 @@ async def index(
 
     log.info("decompressed reference otus json")
 
-    async with aiofiles.open(json_path) as f:
-        data = json.loads(await f.read())
+    data = await asyncio.to_thread(lambda: json.loads(json_path.read_text()))
 
     sequence_lengths = {}
     sequence_otu_map = {}

virtool_workflow/runtime/discover.py

@@ -88,7 +88,11 @@ def import_module_from_file(module_name: str, path: Path) -> ModuleType:
     spec = spec_from_file_location(module_name, path)
     if spec is None:
         raise ImportError(f"could not import {path}")
-    module = spec.loader.load_module(module_from_spec(spec).__name__)
+
+    module = module_from_spec(spec)
+    if spec.loader is None:
+        raise ImportError(f"could not load {path}")
+    spec.loader.exec_module(module)
 
     sys.path.remove(module_parent)
 

virtool_workflow/runtime/run.py

@@ -38,7 +38,7 @@ from virtool_workflow.runtime.redis import (
     get_next_job_with_timeout,
     wait_for_cancellation,
 )
-from virtool_workflow.runtime.sentry import configure_sentry
+from virtool_workflow.runtime.sentry import configure_sentry, set_workflow_context
 from virtool_workflow.utils import configure_logs, get_virtool_workflow_version
 from virtool_workflow.workflow import Workflow
 
@@ -182,6 +182,9 @@ async def run_workflow(
         scope["proc"] = config.proc
         scope["results"] = {}
 
+        # Set Sentry context with workflow metadata
+        set_workflow_context(job.workflow, job.id)
+
         async with create_work_path(config) as work_path:
             scope["work_path"] = work_path
 

virtool_workflow/runtime/sentry.py

@@ -9,7 +9,7 @@ from virtool_workflow.utils import get_virtool_workflow_version
 logger = get_logger("runtime")
 
 
-def configure_sentry(dsn: str):
+def configure_sentry(dsn: str | None) -> None:
     """Initialize Sentry for log aggregation."""
     if dsn:
         logger.info("initializing sentry", dsn=f"{dsn[:15]}...")
@@ -26,3 +26,25 @@ def configure_sentry(dsn: str):
         )
     else:
         logger.info("sentry disabled because no dsn was provided")
+
+
+def set_workflow_context(
+    workflow_name: str,
+    job_id: str,
+):
+    """Set workflow context for Sentry reporting."""
+    try:
+        with open("VERSION") as f:
+            workflow_version = f.read().strip()
+    except FileNotFoundError:
+        workflow_version = "UNKNOWN"
+
+    sentry_sdk.set_context(
+        "workflow",
+        {
+            "workflow_name": workflow_name,
+            "workflow_version": workflow_version,
+            "virtool_workflow_version": get_virtool_workflow_version(),
+            "job_id": job_id,
+        },
+    )

virtool_workflow/utils.py

@@ -82,7 +82,7 @@ async def make_directory(path: Path):
 
 def untar(path: Path, target_path: Path):
     with tarfile.open(path, "r:gz") as tar:
-        tar.extractall(target_path)
+        tar.extractall(target_path, filter="data")
 
 
 def move_all_model_files(source_path: Path, target_path: Path):

{virtool_workflow-7.1.5.dist-info → virtool_workflow-7.2.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: virtool-workflow
-Version: 7.1.5
+Version: 7.2.1
 Summary: A framework for developing bioinformatics workflows for Virtool.
 License: MIT
 Author: Ian Boyes
@@ -8,10 +8,10 @@ Maintainer: Ian Boyes
 Requires-Python: >=3.12.3,<3.13.0
 Classifier: License :: OSI Approved :: MIT License
 Classifier: Programming Language :: Python :: 3
-Requires-Dist: aiofiles (>=0.7.0,<0.8.0)
+Requires-Dist: aiofiles (>=24.1.0,<25.0.0)
 Requires-Dist: aiohttp (>=3.8.1,<4.0.0)
 Requires-Dist: biopython (>=1.81,<2.0)
-Requires-Dist: click (>=8.1.7,<9.0.0)
+Requires-Dist: click (>=8.2.1,<9.0.0)
 Requires-Dist: orjson (>=3.9.9,<4.0.0)
 Requires-Dist: pydantic-factories (>=1.17.3,<2.0.0)
 Requires-Dist: pyfixtures (>=1.0.0,<2.0.0)

{virtool_workflow-7.1.5.dist-info → virtool_workflow-7.2.1.dist-info}/RECORD

@@ -6,13 +6,13 @@ virtool_workflow/analysis/trimming.py,sha256=3Dk0J322ZhBzhuplaVgxZv4l7MLu7ZhqNQH
 virtool_workflow/analysis/utils.py,sha256=YU1_yInZzTNl9nKQTebUz47kUEqZ__d0k-RMLX8DWOA,1108
 virtool_workflow/api/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 virtool_workflow/api/acquire.py,sha256=M8Wf6ck6YOPQU_0dalAYxuQbukj6c4Uat1OPhV7Gj2w,2157
-virtool_workflow/api/client.py,sha256=wOAAX0Na1IrULWUTGq0wzbjquZqmZqCnc0gb2A2JvLE,4604
-virtool_workflow/api/utils.py,sha256=UrOsrGMMJ6PbWcl84A4tKD6Lq7C8hVQIZ_TpQR75KHU,3123
+virtool_workflow/api/client.py,sha256=fXCriMv0s2L5O78fn8uIBWWPtxtPz0tUk1XHb1L7RGU,4766
+virtool_workflow/api/utils.py,sha256=Wzf1YseJKUBnUk0Ge8DMwgigWMS7Dn-VFNLlBFo6SW4,4341
 virtool_workflow/cli.py,sha256=yEl1LziABKbjc5MCOoUGLy-iuqehOou37x9ox4LA92M,1441
 virtool_workflow/data/__init__.py,sha256=L2bQ5rMpnCMuLWDMwel4-wPchfQCTHH3EEwxDqwGOsY,610
 virtool_workflow/data/analyses.py,sha256=tGUU3gLyk_4vn1Xb7FccvBc_HQzYW_AoU7Dh1OkWhAA,3105
-virtool_workflow/data/hmms.py,sha256=pntKYyWeipApMNUGYh90ReyTbJgbCPlE-cm66Eo6teA,3225
-virtool_workflow/data/indexes.py,sha256=hxqxJBTt2VnclR8YjqAAEtk_Ig7OEPBx0sjxf4Puy9Y,9007
+virtool_workflow/data/hmms.py,sha256=qHtkQO0bF6V7GGVoGDpdhXXR9RxaetWiTL3of6Pd88Y,3210
+virtool_workflow/data/indexes.py,sha256=vmrqZjndrTV9tUfVjudIzPa0clJHy1kWrxzRQH81Ess,8981
 virtool_workflow/data/jobs.py,sha256=YYWxWoiVtHIBKkpOTpZl-Ute5rg-RAnMGAzBoixDNQo,1663
 virtool_workflow/data/ml.py,sha256=haYHZfDbYOk0ftg6MCXbQpoDFl1VbdnQ_TuNvI9pmw0,2186
 virtool_workflow/data/samples.py,sha256=fgyuQSavZiqwrqEH7Y_CoAR5gJIcO1ujrIyO5tDrXeI,4971
@@ -27,19 +27,19 @@ virtool_workflow/pytest_plugin/data.py,sha256=objP7cn_4u2MtsDaf0amG0EtBv3-jthj9l
 virtool_workflow/pytest_plugin/utils.py,sha256=lxyWqHKWXGxQXBIbcYUCC6PvjSeruhSDnD-dPX1mL5Y,211
 virtool_workflow/runtime/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 virtool_workflow/runtime/config.py,sha256=fHkprWxxqWeEWKOobxpVvSOS5Njk5Mq32sH-p5o_d8g,546
-virtool_workflow/runtime/discover.py,sha256=DsXII4f1cO0JT-csardzN-PfcuJ09SAFBgfDydZikFo,2933
+virtool_workflow/runtime/discover.py,sha256=-ki36Z4tV0uH_HDP4iD94KT0-JIZ8kHGciuH6wxYDdg,3016
 virtool_workflow/runtime/events.py,sha256=ZIX3veBfC_2yNTdClHJaYMPn4KAF9JZqpQ3qDq4ox_E,138
 virtool_workflow/runtime/hook.py,sha256=ZfmvDHoSE7Qwk4Rpcjyd-pYuwP55mzEddec2ngRZpsg,4328
 virtool_workflow/runtime/path.py,sha256=J8CsNMTg4XgDtib0gVSsLNvv17q793M-ydsxN6pkHrI,562
 virtool_workflow/runtime/ping.py,sha256=Xm4udRCcldHfSV1Rjpiqr-6vE7G6kaVaZgDlzNSxdao,1434
 virtool_workflow/runtime/redis.py,sha256=m-Dtdpbho-Qa9W5IYJCeEEZ7vv04hYu5yALsxOJr0FY,1813
-virtool_workflow/runtime/run.py,sha256=l0HguLLdYRVKvYTnIjoUB-9h6oDxffK-xzZYqO9lSIM,7599
+virtool_workflow/runtime/run.py,sha256=pkfNyjaDBsVHrEGqmz5Ah5FkHJN4SkCHKkEmfDO1QRU,7725
 virtool_workflow/runtime/run_subprocess.py,sha256=OFHVQ2ao16X8I9nl6Nm_f3IL7rHnmLH8ixtlAFKlOLk,4790
-virtool_workflow/runtime/sentry.py,sha256=y4mh1sT-hoijRtKr4meE7yn1E-alxRKDylzGxmN6ndU,755
-virtool_workflow/utils.py,sha256=jLsOtwh3RP7BCdgcr_rQr3DqzK1nLP2NILFXxbgAyGk,2530
+virtool_workflow/runtime/sentry.py,sha256=WJ9CI2fshJD5As6s3umUtJrUq4Z-hQRrSxAgbjceXrk,1326
+virtool_workflow/utils.py,sha256=9rHSlcvZsV2qE1Wo-cUDqESIATswr1zR_P8V-wP6jdk,2545
 virtool_workflow/workflow.py,sha256=W8IEFzd28wjdNNlqRrPDKyX9LeQpR6Vxy7zKeYEMQEc,2655
-virtool_workflow-7.1.5.dist-info/LICENSE,sha256=nkoVQw9W4aoQM9zgtNzHDmztap5TuXZ1L2-87vNr3w8,1097
-virtool_workflow-7.1.5.dist-info/METADATA,sha256=2noOApATUCmR9G0i7FCgHWjzATEaC1KkV_90p9Aso_A,1757
-virtool_workflow-7.1.5.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-virtool_workflow-7.1.5.dist-info/entry_points.txt,sha256=d4MA8ZDTJOU0jKZ3ymtHZbfLVoRPgItEIN5U4uIqay8,62
-virtool_workflow-7.1.5.dist-info/RECORD,,
+virtool_workflow-7.2.1.dist-info/LICENSE,sha256=nkoVQw9W4aoQM9zgtNzHDmztap5TuXZ1L2-87vNr3w8,1097
+virtool_workflow-7.2.1.dist-info/METADATA,sha256=NoJWREvu-JpIz3LcAUZscYbXA-srOf--v9X94A3EygI,1759
+virtool_workflow-7.2.1.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+virtool_workflow-7.2.1.dist-info/entry_points.txt,sha256=d4MA8ZDTJOU0jKZ3ymtHZbfLVoRPgItEIN5U4uIqay8,62
+virtool_workflow-7.2.1.dist-info/RECORD,,