From d3defb4da92225ea47bc8933f99014fdc0fce40f Mon Sep 17 00:00:00 2001 From: Paul Fidika Date: Tue, 14 Jul 2026 12:01:46 -0600 Subject: [PATCH] fix: probe CUDA per concrete worker image --- CHANGELOG.md | 10 ++++++ pyproject.toml | 2 +- src/gen_worker/cuda_probe.py | 28 ++++++++++++++++ src/gen_worker/entrypoint.py | 4 +-- tests/test_cuda_probe.py | 65 +++++++++++++++++++++++++++++++++++- uv.lock | 2 +- 6 files changed, 106 insertions(+), 5 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 0b15ef9..7b48704 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,15 @@ # Changelog +## 0.25.2 (2026-07-14) + +- **Mixed CPU/GPU releases probe the device owned by their concrete image.** + A release-level discovery manifest contains both lanes, so the prior + any-GPU-function check killed the CPU conversion image before worker hello. + Mixed manifests now use the installed Torch build as the lane signal: CUDA + images retain the bad-host health probe, while CPU-only images start their + CPU functions without an environment-variable override. GPU-only manifests + still fail closed when CUDA is absent. + ## 0.25.0 (2026-07-14) - **cl#27: local-only GGUF fit rung.** A bare Tensorhub binding can select the diff --git a/pyproject.toml b/pyproject.toml index 710a97a..5bc0736 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "gen-worker" -version = "0.25.1" +version = "0.25.2" description = "A library used to build custom functions in Cozy Creator's serverless function platform." readme = "README.md" license = "MIT" diff --git a/src/gen_worker/cuda_probe.py b/src/gen_worker/cuda_probe.py index 7d51cf9..0169e16 100644 --- a/src/gen_worker/cuda_probe.py +++ b/src/gen_worker/cuda_probe.py @@ -56,3 +56,31 @@ def manifest_needs_cuda(manifest: Optional[dict[str, Any]]) -> bool: if bool((fn.get("resources") or {}).get("gpu")): return True return False + + +def should_probe_cuda( + manifest: Optional[dict[str, Any]], *, cuda_build: Optional[bool] = None +) -> bool: + """Whether this concrete worker image must pass the CUDA health probe. + + A manifest may contain both CPU and GPU functions because one endpoint + release can publish separate ``accelerator=none`` and ``accelerator=cuda`` + images. In that mixed case the installed torch build is the authoritative + signal: probe CUDA images and let CPU-only images serve the CPU lane. A + GPU-only manifest is always probed so an accidentally CPU-built image + fails before it can register. + """ + functions = (manifest or {}).get("functions", []) or [] + gpu_requirements = [bool((fn.get("resources") or {}).get("gpu")) for fn in functions] + if not any(gpu_requirements): + return False + if all(gpu_requirements): + return True + if cuda_build is None: + try: + import torch + + cuda_build = bool(torch.version.cuda) + except Exception: + return True + return cuda_build diff --git a/src/gen_worker/entrypoint.py b/src/gen_worker/entrypoint.py index 203adc0..46555ce 100644 --- a/src/gen_worker/entrypoint.py +++ b/src/gen_worker/entrypoint.py @@ -33,7 +33,7 @@ import msgspec from .config import get_settings -from .cuda_probe import CUDA_PROBE_FAILED_MARKER, manifest_needs_cuda, probe_cuda +from .cuda_probe import CUDA_PROBE_FAILED_MARKER, probe_cuda, should_probe_cuda from .models.cache_paths import tensorhub_cas_dir try: from .worker import Worker @@ -215,7 +215,7 @@ def _run_main() -> int: # device actually works BEFORE we hello the orchestrator and accept a # job — a busy/unavailable GPU (RunPod bad-host fault) must kill this # pod now, not terminal-fail a real request at model load. - if manifest_needs_cuda(manifest): + if should_probe_cuda(manifest): probe = probe_cuda() if not probe.ok: logger.error("%s: %s", CUDA_PROBE_FAILED_MARKER, probe.reason) diff --git a/tests/test_cuda_probe.py b/tests/test_cuda_probe.py index ba31ac6..6de017d 100644 --- a/tests/test_cuda_probe.py +++ b/tests/test_cuda_probe.py @@ -13,7 +13,12 @@ import torch from gen_worker import entrypoint -from gen_worker.cuda_probe import CUDA_PROBE_FAILED_MARKER, manifest_needs_cuda, probe_cuda +from gen_worker.cuda_probe import ( + CUDA_PROBE_FAILED_MARKER, + manifest_needs_cuda, + probe_cuda, + should_probe_cuda, +) class _FakeTensor: @@ -77,6 +82,30 @@ def test_manifest_needs_cuda_false_for_missing_manifest() -> None: assert manifest_needs_cuda({}) is False +@pytest.mark.parametrize( + ("gpu_flags", "cuda_build", "expected"), + [ + ([False], False, False), + ([False], True, False), + ([True], False, True), + ([True], True, True), + ([False, True], False, False), + ([False, True], True, True), + ], +) +def test_should_probe_cuda_uses_torch_build_only_for_mixed_manifests( + gpu_flags: list[bool], cuda_build: bool, expected: bool +) -> None: + manifest = { + "functions": [ + {"name": str(i), "resources": {"gpu": True} if gpu else {}} + for i, gpu in enumerate(gpu_flags) + ] + } + + assert should_probe_cuda(manifest, cuda_build=cuda_build) is expected + + def _write_manifest(path: Path, *, gpu: bool) -> None: manifest = { "functions": [ @@ -146,3 +175,37 @@ def run(self) -> int: code = entrypoint._run_main() assert code == 0 + + +def test_entrypoint_skips_probe_for_mixed_manifest_in_cpu_torch_image( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + manifest_path = tmp_path / "endpoint.lock" + manifest_path.write_bytes( + msgspec.toml.encode( + { + "functions": [ + {"name": "cpu", "module": "fake_mod", "resources": {}}, + {"name": "gpu", "module": "fake_mod", "resources": {"gpu": True}}, + ] + } + ) + ) + _base_env(monkeypatch, tmp_path, manifest_path) + monkeypatch.setattr(torch.version, "cuda", None) + + def _unexpected_probe(*a: Any, **k: Any) -> None: + raise AssertionError("a mixed release's CPU image must not probe CUDA") + + monkeypatch.setattr(entrypoint, "probe_cuda", _unexpected_probe) + + class _StubWorker: + def __init__(self, *a: Any, **k: Any) -> None: + pass + + def run(self) -> int: + return 0 + + monkeypatch.setattr(entrypoint, "Worker", _StubWorker) + + assert entrypoint._run_main() == 0 diff --git a/uv.lock b/uv.lock index facd7a1..87da043 100644 --- a/uv.lock +++ b/uv.lock @@ -583,7 +583,7 @@ wheels = [ [[package]] name = "gen-worker" -version = "0.25.1" +version = "0.25.2" source = { editable = "." } dependencies = [ { name = "blake3" },