flashruntime 0.5.0__tar.gz → 0.6.0__tar.gz

{flashruntime-0.5.0/flashruntime.egg-info → flashruntime-0.6.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: flashruntime
-Version: 0.5.0
+Version: 0.6.0
 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.5.0 → flashruntime-0.6.0}/flashruntime/protocol/v1alpha1.py

@@ -332,6 +332,17 @@ class NodeCapabilities(BaseModel):
     #: 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)
+    #: Bytes this host will lend to cached dataset shards. Advertised so the
+    #: control plane can refuse a job at SUBMIT time when no host in the pool
+    #: can hold a slice — twenty machines each downloading for forty minutes
+    #: and then refusing is the failure this exists to replace.
+    #:
+    #: Defaults to 0, and 0 means "send me no dataset work". Fail closed, the
+    #: same polarity as `can_install_dependencies`: an agent whose
+    #: registration predates this field has nowhere to put a shard and must
+    #: never be handed one. Bytes rather than GB to match `memory_bytes` —
+    #: one unit for every capacity on this model.
+    dataset_cache_bytes: int = 0
     #: Team pools this node serves, as pool ids minted by the cloud control
     #: plane. Stamped SERVER-SIDE by the cloud API's agent proxy from the
     #: machine owner's memberships — an agent's self-reported value is

{flashruntime-0.5.0 → flashruntime-0.6.0}/flashruntime/recipes/command.py

@@ -66,6 +66,18 @@ class CommandRecipe(WorkloadRecipe):
                 )
             elif not all(isinstance(d, str) for d in extra_deps):
                 problems.append("'extra_dependencies' must be a list of strings")
+        slices = params.get("dataset_slices")
+        if slices is not None:
+            if not isinstance(slices, list) or isinstance(slices, (str, bytes)):
+                problems.append(
+                    "'dataset_slices' must be a list with one entry per task "
+                    "(a list, not a single string)"
+                )
+            elif not all(isinstance(s, list) for s in slices):
+                problems.append(
+                    "'dataset_slices' must be a list of per-task lists, one "
+                    "list of dataset slices for each task"
+                )
         problems.extend(self._unpack_problems(params))
         problems.extend(partition_problems(params))
         problems.extend(declaration_problems(params.get("validators")))
@@ -239,6 +251,31 @@ class CommandRecipe(WorkloadRecipe):
                 # hop have tests that pass while it is broken, because each
                 # constructs the payload directly.
                 payload["extra_dependencies"] = list(p["extra_dependencies"])
+            if p.get("dataset_slices") is not None:
+                # The ONLY per-task payload key in this method. Every other
+                # forward above copies one value to every task; each task
+                # fetches a DIFFERENT slice of the manifest, so this one is
+                # indexed by task position exactly as `task_params` is.
+                #
+                # Computed by the flashml-cloud compiler and opaque here:
+                # the byte-weighted mapping lives beside `fleet_shape` in
+                # that repo's `elastic.py`, and a second implementation here
+                # is the drift that module's docstring warns about.
+                #
+                # Same failure shape as `local_inputs` and `dependencies`:
+                # drop this line and every host fetches nothing, the workload
+                # dies on a missing path, and both ends' unit tests still
+                # pass because each constructs the payload directly.
+                slices = p["dataset_slices"]
+                if i >= len(slices):
+                    raise ValueError(
+                        f"task {i}: 'dataset_slices' has {len(slices)} "
+                        f"entr{'y' if len(slices) == 1 else 'ies'} but this "
+                        f"job expanded to more tasks — every task needs a "
+                        f"slice, even an empty one"
+                    )
+                # `list(...)` so a payload never aliases the caller's spec.
+                payload["datasets"] = list(slices[i])
             gpus = spec.spec.resources.gpuPerTask
             if gpus:
                 # The GPU requirement is the ONLY payload key sourced from

{flashruntime-0.5.0 → flashruntime-0.6.0}/flashruntime/scheduler/__init__.py

@@ -386,6 +386,56 @@ class IsolationAwarePlacement(PlacementPolicy):
     waiver relaxes the sandbox tier and nothing else. No waiver makes a
     read-only container with no network able to run `pip install`.
 
+    A ninth gate applies to tasks whose payload carries `datasets`: the sum
+    of every slice entry's `size` must fit inside the claiming node's
+    `capabilities.dataset_cache_bytes`. This is CAPACITY, not capability —
+    every tier can fetch, because the AGENT fetches before the sandbox
+    closes, so unlike `extra_dependencies` there is no tier that physically
+    cannot. The only question is whether the bytes fit.
+
+    Not to be confused with the fourth gate. `local_inputs`/`local_datasets`
+    asks whether a host already HOLDS named data it will never upload; this
+    one asks whether a host has ROOM for data it is about to fetch from a
+    public origin. Different payload key, different capability, different
+    question.
+
+    It takes the argv/local-data/gpu/exclude_nodes/pool/dependencies
+    polarity — **fail closed** — and every leg of it:
+
+    - A node advertising nothing (an agent whose registration predates the
+      field) is refused, and so is one advertising `0`, which is the model
+      default and means "send me no dataset work". A shard with nowhere to
+      go is a task that dies after a forty-minute download rather than one
+      that never started, and twenty machines doing that at once is the
+      failure the advertisement exists to prevent.
+    - The advertisement must be a genuine `int`. `bool` is excluded
+      explicitly for the reason the `gpus` gate excludes it: it is an `int`
+      subclass, so a `dataset_cache_bytes: true` typo would read as the
+      capacity 1.
+    - `capabilities` itself may be absent or type-confused; that is read as
+      advertising nothing rather than allowed to raise. `isinstance`, never
+      `(node.get("capabilities") or {})` — a string `capabilities` value has
+      no `.get` and would crash the predicate inside the claim path, taking
+      the whole claim down rather than one task.
+    - A non-dict slice, a non-dict entry, or an entry whose `size` is not a
+      non-negative `int` makes the task ineligible everywhere. The manifest
+      is submitter-controlled, so an unreadable size must NOT contribute 0 to
+      the total: that is precisely how a hostile slice would buy itself an
+      unbounded fetch on a node that advertised a modest cache.
+    - The requirement is the sum over every entry of every slice. Checking
+      per-dataset, or against the largest entry, would place a task that then
+      fills a volunteer's disk halfway through the fetch.
+    - Empty and absent both mean "declares no data": `if
+      task.payload.get("datasets")` treats `None` and `[]` alike, so every
+      job deployed today is completely unaffected. A slice that is PRESENT
+      but happens to carry no entries is still dataset work and still needs
+      a readable advertisement — the fail-closed leg is the advertisement,
+      not the total.
+
+    `allowFallback` does not waive this gate either. It waives the sandbox
+    tier and nothing else; a submitter's isolation posture has nothing to say
+    about how much disk a stranger's machine offered.
+
     Everything genuinely standard keeps the fail-open placement default."""
 
     def eligible(self, task: TaskSpec, node: NodeView) -> bool:
@@ -516,6 +566,48 @@ class IsolationAwarePlacement(PlacementPolicy):
         if task.payload.get("extra_dependencies"):
             if node.get("can_install_dependencies") is not True:
                 return False
+        # Capacity, not capability, and the only gate in this class that is.
+        # Every tier's AGENT has network and a writable disk before the
+        # sandbox closes, so unlike `extra_dependencies` there is no tier
+        # that physically cannot fetch — the only question is whether the
+        # bytes fit. Distinct from the `local_inputs` gate above: that one
+        # asks whether a host already HOLDS named data, this one whether it
+        # has ROOM for data it is about to fetch.
+        #
+        # Fail-closed on an absent or type-confused advertisement, and
+        # checked before the allowFallback waiver below for the same reason
+        # as every gate above: the waiver relaxes the sandbox tier and has
+        # nothing to say about how much disk a stranger's machine offered.
+        if task.payload.get("datasets"):
+            needed = 0
+            for dataset in task.payload["datasets"]:
+                if not isinstance(dataset, dict):
+                    return False  # type-confused slice ⇒ fail closed, no crash
+                for entry in dataset.get("entries") or []:
+                    size = entry.get("size") if isinstance(entry, dict) else None
+                    # A non-int size is a malformed slice, not a free one:
+                    # letting it contribute 0 is how a hostile manifest buys
+                    # an unbounded fetch on a modest cache. `bool` is an
+                    # `int` subclass, so `size: true` is excluded explicitly.
+                    if (
+                        not isinstance(size, int)
+                        or isinstance(size, bool)
+                        or size < 0
+                    ):
+                        return False
+                    needed += size
+            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("dataset_cache_bytes")
+                if isinstance(capabilities, dict)
+                else None
+            )
+            if not isinstance(advertised, int) or isinstance(advertised, bool):
+                return False  # absent/type-confused ⇒ advertises nothing
+            if needed > advertised:
+                return False
         isolation = task.payload.get("isolation")
         if isolation is None:
             return True  # no isolation payload ⇒ standard, runs anywhere

{flashruntime-0.5.0 → flashruntime-0.6.0/flashruntime.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: flashruntime
-Version: 0.5.0
+Version: 0.6.0
 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.5.0 → flashruntime-0.6.0}/flashruntime.egg-info/SOURCES.txt

@@ -138,6 +138,7 @@ tests/test_leases.py
 tests/test_leases_scope.py
 tests/test_leases_sqlite.py
 tests/test_monitor_sampler.py
+tests/test_placement_datasets.py
 tests/test_placement_dependencies.py
 tests/test_placement_exclusion.py
 tests/test_placement_gpu.py
@@ -146,11 +147,13 @@ tests/test_placement_pool.py
 tests/test_placement_reliability.py
 tests/test_placement_trusted_argv.py
 tests/test_planner.py
+tests/test_protocol_datasets.py
 tests/test_protocol_evidence.py
 tests/test_protocol_gpu.py
 tests/test_protocol_local_datasets.py
 tests/test_protocol_pools.py
 tests/test_protocol_v1alpha1.py
+tests/test_recipe_datasets.py
 tests/test_recipe_gpu.py
 tests/test_recipe_local_inputs.py
 tests/test_recipe_partition.py

{flashruntime-0.5.0 → flashruntime-0.6.0}/pyproject.toml

@@ -46,7 +46,7 @@ name = "flashruntime"
 # 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.5.0"
+version = "0.6.0"
 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.6.0/tests/test_placement_datasets.py

@@ -0,0 +1,354 @@
+"""The ninth placement gate: a node must have room for the slice it is
+offered — `capabilities.dataset_cache_bytes` must cover the sum of every
+declared entry's `size` (AGENTS.md rule 3).
+
+Capacity, not capability. Unlike installing dependencies — which the
+container tiers physically cannot do — every tier's AGENT fetches before the
+sandbox closes, so there is no tier that *cannot* hold a shard. The only
+question this gate asks is whether the bytes fit.
+
+Fail closed, taking the argv/local-data/gpu polarity: a node advertising
+nothing (an agent whose registration predates the field) holds no dataset
+work. A shard with nowhere to go is a task that dies after a forty-minute
+download rather than one that never started.
+
+Not to be confused with `local_inputs`/`local_datasets`
+(`test_placement_local_data.py`), which is a different feature with its own
+gate: that one asks whether a host already HOLDS named data it will never
+upload; this one asks whether a host has room for data it is about to fetch.
+"""
+
+from __future__ import annotations
+
+import pytest
+
+
+def _entry(size=100, **over):
+    entry = {
+        "path": "a.parquet",
+        "url": "https://origin.invalid/a.parquet",
+        "size": size,
+        "integrity": {"kind": "sha256", "value": "aa"},
+    }
+    entry.update(over)
+    return entry
+
+
+def _dataset(*sizes, name="d", **over):
+    dataset = {"name": name, "split": "shard",
+               "entries": [_entry(size) for size in sizes]}
+    dataset.update(over)
+    return dataset
+
+
+def _task(datasets, **payload_extra):
+    """A task carrying `datasets` RAW, so tests can poison it with values the
+    compiler would never emit (the trick `_local_task` plays on `local_inputs`
+    and `_task_raw` plays on the isolation payload)."""
+    from flashruntime.protocol.v1alpha1 import TaskSpec
+
+    payload = {"module": "flashml_workloads.sklearn_trial"}
+    if datasets is not None:
+        payload["datasets"] = datasets
+    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(cache_bytes=None, **over):
+    """A node view shaped like the claim endpoint's. `cache_bytes` defaults to
+    ABSENT (not 0) so a bare `_node()` models an agent whose registration
+    predates the field — the fail-closed case."""
+    node = {"node_id": "n1", "sandbox_capable": True, "capabilities": {}}
+    if cache_bytes is not None:
+        node["capabilities"] = {"dataset_cache_bytes": cache_bytes}
+    node.update(over)
+    return node
+
+
+def test_a_node_with_room_is_eligible():
+    from flashruntime.scheduler import IsolationAwarePlacement
+
+    assert IsolationAwarePlacement().eligible(
+        _task([_dataset(100)]), _node(1000)) is True
+
+
+def test_a_node_without_room_is_not():
+    from flashruntime.scheduler import IsolationAwarePlacement
+
+    assert IsolationAwarePlacement().eligible(
+        _task([_dataset(2000)]), _node(1000)) is False
+
+
+def test_an_exact_fit_is_eligible():
+    """The boundary is `needed > advertised`, not `>=`: a node that advertised
+    exactly this many bytes said yes to exactly this many bytes."""
+    from flashruntime.scheduler import IsolationAwarePlacement
+
+    assert IsolationAwarePlacement().eligible(
+        _task([_dataset(1000)]), _node(1000)) is True
+
+
+def test_the_requirement_is_the_sum_over_every_entry_of_every_dataset():
+    """One entry fits; the whole slice does not. Summing per-dataset or
+    checking only the largest entry would place a task that then fills a
+    volunteer's disk halfway through the fetch."""
+    from flashruntime.scheduler import IsolationAwarePlacement
+
+    policy = IsolationAwarePlacement()
+    task = _task([_dataset(400, 400, name="d1"), _dataset(400, name="d2")])
+    assert policy.eligible(task, _node(1000)) is False
+    assert policy.eligible(task, _node(1200)) is True
+
+
+def test_a_node_predating_the_field_holds_no_dataset_work():
+    """Fail closed — an absent capability reads as 'advertises nothing', never
+    as 'advertises everything'."""
+    from flashruntime.scheduler import IsolationAwarePlacement
+
+    assert IsolationAwarePlacement().eligible(_task([_dataset(1)]), _node()) is False
+
+
+def test_a_zero_advertisement_holds_no_dataset_work():
+    """`0` is the model default and it means 'send me no dataset work' — the
+    key being PRESENT must not read as consent."""
+    from flashruntime.scheduler import IsolationAwarePlacement
+
+    assert IsolationAwarePlacement().eligible(_task([_dataset(1)]), _node(0)) is False
+
+
+@pytest.mark.parametrize(
+    "capabilities",
+    [
+        pytest.param("not-a-dict", id="bare-string"),
+        pytest.param(None, id="explicit-none"),
+        pytest.param(["dataset_cache_bytes"], id="list"),
+        pytest.param(17179869184, id="bare-int"),
+    ],
+)
+def test_a_type_confused_capabilities_value_does_not_raise(capabilities):
+    """`(node.get("capabilities") or {})` would crash the predicate on a
+    string — it has no `.get`, and a crash inside the claim path takes the
+    whole claim down rather than one task. The house rule is `isinstance`
+    (documented at `scheduler/__init__.py:234-236`), and the reading is
+    'advertises nothing'."""
+    from flashruntime.scheduler import IsolationAwarePlacement
+
+    node = _node()
+    node["capabilities"] = capabilities
+    assert IsolationAwarePlacement().eligible(_task([_dataset(1)]), node) is False
+
+
+@pytest.mark.parametrize(
+    "advertised",
+    [
+        pytest.param("17179869184", id="numeric-string"),
+        pytest.param(16.0, id="float"),
+        pytest.param(None, id="explicit-none"),
+        pytest.param([16], id="list"),
+        pytest.param(True, id="truthy-boolean"),
+    ],
+)
+def test_a_non_int_advertisement_holds_no_dataset_work(advertised):
+    """A capacity that is not a genuine `int` is an advertisement nobody can
+    act on. `bool` is named explicitly because it is an `int` subclass: a
+    `dataset_cache_bytes: true` typo must not read as the capacity 1."""
+    from flashruntime.scheduler import IsolationAwarePlacement
+
+    assert IsolationAwarePlacement().eligible(
+        _task([_dataset(1)]), _node(advertised)) is False
+
+
+@pytest.mark.parametrize(
+    "size",
+    [
+        pytest.param("100", id="numeric-string"),
+        pytest.param(100.5, id="float"),
+        pytest.param(None, id="explicit-none"),
+        pytest.param(True, id="truthy-boolean"),
+        pytest.param(-1, id="negative"),
+    ],
+)
+def test_a_non_int_entry_size_is_a_malformed_slice_not_a_free_one(size):
+    """The manifest is submitter-controlled. An unreadable size must make the
+    task unplaceable, never contribute 0 to the total — the latter is how a
+    hostile slice buys itself an unbounded fetch on a node that advertised a
+    modest cache."""
+    from flashruntime.scheduler import IsolationAwarePlacement
+
+    task = _task([{"name": "d", "split": "shard", "entries": [_entry(size)]}])
+    assert IsolationAwarePlacement().eligible(task, _node(1_000_000)) is False
+
+
+def test_an_entry_missing_its_size_is_refused():
+    from flashruntime.scheduler import IsolationAwarePlacement
+
+    entry = _entry()
+    del entry["size"]
+    task = _task([{"name": "d", "split": "shard", "entries": [entry]}])
+    assert IsolationAwarePlacement().eligible(task, _node(1_000_000)) is False
+
+
+@pytest.mark.parametrize(
+    "entry",
+    [
+        pytest.param("a.parquet", id="bare-string"),
+        pytest.param(None, id="explicit-none"),
+        pytest.param(100, id="bare-int"),
+    ],
+)
+def test_a_non_dict_entry_is_refused(entry):
+    from flashruntime.scheduler import IsolationAwarePlacement
+
+    task = _task([{"name": "d", "split": "shard", "entries": [entry]}])
+    assert IsolationAwarePlacement().eligible(task, _node(1_000_000)) is False
+
+
+@pytest.mark.parametrize(
+    "dataset",
+    [
+        pytest.param("imdb", id="bare-string"),
+        pytest.param(None, id="explicit-none"),
+        pytest.param([{"path": "a.parquet"}], id="nested-list"),
+    ],
+)
+def test_a_non_dict_dataset_is_refused(dataset):
+    """A slice that is not a mapping was built wrong, and the gate must say so
+    by refusing rather than by raising inside the claim path."""
+    from flashruntime.scheduler import IsolationAwarePlacement
+
+    assert IsolationAwarePlacement().eligible(
+        _task([dataset]), _node(1_000_000)) is False
+
+
+@pytest.mark.parametrize(
+    "datasets",
+    [
+        pytest.param("imdb", id="bare-string"),
+        pytest.param({"name": "imdb"}, id="single-dict-not-a-list"),
+    ],
+)
+def test_a_type_confused_datasets_payload_is_ineligible_everywhere(datasets):
+    """Present but not a list of slices: fail closed on every node, including
+    one with room to spare, and never crash the predicate."""
+    from flashruntime.scheduler import IsolationAwarePlacement
+
+    assert IsolationAwarePlacement().eligible(
+        _task(datasets), _node(1_000_000)) is False
+
+
+def test_a_task_without_datasets_is_unaffected():
+    """Every job deployed today. The gate must be invisible to them — including
+    on a node that advertises no cache at all."""
+    from flashruntime.scheduler import IsolationAwarePlacement
+
+    policy = IsolationAwarePlacement()
+    assert policy.eligible(_task(None), _node()) is True
+    assert policy.eligible(_task(None), _node(1000)) is True
+
+
+def test_an_empty_datasets_list_requires_nothing():
+    """Absent and empty mean the same thing here, exactly as they do for
+    `extra_dependencies`: a job that declares no data runs anywhere."""
+    from flashruntime.scheduler import IsolationAwarePlacement
+
+    assert IsolationAwarePlacement().eligible(_task([]), _node()) is True
+
+
+def test_a_slice_with_no_entries_still_needs_an_advertisement():
+    """`needed == 0` is not a licence: a node that advertised nothing has said
+    it wants no dataset work, and a slice that happens to be empty this time is
+    still dataset work. Fail closed on the advertisement, not on the total."""
+    from flashruntime.scheduler import IsolationAwarePlacement
+
+    policy = IsolationAwarePlacement()
+    empty = _task([{"name": "d", "split": "shard", "entries": []}])
+    assert policy.eligible(empty, _node()) is False
+    assert policy.eligible(empty, _node(1000)) is True
+
+
+def test_allow_fallback_does_not_waive_this_gate():
+    """The waiver covers the sandbox tier and nothing else. A submitter's
+    isolation posture has nothing to say about how much disk a stranger's
+    machine offered."""
+    from flashruntime.scheduler import IsolationAwarePlacement
+
+    task = _task(
+        [_dataset(2000)],
+        pool="p1",
+        isolation={"tier": "sandboxed", "allowFallback": True},
+    )
+    node = _node(1000)
+    node["capabilities"]["pools"] = ["p1"]
+    assert IsolationAwarePlacement().eligible(task, node) is False
+
+
+def test_claim_over_a_dataset_queue_serves_the_node_with_room():
+    """A slice too large for the asking node must not block the queue: the next
+    clean task is still leased, and the dataset task waits for a host that fits
+    it."""
+    from flashruntime.leases import LeaseManager
+    from flashruntime.protocol.v1alpha1 import TaskSpec
+    from flashruntime.scheduler import IsolationAwarePlacement
+
+    mgr = LeaseManager()
+    mgr.add_task(_task([_dataset(2000)]))  # head of queue, needs room
+    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(1000))
+    assert lease is not None
+    assert lease.task_id == "task-001"
+
+    lease2 = mgr.claim("n2", policy=policy, node=_node(1_000_000, node_id="n2"))
+    assert lease2 is not None
+    assert lease2.task_id == "task-000"
+
+
+def test_claim_endpoint_leases_dataset_work_only_to_a_host_with_room():
+    """End to end over the REAL claim endpoint, not a hand-built node dict.
+
+    Every other test in this file hands `eligible()` a dictionary shaped like
+    the claim endpoint's node view. None of them touches the code that
+    actually BUILDS that view (`service/modea.py`), so all of them would pass
+    while the endpoint omitted `dataset_cache_bytes` — the gate would read an
+    absent capability on EVERY node and refuse every dataset task everywhere.
+    Fail closed, and also permanently stuck: inert in production, green in CI.
+    """
+    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(_task([_dataset(4096)]))
+
+    def register(node_id: str, cache_bytes: int):
+        r = client.post(
+            "/v1alpha1/nodes/register",
+            json={
+                "node_id": node_id, "kubernetes_node": "", "hostname": node_id,
+                "capabilities": {"dataset_cache_bytes": cache_bytes},
+            },
+        )
+        assert r.status_code == 200
+
+    register("no-room-node", 0)
+    register("roomy-node", 1024 ** 3)
+
+    assert client.post(
+        "/v1alpha1/leases/claim", json={"node_id": "no-room-node"}
+    ).status_code == 204
+    r = client.post("/v1alpha1/leases/claim", json={"node_id": "roomy-node"})
+    assert r.status_code == 200, r.text
+    assert r.json()["task_id"] == "task-000"

flashruntime-0.6.0/tests/test_protocol_datasets.py

@@ -0,0 +1,28 @@
+"""Wire field for dataset cache capacity (AGENTS.md rule 3: fail closed).
+
+Bytes, not GB, to match `memory_bytes` — one unit for every capacity on
+this model. Zero means "send me no dataset work", which is what an agent
+predating the field must be read as saying.
+"""
+
+from __future__ import annotations
+
+from flashruntime.protocol.v1alpha1 import NodeCapabilities, NodeRegistration
+
+
+def test_dataset_cache_bytes_defaults_to_zero():
+    assert NodeCapabilities().dataset_cache_bytes == 0
+
+
+def test_dataset_cache_bytes_round_trips():
+    caps = NodeCapabilities.model_validate({"dataset_cache_bytes": 17179869184})
+    assert caps.model_dump()["dataset_cache_bytes"] == 17179869184
+
+
+def test_old_wire_shapes_still_validate():
+    """An agent on 0.5.0 sends none of this; nothing may break."""
+    reg = NodeRegistration.model_validate(
+        {"node_id": "n", "kubernetes_node": "", "hostname": "h",
+         "capabilities": {"cpu_cores": 4}}
+    )
+    assert reg.capabilities.dataset_cache_bytes == 0