flashruntime 0.4.0__tar.gz → 0.4.1__tar.gz

{flashruntime-0.4.0/flashruntime.egg-info → flashruntime-0.4.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: flashruntime
-Version: 0.4.0
+Version: 0.4.1
 Summary: Open fault-tolerant distributed ML runtime: strategy planning, job protocol, leases, checkpointing, and recovery across heterogeneous compute.
 Author: Zolli Labs
 License-Expression: Apache-2.0

{flashruntime-0.4.0 → flashruntime-0.4.1}/flashruntime/protocol/v1alpha1.py

@@ -78,6 +78,13 @@ class ResourcesSpec(BaseModel):
     maximumWorkers: int = Field(ge=1, default=3)
     cpuPerTask: float = Field(gt=0, default=1)
     memoryPerTask: str = "512Mi"
+    #: GPUs each task needs. 0 — the default, and every job that exists
+    #: today — means "no GPU required" and places anywhere. A non-zero value
+    #: is a hardware *requirement*, enforced fail-closed at placement
+    #: (`scheduler.IsolationAwarePlacement`): a CUDA task on a CPU-only host
+    #: does not politely requeue, it crashes or silently runs on the CPU two
+    #: orders of magnitude slower while reporting success.
+    gpuPerTask: int = Field(ge=0, default=0)
 
     @field_validator("maximumWorkers")
     @classmethod
@@ -269,10 +276,39 @@ class NodeEnvironment(str, Enum):
     EDGE = "edge"
 
 
+class GpuInfo(BaseModel):
+    """One GPU as the host's driver reports it.
+
+    Every field but `index` is optional: a probe that cannot read a value says
+    nothing rather than guessing, so a partial reading is still reportable and
+    an unreadable one is still `[]` rather than a fabricated device.
+
+    Typed now although placement matches only on COUNT, because the wire
+    format is the expensive thing to change later — these agents run on
+    machines we cannot reach, and a field added to an existing model is cheap
+    where turning a `dict` into a model is not. `memory_total_mb` and
+    `compute_capability` are collected from the start so the data is already
+    flowing when matching rules arrive (spec §10.1).
+
+    No `schema_version`: like `CheckpointPart`, this is a nested element of a
+    versioned message (`NodeRegistration`), not a wire message of its own.
+    """
+
+    index: int
+    name: str = ""
+    memory_total_mb: int | None = None
+    driver_version: str = ""
+    compute_capability: str = ""
+
+
 class NodeCapabilities(BaseModel):
     cpu_cores: float | None = None
     memory_bytes: int | None = None
-    gpus: list[dict[str, Any]] = Field(default_factory=list)
+    #: Devices this host's driver reports, one entry per GPU. Empty is the
+    #: honest answer for a host with no driver, no `nvidia-smi`, or output
+    #: this runtime cannot parse — never a guess. Placement reads the LENGTH
+    #: of this list and nothing else in v1.
+    gpus: list[GpuInfo] = Field(default_factory=list)
     os: str = ""
     architecture: str = ""
 

{flashruntime-0.4.0 → flashruntime-0.4.1}/flashruntime/recipes/command.py

@@ -160,6 +160,26 @@ class CommandRecipe(WorkloadRecipe):
                 #
                 # `list(...)` so a payload never aliases the caller's spec.
                 payload["local_inputs"] = list(p["local_inputs"])
+            gpus = spec.spec.resources.gpuPerTask
+            if gpus:
+                # The GPU requirement is the ONLY payload key sourced from
+                # `spec.spec.resources` rather than `workload.parameters` —
+                # a branch of the spec nothing else in this method reads.
+                # That makes it the easiest forward in the file to leave out
+                # and the hardest to notice missing: there is no unrecognised
+                # parameter sitting in `p` to go looking for.
+                #
+                # Dropping it does NOT fail closed. `IsolationAwarePlacement`
+                # reads `task.payload["gpus"]`; with the key absent it sees a
+                # task requiring nothing, places it on any node, and a CUDA
+                # job lands on a CPU-only volunteer. Both ends of this hop
+                # have tests that pass while it is broken, because each
+                # constructs the payload directly.
+                #
+                # Absent stays absent, never 0 — the no-GPU path must keep
+                # exercising the key-missing branch, as unpack_inputs and
+                # local_inputs do.
+                payload["gpus"] = int(gpus)
             tasks.append(
                 TaskSpec(
                     task_id=task_id,

{flashruntime-0.4.0 → flashruntime-0.4.1}/flashruntime/scheduler/__init__.py

@@ -172,6 +172,43 @@ class IsolationAwarePlacement(PlacementPolicy):
     is not in its map) — failing closed here makes it fail before anything
     touches the data, and before an attempt is burned.
 
+    A fifth gate applies to tasks whose payload asks for `gpus: N`: the
+    claiming node's `capabilities.gpus` must be a list of at least N entries.
+    It takes the argv/local-data polarity — **fail closed** — and deliberately
+    NOT the module gate's fail-open one, even though "the node cannot run
+    this" sounds like the same availability concern:
+
+    - A misplaced module task wastes retry attempts and nothing else. A CUDA
+      job on a CPU-only box does not politely fail and requeue. It either
+      crashes on `torch.cuda.is_available()` or, worse, silently falls back
+      to the CPU and runs two orders of magnitude slower while reporting
+      success. The second outcome is not a failure anything here can detect;
+      it is a bill.
+    - The capability must be a genuine *list*, because its LENGTH is the
+      whole matching rule in v1. Absent, `None`, a bare string, a dict, or a
+      bare `int` count as NO GPUs. The bare `int` deserves naming: `1` reads
+      like "one GPU" and is exactly what a hand-written node view would put
+      there, but accepting it would mean a second, looser matching rule
+      beside the one `NodeCapabilities.gpus` actually feeds.
+    - `capabilities` itself may be absent or type-confused; that is read as
+      no GPUs rather than allowed to raise. `(node.get("capabilities") or
+      {}).get(...)` is NOT sufficient for this — a string capabilities value
+      has no `.get` and would crash the predicate.
+    - The requirement itself must be a non-negative `int`. Anything else
+      (`"1"`, `-1`, `1.5`, a list) makes the task ineligible everywhere
+      rather than crashing, exactly as a non-list `local_inputs` does.
+      **`bool` is a subclass of `int`**, so `True` is excluded explicitly: a
+      `gpus: true` typo must not silently mean "one GPU" and place real work.
+    - `gpus: 0` requires nothing and runs anywhere, exactly like tier
+      `standard` and an empty `local_inputs`.
+
+    The gate is ONE-DIRECTIONAL. A node with GPUs still receives CPU work;
+    reserving GPU hosts for GPU jobs is a scheduling optimisation and a
+    separate decision, and making it a gate here would idle the scarcest
+    hardware on the network. `allowFallback` does not waive this gate either
+    — hardware either exists on a host or it does not, and the submitter's
+    isolation posture has nothing to say about it.
+
     Everything genuinely standard keeps the fail-open placement default."""
 
     def eligible(self, task: TaskSpec, node: NodeView) -> bool:
@@ -199,6 +236,28 @@ class IsolationAwarePlacement(PlacementPolicy):
                 return False  # absent/None/type-confused capability ⇒ not capable
             if any(name not in advertised for name in local_inputs):
                 return False
+        # Fail-closed like the argv and local-data gates, and checked before
+        # the allowFallback waiver below for the same reason: a submitter
+        # cannot waive their way onto hardware a host does not have.
+        required_gpus = task.payload.get("gpus")
+        if required_gpus is not None:
+            # `bool` is a subclass of `int`: without the explicit exclusion,
+            # a `gpus: true` typo would read as "1 GPU" and place real work.
+            if (
+                not isinstance(required_gpus, int)
+                or isinstance(required_gpus, bool)
+                or required_gpus < 0
+            ):
+                return False  # type-confused requirement ⇒ fail closed, no crash
+            if required_gpus > 0:
+                capabilities = node.get("capabilities")
+                # isinstance, not `or {}` — a string capabilities value has no
+                # `.get` and must fail closed rather than crash the predicate.
+                advertised = (
+                    capabilities.get("gpus") if isinstance(capabilities, dict) else None
+                )
+                if not isinstance(advertised, list) or len(advertised) < required_gpus:
+                    return False  # absent/short/type-confused ⇒ not capable
         isolation = task.payload.get("isolation")
         if isolation is None:
             return True  # no isolation payload ⇒ standard, runs anywhere

{flashruntime-0.4.0 → flashruntime-0.4.1/flashruntime.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: flashruntime
-Version: 0.4.0
+Version: 0.4.1
 Summary: Open fault-tolerant distributed ML runtime: strategy planning, job protocol, leases, checkpointing, and recovery across heterogeneous compute.
 Author: Zolli Labs
 License-Expression: Apache-2.0

{flashruntime-0.4.0 → flashruntime-0.4.1}/flashruntime.egg-info/SOURCES.txt

@@ -118,10 +118,13 @@ tests/test_leases.py
 tests/test_leases_scope.py
 tests/test_leases_sqlite.py
 tests/test_monitor_sampler.py
+tests/test_placement_gpu.py
 tests/test_placement_local_data.py
 tests/test_planner.py
+tests/test_protocol_gpu.py
 tests/test_protocol_local_datasets.py
 tests/test_protocol_v1alpha1.py
+tests/test_recipe_gpu.py
 tests/test_recipe_local_inputs.py
 tests/test_recipes_command.py
 tests/test_recovery.py

{flashruntime-0.4.0 → flashruntime-0.4.1}/pyproject.toml

@@ -5,12 +5,23 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "flashruntime"
-# 0.4.0 adds NodeRegistration.local_datasets — a new wire field, so consumers
-# must pin a range that includes it. flashnode’s floor moves to >=0.4,<0.5.
-# (0.3.0 was the first published release; 0.1.0 and 0.2.0 were declared here
-# but never reached PyPI.) Released from the Zolli-Labs/flashml monorepo by
-# pushing a `flashruntime-v0.4.0` tag.
-version = "0.4.0"
+# 0.4.1 adds GpuInfo and ResourcesSpec.gpuPerTask. Both are additive wire
+# fields, so the compatibility range does not move — but the PATCH bump is
+# load-bearing, not bookkeeping. The `flashruntime-v0.4.0` tag was cut BEFORE
+# the protocol commit that added them, and this tree went on calling itself
+# 0.4.0 afterwards: one version string named two different protocols, which is
+# exactly the drift the monorepo consolidation removed. Bump on any
+# wire-visible change, even an additive one.
+#
+# flashnode’s floor moves to >=0.4.1,<0.5 in the same change — it imports
+# GpuInfo at module scope, so for the agent this floor is an import-time hard
+# requirement, not a silent-degradation one.
+#
+# 0.4.0 added NodeRegistration.local_datasets — a new wire field, so consumers
+# must pin a range that includes it. (0.3.0 was the first published release;
+# 0.1.0 and 0.2.0 were declared here but never reached PyPI.) Released from
+# the Zolli-Labs/flashml monorepo by pushing a `flashruntime-vX.Y.Z` tag.
+version = "0.4.1"
 description = "Open fault-tolerant distributed ML runtime: strategy planning, job protocol, leases, checkpointing, and recovery across heterogeneous compute."
 readme = "README.md"
 # PEP 639 SPDX expression (string form) — replaces the deprecated

flashruntime-0.4.1/tests/test_placement_gpu.py

@@ -0,0 +1,301 @@
+"""The fifth fail-closed placement gate: a task whose payload asks for `gpus: N`
+may only be leased to a node whose `capabilities.gpus` is a list of at least N
+entries (AGENTS.md rule 3).
+
+Polarity follows `argv_capable` / `local_datasets`, not `module_capable`. A
+misplaced module task only wastes retry attempts; a CUDA job on a CPU-only box
+either crashes on `torch.cuda.is_available()` or silently falls back to the CPU
+and runs two orders of magnitude slower while reporting success. Neither is
+something to discover from a bill.
+
+The gate is one-directional: a GPU node still receives CPU work.
+"""
+
+from __future__ import annotations
+
+import pytest
+
+
+def _gpu_task(gpus=1, **payload_extra):
+    """A task requiring GPUs. `gpus` is passed through RAW so tests can poison
+    it with type-confused values the protocol would never emit (the same trick
+    `_local_task` plays on `local_inputs`)."""
+    from flashruntime.protocol.v1alpha1 import TaskSpec
+
+    payload = {"module": "flashml_workloads.sklearn_trial"}
+    if gpus is not None:
+        payload["gpus"] = gpus
+    payload.update(payload_extra)
+    return TaskSpec(
+        task_id="task-000", job_id="job-a", commit_key="job-a/task-000/m.json",
+        payload=payload,
+    )
+
+
+def _node(gpu_count=0, **extra):
+    """A node view shaped like the claim endpoint's: the GPU list lives under
+    `capabilities`, which is `NodeCapabilities.model_dump()` — plain dicts by
+    the time placement sees them, not GpuInfo instances."""
+    node = {
+        "node_id": "n1",
+        "capabilities": {
+            "cpu_cores": 8,
+            "gpus": [{"index": i, "name": "NVIDIA A10G"} for i in range(gpu_count)],
+        },
+    }
+    node.update(extra)
+    return node
+
+
+# ---------------------------------------------------------------------------
+# DoD 3–6: the four placement outcomes
+# ---------------------------------------------------------------------------
+
+
+def test_one_gpu_job_is_refused_on_a_node_advertising_none():
+    from flashruntime.scheduler import IsolationAwarePlacement
+
+    assert IsolationAwarePlacement().eligible(_gpu_task(1), _node(0)) is False
+
+
+def test_one_gpu_job_is_placed_on_a_node_advertising_one():
+    from flashruntime.scheduler import IsolationAwarePlacement
+
+    assert IsolationAwarePlacement().eligible(_gpu_task(1), _node(1)) is True
+
+
+def test_two_gpu_job_is_refused_on_a_one_gpu_node():
+    from flashruntime.scheduler import IsolationAwarePlacement
+
+    assert IsolationAwarePlacement().eligible(_gpu_task(2), _node(1)) is False
+
+
+def test_two_gpu_job_is_placed_on_a_two_gpu_node():
+    from flashruntime.scheduler import IsolationAwarePlacement
+
+    assert IsolationAwarePlacement().eligible(_gpu_task(2), _node(2)) is True
+
+
+def test_a_job_with_no_gpu_requirement_is_placed_on_both():
+    """The key must be ABSENT, not 0 — this is the branch every job that
+    exists today takes."""
+    from flashruntime.scheduler import IsolationAwarePlacement
+
+    policy = IsolationAwarePlacement()
+    task = _gpu_task(None)  # no "gpus" key in the payload at all
+    assert policy.eligible(task, _node(0)) is True
+    assert policy.eligible(task, _node(2)) is True
+
+
+def test_gpus_zero_requires_nothing_and_runs_anywhere():
+    """`gpus: 0` is what an explicit opt-out compiles to. It demands nothing,
+    like tier "standard" and an empty `local_inputs`."""
+    from flashruntime.scheduler import IsolationAwarePlacement
+
+    policy = IsolationAwarePlacement()
+    assert policy.eligible(_gpu_task(0), _node(0)) is True
+    assert policy.eligible(_gpu_task(0), {"node_id": "n1"}) is True
+
+
+def test_a_gpu_node_still_receives_cpu_work():
+    """The gate is one-directional. Reserving GPU hosts for GPU jobs is a
+    separate scheduling decision; making it a gate here would idle the
+    scarcest hardware on the network."""
+    from flashruntime.scheduler import IsolationAwarePlacement
+
+    assert IsolationAwarePlacement().eligible(_gpu_task(None), _node(4)) is True
+
+
+# ---------------------------------------------------------------------------
+# Fail-closed: the advertised capability
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+    "advertised",
+    [
+        pytest.param([], id="empty-list"),
+        pytest.param(None, id="explicit-none"),
+        pytest.param("gpu", id="bare-string-type-confusion"),
+        pytest.param({}, id="empty-dict-type-confusion"),
+        pytest.param({"0": "NVIDIA A10G"}, id="dict-type-confusion"),
+        pytest.param(({"index": 0},), id="tuple-type-confusion"),
+        pytest.param(True, id="truthy-boolean"),
+        pytest.param(1, id="bare-count-int"),
+    ],
+)
+def test_ineligible_unless_the_capability_is_a_list_long_enough(advertised):
+    """Fail closed on absence AND on type confusion. A bare `1` is the most
+    tempting wrong answer — it reads like "one GPU" and has no length, so
+    accepting it would mean writing a second, looser matching rule beside the
+    one the probe actually feeds."""
+    from flashruntime.scheduler import IsolationAwarePlacement
+
+    node = {"node_id": "n1", "capabilities": {"gpus": advertised}}
+    assert IsolationAwarePlacement().eligible(_gpu_task(1), node) is False
+
+
+def test_node_with_no_capabilities_key_at_all_is_ineligible():
+    """An already-deployed agent whose registration predates GPU probing
+    advertises nothing — absent must read as "has none", never "unknown, try
+    it and see"."""
+    from flashruntime.scheduler import IsolationAwarePlacement
+
+    assert IsolationAwarePlacement().eligible(_gpu_task(1), {"node_id": "n1"}) is False
+
+
+def test_node_with_capabilities_none_is_ineligible():
+    from flashruntime.scheduler import IsolationAwarePlacement
+
+    node = {"node_id": "n1", "capabilities": None}
+    assert IsolationAwarePlacement().eligible(_gpu_task(1), node) is False
+
+
+def test_node_with_type_confused_capabilities_is_ineligible_without_crashing():
+    from flashruntime.scheduler import IsolationAwarePlacement
+
+    node = {"node_id": "n1", "capabilities": "8 cores, 1 gpu"}
+    assert IsolationAwarePlacement().eligible(_gpu_task(1), node) is False
+
+
+def test_node_capabilities_without_a_gpus_key_is_ineligible():
+    from flashruntime.scheduler import IsolationAwarePlacement
+
+    node = {"node_id": "n1", "capabilities": {"cpu_cores": 8}}
+    assert IsolationAwarePlacement().eligible(_gpu_task(1), node) is False
+
+
+# ---------------------------------------------------------------------------
+# Fail-closed: the requirement itself
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+    "required",
+    [
+        pytest.param("1", id="numeric-string"),
+        pytest.param(-1, id="negative"),
+        pytest.param(1.5, id="fractional"),
+        pytest.param(1.0, id="whole-float"),
+        pytest.param([1], id="list"),
+        pytest.param({"count": 1}, id="dict"),
+        pytest.param(True, id="boolean-true"),
+    ],
+)
+def test_type_confused_requirement_is_ineligible_everywhere(required):
+    """A poisoned requirement must fail closed on EVERY node — including one
+    with plenty of GPUs — and must never crash the predicate.
+
+    `True` is the one that needs saying out loud: `bool` is a subclass of
+    `int` in Python, so `isinstance(True, int)` is True and `True >= 1`.
+    Without the explicit bool exclusion, `gpus: true` would silently mean
+    "one GPU" — a JSON typo that placed real work.
+    """
+    from flashruntime.scheduler import IsolationAwarePlacement
+
+    policy = IsolationAwarePlacement()
+    assert policy.eligible(_gpu_task(required), _node(0)) is False
+    assert policy.eligible(_gpu_task(required), _node(8)) is False
+
+
+def test_boolean_false_is_not_read_as_zero_gpus():
+    """The mirror of the above: `False` is not "no GPU required" either. It is
+    a type-confused requirement and fails closed like every other."""
+    from flashruntime.scheduler import IsolationAwarePlacement
+
+    assert IsolationAwarePlacement().eligible(_gpu_task(False), _node(8)) is False
+
+
+# ---------------------------------------------------------------------------
+# Interaction with the existing gates
+# ---------------------------------------------------------------------------
+
+
+def test_allow_fallback_cannot_bypass_the_gpu_gate():
+    """allowFallback waives the sandbox-tier requirement only. It is the
+    submitter's statement about their own isolation posture, and has nothing
+    to say about hardware that either exists or does not."""
+    from flashruntime.scheduler import IsolationAwarePlacement
+
+    task = _gpu_task(1, isolation={"tier": "sandboxed", "allowFallback": True})
+    node = _node(0, sandbox_capable=True)
+    assert IsolationAwarePlacement().eligible(task, node) is False
+
+
+def test_a_gpu_node_still_fails_the_argv_gate_it_does_not_pass():
+    """The GPU gate adds a requirement; it never satisfies another one."""
+    from flashruntime.protocol.v1alpha1 import TaskSpec
+    from flashruntime.scheduler import IsolationAwarePlacement
+
+    task = TaskSpec(
+        task_id="task-000", job_id="job-a", commit_key="job-a/task-000/m.json",
+        payload={"argv": ["python", "train.py"], "gpus": 1},
+    )
+    assert IsolationAwarePlacement().eligible(task, _node(1)) is False
+    assert IsolationAwarePlacement().eligible(task, _node(1, argv_capable=True)) is True
+
+
+def test_claim_over_a_gpu_queue_serves_the_clean_task():
+    """A GPU task at the head of the queue must not block a CPU-only node —
+    the next clean task is still leased, and the GPU task is never
+    mis-leased."""
+    from flashruntime.leases import LeaseManager
+    from flashruntime.protocol.v1alpha1 import TaskSpec
+    from flashruntime.scheduler import IsolationAwarePlacement
+
+    mgr = LeaseManager()
+    mgr.add_task(_gpu_task(1))  # head of queue, needs a GPU
+    mgr.add_task(
+        TaskSpec(task_id="task-001", job_id="job-a", commit_key="job-a/task-001/m.json")
+    )
+    policy = IsolationAwarePlacement()
+
+    lease = mgr.claim("n1", policy=policy, node=_node(0))
+    assert lease is not None
+    assert lease.task_id == "task-001"
+
+    lease2 = mgr.claim("n2", policy=policy, node=_node(1, node_id="n2"))
+    assert lease2 is not None
+    assert lease2.task_id == "task-000"
+
+
+def test_claim_endpoint_leases_gpu_work_only_to_a_gpu_host():
+    """End to end over the claim endpoint: the registry's claim-time node view
+    must forward `capabilities.gpus`, or the gate reads an absent capability on
+    EVERY node and a GPU task becomes unplaceable everywhere.
+
+    This pins hop 1 of the three that broke the `local_datasets` work — the
+    node view. It is believed correct today (`capabilities` is forwarded whole
+    as `model_dump()`), and this test is what keeps it that way.
+    """
+    import pathlib
+
+    import fastapi
+    from fastapi.testclient import TestClient
+
+    from flashruntime.leases import LeaseManager
+    from flashruntime.service.modea import ModeAState, build_router
+
+    state = ModeAState(LeaseManager(), artifacts_dir=pathlib.Path("/tmp"))
+    app = fastapi.FastAPI()
+    app.include_router(build_router(state))
+    client = TestClient(app)
+    state.manager.add_task(_gpu_task(1))
+
+    def register(node_id: str, gpus: list[dict]):
+        r = client.post(
+            "/v1alpha1/nodes/register",
+            json={
+                "node_id": node_id, "kubernetes_node": "", "hostname": node_id,
+                "capabilities": {"cpu_cores": 8, "gpus": gpus},
+            },
+        )
+        assert r.status_code == 200
+
+    register("cpu-node", [])
+    register("gpu-node", [{"index": 0, "name": "NVIDIA A10G", "memory_total_mb": 22731}])
+
+    assert client.post("/v1alpha1/leases/claim", json={"node_id": "cpu-node"}).status_code == 204
+    r = client.post("/v1alpha1/leases/claim", json={"node_id": "gpu-node"})
+    assert r.status_code == 200
+    assert r.json()["task_id"] == "task-000"

flashruntime-0.4.1/tests/test_protocol_gpu.py

@@ -0,0 +1,152 @@
+"""`GpuInfo` and `ResourcesSpec.gpuPerTask` — the two wire additions GPU
+placement needs.
+
+These are protocol tests, not scheduler tests: they pin the *wire* behaviour
+the fifth placement gate is later allowed to trust, exactly as
+`test_protocol_local_datasets.py` does for `local_datasets`.
+
+The typing change matters more than it looks. `NodeCapabilities.gpus` was
+`list[dict[str, Any]]`; it becomes `list[GpuInfo]`. Agents already deployed on
+machines we cannot reach send plain JSON objects, so a raw dict must still
+coerce — a model that only accepted `GpuInfo` instances would be a breaking
+change dressed up as a type annotation.
+"""
+
+from __future__ import annotations
+
+import pytest
+from pydantic import ValidationError
+
+from flashruntime.protocol.v1alpha1 import GpuInfo, NodeCapabilities, ResourcesSpec
+
+
+# ---------------------------------------------------------------------------
+# GpuInfo
+# ---------------------------------------------------------------------------
+
+
+def test_gpu_info_round_trips_with_only_index_set():
+    """Every field but `index` is optional: a probe that cannot read a value
+    says nothing rather than guessing, and the omission must survive JSON."""
+    gpu = GpuInfo(index=0)
+    restored = GpuInfo.model_validate_json(gpu.model_dump_json())
+
+    assert restored.index == 0
+    assert restored.name == ""
+    assert restored.memory_total_mb is None
+    assert restored.driver_version == ""
+    assert restored.compute_capability == ""
+
+
+def test_gpu_info_round_trips_a_fully_populated_device():
+    gpu = GpuInfo(
+        index=1,
+        name="NVIDIA GeForce RTX 4090",
+        memory_total_mb=24564,
+        driver_version="550.54.14",
+        compute_capability="8.9",
+    )
+    restored = GpuInfo.model_validate_json(gpu.model_dump_json())
+
+    assert restored == gpu
+
+
+def test_gpu_info_requires_an_index():
+    """The one field a probe can always report. Without it there is no way to
+    tell two devices apart, so it is not defaulted."""
+    with pytest.raises(ValidationError):
+        GpuInfo.model_validate({"name": "NVIDIA GeForce RTX 4090"})
+
+
+# ---------------------------------------------------------------------------
+# NodeCapabilities.gpus
+# ---------------------------------------------------------------------------
+
+
+def test_node_capabilities_accepts_gpu_info_instances():
+    caps = NodeCapabilities(gpus=[GpuInfo(index=0)])
+
+    assert len(caps.gpus) == 1
+    assert caps.gpus[0].index == 0
+
+
+def test_node_capabilities_coerces_raw_dicts_from_the_wire():
+    """Deployed agents send JSON objects, not `GpuInfo` instances. The typing
+    change must be additive on the wire or every already-enrolled machine
+    fails to register the moment the coordinator upgrades."""
+    caps = NodeCapabilities.model_validate(
+        {
+            "cpu_cores": 8,
+            "gpus": [
+                {"index": 0, "name": "NVIDIA A10G", "memory_total_mb": 22731},
+                {"index": 1},
+            ],
+        }
+    )
+
+    assert [g.index for g in caps.gpus] == [0, 1]
+    assert isinstance(caps.gpus[0], GpuInfo)
+    assert caps.gpus[0].name == "NVIDIA A10G"
+    assert caps.gpus[1].name == ""
+
+
+def test_node_capabilities_ignores_fields_a_newer_probe_might_add():
+    """Forward compatibility in the other direction: a newer agent reporting a
+    field this coordinator does not know must not fail registration."""
+    caps = NodeCapabilities.model_validate(
+        {"gpus": [{"index": 0, "some_future_field": "ignored"}]}
+    )
+
+    assert caps.gpus[0].index == 0
+
+
+def test_node_capabilities_defaults_to_no_gpus():
+    """Absent means "no GPU", never "unknown, assume yes" — the gate reads
+    this list and its length is the whole matching rule in v1."""
+    assert NodeCapabilities().gpus == []
+    assert NodeCapabilities.model_validate({"cpu_cores": 4}).gpus == []
+
+
+def test_node_capabilities_gpu_default_is_not_shared_between_instances():
+    # `default_factory` and not a bare `[]`: one shared list would let a
+    # single node's devices leak into every other node's capabilities.
+    first = NodeCapabilities()
+    second = NodeCapabilities()
+
+    assert first.gpus is not second.gpus
+    first.gpus.append(GpuInfo(index=0))
+    assert second.gpus == []
+
+
+# ---------------------------------------------------------------------------
+# ResourcesSpec.gpuPerTask
+# ---------------------------------------------------------------------------
+
+
+def test_gpu_per_task_defaults_to_zero():
+    """0 means "no GPU required", which is every job that exists today."""
+    assert ResourcesSpec().gpuPerTask == 0
+
+
+def test_gpu_per_task_accepts_a_positive_count():
+    assert ResourcesSpec(gpuPerTask=2).gpuPerTask == 2
+
+
+@pytest.mark.parametrize(
+    "value",
+    [
+        pytest.param(-1, id="negative"),
+        pytest.param(1.5, id="fractional"),
+        pytest.param("one", id="not-a-number"),
+    ],
+)
+def test_gpu_per_task_rejects_values_that_are_not_a_device_count(value):
+    with pytest.raises(ValidationError):
+        ResourcesSpec(gpuPerTask=value)
+
+
+def test_gpu_per_task_round_trips_through_json():
+    spec = ResourcesSpec(gpuPerTask=1)
+    restored = ResourcesSpec.model_validate_json(spec.model_dump_json())
+
+    assert restored.gpuPerTask == 1