diff --git a/.github/workflows/quality.yml b/.github/workflows/quality.yml new file mode 100644 index 0000000..2771119 --- /dev/null +++ b/.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 diff --git a/README.md b/README.md index 52558c4..d04a273 100644 --- a/README.md +++ b/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 ``` diff --git a/maestro_worker_python/cli.py b/maestro_worker_python/cli.py index 0d043c8..7fd4544 100755 --- a/maestro_worker_python/cli.py +++ b/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) diff --git a/maestro_worker_python/config.py b/maestro_worker_python/config.py index 79fee10..cd58345 100644 --- a/maestro_worker_python/config.py +++ b/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() diff --git a/maestro_worker_python/convert_files.py b/maestro_worker_python/convert_files.py index 39d84fc..e62fd8e 100644 --- a/maestro_worker_python/convert_files.py +++ b/maestro_worker_python/convert_files.py @@ -4,9 +4,10 @@ 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 @@ -14,30 +15,32 @@ 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,7 +58,7 @@ 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 = [] @@ -63,11 +66,7 @@ def convert_files_manager(*convert_files: FileToConvert) -> None | str | list[st 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( diff --git a/maestro_worker_python/download_file.py b/maestro_worker_python/download_file.py index a289fb8..2b3dfc6 100644 --- a/maestro_worker_python/download_file.py +++ b/maestro_worker_python/download_file.py @@ -1,27 +1,31 @@ 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}") @@ -29,7 +33,7 @@ def download_file(url: str, filename: str = None): @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 = [] diff --git a/maestro_worker_python/get_duration.py b/maestro_worker_python/get_duration.py index 9457c94..57dcab7 100644 --- a/maestro_worker_python/get_duration.py +++ b/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 diff --git a/maestro_worker_python/health.py b/maestro_worker_python/health.py index 177024d..a5225ba 100644 --- a/maestro_worker_python/health.py +++ b/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 diff --git a/maestro_worker_python/init.py b/maestro_worker_python/init.py index 2793efb..875c198 100644 --- a/maestro_worker_python/init.py +++ b/maestro_worker_python/init.py @@ -5,13 +5,10 @@ from importlib import metadata from tempfile import TemporaryDirectory - _VERSION_PLACEHOLDER = "MAESTRO_WORKER_PYTHON_VERSION" -def _find_conflicts( - scaffold_dir: pathlib.Path, target_dir: pathlib.Path -) -> list[pathlib.Path]: +def _find_conflicts(scaffold_dir: pathlib.Path, target_dir: pathlib.Path) -> list[pathlib.Path]: conflicts = [] for scaffold_path in sorted(scaffold_dir.rglob("*")): relative_path = scaffold_path.relative_to(scaffold_dir) @@ -24,24 +21,20 @@ def _find_conflicts( elif scaffold_path.is_dir(): if not target_path.is_dir(): conflicts.append(relative_path) - elif not target_path.is_file() or ( - scaffold_path.read_bytes() != target_path.read_bytes() - ): + elif not target_path.is_file() or (scaffold_path.read_bytes() != target_path.read_bytes()): conflicts.append(relative_path) return conflicts def main(): - parser = argparse.ArgumentParser( - formatter_class=argparse.ArgumentDefaultsHelpFormatter) - parser.add_argument("--folder", default="./", - help="Folder to create the worker in") + parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter) + parser.add_argument("--folder", default="./", help="Folder to create the worker in") - args = parser.parse_args().__dict__ + args = parser.parse_args() scaffold_dir = pathlib.Path(__file__).parent.resolve() / "scaffold" - target_dir = pathlib.Path(args.get("folder")) + target_dir = pathlib.Path(args.folder) if target_dir.is_symlink() or (target_dir.exists() and not target_dir.is_dir()): parser.error(f"target path is not a directory: {target_dir}") @@ -59,9 +52,7 @@ def main(): conflicts = _find_conflicts(rendered_scaffold_dir, target_dir) if conflicts: conflict_list = ", ".join(str(path) for path in conflicts) - parser.error( - f"refusing to overwrite modified scaffold files: {conflict_list}" - ) + parser.error(f"refusing to overwrite modified scaffold files: {conflict_list}") print(r"""___ ___ _ | \/ | | | diff --git a/maestro_worker_python/kill_process.py b/maestro_worker_python/kill_process.py index 0e214f4..f5d38d7 100644 --- a/maestro_worker_python/kill_process.py +++ b/maestro_worker_python/kill_process.py @@ -1,27 +1,28 @@ -import psutil import logging from contextlib import suppress +import psutil + def kill_child_processes(): current_process = psutil.Process() - logging.info(f'Terminating everything from parent process PID: {current_process.pid}') - + logging.info(f"Terminating everything from parent process PID: {current_process.pid}") + children = current_process.children(recursive=True) for child in children: - logging.info(f'Sending SIGTERM to child process PID: {child.pid}') + logging.info(f"Sending SIGTERM to child process PID: {child.pid}") with suppress(psutil.NoSuchProcess): child.terminate() _, still_alive = psutil.wait_procs(children, timeout=3) for p in still_alive: - logging.warning(f'Child process PID: {p.pid} did not terminate, killing forcefully') + logging.warning(f"Child process PID: {p.pid} did not terminate, killing forcefully") with suppress(psutil.NoSuchProcess): p.kill() def terminate_current_process(): current_process = psutil.Process() - logging.info(f'Sending SIGTERM to current process PID: {current_process.pid}') + logging.info(f"Sending SIGTERM to current process PID: {current_process.pid}") with suppress(psutil.NoSuchProcess): current_process.terminate() diff --git a/maestro_worker_python/response.py b/maestro_worker_python/response.py index 32e954f..d65a0ec 100644 --- a/maestro_worker_python/response.py +++ b/maestro_worker_python/response.py @@ -1,10 +1,12 @@ -from typing import Dict, Optional +from typing import Any + from pydantic import BaseModel + class WorkerResponse(BaseModel): - billable_seconds: Optional[int] = ... - stats: Dict[str, float] - result: Dict + billable_seconds: float | None + stats: dict[str, float] + result: dict[str, Any] class ValidationError(Exception): diff --git a/maestro_worker_python/run_upload_server.py b/maestro_worker_python/run_upload_server.py index e77c42c..5525088 100644 --- a/maestro_worker_python/run_upload_server.py +++ b/maestro_worker_python/run_upload_server.py @@ -1,6 +1,7 @@ -import uvicorn -import logging import argparse +import logging + +import uvicorn def main(): @@ -8,6 +9,4 @@ def main(): parser.add_argument("--port", type=int, default=9090) args = parser.parse_args() logging.info("Starting upload server on folder ./uploads using port 9090") - uvicorn.run("maestro_worker_python.upload_server:app", - host="0.0.0.0", port=args.port, log_level='info' - ) + uvicorn.run("maestro_worker_python.upload_server:app", host="0.0.0.0", port=args.port, log_level="info") diff --git a/maestro_worker_python/scaffold/worker.py b/maestro_worker_python/scaffold/worker.py index 10572e8..5ec4ceb 100644 --- a/maestro_worker_python/scaffold/worker.py +++ b/maestro_worker_python/scaffold/worker.py @@ -1,13 +1,14 @@ import logging from time import time -from maestro_worker_python.response import WorkerResponse, ValidationError + +from maestro_worker_python.response import ValidationError, WorkerResponse def your_model(input): return f"{input} World" -class MoisesWorker(object): +class MoisesWorker: def __init__(self): logging.info("Loading model...") self.model = your_model @@ -19,13 +20,13 @@ def inference(self, input_data): if len(input_example) > 25: raise ValidationError("input is too big") time_start = time() - result = self.model(input_example) + model_result = self.model(input_example) time_end = time() # Send response with the result and the time it took to process the request return WorkerResponse( - billable_seconds=0, + billable_seconds=0.0, stats={"duration": time_end - time_start}, - result={}, + result={"output": model_result}, ) finally: logging.info("cleaning up") diff --git a/maestro_worker_python/serve.py b/maestro_worker_python/serve.py index 220dc26..65a3581 100644 --- a/maestro_worker_python/serve.py +++ b/maestro_worker_python/serve.py @@ -1,24 +1,27 @@ -import os -import socket +import asyncio import datetime import logging -import json_logging +import os +import socket import traceback -import asyncio -import sentry_sdk -import pydantic +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager, suppress from urllib.parse import urlparse + +import json_logging +import pydantic +import sentry_sdk from fastapi import FastAPI, Request from fastapi.encoders import jsonable_encoder +from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import JSONResponse from starlette.concurrency import run_in_threadpool -from fastapi.middleware.cors import CORSMiddleware from .config import settings from .health import get_health_metadata +from .kill_process import kill_child_processes, terminate_current_process from .load_worker import load_worker from .response import ValidationError, WorkerResponse -from .kill_process import terminate_current_process, kill_child_processes def filter_transactions(event, hint): @@ -38,12 +41,33 @@ def filter_transactions(event, hint): before_send_transaction=filter_transactions, ) -app = FastAPI() + +HTTP_READY_PATH = "/tmp/http_ready" + + +@asynccontextmanager +async def lifespan(_app: FastAPI) -> AsyncIterator[None]: + # Prime the cache before readiness probes start hitting /health. + get_health_metadata() + with open(HTTP_READY_PATH, "a") as ready_file: + ready_file.write("Ready to serve") + + try: + yield + finally: + logging.info("Shutting down, bye bye") + kill_child_processes() + with suppress(FileNotFoundError): + os.remove(HTTP_READY_PATH) + + +app = FastAPI(lifespan=lifespan) origins = ["*"] app.add_middleware( - CORSMiddleware, + # ty does not yet model Starlette's class-based middleware factory protocol. + CORSMiddleware, # ty: ignore[invalid-argument-type] allow_origins=origins, allow_credentials=True, allow_methods=["*"], @@ -68,9 +92,9 @@ def filter_transactions(event, hint): @app.exception_handler(500) async def internal_exception_handler(request: Request, exc: Exception): global error_counter - tb = ''.join(traceback.format_exception(None, exc, exc.__traceback__)) + tb = "".join(traceback.format_exception(None, exc, exc.__traceback__)) try: - return JSONResponse(status_code=500, content=jsonable_encoder({"error": tb})) + return JSONResponse(status_code=500, content=jsonable_encoder({"error": tb})) finally: if error_counter > 10: logging.error("Too many consecutive errors, shutting down worker") @@ -82,7 +106,7 @@ async def internal_exception_handler(request: Request, exc: Exception): @app.exception_handler(ValidationError) async def validation_error_handler(request: Request, exc: ValidationError): - return JSONResponse(status_code=400, content=jsonable_encoder({"error": exc.reason})) + return JSONResponse(status_code=400, content=jsonable_encoder({"error": exc.reason})) @app.exception_handler(pydantic.ValidationError) @@ -111,19 +135,3 @@ async def index(request: Request): @app.get("/health") async def health(request: Request): return {"ok": True, **get_health_metadata()} - - -@app.on_event("shutdown") -def shutdown_event(): - logging.info("Shutting down, bye bye") - kill_child_processes() - os.remove("/tmp/http_ready") - - -@app.on_event("startup") -def startup_event(): - # Prime the cache before readiness probes start hitting /health. - get_health_metadata() - # Create a file to indicate that the server is running - with open("/tmp/http_ready", "a") as f: - f.write("Ready to serve") diff --git a/maestro_worker_python/server.py b/maestro_worker_python/server.py index d424d6f..2487afb 100644 --- a/maestro_worker_python/server.py +++ b/maestro_worker_python/server.py @@ -1,44 +1,56 @@ -import uvicorn import argparse -import os import logging +import os + +import uvicorn + + +def _parse_bool(value: str) -> bool: + normalized = value.casefold() + if normalized in {"1", "true", "yes", "on"}: + return True + if normalized in {"0", "false", "no", "off"}: + return False + raise argparse.ArgumentTypeError(f"expected a boolean value, got {value!r}") def main(): - parser = argparse.ArgumentParser( - formatter_class=argparse.ArgumentDefaultsHelpFormatter) - parser.add_argument("--worker", default="./worker.py", - help="Python file or importable module that contains a MoisesWorker class") - parser.add_argument("--base_path", default="/", - help="DEPRECATED") - parser.add_argument("--port", default=8000, help="Port to run uvicorn on") - parser.add_argument("--reload", default=False, - help="Reload the server on code changes") + parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter) + parser.add_argument( + "--worker", default="./worker.py", help="Python file or importable module that contains a MoisesWorker class" + ) + parser.add_argument("--base_path", default="/", help="DEPRECATED") + parser.add_argument("--port", type=int, default=8000, help="Port to run uvicorn on") + parser.add_argument("--reload", type=_parse_bool, default=False, help="Reload the server on code changes") - args = parser.parse_args().__dict__ + args = parser.parse_args() logging.info(f"Running maestro server with {str(args)}") - os.environ["MODEL_PATH"] = args.get("worker") + os.environ["MODEL_PATH"] = args.worker - log_level = os.environ.get('LOG_LEVEL', 'INFO').upper() + log_level = os.environ.get("LOG_LEVEL", "INFO").upper() logging_config = { - 'version': 1, - 'disable_existing_loggers': False, - 'handlers': { - 'default_handler': { - 'class': 'logging.StreamHandler', - 'level': log_level, + "version": 1, + "disable_existing_loggers": False, + "handlers": { + "default_handler": { + "class": "logging.StreamHandler", + "level": log_level, }, }, - 'loggers': { - '': { - '#handlers': ['default_handler'], + "loggers": { + "": { + "#handlers": ["default_handler"], }, - 'root': { - '#handlers': ['default_handler'], - } - } + "root": { + "#handlers": ["default_handler"], + }, + }, } uvicorn.run( - "maestro_worker_python.serve:app", host='0.0.0.0', port=int(args.get("port")), reload=args.get("reload"), - log_level=log_level.lower(), log_config=logging_config, + "maestro_worker_python.serve:app", + host="0.0.0.0", + port=args.port, + reload=args.reload, + log_level=log_level.lower(), + log_config=logging_config, ) diff --git a/maestro_worker_python/timer.py b/maestro_worker_python/timer.py index c6105c5..5747f19 100644 --- a/maestro_worker_python/timer.py +++ b/maestro_worker_python/timer.py @@ -1,4 +1,5 @@ from __future__ import annotations + import time diff --git a/maestro_worker_python/upload_files.py b/maestro_worker_python/upload_files.py index e329a95..b0f6bba 100644 --- a/maestro_worker_python/upload_files.py +++ b/maestro_worker_python/upload_files.py @@ -1,13 +1,13 @@ from __future__ import annotations +import concurrent.futures import json import logging import tempfile import threading -import concurrent.futures from dataclasses import dataclass from os.path import join -from typing import List, Any +from typing import Any import requests @@ -19,16 +19,21 @@ class UploadFile: signed_url: str -def upload_files(upload_files: List[UploadFile]): +def upload_files(upload_files: list[UploadFile]): logging.info(f"Uploading {len(upload_files)} files") threads_upload = [] did_raise_exception = threading.Event() for file_ in upload_files: - t = threading.Thread(target=_upload, args=( - file_, did_raise_exception,)) + t = threading.Thread( + target=_upload, + args=( + file_, + did_raise_exception, + ), + ) threads_upload.append(t) t.start() - + for t in threads_upload: t.join() @@ -40,7 +45,9 @@ def _upload(upload_file: UploadFile, did_raise_exception): logging.info(f"Uploading:{upload_file.file_path}") try: with open(upload_file.file_path, "rb") as data: - response = requests.put(upload_file.signed_url, data=data, headers={"Content-Type": upload_file.file_type}, timeout=300) + response = requests.put( + upload_file.signed_url, data=data, headers={"Content-Type": upload_file.file_type}, timeout=300 + ) response.raise_for_status() logging.info(f"Uploaded {upload_file.file_path}") except Exception as e: @@ -87,23 +94,32 @@ class AsyncUploader: def __init__(self, max_workers: int | None = None): self._max_workers = max_workers - self._exectuor = None - self._futures = [] + self._executor: concurrent.futures.ThreadPoolExecutor | None = None + self._futures: list[concurrent.futures.Future[None]] = [] def __enter__(self): - self._exectuor = concurrent.futures.ThreadPoolExecutor(max_workers=self._max_workers) + self._executor = concurrent.futures.ThreadPoolExecutor(max_workers=self._max_workers) return self def __exit__(self, exc_type, exc_value, exc_traceback): - self._exectuor.shutdown() + if self._executor is None: + raise RuntimeError("AsyncUploader must be entered before it can be exited") + self._executor.shutdown() for future in self._futures: - if future.exception(): - raise future.exception() - - def upload_file(self, file_path: str, file_type: str, signed_url: str | None) -> concurrent.futures.Future: - self._futures.append(self._exectuor.submit(upload_files, [UploadFile(file_path, file_type, signed_url)])) - return self._futures[-1] - - def upload_json(self, data: Any, signed_url: str | None) -> concurrent.futures.Future: - self._futures.append(self._exectuor.submit(upload_json_data, [UploadJsonData(data, signed_url)])) - return self._futures[-1] + exception = future.exception() + if exception is not None: + raise exception + + def upload_file(self, file_path: str, file_type: str, signed_url: str) -> concurrent.futures.Future[None]: + if self._executor is None: + raise RuntimeError("AsyncUploader must be used as a context manager") + future = self._executor.submit(upload_files, [UploadFile(file_path, file_type, signed_url)]) + self._futures.append(future) + return future + + def upload_json(self, data: Any, signed_url: str) -> concurrent.futures.Future[None]: + if self._executor is None: + raise RuntimeError("AsyncUploader must be used as a context manager") + future = self._executor.submit(upload_json_data, [UploadJsonData(data, signed_url)]) + self._futures.append(future) + return future diff --git a/maestro_worker_python/upload_server.py b/maestro_worker_python/upload_server.py index 717359f..318e6a4 100644 --- a/maestro_worker_python/upload_server.py +++ b/maestro_worker_python/upload_server.py @@ -1,10 +1,12 @@ import os from glob import glob + from fastapi import FastAPI, Request from fastapi.responses import FileResponse app = FastAPI() + @app.put("/upload-file/{filename}") async def put_file(filename: str, request: Request): data = await request.body() @@ -41,4 +43,4 @@ async def clean(): @app.get("/list-files") async def list_files(): - return {"files": glob("uploads/*")} \ No newline at end of file + return {"files": glob("uploads/*")} diff --git a/pyproject.toml b/pyproject.toml index 94ce0df..86c004b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -28,11 +28,20 @@ maestro-upload-server = "maestro_worker_python.run_upload_server:main" dev = [ "pytest>=9.0.3,<10", "pytest-httpserver>=1.0.6,<2", + "ruff>=0.14.9,<0.15", + "ty==0.0.14", ] [build-system] requires = ["uv_build>=0.11.25,<0.12"] build-backend = "uv_build" +[tool.ruff] +line-length = 120 +target-version = "py310" + +[tool.ruff.lint] +select = ["E4", "E7", "E9", "F", "I", "UP", "B"] + [tool.uv.build-backend] module-root = "" diff --git a/tests/silent with space.ogg b/tests/fixtures/silent with space.ogg similarity index 100% rename from tests/silent with space.ogg rename to tests/fixtures/silent with space.ogg diff --git a/tests/silent with space.wav b/tests/fixtures/silent with space.wav similarity index 100% rename from tests/silent with space.wav rename to tests/fixtures/silent with space.wav diff --git a/tests/silent.ogg b/tests/fixtures/silent.ogg similarity index 100% rename from tests/silent.ogg rename to tests/fixtures/silent.ogg diff --git a/tests/video-no-audio.mp4 b/tests/fixtures/video-no-audio.mp4 similarity index 100% rename from tests/video-no-audio.mp4 rename to tests/fixtures/video-no-audio.mp4 diff --git a/tests/test_convert_files.py b/tests/test_convert_files.py index 8ab300a..b90590a 100644 --- a/tests/test_convert_files.py +++ b/tests/test_convert_files.py @@ -12,7 +12,7 @@ ) from maestro_worker_python.response import ValidationError -TEST_PATH = Path(__file__).resolve().parent +FIXTURES = Path(__file__).resolve().parent / "fixtures" @pytest.fixture(scope="session") @@ -31,11 +31,11 @@ def corrupt_audio_file(tmp_path_factory): @pytest.mark.parametrize("file_format", ["m4a", "wav"]) def test_should_re_raise_exceptions_in_thread(invalid_audio_file, file_format): - with pytest.raises(FileConversionError) as exc: + with pytest.raises(FileConversionError): convert_files( [ FileToConvert( - input_file_path=TEST_PATH / "foobar.mp3", + input_file_path=FIXTURES / "foobar.mp3", output_file_path=f"{invalid_audio_file}.wav", file_format=file_format, ) @@ -43,10 +43,20 @@ def test_should_re_raise_exceptions_in_thread(invalid_audio_file, file_format): ) +def test_convert_files_requires_output_path(): + with pytest.raises(ValueError, match="output_file_path is required"): + convert_files( + [ + FileToConvert( + input_file_path=FIXTURES / "silent.ogg", + file_format="wav", + ) + ] + ) + + @pytest.mark.parametrize("file_format", ["m4a", "wav"]) -def test_should_raise_validation_error_if_audio_file_is_invalid( - invalid_audio_file, file_format -): +def test_should_raise_validation_error_if_audio_file_is_invalid(invalid_audio_file, file_format): with pytest.raises(ValidationError) as exc: convert_files( [ @@ -62,9 +72,7 @@ def test_should_raise_validation_error_if_audio_file_is_invalid( @pytest.mark.parametrize("file_format", ["m4a", "wav"]) -def test_should_raise_validation_error_if_audio_file_is_corrupt( - corrupt_audio_file, file_format -): +def test_should_raise_validation_error_if_audio_file_is_corrupt(corrupt_audio_file, file_format): with pytest.raises(ValidationError) as exc: convert_files( [ @@ -82,8 +90,8 @@ def test_should_raise_validation_error_if_audio_file_is_corrupt( @pytest.mark.parametrize("file_format", ["m4a", "wav"]) def test_should_raise_validation_error_if_source_has_no_audio(file_format, caplog): input_file_path, output_file_path = ( - TEST_PATH / "video-no-audio.mp4", - TEST_PATH / "output.wav", + FIXTURES / "video-no-audio.mp4", + FIXTURES / "output.wav", ) with pytest.raises(ValidationError) as exc: convert_files( @@ -111,8 +119,8 @@ def test_should_raise_validation_error_if_source_has_no_audio(file_format, caplo ) def test_should_convert_valid_wav_audio_file(input_name, output_name, format, sample_rate): input_file_path, output_file_path = ( - TEST_PATH / input_name, - TEST_PATH / output_name, + FIXTURES / input_name, + FIXTURES / output_name, ) convert_files( [ @@ -140,8 +148,8 @@ def test_should_convert_valid_wav_audio_file(input_name, output_name, format, sa ) def test_should_convert_valid_m4a_audio_file(input_name, output_name, format, sample_rate): input_file_path, output_file_path = ( - TEST_PATH / input_name, - TEST_PATH / output_name, + FIXTURES / input_name, + FIXTURES / output_name, ) convert_files( [ @@ -158,8 +166,8 @@ def test_should_convert_valid_m4a_audio_file(input_name, output_name, format, sa def test_should_convert_multiple_valid_audio_files_and_delete_after_context(): input_file_path, output_file_path = ( - TEST_PATH / "silent.ogg", - TEST_PATH / "silent.wav", + FIXTURES / "silent.ogg", + FIXTURES / "silent.wav", ) converted_files_list = [] with convert_files_manager( @@ -174,9 +182,10 @@ def test_should_convert_multiple_valid_audio_files_and_delete_after_context(): file_format="wav", ), ) as converted_files: + assert isinstance(converted_files, list) converted_files_list = converted_files result = [os.path.exists(path) for path in converted_files_list] - assert all(result) == False + assert not any(result) def _get_hash(file_name, sample_rate): @@ -188,17 +197,21 @@ def _get_hash(file_name, sample_rate): str(file_name), ] if sample_rate: - command.extend([ - "-ar", - str(sample_rate), - ]) - command.extend([ - "-map", - "0", - "-f", - "hash", - "-", - ]) + command.extend( + [ + "-ar", + str(sample_rate), + ] + ) + command.extend( + [ + "-map", + "0", + "-f", + "hash", + "-", + ] + ) process = subprocess.run(command, shell=False, capture_output=True, check=True) return process.stdout.split(b"=")[1].strip() diff --git a/tests/test_downlaod_file.py b/tests/test_download_file.py similarity index 89% rename from tests/test_downlaod_file.py rename to tests/test_download_file.py index 133fa60..6be7dcc 100644 --- a/tests/test_downlaod_file.py +++ b/tests/test_download_file.py @@ -1,5 +1,7 @@ import os + import pytest + from maestro_worker_python.download_file import download_file, download_files_manager from maestro_worker_python.response import ValidationError @@ -29,10 +31,11 @@ def test_download_files_manager(httpserver): files_content = [] with download_files_manager(url, url) as downloaded_files: + assert isinstance(downloaded_files, list) for file in downloaded_files: with open(file) as f: files_content.append(f.read() == "hello") - assert all(files_content) == True + assert all(files_content) def test_download_files_manager_delete(httpserver): @@ -41,7 +44,8 @@ def test_download_files_manager_delete(httpserver): files_path_exists = [] with download_files_manager(url, url) as downloaded_files: + assert isinstance(downloaded_files, list) files_path = downloaded_files for path in files_path: files_path_exists.append(os.path.exists(path)) - assert all(files_path_exists) == False + assert not any(files_path_exists) diff --git a/tests/test_entry_points.py b/tests/test_entry_points.py index ff11ba7..12861fc 100644 --- a/tests/test_entry_points.py +++ b/tests/test_entry_points.py @@ -13,9 +13,7 @@ @pytest.mark.parametrize("name,value", EXPECTED_SCRIPTS.items()) def test_console_script_is_registered(name, value): scripts = { - ep.name: ep - for ep in distribution("maestro-worker-python").entry_points - if ep.group == "console_scripts" + ep.name: ep for ep in distribution("maestro-worker-python").entry_points if ep.group == "console_scripts" } assert name in scripts assert scripts[name].value == value diff --git a/tests/test_get_duration.py b/tests/test_get_duration.py new file mode 100644 index 0000000..c7f7f61 --- /dev/null +++ b/tests/test_get_duration.py @@ -0,0 +1,7 @@ +import maestro_worker_python.get_duration as duration_module + + +def test_get_duration_preserves_fractional_seconds(monkeypatch): + monkeypatch.setattr(duration_module, "check_output", lambda _command: b"42.711293\n") + + assert duration_module.get_duration("audio.wav") == 42.711293 diff --git a/tests/test_health.py b/tests/test_health.py index 364819c..9f24169 100644 --- a/tests/test_health.py +++ b/tests/test_health.py @@ -12,9 +12,7 @@ def _initialized_torch(cuda_version="12.4", sm_count=30): cuda=SimpleNamespace( is_initialized=lambda: True, current_device=lambda: 0, - get_device_properties=lambda _device: SimpleNamespace( - multi_processor_count=sm_count - ), + get_device_properties=lambda _device: SimpleNamespace(multi_processor_count=sm_count), ), ) @@ -23,17 +21,11 @@ def _initialized_torch(cuda_version="12.4", sm_count=30): def nvml_host(monkeypatch): lifecycle = [] monkeypatch.setattr(health.pynvml, "nvmlInit", lambda: lifecycle.append("init")) - monkeypatch.setattr( - health.pynvml, "nvmlShutdown", lambda: lifecycle.append("shutdown") - ) + monkeypatch.setattr(health.pynvml, "nvmlShutdown", lambda: lifecycle.append("shutdown")) monkeypatch.setattr(health.pynvml, "nvmlSystemGetDriverVersion", lambda: "610.12") - monkeypatch.setattr( - health.pynvml, "nvmlSystemGetCudaDriverVersion_v2", lambda: 13030 - ) + monkeypatch.setattr(health.pynvml, "nvmlSystemGetCudaDriverVersion_v2", lambda: 13030) monkeypatch.setattr(health.pynvml, "nvmlDeviceGetCount", lambda: 0) - monkeypatch.setattr( - health.pynvml, "nvmlDeviceIsMigDeviceHandle", lambda _device: False - ) + monkeypatch.setattr(health.pynvml, "nvmlDeviceIsMigDeviceHandle", lambda _device: False) monkeypatch.setattr(health.pynvml, "nvmlDeviceGetUUID", lambda device: device) monkeypatch.setattr( health, @@ -47,32 +39,22 @@ def nvml_host(monkeypatch): return lifecycle -def test_collect_health_metadata_reports_nvml_gpu_and_visible_mig( - monkeypatch, nvml_host -): +def test_collect_health_metadata_reports_nvml_gpu_and_visible_mig(monkeypatch, nvml_host): def mig_device(_device, index): if index in {0, 2}: return f"mig-{index}" raise health.pynvml.NVMLError(health.pynvml.NVML_ERROR_NOT_FOUND) monkeypatch.setattr(health.pynvml, "nvmlDeviceGetCount", lambda: 1) - monkeypatch.setattr( - health.pynvml, "nvmlDeviceGetHandleByIndex", lambda _index: "gpu-0" - ) - monkeypatch.setattr( - health.pynvml, "nvmlDeviceGetName", lambda _device: "NVIDIA H100" - ) + monkeypatch.setattr(health.pynvml, "nvmlDeviceGetHandleByIndex", lambda _index: "gpu-0") + monkeypatch.setattr(health.pynvml, "nvmlDeviceGetName", lambda _device: "NVIDIA H100") monkeypatch.setattr( health.pynvml, "nvmlDeviceGetCudaComputeCapability", lambda _device: (9, 0), ) - monkeypatch.setattr( - health.pynvml, "nvmlDeviceGetMaxMigDeviceCount", lambda _device: 3 - ) - monkeypatch.setattr( - health.pynvml, "nvmlDeviceGetMigDeviceHandleByIndex", mig_device - ) + monkeypatch.setattr(health.pynvml, "nvmlDeviceGetMaxMigDeviceCount", lambda _device: 3) + monkeypatch.setattr(health.pynvml, "nvmlDeviceGetMigDeviceHandleByIndex", mig_device) monkeypatch.setattr( health.pynvml, "nvmlDeviceGetAttributes", @@ -90,9 +72,7 @@ def mig_device(_device, index): }[device], ) monkeypatch.setenv("WORKER_VERSION", "git-abc123") - monkeypatch.setitem( - sys.modules, "torch", SimpleNamespace(version=SimpleNamespace(cuda="12.4")) - ) + monkeypatch.setitem(sys.modules, "torch", SimpleNamespace(version=SimpleNamespace(cuda="12.4"))) assert health.collect_health_metadata() == { "worker_version": "git-abc123", @@ -124,19 +104,11 @@ def mig_device(_device, index): assert nvml_host == ["init", "shutdown"] -def test_collect_health_metadata_detects_directly_visible_mig_device( - monkeypatch, nvml_host -): +def test_collect_health_metadata_detects_directly_visible_mig_device(monkeypatch, nvml_host): monkeypatch.setattr(health.pynvml, "nvmlDeviceGetCount", lambda: 1) - monkeypatch.setattr( - health.pynvml, "nvmlDeviceGetHandleByIndex", lambda _index: "mig-0" - ) - monkeypatch.setattr( - health.pynvml, "nvmlDeviceIsMigDeviceHandle", lambda _device: True - ) - monkeypatch.setattr( - health.pynvml, "nvmlDeviceGetName", lambda _device: "NVIDIA H100 MIG" - ) + monkeypatch.setattr(health.pynvml, "nvmlDeviceGetHandleByIndex", lambda _index: "mig-0") + monkeypatch.setattr(health.pynvml, "nvmlDeviceIsMigDeviceHandle", lambda _device: True) + monkeypatch.setattr(health.pynvml, "nvmlDeviceGetName", lambda _device: "NVIDIA H100 MIG") monkeypatch.setattr( health.pynvml, "nvmlDeviceGetCudaComputeCapability", @@ -173,19 +145,13 @@ def test_collect_health_metadata_detects_directly_visible_mig_device( assert nvml_host == ["init", "shutdown"] -def test_collect_health_metadata_keeps_mig_detection_when_attributes_fail( - monkeypatch, nvml_host -): +def test_collect_health_metadata_keeps_mig_detection_when_attributes_fail(monkeypatch, nvml_host): def attributes(_device): raise health.pynvml.NVMLError(health.pynvml.NVML_ERROR_NOT_SUPPORTED) monkeypatch.setattr(health.pynvml, "nvmlDeviceGetCount", lambda: 1) - monkeypatch.setattr( - health.pynvml, "nvmlDeviceGetHandleByIndex", lambda _index: "mig-0" - ) - monkeypatch.setattr( - health.pynvml, "nvmlDeviceIsMigDeviceHandle", lambda _device: True - ) + monkeypatch.setattr(health.pynvml, "nvmlDeviceGetHandleByIndex", lambda _index: "mig-0") + monkeypatch.setattr(health.pynvml, "nvmlDeviceIsMigDeviceHandle", lambda _device: True) monkeypatch.setattr(health.pynvml, "nvmlDeviceGetName", lambda _device: None) monkeypatch.setattr( health.pynvml, @@ -211,15 +177,11 @@ def attributes(_device): @pytest.mark.parametrize("label", ["CUDA Version", "CUDA UMD Version"]) -def test_collect_health_metadata_falls_back_for_cuda_driver_version( - monkeypatch, nvml_host, label -): +def test_collect_health_metadata_falls_back_for_cuda_driver_version(monkeypatch, nvml_host, label): def cuda_driver_version(): raise health.pynvml.NVMLError(health.pynvml.NVML_ERROR_FUNCTION_NOT_FOUND) - monkeypatch.setattr( - health.pynvml, "nvmlSystemGetCudaDriverVersion_v2", cuda_driver_version - ) + monkeypatch.setattr(health.pynvml, "nvmlSystemGetCudaDriverVersion_v2", cuda_driver_version) monkeypatch.setattr( health, "_run_nvidia_smi", @@ -248,9 +210,7 @@ def mig_count(_device): monkeypatch.setattr(health.pynvml, "nvmlSystemGetDriverVersion", driver_version) monkeypatch.setattr(health.pynvml, "nvmlDeviceGetCount", lambda: 1) - monkeypatch.setattr( - health.pynvml, "nvmlDeviceGetHandleByIndex", lambda _index: "gpu-0" - ) + monkeypatch.setattr(health.pynvml, "nvmlDeviceGetHandleByIndex", lambda _index: "gpu-0") monkeypatch.setattr(health.pynvml, "nvmlDeviceGetName", model) monkeypatch.setattr( health.pynvml, @@ -332,9 +292,7 @@ def test_collect_health_metadata_detects_mps_from_memory_limit(monkeypatch, nvml assert nvml_host == ["init", "shutdown"] -def test_cached_health_refreshes_torch_metadata_without_reprobing_host( - monkeypatch, nvml_host -): +def test_cached_health_refreshes_torch_metadata_without_reprobing_host(monkeypatch, nvml_host): health._get_host_metadata.cache_clear() monkeypatch.setenv("CUDA_MPS_ACTIVE_THREAD_PERCENTAGE", "50") diff --git a/tests/test_init.py b/tests/test_init.py index 97ce382..9847a20 100644 --- a/tests/test_init.py +++ b/tests/test_init.py @@ -3,6 +3,7 @@ import pytest from maestro_worker_python import init +from maestro_worker_python.load_worker import load_worker def test_init_creates_complete_worker_scaffold(tmp_path, monkeypatch): @@ -27,25 +28,18 @@ def test_init_creates_complete_worker_scaffold(tmp_path, monkeypatch): pyproject = (target / "pyproject.toml").read_text() assert ( '"maestro-worker-python @ ' - 'https://github.com/moises-ai/maestro-worker-python/' - 'archive/refs/tags/4.2.0.tar.gz"' - in pyproject + "https://github.com/moises-ai/maestro-worker-python/" + 'archive/refs/tags/4.2.0.tar.gz"' in pyproject ) dockerfile = (target / "Dockerfile").read_text() assert "ARG BASE_IMAGE=python:3.12-slim-trixie" in dockerfile assert "FROM ${BASE_IMAGE}" in dockerfile - assert ( - "COPY --from=ghcr.io/astral-sh/uv:0.11.25 /uv /uvx /bin/" - in dockerfile - ) + assert "COPY --from=ghcr.io/astral-sh/uv:0.11.25 /uv /uvx /bin/" in dockerfile assert "UV_LINK_MODE=copy" in dockerfile assert "COPY pyproject.toml uv.lock ./" in dockerfile assert "--mount=type=cache,target=/root/.cache/uv" in dockerfile - assert ( - 'PYTHON_BIN="$(uv python find --no-project --no-python-downloads)"' - in dockerfile - ) + assert 'PYTHON_BIN="$(uv python find --no-project --no-python-downloads)"' in dockerfile assert 'UV_PROJECT_ENVIRONMENT="$("$PYTHON_BIN" -c' in dockerfile assert 'sysconfig.get_config_var("prefix")' in dockerfile assert "uv sync --locked --no-dev --no-install-project --inexact" in dockerfile @@ -65,10 +59,12 @@ def test_init_creates_complete_worker_scaffold(tmp_path, monkeypatch): assert "BASE_IMAGE" in readme assert "must provide a Python interpreter" in readme + generated_worker = load_worker(str(target / "worker.py")) + response = generated_worker.MoisesWorker().inference({"input_1": "Hello"}) + assert response.result == {"output": "Hello World"} + -def test_init_allows_identical_rerun_and_preserves_unrelated_files( - tmp_path, monkeypatch -): +def test_init_allows_identical_rerun_and_preserves_unrelated_files(tmp_path, monkeypatch): target = tmp_path / "existing-directory" target.mkdir() unrelated_file = target / "notes.txt" @@ -82,9 +78,7 @@ def test_init_allows_identical_rerun_and_preserves_unrelated_files( assert unrelated_file.read_text() == "keep me" -def test_init_aborts_before_overwriting_modified_scaffold_file( - tmp_path, monkeypatch, capsys -): +def test_init_aborts_before_overwriting_modified_scaffold_file(tmp_path, monkeypatch, capsys): target = tmp_path / "existing-worker" target.mkdir() worker = target / "worker.py" diff --git a/tests/test_worker_loading.py b/tests/test_load_worker.py similarity index 100% rename from tests/test_worker_loading.py rename to tests/test_load_worker.py diff --git a/tests/test_response.py b/tests/test_response.py new file mode 100644 index 0000000..699ffd5 --- /dev/null +++ b/tests/test_response.py @@ -0,0 +1,21 @@ +from maestro_worker_python.response import WorkerResponse + + +def test_worker_response_preserves_fractional_billable_seconds(): + response = WorkerResponse( + billable_seconds=42.711293, + stats={"duration": 1.25}, + result={}, + ) + + assert response.billable_seconds == 42.711293 + + +def test_worker_response_accepts_integer_billable_seconds(): + response = WorkerResponse( + billable_seconds=42, + stats={"duration": 1.25}, + result={}, + ) + + assert response.billable_seconds == 42.0 diff --git a/tests/test_serve.py b/tests/test_serve.py new file mode 100644 index 0000000..32002e5 --- /dev/null +++ b/tests/test_serve.py @@ -0,0 +1,57 @@ +import asyncio +import importlib +import sys +from pathlib import Path + +import pytest + +from maestro_worker_python.config import settings + + +def _import_serve(monkeypatch: pytest.MonkeyPatch, worker_path: Path): + monkeypatch.setattr(settings, "model_path", str(worker_path)) + monkeypatch.setattr(settings, "sentry_dsn", None) + monkeypatch.setattr(settings, "enable_json_logging", False) + sys.modules.pop("maestro_worker_python.serve", None) + return importlib.import_module("maestro_worker_python.serve") + + +@pytest.fixture +def serve_module(tmp_path, monkeypatch): + worker_path = tmp_path / "worker.py" + worker_path.write_text("class MoisesWorker:\n pass\n") + return _import_serve(monkeypatch, worker_path) + + +def test_lifespan_creates_ready_file_and_cleans_up_on_shutdown(serve_module, tmp_path, monkeypatch): + ready_path = tmp_path / "http_ready" + events: list[str] = [] + + monkeypatch.setattr(serve_module, "HTTP_READY_PATH", str(ready_path)) + monkeypatch.setattr(serve_module, "get_health_metadata", lambda: events.append("health") or {}) + monkeypatch.setattr(serve_module, "kill_child_processes", lambda: events.append("kill")) + + async def run(): + assert not ready_path.exists() + async with serve_module.lifespan(serve_module.app): + assert events == ["health"] + assert ready_path.read_text() == "Ready to serve" + assert events == ["health", "kill"] + assert not ready_path.exists() + + asyncio.run(run()) + + +def test_lifespan_shutdown_tolerates_missing_ready_file(serve_module, tmp_path, monkeypatch): + ready_path = tmp_path / "http_ready" + monkeypatch.setattr(serve_module, "HTTP_READY_PATH", str(ready_path)) + monkeypatch.setattr(serve_module, "get_health_metadata", lambda: {}) + monkeypatch.setattr(serve_module, "kill_child_processes", lambda: None) + + async def run(): + async with serve_module.lifespan(serve_module.app): + ready_path.unlink() + assert not ready_path.exists() + + asyncio.run(run()) + assert not ready_path.exists() diff --git a/tests/test_server.py b/tests/test_server.py new file mode 100644 index 0000000..50ae8d8 --- /dev/null +++ b/tests/test_server.py @@ -0,0 +1,20 @@ +import argparse + +import pytest + +from maestro_worker_python.server import _parse_bool + + +@pytest.mark.parametrize("value", ["1", "true", "TRUE", "yes", "on"]) +def test_parse_bool_accepts_true_values(value): + assert _parse_bool(value) is True + + +@pytest.mark.parametrize("value", ["0", "false", "FALSE", "no", "off"]) +def test_parse_bool_accepts_false_values(value): + assert _parse_bool(value) is False + + +def test_parse_bool_rejects_unknown_value(): + with pytest.raises(argparse.ArgumentTypeError, match="expected a boolean value"): + _parse_bool("sometimes") diff --git a/uv.lock b/uv.lock index fd276f7..70de308 100644 --- a/uv.lock +++ b/uv.lock @@ -209,6 +209,8 @@ dependencies = [ dev = [ { name = "pytest" }, { name = "pytest-httpserver" }, + { name = "ruff" }, + { name = "ty" }, ] [package.metadata] @@ -228,6 +230,8 @@ requires-dist = [ dev = [ { name = "pytest", specifier = ">=9.0.3,<10" }, { name = "pytest-httpserver", specifier = ">=1.0.6,<2" }, + { name = "ruff", specifier = ">=0.14.9,<0.15" }, + { name = "ty", specifier = "==0.0.14" }, ] [[package]] @@ -513,6 +517,32 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/a0/f4/c67b0b3f1b9245e8d266f0f112c500d50e5b4e83cb6f3b71b6528104182a/requests-2.34.2-py3-none-any.whl", hash = "sha256:2a0d60c172f83ac6ab31e4554906c0f3b3588d37b5cb939b1c061f4907e278e0", size = 73075, upload-time = "2026-05-14T19:25:26.443Z" }, ] +[[package]] +name = "ruff" +version = "0.14.14" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/2e/06/f71e3a86b2df0dfa2d2f72195941cd09b44f87711cb7fa5193732cb9a5fc/ruff-0.14.14.tar.gz", hash = "sha256:2d0f819c9a90205f3a867dbbd0be083bee9912e170fd7d9704cc8ae45824896b", size = 4515732, upload-time = "2026-01-22T22:30:17.527Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d2/89/20a12e97bc6b9f9f68343952da08a8099c57237aef953a56b82711d55edd/ruff-0.14.14-py3-none-linux_armv6l.whl", hash = "sha256:7cfe36b56e8489dee8fbc777c61959f60ec0f1f11817e8f2415f429552846aed", size = 10467650, upload-time = "2026-01-22T22:30:08.578Z" }, + { url = "https://files.pythonhosted.org/packages/a3/b1/c5de3fd2d5a831fcae21beda5e3589c0ba67eec8202e992388e4b17a6040/ruff-0.14.14-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:6006a0082336e7920b9573ef8a7f52eec837add1265cc74e04ea8a4368cd704c", size = 10883245, upload-time = "2026-01-22T22:30:04.155Z" }, + { url = "https://files.pythonhosted.org/packages/b8/7c/3c1db59a10e7490f8f6f8559d1db8636cbb13dccebf18686f4e3c9d7c772/ruff-0.14.14-py3-none-macosx_11_0_arm64.whl", hash = "sha256:026c1d25996818f0bf498636686199d9bd0d9d6341c9c2c3b62e2a0198b758de", size = 10231273, upload-time = "2026-01-22T22:30:34.642Z" }, + { url = "https://files.pythonhosted.org/packages/a1/6e/5e0e0d9674be0f8581d1f5e0f0a04761203affce3232c1a1189d0e3b4dad/ruff-0.14.14-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f666445819d31210b71e0a6d1c01e24447a20b85458eea25a25fe8142210ae0e", size = 10585753, upload-time = "2026-01-22T22:30:31.781Z" }, + { url = "https://files.pythonhosted.org/packages/23/09/754ab09f46ff1884d422dc26d59ba18b4e5d355be147721bb2518aa2a014/ruff-0.14.14-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:3c0f18b922c6d2ff9a5e6c3ee16259adc513ca775bcf82c67ebab7cbd9da5bc8", size = 10286052, upload-time = "2026-01-22T22:30:24.827Z" }, + { url = "https://files.pythonhosted.org/packages/c8/cc/e71f88dd2a12afb5f50733851729d6b571a7c3a35bfdb16c3035132675a0/ruff-0.14.14-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:1629e67489c2dea43e8658c3dba659edbfd87361624b4040d1df04c9740ae906", size = 11043637, upload-time = "2026-01-22T22:30:13.239Z" }, + { url = "https://files.pythonhosted.org/packages/67/b2/397245026352494497dac935d7f00f1468c03a23a0c5db6ad8fc49ca3fb2/ruff-0.14.14-py3-none-manylinux_2_17_ppc64.manylinux2014_ppc64.whl", hash = "sha256:27493a2131ea0f899057d49d303e4292b2cae2bb57253c1ed1f256fbcd1da480", size = 12194761, upload-time = "2026-01-22T22:30:22.542Z" }, + { url = "https://files.pythonhosted.org/packages/5b/06/06ef271459f778323112c51b7587ce85230785cd64e91772034ddb88f200/ruff-0.14.14-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:01ff589aab3f5b539e35db38425da31a57521efd1e4ad1ae08fc34dbe30bd7df", size = 12005701, upload-time = "2026-01-22T22:30:20.499Z" }, + { url = "https://files.pythonhosted.org/packages/41/d6/99364514541cf811ccc5ac44362f88df66373e9fec1b9d1c4cc830593fe7/ruff-0.14.14-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:1cc12d74eef0f29f51775f5b755913eb523546b88e2d733e1d701fe65144e89b", size = 11282455, upload-time = "2026-01-22T22:29:59.679Z" }, + { url = "https://files.pythonhosted.org/packages/ca/71/37daa46f89475f8582b7762ecd2722492df26421714a33e72ccc9a84d7a5/ruff-0.14.14-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:bb8481604b7a9e75eff53772496201690ce2687067e038b3cc31aaf16aa0b974", size = 11215882, upload-time = "2026-01-22T22:29:57.032Z" }, + { url = "https://files.pythonhosted.org/packages/2c/10/a31f86169ec91c0705e618443ee74ede0bdd94da0a57b28e72db68b2dbac/ruff-0.14.14-py3-none-manylinux_2_31_riscv64.whl", hash = "sha256:14649acb1cf7b5d2d283ebd2f58d56b75836ed8c6f329664fa91cdea19e76e66", size = 11180549, upload-time = "2026-01-22T22:30:27.175Z" }, + { url = "https://files.pythonhosted.org/packages/fd/1e/c723f20536b5163adf79bdd10c5f093414293cdf567eed9bdb7b83940f3f/ruff-0.14.14-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:e8058d2145566510790eab4e2fad186002e288dec5e0d343a92fe7b0bc1b3e13", size = 10543416, upload-time = "2026-01-22T22:30:01.964Z" }, + { url = "https://files.pythonhosted.org/packages/3e/34/8a84cea7e42c2d94ba5bde1d7a4fae164d6318f13f933d92da6d7c2041ff/ruff-0.14.14-py3-none-musllinux_1_2_armv7l.whl", hash = "sha256:e651e977a79e4c758eb807f0481d673a67ffe53cfa92209781dfa3a996cf8412", size = 10285491, upload-time = "2026-01-22T22:30:29.51Z" }, + { url = "https://files.pythonhosted.org/packages/55/ef/b7c5ea0be82518906c978e365e56a77f8de7678c8bb6651ccfbdc178c29f/ruff-0.14.14-py3-none-musllinux_1_2_i686.whl", hash = "sha256:cc8b22da8d9d6fdd844a68ae937e2a0adf9b16514e9a97cc60355e2d4b219fc3", size = 10733525, upload-time = "2026-01-22T22:30:06.499Z" }, + { url = "https://files.pythonhosted.org/packages/6a/5b/aaf1dfbcc53a2811f6cc0a1759de24e4b03e02ba8762daabd9b6bd8c59e3/ruff-0.14.14-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:16bc890fb4cc9781bb05beb5ab4cd51be9e7cb376bf1dd3580512b24eb3fda2b", size = 11315626, upload-time = "2026-01-22T22:30:36.848Z" }, + { url = "https://files.pythonhosted.org/packages/2c/aa/9f89c719c467dfaf8ad799b9bae0df494513fb21d31a6059cb5870e57e74/ruff-0.14.14-py3-none-win32.whl", hash = "sha256:b530c191970b143375b6a68e6f743800b2b786bbcf03a7965b06c4bf04568167", size = 10502442, upload-time = "2026-01-22T22:30:38.93Z" }, + { url = "https://files.pythonhosted.org/packages/87/44/90fa543014c45560cae1fffc63ea059fb3575ee6e1cb654562197e5d16fb/ruff-0.14.14-py3-none-win_amd64.whl", hash = "sha256:3dde1435e6b6fe5b66506c1dff67a421d0b7f6488d466f651c07f4cab3bf20fd", size = 11630486, upload-time = "2026-01-22T22:30:10.852Z" }, + { url = "https://files.pythonhosted.org/packages/9e/6a/40fee331a52339926a92e17ae748827270b288a35ef4a15c9c8f2ec54715/ruff-0.14.14-py3-none-win_arm64.whl", hash = "sha256:56e6981a98b13a32236a72a8da421d7839221fa308b223b9283312312e5ac76c", size = 10920448, upload-time = "2026-01-22T22:30:15.417Z" }, +] + [[package]] name = "sentry-sdk" version = "2.66.0" @@ -580,6 +610,30 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/7b/61/cceae43728b7de99d9b847560c262873a1f6c98202171fd5ed62640b494b/tomli-2.4.1-py3-none-any.whl", hash = "sha256:0d85819802132122da43cb86656f8d1f8c6587d54ae7dcaf30e90533028b49fe", size = 14583, upload-time = "2026-03-25T20:22:03.012Z" }, ] +[[package]] +name = "ty" +version = "0.0.14" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/af/57/22c3d6bf95c2229120c49ffc2f0da8d9e8823755a1c3194da56e51f1cc31/ty-0.0.14.tar.gz", hash = "sha256:a691010565f59dd7f15cf324cdcd1d9065e010c77a04f887e1ea070ba34a7de2", size = 5036573, upload-time = "2026-01-27T00:57:31.427Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/99/cb/cc6d1d8de59beb17a41f9a614585f884ec2d95450306c173b3b7cc090d2e/ty-0.0.14-py3-none-linux_armv6l.whl", hash = "sha256:32cf2a7596e693094621d3ae568d7ee16707dce28c34d1762947874060fdddaa", size = 10034228, upload-time = "2026-01-27T00:57:53.133Z" }, + { url = "https://files.pythonhosted.org/packages/f3/96/dd42816a2075a8f31542296ae687483a8d047f86a6538dfba573223eaf9a/ty-0.0.14-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:f971bf9805f49ce8c0968ad53e29624d80b970b9eb597b7cbaba25d8a18ce9a2", size = 9939162, upload-time = "2026-01-27T00:57:43.857Z" }, + { url = "https://files.pythonhosted.org/packages/ff/b4/73c4859004e0f0a9eead9ecb67021438b2e8e5fdd8d03e7f5aca77623992/ty-0.0.14-py3-none-macosx_11_0_arm64.whl", hash = "sha256:45448b9e4806423523268bc15e9208c4f3f2ead7c344f615549d2e2354d6e924", size = 9418661, upload-time = "2026-01-27T00:58:03.411Z" }, + { url = "https://files.pythonhosted.org/packages/58/35/839c4551b94613db4afa20ee555dd4f33bfa7352d5da74c5fa416ffa0fd2/ty-0.0.14-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:ee94a9b747ff40114085206bdb3205a631ef19a4d3fb89e302a88754cbbae54c", size = 9837872, upload-time = "2026-01-27T00:57:23.718Z" }, + { url = "https://files.pythonhosted.org/packages/41/2b/bbecf7e2faa20c04bebd35fc478668953ca50ee5847ce23e08acf20ea119/ty-0.0.14-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:6756715a3c33182e9ab8ffca2bb314d3c99b9c410b171736e145773ee0ae41c3", size = 9848819, upload-time = "2026-01-27T00:57:58.501Z" }, + { url = "https://files.pythonhosted.org/packages/be/60/3c0ba0f19c0f647ad9d2b5b5ac68c0f0b4dc899001bd53b3a7537fb247a2/ty-0.0.14-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:89d0038a2f698ba8b6fec5cf216a4e44e2f95e4a5095a8c0f57fe549f87087c2", size = 10324371, upload-time = "2026-01-27T00:57:29.291Z" }, + { url = "https://files.pythonhosted.org/packages/24/32/99d0a0b37d0397b0a989ffc2682493286aa3bc252b24004a6714368c2c3d/ty-0.0.14-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:2c64a83a2d669b77f50a4957039ca1450626fb474619f18f6f8a3eb885bf7544", size = 10865898, upload-time = "2026-01-27T00:57:33.542Z" }, + { url = "https://files.pythonhosted.org/packages/1a/88/30b583a9e0311bb474269cfa91db53350557ebec09002bfc3fb3fc364e8c/ty-0.0.14-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:242488bfb547ef080199f6fd81369ab9cb638a778bb161511d091ffd49c12129", size = 10555777, upload-time = "2026-01-27T00:58:05.853Z" }, + { url = "https://files.pythonhosted.org/packages/cd/a2/cb53fb6325dcf3d40f2b1d0457a25d55bfbae633c8e337bde8ec01a190eb/ty-0.0.14-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:4790c3866f6c83a4f424fc7d09ebdb225c1f1131647ba8bdc6fcdc28f09ed0ff", size = 10412913, upload-time = "2026-01-27T00:57:38.834Z" }, + { url = "https://files.pythonhosted.org/packages/42/8f/f2f5202d725ed1e6a4e5ffaa32b190a1fe70c0b1a2503d38515da4130b4c/ty-0.0.14-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:950f320437f96d4ea9a2332bbfb5b68f1c1acd269ebfa4c09b6970cc1565bd9d", size = 9837608, upload-time = "2026-01-27T00:57:55.898Z" }, + { url = "https://files.pythonhosted.org/packages/f7/ba/59a2a0521640c489dafa2c546ae1f8465f92956fede18660653cce73b4c5/ty-0.0.14-py3-none-musllinux_1_2_armv7l.whl", hash = "sha256:4a0ec3ee70d83887f86925bbc1c56f4628bd58a0f47f6f32ddfe04e1f05466df", size = 9884324, upload-time = "2026-01-27T00:57:46.786Z" }, + { url = "https://files.pythonhosted.org/packages/03/95/8d2a49880f47b638743212f011088552ecc454dd7a665ddcbdabea25772a/ty-0.0.14-py3-none-musllinux_1_2_i686.whl", hash = "sha256:a1a4e6b6da0c58b34415955279eff754d6206b35af56a18bb70eb519d8d139ef", size = 10033537, upload-time = "2026-01-27T00:58:01.149Z" }, + { url = "https://files.pythonhosted.org/packages/e9/40/4523b36f2ce69f92ccf783855a9e0ebbbd0f0bb5cdce6211ee1737159ed3/ty-0.0.14-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:dc04384e874c5de4c5d743369c277c8aa73d1edea3c7fc646b2064b637db4db3", size = 10495910, upload-time = "2026-01-27T00:57:26.691Z" }, + { url = "https://files.pythonhosted.org/packages/08/d5/655beb51224d1bfd4f9ddc0bb209659bfe71ff141bcf05c418ab670698f0/ty-0.0.14-py3-none-win32.whl", hash = "sha256:b20e22cf54c66b3e37e87377635da412d9a552c9bf4ad9fc449fed8b2e19dad2", size = 9507626, upload-time = "2026-01-27T00:57:41.43Z" }, + { url = "https://files.pythonhosted.org/packages/b6/d9/c569c9961760e20e0a4bc008eeb1415754564304fd53997a371b7cf3f864/ty-0.0.14-py3-none-win_amd64.whl", hash = "sha256:e312ff9475522d1a33186657fe74d1ec98e4a13e016d66f5758a452c90ff6409", size = 10437980, upload-time = "2026-01-27T00:57:36.422Z" }, + { url = "https://files.pythonhosted.org/packages/ad/0c/186829654f5bfd9a028f6648e9caeb11271960a61de97484627d24443f91/ty-0.0.14-py3-none-win_arm64.whl", hash = "sha256:b6facdbe9b740cb2c15293a1d178e22ffc600653646452632541d01c36d5e378", size = 9885831, upload-time = "2026-01-27T00:57:49.747Z" }, +] + [[package]] name = "typing-extensions" version = "4.16.0"