Fix CUDA probing for mixed CPU/GPU releases

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

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"

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

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)

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