Write the card a colocated engine was given onto its pod (#2992)

miles/utils/external_utils/colocate_pairing/controller.py

@@ -1,6 +1,8 @@
 from __future__ import annotations
 
+import asyncio
 import logging
+from typing import NamedTuple
 
 from kubernetes_asyncio import client
 
@@ -20,16 +22,21 @@ logger = logging.getLogger(__name__)
 _UNRELATED_KEY_PREFIX = "__unrelated__/"
 
 
+class InferencePlacement(NamedTuple):
+    trainer_coord: PodCoordinate
+    base_gpu_id: int
+
+
 class PairingController:
     _loop: ReconcileLoop
 
     def __init__(self, *, config: PairingConfig, core_v1: client.CoreV1Api) -> None:
         self._config = config
         self._core_v1 = core_v1
-        trainer_of_inference = {
+        placement_of_inference = {
             PodCoordinate(
                 pool_id=pool.pool_id, cell_index=cell_index, pod_in_cell_index=pod_index
-            ): _target_trainer_pod(
+            ): _place_inference_pod(
                 inference_cell_index=cell_index,
                 inference_pod_index=pod_index,
                 layout=pool.layout,
@@ -39,12 +46,14 @@ class PairingController:
             for cell_index in range(pool.layout.num_inference_cells)
             for pod_index in range(pool.layout.num_pods_per_inference_cell)
         }
-        self._inferences_of_trainer: dict[PodCoordinate, list[PodCoordinate]] = {}
-        for inference, trainer in trainer_of_inference.items():
-            self._inferences_of_trainer.setdefault(trainer, []).append(inference)
+        self._inferences_of_trainer: dict[PodCoordinate, list[tuple[PodCoordinate, int]]] = {}
+        for inference, placement in placement_of_inference.items():
+            self._inferences_of_trainer.setdefault(placement.trainer_coord, []).append(
+                (inference, placement.base_gpu_id)
+            )
 
         self._trainer_key_of_coord = {
-            inference: trainer.key for inference, trainer in trainer_of_inference.items()
+            inference: placement.trainer_coord.key for inference, placement in placement_of_inference.items()
         } | {trainer: trainer.key for trainer in self._inferences_of_trainer}
 
     def set_loop(self, loop: ReconcileLoop) -> None:
@@ -58,12 +67,12 @@ class PairingController:
         trainer_coord = next((coord for coord in pods_by_coord if coord.key == pair_key), None)
         if trainer_coord is None:
             return
-        gated = [
-            pod
-            for inference_coord in self._inferences_of_trainer.get(trainer_coord, [])
+        gated_pods_and_base_gpu_ids = [
+            (pod, base_gpu_id)
+            for inference_coord, base_gpu_id in self._inferences_of_trainer.get(trainer_coord, [])
             if (pod := pods_by_coord.get(inference_coord)) is not None and is_gated(pod)
         ]
-        if not gated:
+        if not gated_pods_and_base_gpu_ids:
             return
 
         trainer_node_name = pods_by_coord[trainer_coord].spec.node_name
@@ -71,21 +80,47 @@ class PairingController:
             logger.info(
                 "Waiting for %s to be scheduled before releasing %s",
                 trainer_coord.key,
-                [pod.metadata.name for pod in gated],
+                [pod.metadata.name for pod, _ in gated_pods_and_base_gpu_ids],
             )
             return
 
-        for inference_pod in gated:
-            await self._release(inference_pod, node_name=trainer_node_name, trainer_key=trainer_coord.key)
+        outcomes = await asyncio.gather(
+            *(
+                self._release(
+                    inference_pod,
+                    node_name=trainer_node_name,
+                    base_gpu_id=base_gpu_id,
+                    trainer_key=trainer_coord.key,
+                )
+                for inference_pod, base_gpu_id in gated_pods_and_base_gpu_ids
+            ),
+            return_exceptions=True,
+        )
+        for (inference_pod, _), outcome in zip(gated_pods_and_base_gpu_ids, outcomes, strict=True):
+            if isinstance(outcome, BaseException):
+                logger.error(
+                    "Releasing %s failed; the next reconcile of %s tries it again",
+                    inference_pod.metadata.name,
+                    trainer_coord.key,
+                    exc_info=outcome,
+                )
 
-    async def _release(self, inference_pod: Pod, *, node_name: str, trainer_key: str) -> None:
-        logger.info("Releasing %s onto %s, where %s runs", inference_pod.metadata.name, node_name, trainer_key)
+    async def _release(self, inference_pod: Pod, *, node_name: str, base_gpu_id: int, trainer_key: str) -> None:
+        logger.info(
+            "Releasing %s onto gpu %s of %s, where %s runs",
+            inference_pod.metadata.name,
+            base_gpu_id,
+            node_name,
+            trainer_key,
+        )
         await self._core_v1.patch_namespaced_pod(
             name=inference_pod.metadata.name,
             namespace=self._config.namespace,
             body=release_patch(
                 node_name=node_name,
+                base_gpu_id=base_gpu_id,
                 gates=gate_names(inference_pod),
+                annotations=inference_pod.metadata.annotations,
             ),
         )
 
@@ -95,9 +130,9 @@ class PairingController:
         return f"{_UNRELATED_KEY_PREFIX}{pod.metadata.name}"
 
 
-def _target_trainer_pod(
+def _place_inference_pod(
     *, inference_cell_index: int, inference_pod_index: int, layout: PairingLayout, trainer_pool_id: str
-) -> PodCoordinate:
+) -> InferencePlacement:
     assert 0 <= inference_cell_index < layout.num_inference_cells, f"{inference_cell_index=} outside {layout}"
     assert 0 <= inference_pod_index < layout.num_pods_per_inference_cell, f"{inference_pod_index=} outside {layout}"
 
@@ -109,4 +144,9 @@ def _target_trainer_pod(
     trainer_cell_index, trainer_pod_index = divmod(
         absolute_gpu // layout.num_gpus_per_node, layout.num_pods_per_trainer_cell
     )
-    return PodCoordinate(pool_id=trainer_pool_id, cell_index=trainer_cell_index, pod_in_cell_index=trainer_pod_index)
+    return InferencePlacement(
+        trainer_coord=PodCoordinate(
+            pool_id=trainer_pool_id, cell_index=trainer_cell_index, pod_in_cell_index=trainer_pod_index
+        ),
+        base_gpu_id=absolute_gpu % layout.num_gpus_per_node,
+    )

miles/utils/external_utils/colocate_pairing/pods.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+from collections.abc import Mapping
+
 from miles.utils.pydantic_utils import FrozenStrictBaseModel
 from miles.utils.workers.k8s_types import Pod
 from miles.utils.workers.worker_provider.kubernetes.core import pod_view
@@ -7,7 +9,7 @@ from miles.utils.workers.worker_provider.kubernetes.helm.env import DEFAULT_LABE
 
 _GATE_NAME = "miles.radixark.io/colocate-pairing"
 
-_HOSTNAME_LABEL = "kubernetes.io/hostname"
+_NODE_NAME_FIELD = "metadata.name"
 
 
 class PodCoordinate(FrozenStrictBaseModel):
@@ -29,15 +31,36 @@ def coordinate_of(pod: Pod) -> PodCoordinate | None:
     )
 
 
-def release_patch(*, node_name: str, gates: list[str], has_node_selector: bool) -> list[dict[str, object]]:
-    index = gates.index(_GATE_NAME)
-    pin = (
-        {"op": "add", "path": f"/spec/nodeSelector/{_escape_pointer(_HOSTNAME_LABEL)}", "value": node_name}
-        if has_node_selector
-        else {"op": "add", "path": "/spec/nodeSelector", "value": {_HOSTNAME_LABEL: node_name}}
+def release_patch(
+    *,
+    node_name: str,
+    base_gpu_id: int,
+    gates: list[str],
+    annotations: Mapping[str, str],
+) -> list[dict[str, object]]:
+    key = DEFAULT_LABEL_KEYS.base_gpu_id_annotation
+    assert annotations, (
+        f"the pod carries no annotations, so adding {key} under a map that does not exist is a patch "
+        f"the apiserver refuses"
     )
+
+    index = gates.index(_GATE_NAME)
+    pin = {
+        "op": "add",
+        "path": "/spec/affinity",
+        "value": {
+            "nodeAffinity": {
+                "requiredDuringSchedulingIgnoredDuringExecution": {
+                    "nodeSelectorTerms": [
+                        {"matchFields": [{"key": _NODE_NAME_FIELD, "operator": "In", "values": [node_name]}]}
+                    ]
+                }
+            }
+        },
+    }
     return [
         pin,
+        {"op": "add", "path": f"/metadata/annotations/{_escape_pointer(key)}", "value": str(base_gpu_id)},
         {"op": "test", "path": f"/spec/schedulingGates/{index}/name", "value": _GATE_NAME},
         {"op": "remove", "path": f"/spec/schedulingGates/{index}"},
     ]

miles/utils/workers/worker_provider/kubernetes/core/pod_view.py

@@ -14,6 +14,7 @@ class CellLabelKeys(FrozenStrictBaseModel):
     cell_size_annotation: str
     meta_annotation_prefix: str
     gpu_ids_meta: str
+    base_gpu_id_annotation: str
 
 
 class ParsedPod(FrozenStrictBaseModel):

miles/utils/workers/worker_provider/kubernetes/helm/env.py

@@ -14,6 +14,7 @@ DEFAULT_LABEL_KEYS = CellLabelKeys(
     cell_size_annotation="leaderworkerset.sigs.k8s.io/size",
     meta_annotation_prefix="miles.radixark.io/meta-",
     gpu_ids_meta="gpu_ids",
+    base_gpu_id_annotation="miles.radixark.io/base-gpu-id",
 )
 
 NAMESPACE_ENV_VAR = "MILES_K8S_NAMESPACE"

tests/fast/utils/external_utils/test_colocate_pairing.py

@@ -15,7 +15,11 @@ from tests.fast.utils.workers.reconcile.utils import FakeSource, replace_of, set
 
 from miles.utils.external_utils.colocate_pairing import pods as pairing_pods
 from miles.utils.external_utils.colocate_pairing.config import InferencePool, PairingConfig, PairingLayout
-from miles.utils.external_utils.colocate_pairing.controller import PairingController, _target_trainer_pod
+from miles.utils.external_utils.colocate_pairing.controller import (
+    InferencePlacement,
+    PairingController,
+    _place_inference_pod,
+)
 from miles.utils.external_utils.colocate_pairing.pods import PodCoordinate
 from miles.utils.external_utils.command_utils.helm_backend.launcher.values.colocate import _assert_colocate_supported
 from miles.utils.test_utils.clock import FakeClock
@@ -64,8 +68,8 @@ def _sub_node_layout(gpu_offset: int = 0, num_inference_cells: int = 4) -> Pairi
     )
 
 
-def _target(inference_cell_index: int, layout: PairingLayout, inference_pod_index: int = 0) -> PodCoordinate:
-    return _target_trainer_pod(
+def _place(inference_cell_index: int, layout: PairingLayout, inference_pod_index: int = 0) -> InferencePlacement:
+    return _place_inference_pod(
         inference_cell_index=inference_cell_index,
         inference_pod_index=inference_pod_index,
         layout=layout,
@@ -73,6 +77,14 @@ def _target(inference_cell_index: int, layout: PairingLayout, inference_pod_inde
     )
 
 
+def _target(inference_cell_index: int, layout: PairingLayout, inference_pod_index: int = 0) -> PodCoordinate:
+    return _place(inference_cell_index, layout, inference_pod_index).trainer_coord
+
+
+def _base_gpu(inference_cell_index: int, layout: PairingLayout, inference_pod_index: int = 0) -> int:
+    return _place(inference_cell_index, layout, inference_pod_index).base_gpu_id
+
+
 def _coordinate(cell_index: int, pod_index: int = 0, pool_id: str = TRAINER_POOL_ID) -> PodCoordinate:
     return PodCoordinate(pool_id=pool_id, cell_index=cell_index, pod_in_cell_index=pod_index)
 
@@ -319,6 +331,164 @@ def _all_targets(layout: PairingLayout) -> list[PodCoordinate]:
     ]
 
 
+class TestBaseGpuIdOfAnInferencePod:
+    def test_a_whole_node_pod_starts_at_the_first_card(self):
+        """Its pod holds every card of the node, so the node-local numbering starts where the node does."""
+        layout = _layout(num_inference_cells=2, num_trainer_cells=2)
+
+        assert [_base_gpu(index, layout) for index in range(2)] == [0, 0]
+
+    def test_sub_node_pods_tile_the_cards_of_the_node_they_share(self):
+        """Two half-node pods on one trainer node must not both claim the cards the node numbers from zero."""
+        layout = _sub_node_layout()
+
+        assert [_base_gpu(index, layout) for index in range(4)] == [0, 4, 0, 4]
+
+    def test_one_gpu_pods_walk_the_whole_node(self):
+        """The recipe default is one gpu per engine, and eight of them cover a node card by card."""
+        layout = _layout(
+            num_inference_cells=8, num_trainer_cells=1, num_pods_per_trainer_cell=1, num_gpus_per_inference_pod=1
+        )
+
+        assert [_base_gpu(index, layout) for index in range(8)] == list(range(8))
+
+    def test_the_pods_of_a_wide_engine_each_start_at_their_own_node(self):
+        """An engine wider than a node is one whole-node pod per node, and each is handed its cards from zero."""
+        layout = _layout(
+            num_inference_cells=1,
+            num_trainer_cells=1,
+            num_pods_per_inference_cell=2,
+            num_pods_per_trainer_cell=2,
+        )
+
+        assert [_base_gpu(0, layout, inference_pod_index=index) for index in range(2)] == [0, 0]
+
+    def test_the_pods_of_a_wide_engine_pair_with_adjacent_trainer_pods(self):
+        """Its two pods sit on two nodes, so they must wait on the two trainer pods holding those nodes."""
+        layout = _layout(
+            num_inference_cells=1,
+            num_trainer_cells=1,
+            num_pods_per_inference_cell=2,
+            num_pods_per_trainer_cell=2,
+        )
+
+        assert [_target(0, layout, inference_pod_index=index) for index in range(2)] == [
+            _coordinate(0, 0),
+            _coordinate(0, 1),
+        ]
+
+    def test_the_pool_offset_moves_the_first_card_along(self):
+        """A second pool starts part way into the node its neighbour shares, and its offset says how far."""
+        layout = _sub_node_layout(gpu_offset=4, num_inference_cells=3)
+
+        assert [_base_gpu(index, layout) for index in range(3)] == [4, 0, 4]
+
+    def test_a_multi_pod_cell_of_sub_node_pods_walks_across_the_trainer_pods(self):
+        """Cells and pods per cell were only ever varied one at a time; together they index one flat gpu line."""
+        layout = _layout(
+            num_inference_cells=2,
+            num_trainer_cells=1,
+            num_pods_per_inference_cell=2,
+            num_pods_per_trainer_cell=2,
+            num_gpus_per_inference_pod=4,
+        )
+        placements = [_place(cell, layout, pod) for cell in range(2) for pod in range(2)]
+
+        assert [(placement.trainer_coord, placement.base_gpu_id) for placement in placements] == [
+            (_coordinate(0, 0), 0),
+            (_coordinate(0, 0), 4),
+            (_coordinate(0, 1), 0),
+            (_coordinate(0, 1), 4),
+        ]
+
+    def test_a_multi_pod_cell_of_quarter_node_pods_stays_on_one_trainer_pod(self):
+        """Four quarter-node pods spread over two cells still share one node, a pair of its cards each."""
+        layout = _layout(
+            num_inference_cells=2,
+            num_trainer_cells=1,
+            num_pods_per_inference_cell=2,
+            num_pods_per_trainer_cell=2,
+            num_gpus_per_inference_pod=2,
+        )
+        placements = [_place(cell, layout, pod) for cell in range(2) for pod in range(2)]
+
+        assert {placement.trainer_coord for placement in placements} == {_coordinate(0, 0)}
+        assert [placement.base_gpu_id for placement in placements] == [0, 2, 4, 6]
+
+
+_DISJOINT_CARD_LAYOUTS = {
+    "half node pods, two cells of two": _layout(
+        num_inference_cells=2,
+        num_trainer_cells=1,
+        num_pods_per_inference_cell=2,
+        num_pods_per_trainer_cell=2,
+        num_gpus_per_inference_pod=4,
+    ),
+    "quarter node pods on one trainer pod": _layout(
+        num_inference_cells=2,
+        num_trainer_cells=1,
+        num_pods_per_inference_cell=2,
+        num_pods_per_trainer_cell=2,
+        num_gpus_per_inference_pod=2,
+    ),
+    "one gpu pods across two nodes": _layout(
+        num_inference_cells=16,
+        num_trainer_cells=2,
+        num_pods_per_trainer_cell=1,
+        num_gpus_per_inference_pod=1,
+    ),
+    "half node pods offset by one pod": _sub_node_layout(gpu_offset=4, num_inference_cells=3),
+}
+
+
+def _cards_by_trainer_pod(layouts: dict[str, PairingLayout]) -> dict[PodCoordinate, list[tuple[str, range]]]:
+    cards: dict[PodCoordinate, list[tuple[str, range]]] = {}
+    for pool_id, layout in layouts.items():
+        for cell_index in range(layout.num_inference_cells):
+            for pod_index in range(layout.num_pods_per_inference_cell):
+                placement = _place(cell_index, layout, pod_index)
+                span = range(placement.base_gpu_id, placement.base_gpu_id + layout.num_gpus_per_inference_pod)
+                cards.setdefault(placement.trainer_coord, []).append((pool_id, span))
+    return cards
+
+
+def _cards_claimed_by_two_trainer_pods(cards: dict[PodCoordinate, list[tuple[str, range]]]) -> list[Any]:
+    return [
+        (trainer, first, second)
+        for trainer, entries in cards.items()
+        for first, second in itertools.combinations(entries, 2)
+        if set(first[1]) & set(second[1])
+    ]
+
+
+class TestTheEnginesOfOneTrainerPodHoldDifferentCards:
+    @pytest.mark.parametrize("name", sorted(_DISJOINT_CARD_LAYOUTS))
+    def test_no_two_pods_of_a_pool_are_given_an_overlapping_stretch(self, name: str):
+        """Two engines on one node sharing a card is the failure this whole mechanism exists to avoid."""
+        cards = _cards_by_trainer_pod({name: _DISJOINT_CARD_LAYOUTS[name]})
+
+        assert _cards_claimed_by_two_trainer_pods(cards) == []
+
+    @pytest.mark.parametrize("name", sorted(_DISJOINT_CARD_LAYOUTS))
+    def test_at_least_one_trainer_pod_really_seats_several_of_them(self, name: str):
+        """A layout whose trainer pods seat one engine each would satisfy the test above without testing it."""
+        cards = _cards_by_trainer_pod({name: _DISJOINT_CARD_LAYOUTS[name]})
+
+        assert max(len(entries) for entries in cards.values()) >= 2
+
+    def test_two_pools_splitting_a_node_are_given_different_halves_of_it(self):
+        """Prefill and decode share a trainer node on purpose, and only their offsets keep them apart."""
+        cards = _cards_by_trainer_pod(
+            {
+                INFERENCE_POOL_ID: _sub_node_layout(num_inference_cells=1),
+                DECODE_POOL_ID: _sub_node_layout(gpu_offset=4, num_inference_cells=1),
+            }
+        )
+
+        assert _cards_claimed_by_two_trainer_pods(cards) == []
+        assert len(cards[_coordinate(0, 0)]) == 2
+
+
 class TestGpuOffsetPairing:
     @pytest.mark.parametrize(
         "layout",
@@ -530,44 +700,109 @@ class TestPoolsClaimDistinctGpus:
         )
 
 
+BASE_GPU_ID_ANNOTATION = DEFAULT_LABEL_KEYS.base_gpu_id_annotation
+_ESCAPED_BASE_GPU_ID_ANNOTATION = "miles.radixark.io~1base-gpu-id"
+_CHART_META_ANNOTATION = f"{DEFAULT_LABEL_KEYS.meta_annotation_prefix}{DEFAULT_LABEL_KEYS.gpu_ids_meta}"
+_POD_ANNOTATIONS = {_CHART_META_ANNOTATION: "0", "leaderworkerset.sigs.k8s.io/size": "2"}
+
+
+def _release_patch(
+    *,
+    node_name: str = "gpu-7",
+    base_gpu_id: int = 0,
+    gates: list[str] | None = None,
+    annotations: dict[str, str] | None = None,
+) -> list[dict[str, Any]]:
+    return pairing_pods.release_patch(
+        node_name=node_name,
+        base_gpu_id=base_gpu_id,
+        gates=[pairing_pods._GATE_NAME] if gates is None else gates,
+        annotations=_POD_ANNOTATIONS if annotations is None else annotations,
+    )
+
+
 class TestReleasePatch:
     def test_pins_the_pod_to_one_node_and_removes_the_gate(self):
         """Both in one patch, so a controller restart cannot leave a pinned pod still gated."""
-        patch = pairing_pods.release_patch(node_name="gpu-7", gates=[pairing_pods._GATE_NAME], has_node_selector=False)
+        patch = _release_patch()
 
-        assert patch[0]["value"] == {"kubernetes.io/hostname": "gpu-7"}
-        assert patch[1:] == [
+        assert patch[0] == {
+            "op": "add",
+            "path": "/spec/affinity",
+            "value": {
+                "nodeAffinity": {
+                    "requiredDuringSchedulingIgnoredDuringExecution": {
+                        "nodeSelectorTerms": [
+                            {"matchFields": [{"key": "metadata.name", "operator": "In", "values": ["gpu-7"]}]}
+                        ]
+                    }
+                }
+            },
+        }
+        assert patch[2:] == [
             {"op": "test", "path": "/spec/schedulingGates/0/name", "value": pairing_pods._GATE_NAME},
             {"op": "remove", "path": "/spec/schedulingGates/0"},
         ]
 
-    def test_adds_one_key_when_the_pod_already_has_a_selector(self):
-        """Replacing the map would drop the run's own nodeSelector, and a gated pod may only gain keys."""
-        patch = pairing_pods.release_patch(node_name="gpu-7", gates=[pairing_pods._GATE_NAME], has_node_selector=True)
+    def test_tells_the_pod_which_card_of_the_node_is_its_own(self):
+        """The pod cannot work this out itself, and the same patch that seats it is what hands it over."""
+        patch = _release_patch(base_gpu_id=5)
 
-        assert patch[0] == {
+        assert patch[1] == {
             "op": "add",
-            "path": "/spec/nodeSelector/kubernetes.io~1hostname",
-            "value": "gpu-7",
+            "path": f"/metadata/annotations/{_ESCAPED_BASE_GPU_ID_ANNOTATION}",
+            "value": "5",
         }
 
+    def test_writes_the_card_as_a_string_the_downward_api_can_serve(self):
+        """A fieldRef reads an annotation, and an annotation value that is not a string is not valid."""
+        [operation] = [op for op in _release_patch(base_gpu_id=7) if str(op["path"]).startswith("/metadata")]
+
+        assert operation["value"] == "7"
+
+    def test_never_replaces_the_map_of_a_pod_that_carries_annotations(self):
+        """A whole-map add replaces it, dropping the chart's gpu meta and the platform's own bookkeeping."""
+        patch = _release_patch(base_gpu_id=3, annotations=_POD_ANNOTATIONS)
+
+        assert patch[1]["path"] == f"/metadata/annotations/{_ESCAPED_BASE_GPU_ID_ANNOTATION}"
+        assert not any(operation["path"] == "/metadata/annotations" for operation in patch)
+
+    def test_refuses_to_build_a_patch_for_a_pod_that_carries_no_annotations(self):
+        """The apiserver would reject the add anyway; failing here names the pod state that caused it."""
+        with pytest.raises(AssertionError, match="carries no annotations"):
+            _release_patch(base_gpu_id=3, annotations={})
+
+    def test_pins_by_node_name_rather_than_by_hostname_label(self):
+        """A node whose hostname label differs from its name would match no node at all."""
+        [term] = _release_patch()[0]["value"]["nodeAffinity"]["requiredDuringSchedulingIgnoredDuringExecution"][
+            "nodeSelectorTerms"
+        ]
+
+        assert term == {"matchFields": [{"key": "metadata.name", "operator": "In", "values": ["gpu-7"]}]}
+
     def test_removes_only_its_own_gate(self):
         """Dropping the whole list would release a pod another controller is still deliberately holding back."""
         gates = ["other.io/first", pairing_pods._GATE_NAME, "other.io/last"]
 
-        patch = pairing_pods.release_patch(node_name="gpu-7", gates=gates, has_node_selector=False)
+        patch = _release_patch(gates=gates)
 
-        assert patch[1:] == [
+        assert patch[2:] == [
             {"op": "test", "path": "/spec/schedulingGates/1/name", "value": pairing_pods._GATE_NAME},
             {"op": "remove", "path": "/spec/schedulingGates/1"},
         ]
 
     def test_is_a_json_patch_rather_than_a_merge(self):
         """A merge patch setting the gates to an empty list is silently ignored: the list merges by name."""
-        patch = pairing_pods.release_patch(node_name="gpu-7", gates=[pairing_pods._GATE_NAME], has_node_selector=False)
+        patch = _release_patch()
 
         assert all("op" in operation for operation in patch)
 
+    def test_hands_over_the_card_only_if_the_gate_is_still_there(self):
+        """A test op guards the whole patch, so a pod another pass released is never re-annotated."""
+        patch = _release_patch(base_gpu_id=5)
+
+        assert [operation["op"] for operation in patch] == ["add", "add", "test", "remove"]
+
 
 class TestCoordinateOf:
     def test_reads_the_pool_and_the_indices_the_chart_labelled(self):
@@ -606,6 +841,7 @@ def _pod(
         DEFAULT_LABEL_KEYS.cell_index: str(cell_index),
         DEFAULT_LABEL_KEYS.pod_in_cell_index: str(pod_index),
     }
+    pod.metadata.annotations = {f"{DEFAULT_LABEL_KEYS.meta_annotation_prefix}{DEFAULT_LABEL_KEYS.gpu_ids_meta}": "0"}
     return pod
 
 
@@ -644,6 +880,12 @@ def _controller(core_v1: Any, layout: PairingLayout | None = None) -> PairingCon
     return PairingController(config=_config(pools), core_v1=core_v1)
 
 
+def _base_gpu_id_written(body: list[dict[str, Any]]) -> str:
+    [operation] = [op for op in body if str(op["path"]).startswith("/metadata/annotations")]
+    value = operation["value"]
+    return value if isinstance(value, str) else value[BASE_GPU_ID_ANNOTATION]
+
+
 class TestReconcile:
     def test_releases_a_gated_inference_onto_its_trainer_node(self):
         """This is the whole point: the inference ends up where the trainer that feeds it already runs."""
@@ -655,9 +897,7 @@ class TestReconcile:
         assert core_v1.patched == [
             (
                 _pod_name(INFERENCE_POOL_ID, 0),
-                pairing_pods.release_patch(
-                    node_name="gpu-3", gates=[pairing_pods._GATE_NAME], has_node_selector=False
-                ),
+                _release_patch(node_name="gpu-3"),
             )
         ]
 
@@ -675,13 +915,58 @@ class TestReconcile:
         assert core_v1.patched == [
             (
                 _pod_name(INFERENCE_POOL_ID, index),
-                pairing_pods.release_patch(
-                    node_name="gpu-3", gates=[pairing_pods._GATE_NAME], has_node_selector=False
-                ),
+                _release_patch(node_name="gpu-3", base_gpu_id=index * 4),
             )
             for index in (0, 1)
         ]
 
+    def test_gives_every_engine_of_a_shared_node_a_card_of_its_own(self):
+        """Eight one-gpu engines on one trainer node is the shape this exists for, and 0..7 is the answer."""
+        core_v1 = FakeCoreV1()
+        layout = _layout(
+            num_inference_cells=8,
+            num_trainer_cells=1,
+            num_pods_per_trainer_cell=1,
+            num_gpus_per_inference_pod=1,
+        )
+        pods = [_pod(INFERENCE_POOL_ID, index) for index in range(8)]
+        pods.append(_pod(TRAINER_POOL_ID, 0, node_name="gpu-3", gated=False))
+
+        asyncio.run(_attached(_controller(core_v1, layout), pods).reconcile(_key(TRAINER_POOL_ID, 0)))
+
+        assert [_base_gpu_id_written(body) for _, body in core_v1.patched] == [str(index) for index in range(8)]
+
+    def test_gives_a_whole_node_engine_the_first_card(self):
+        """The engine holds every card of the node, so it is handed them from zero and is told exactly that."""
+        core_v1 = FakeCoreV1()
+        pods = [_pod(INFERENCE_POOL_ID, 1), _pod(TRAINER_POOL_ID, 1, node_name="gpu-3", gated=False)]
+
+        asyncio.run(_attached(_controller(core_v1), pods).reconcile(_key(TRAINER_POOL_ID, 1)))
+
+        assert [_base_gpu_id_written(body) for _, body in core_v1.patched] == ["0"]
+
+    def test_starts_a_second_sub_node_pool_where_its_own_offset_says(self):
+        """Two half-node pools split one trainer node, and only the offset tells the second one apart."""
+        core_v1 = FakeCoreV1()
+        controller = PairingController(
+            config=_config(
+                [
+                    _inference_pool(_sub_node_layout(num_inference_cells=1)),
+                    _inference_pool(_sub_node_layout(gpu_offset=4, num_inference_cells=1), pool_id=DECODE_POOL_ID),
+                ]
+            ),
+            core_v1=core_v1,
+        )
+        pods = [
+            _pod(INFERENCE_POOL_ID, 0),
+            _pod(DECODE_POOL_ID, 0),
+            _pod(TRAINER_POOL_ID, 0, 0, node_name="gpu-3", gated=False),
+        ]
+
+        asyncio.run(_attached(controller, pods).reconcile(_key(TRAINER_POOL_ID, 0, 0)))
+
+        assert [_base_gpu_id_written(body) for _, body in core_v1.patched] == ["0", "4"]
+
     def test_releases_only_the_inference_pods_of_the_trainer_being_reconciled(self):
         """The other trainer pod's node is a different machine, so its inference pods must stay gated."""
         core_v1 = FakeCoreV1()
@@ -742,7 +1027,7 @@ class TestReconcile:
 
         asyncio.run(_attached(_controller(core_v1), pods).reconcile(_key(TRAINER_POOL_ID, 0)))
 
-        assert core_v1.patched[0][1][0]["path"].endswith("kubernetes.io~1hostname")
+        assert core_v1.patched[0][1][0]["path"] == "/spec/affinity"
 
     def test_waits_while_the_trainer_has_no_node(self):
         """Releasing now would let the scheduler put the inference anywhere, which is the bug gates prevent."""
@@ -868,9 +1153,7 @@ class TestSeveralInferencePools:
         assert core_v1.patched == [
             (
                 _pod_name(DECODE_POOL_ID, 1),
-                pairing_pods.release_patch(
-                    node_name="gpu-4", gates=[pairing_pods._GATE_NAME], has_node_selector=False
-                ),
+                _release_patch(node_name="gpu-4"),
             )
         ]
 
@@ -1038,9 +1321,7 @@ class TestEventSequences:
             assert harness.core_v1.patched == [
                 (
                     _pod_name(INFERENCE_POOL_ID, 0),
-                    pairing_pods.release_patch(
-                        node_name="gpu-3", gates=[pairing_pods._GATE_NAME], has_node_selector=False
-                    ),
+                    _release_patch(node_name="gpu-3"),
                 ),
             ]
 
@@ -1054,9 +1335,7 @@ class TestEventSequences:
             assert harness.core_v1.patched == [
                 (
                     _pod_name(INFERENCE_POOL_ID, 1),
-                    pairing_pods.release_patch(
-                        node_name="gpu-9", gates=[pairing_pods._GATE_NAME], has_node_selector=False
-                    ),
+                    _release_patch(node_name="gpu-9"),
                 ),
             ]
 

tests/fast/utils/workers/worker_provider/kubernetes/core/test_cell_view.py

@@ -62,6 +62,7 @@ def make_run(
             cell_size_annotation="cell-size",
             meta_annotation_prefix="meta-",
             gpu_ids_meta="gpu_ids",
+            base_gpu_id_annotation="base-gpu-id",
         ),
     )
 

tests/fast/utils/workers/worker_provider/kubernetes/helm/test_env.py

@@ -55,3 +55,11 @@ class TestReleaseDiscovery:
 
         with pytest.raises(AssertionError, match=env.RELEASE_ENV_VAR):
             env.current_release()
+
+
+class TestTheCardAnnotationStandsApartFromTheMeta:
+    def test_the_base_gpu_id_key_is_not_a_meta_annotation(self) -> None:
+        """Meta annotations are merged across a cell's pods and asserted equal; this one differs per pod."""
+        keys = env.DEFAULT_LABEL_KEYS
+
+        assert not keys.base_gpu_id_annotation.startswith(keys.meta_annotation_prefix)