chore: add code quality tooling and harden worker runtime

.github/workflows/quality.yml

@@ -0,0 +1,34 @@
+name: Quality
+
+on:
+  pull_request:
+  push:
+    branches: [main]
+
+permissions:
+  contents: read
+
+jobs:
+  quality:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v5
+      - uses: astral-sh/setup-uv@v7
+        with:
+          enable-cache: true
+          python-version: "3.12"
+          version: "0.11.25"
+      - name: Install ffmpeg
+        run: |
+          sudo apt-get update
+          sudo apt-get install --yes --no-install-recommends ffmpeg
+      - name: Install dependencies
+        run: uv sync --locked --all-groups
+      - name: Lint
+        run: uv run ruff check .
+      - name: Check formatting
+        run: uv run ruff format --check .
+      - name: Type check
+        run: uv run ty check
+      - name: Test
+        run: uv run pytest

README.md

@@ -74,6 +74,14 @@ curl --request POST --url http://localhost:8000/inference  --header 'Content-Typ
     --data '{"input_1": "Hello"}'
 ```
 
+Workers return a `WorkerResponse` with three fields:
+
+- `billable_seconds`: the billable duration as a number of seconds, including
+  fractional seconds when available, or `null` when it cannot be determined.
+- `stats`: numeric worker measurements. Use `duration` for total worker
+  processing time; additional keys may report individual phases.
+- `result`: the worker-specific JSON object.
+
 The `/health` endpoint reports the worker artifact version supplied through
 `WORKER_VERSION` and available GPU metadata in addition to `ok`. Deployments
 should set it to the exact worker image tag; it is `null` when unset.
@@ -157,6 +165,8 @@ from maestro_worker_python.get_duration import get_duration
 get_duration('./myfile.mp3')
 ```
 
+The returned `float` preserves fractional seconds reported by `ffprobe`.
+
 ## Using Docker Compose
 
 ### Build image
@@ -187,8 +197,11 @@ To bump the package version:
 uv version --bump patch
 ```
 
-Running tests:
+Run the quality checks:
 
 ```bash
 uv run pytest
+uv run ruff check .
+uv run ruff format --check .
+uv run ty check
 ```

maestro_worker_python/cli.py

@@ -10,7 +10,7 @@ def main():
     params = {}
     for arg in sys.argv:
         if arg.startswith("--"):
-            key, value = arg[2:].split("=",1)
+            key, value = arg[2:].split("=", 1)
             params[key] = value
 
     logging.info("Running with", params)

maestro_worker_python/config.py

@@ -1,13 +1,14 @@
-from typing import Optional
 from pydantic_settings import BaseSettings
 
+
 class Settings(BaseSettings):
     log_level: str = "INFO"
     enable_json_logging: bool = False
     model_path: str = "./worker.py"
-    sentry_dsn: Optional[str] = None
+    sentry_dsn: str | None = None
     sentry_traces_sample_rate: float = 1.0
     sentry_errors_sample_rate: float = 1.0
-    environment: str = 'production'
+    environment: str = "production"
+
 
 settings = Settings()

maestro_worker_python/convert_files.py

@@ -4,40 +4,43 @@
 import logging
 import subprocess
 import tempfile
+from collections.abc import Iterator
 from contextlib import contextmanager
 from dataclasses import dataclass
-from typing import List
+from os import PathLike
 
 from .response import ValidationError
 
 logger = logging.getLogger(__name__)
 
 
 class FileConversionError(Exception):
-    def __init__(self, message):
+    def __init__(self, message: str) -> None:
+        super().__init__(message)
         self.message = message
 
 
+StrPath = str | PathLike[str]
+
+
 @dataclass
 class FileToConvert:
-    input_file_path: str
+    input_file_path: StrPath
     file_format: str
-    output_file_path: str | None = None
+    output_file_path: StrPath | None = None
     max_duration: int = 1200
     sample_rate: int | None = 44100
 
 
-def convert_files(convert_files: List[FileToConvert]):
+def convert_files(convert_files: list[FileToConvert]) -> None:
     logger.info(f"Converting {len(convert_files)} files")
 
     futures = []
     with concurrent.futures.ThreadPoolExecutor() as executor:
         for convert_file in convert_files:
-            target_function = (
-                _convert_to_m4a
-                if convert_file.file_format == "m4a"
-                else _convert_to_wav
-            )
+            if convert_file.output_file_path is None:
+                raise ValueError("output_file_path is required when using convert_files")
+            target_function = _convert_to_m4a if convert_file.file_format == "m4a" else _convert_to_wav
             futures.append(
                 executor.submit(
                     target_function,
@@ -55,19 +58,15 @@ def convert_files(convert_files: List[FileToConvert]):
 
 
 @contextmanager
-def convert_files_manager(*convert_files: FileToConvert) -> None | str | list[str]:
+def convert_files_manager(*convert_files: FileToConvert) -> Iterator[None | str | list[str]]:
     try:
         thread_list = []
         list_objects = []
         with concurrent.futures.ThreadPoolExecutor() as executor:
             for convert_file in convert_files:
                 file_format = ".m4a" if convert_file.file_format == "m4a" else ".wav"
                 filename = tempfile.NamedTemporaryFile(suffix=file_format)
-                target_function = (
-                    _convert_to_m4a
-                    if convert_file.file_format == "m4a"
-                    else _convert_to_wav
-                )
+                target_function = _convert_to_m4a if convert_file.file_format == "m4a" else _convert_to_wav
                 thread_list.append(
                     executor.submit(
                         target_function,
@@ -92,7 +91,12 @@ def convert_files_manager(*convert_files: FileToConvert) -> None | str | list[st
             obj.close()
 
 
-def _convert_to_wav(input_file_path: str, output_file_path: str, max_duration: int, sample_rate: int | None = 44100):
+def _convert_to_wav(
+    input_file_path: StrPath,
+    output_file_path: StrPath,
+    max_duration: int,
+    sample_rate: int | None = 44100,
+) -> None:
     command_list = [
         "ffmpeg",
         "-y",
@@ -113,7 +117,12 @@ def _convert_to_wav(input_file_path: str, output_file_path: str, max_duration: i
     _run_subprocess(command_list)
 
 
-def _convert_to_m4a(input_file_path: str, output_file_path: str, max_duration: int, sample_rate: int | None = 44100):
+def _convert_to_m4a(
+    input_file_path: StrPath,
+    output_file_path: StrPath,
+    max_duration: int,
+    sample_rate: int | None = 44100,
+) -> None:
     command_list = [
         "ffmpeg",
         "-y",
@@ -140,7 +149,7 @@ def _convert_to_m4a(input_file_path: str, output_file_path: str, max_duration: i
     _run_subprocess(command_list)
 
 
-def _run_subprocess(command):
+def _run_subprocess(command: list[str]) -> None:
     try:
         process = subprocess.run(command, shell=False, capture_output=True, check=True)
     except subprocess.CalledProcessError as exc:
@@ -164,9 +173,7 @@ def _run_subprocess(command):
                 f"Could not convert because the file is invalid, ffmpeg stderr: {exc.stderr.decode()}"
             ) from exc
 
-        raise FileConversionError(
-            f"Fatal error during conversion, ffmpeg stderr: {exc.stderr.decode()}"
-        ) from exc
+        raise FileConversionError(f"Fatal error during conversion, ffmpeg stderr: {exc.stderr.decode()}") from exc
     else:
         if process.stderr:
             logger.warning(

maestro_worker_python/download_file.py

@@ -1,35 +1,39 @@
 from __future__ import annotations
-import requests
 
 import logging
 import tempfile
-
-from concurrent.futures import ThreadPoolExecutor
-from concurrent.futures import as_completed
+from collections.abc import Iterator
+from concurrent.futures import ThreadPoolExecutor, as_completed
 from contextlib import contextmanager
 
+import requests
+
 from .response import ValidationError
 
 
-def download_file(url: str, filename: str = None):
+def download_file(url: str, filename: str | None = None) -> str:
     logging.info(f"Downloading input: {url}")
     response = requests.get(url, allow_redirects=True, timeout=300)
 
     try:
         response.raise_for_status()
     except requests.exceptions.HTTPError as e:
-        raise ValidationError(f"Bad download input: {response.status_code}: {e}")
-
-    file_name = filename if filename is not None else tempfile.mktemp()
-    with open(file_name, 'wb') as file:
+        raise ValidationError(f"Bad download input: {response.status_code}: {e}") from e
+
+    if filename is None:
+        with tempfile.NamedTemporaryFile(delete=False) as temporary_file:
+            file_name = temporary_file.name
+    else:
+        file_name = filename
+    with open(file_name, "wb") as file:
         file.write(response.content)
 
     logging.info(f"Downloaded input to {file_name}")
     return file_name
 
 
 @contextmanager
-def download_files_manager(*urls: str) -> None | str | list[str]:
+def download_files_manager(*urls: str) -> Iterator[None | str | list[str]]:
     try:
         thread_list = []
         list_objects = []

maestro_worker_python/get_duration.py

@@ -1,17 +1,27 @@
 import logging
 from subprocess import check_output
 
-def get_duration(local_file_path: str) -> int:
+
+def get_duration(local_file_path: str) -> float | None:
     logging.info(f"Getting file duration for {local_file_path}")
     duration = None
     try:
-        duration = check_output(["ffprobe", "-v", "error", "-show_entries",
-                               "format=duration", "-of",
-                               "default=noprint_wrappers=1:nokey=1",
-                               local_file_path])
-        if duration == b'N/A\n':
+        duration = check_output(
+            [
+                "ffprobe",
+                "-v",
+                "error",
+                "-show_entries",
+                "format=duration",
+                "-of",
+                "default=noprint_wrappers=1:nokey=1",
+                local_file_path,
+            ]
+        )
+        if duration == b"N/A\n":
             logging.info(f"File {local_file_path} has no valid duration")
             return None
-        return int(float(duration))
+        return float(duration)
     except Exception as e:
         logging.error(f"Error getting duration: {e}")
+        return None

maestro_worker_python/health.py

@@ -42,13 +42,8 @@ def _partitioning_metadata(
         if percentage is not None and 0 < percentage <= 100:
             configured_active_thread_percentage = percentage
 
-    configured_pinned_device_memory_limit = (
-        os.getenv("CUDA_MPS_PINNED_DEVICE_MEM_LIMIT") or None
-    )
-    if (
-        configured_active_thread_percentage is None
-        and configured_pinned_device_memory_limit is None
-    ):
+    configured_pinned_device_memory_limit = os.getenv("CUDA_MPS_PINNED_DEVICE_MEM_LIMIT") or None
+    if configured_active_thread_percentage is None and configured_pinned_device_memory_limit is None:
         return None
 
     return {
@@ -115,9 +110,7 @@ def _nvml_device_identity(device: Any) -> str:
         return repr(device)
 
 
-def _nvml_gpu_metadata() -> tuple[
-    list[dict[str, str | None]], list[dict[str, int | None]]
-]:
+def _nvml_gpu_metadata() -> tuple[list[dict[str, str | None]], list[dict[str, int | None]]]:
     """Collect visible GPUs and active MIG partitions."""
     try:
         device_count = pynvml.nvmlDeviceGetCount()
@@ -151,9 +144,7 @@ def _nvml_gpu_metadata() -> tuple[
             if identity in seen_mig_devices:
                 continue
             seen_mig_devices.add(identity)
-            visible_mig_partitions.append(
-                _nvml_mig_partition_metadata(mig_device)
-            )
+            visible_mig_partitions.append(_nvml_mig_partition_metadata(mig_device))
 
     return gpus, visible_mig_partitions
 
@@ -215,14 +206,16 @@ def _torch_observed_sm_count() -> int | None:
     torch = sys.modules.get("torch")
     cuda = getattr(torch, "cuda", None)
     is_initialized = getattr(cuda, "is_initialized", None)
-    if not callable(is_initialized):
+    current_device = getattr(cuda, "current_device", None)
+    get_device_properties = getattr(cuda, "get_device_properties", None)
+    if not callable(is_initialized) or not callable(current_device) or not callable(get_device_properties):
         return None
 
     try:
         if not is_initialized():
             return None
-        device = cuda.current_device()
-        sm_count = cuda.get_device_properties(device).multi_processor_count
+        device = current_device()
+        sm_count = get_device_properties(device).multi_processor_count
     except Exception:
         return None