Persist spot and reservation across node pool updates

docs/guides/cost_optimization.md

@@ -111,6 +111,16 @@ kinetic pool add --accelerator a100 --spot
 2. **Use Checkpointing:** Use Kinetic's integration with **Orbax** (see the [Checkpointing Guide](checkpointing.md)) to continuously flush state to Cloud Storage (`gs://`). If a Spot preemption kills your training run midway, you can resume from the last saved checkpoint instead of restarting from scratch.
 3. **Multi-Host TPUs:** While you can provision Spot pools for multi-node TPUs (such as `tpu-v3`, `tpu-v4`, `tpu-v5p`, or `tpu-v6e`), if any single host in the TPU slice is preempted, the entire slice job will fail. Spot pricing is therefore highly effective for single-host jobs (like `tpu-v5litepod-4`, `tpu-v5litepod-8`, or `l4` workloads) where you minimize the probability of aggregate preemption.
 
+### Verify the Provisioning Model
+
+Kinetic records the Spot setting in the cluster state. Later `kinetic pool add`, `kinetic pool remove`, and `kinetic up` commands keep the setting on the pool.
+
+To see the provisioning model of each pool, run `kinetic pool list`. The `Provisioning` row shows `Spot` or `On-demand`.
+
+:::{note}
+Kinetic records the Spot setting when it creates the pool. A pool from an earlier version of Kinetic has no such record, and the `Provisioning` row is absent. The next `kinetic pool add` or `kinetic pool remove` command rebuilds that pool as on-demand. To keep Spot capacity, remove the pool and add it again with `--spot`.
+:::
+
 ---
 
 ## 5. Managing Capacity Reservations

docs/guides/reservations.md

@@ -33,6 +33,22 @@ kinetic pool add \
 
 Kinetic sets `SPECIFIC_RESERVATION` affinity on the node pool so the autoscaler consumes nodes from your reservation instead of competing for on-demand capacity.
 
+## Verify the Reservation
+
+Kinetic records the reservation in the cluster state. Later `kinetic pool add`, `kinetic pool remove`, and `kinetic up` commands keep the reservation on the pool.
+
+To see the reservation of each pool, run:
+
+```bash
+kinetic pool list --project your-project-id
+```
+
+The `Reservation` row shows the reservation name. Kinetic does not show this row for a pool that uses on-demand or Spot capacity.
+
+:::{note}
+Kinetic records the reservation when it creates the pool. A pool from an earlier version of Kinetic has no such record, and the `Reservation` row is absent. The next `kinetic pool add` or `kinetic pool remove` command removes the reservation affinity from that pool. To keep the reservation, remove the pool and add it again with `--reservation`.
+:::
+
 ## Cleaning Up
 
 Remove the reservation when you are done to avoid ongoing charges:

kinetic/cli/infra/program.py

@@ -459,11 +459,14 @@ def _create_accelerator_pools(
   zone: str,
   project_id: str,
   service_account: pulumi.Output[str],
-) -> list[tuple[GpuConfig | TpuConfig, gcp.container.NodePool, int]]:
-  """Create accelerator node pools and return entries for export."""
-  pool_entries: list[
-    tuple[GpuConfig | TpuConfig, gcp.container.NodePool, int]
-  ] = []
+) -> list[tuple[NodePoolConfig, gcp.container.NodePool]]:
+  """Create accelerator node pools and return entries for export.
+
+  Each entry pairs the requesting ``NodePoolConfig`` with the created
+  pool, so the stack export can record every setting that was applied —
+  not just the accelerator shape.
+  """
+  pool_entries: list[tuple[NodePoolConfig, gcp.container.NodePool]] = []
   for np in node_pools:
     accel = np.accelerator
     if isinstance(accel, GpuConfig):
@@ -490,7 +493,7 @@ def _create_accelerator_pools(
       )
     else:
       continue
-    pool_entries.append((accel, pool, np.min_nodes))
+    pool_entries.append((np, pool))
   return pool_entries
 
 
@@ -502,10 +505,16 @@ def _export_stack_outputs(
   repo: gcp.artifactregistry.Repository,
   ar_location: str,
   cluster_name: str,
-  pool_entries: list[tuple[GpuConfig | TpuConfig, gcp.container.NodePool, int]],
+  pool_entries: list[tuple[NodePoolConfig, gcp.container.NodePool]],
   force_destroy: bool,
 ) -> None:
-  """Export all Pulumi stack outputs."""
+  """Export all Pulumi stack outputs.
+
+  The per-pool ``accelerators`` entries are the only record of how each
+  pool was provisioned: ``load_state()`` reads them back and re-declares
+  every pool on the next update. Any setting missing here is silently
+  reset to its default the next time a pool is added or removed.
+  """
   pulumi.export("project", project_id)
   pulumi.export("zone", zone)
   pulumi.export("cluster_name", cluster.name)
@@ -524,30 +533,35 @@ def _export_stack_outputs(
     return
 
   export_outputs = []
-  for accel, pool, min_nodes in pool_entries:
+  for np, pool in pool_entries:
+    accel = np.accelerator
     if isinstance(accel, GpuConfig):
       entry = pool.name.apply(
-        lambda pn, a=accel, mn=min_nodes: {
+        lambda pn, a=accel, cfg=np: {
           "type": "GPU",
           "name": a.name,
           "count": a.count,
           "machine_type": a.machine_type,
           "node_pool": pn,
           "node_count": 1,
-          "min_nodes": mn,
+          "min_nodes": cfg.min_nodes,
+          "spot": a.spot,
+          "reservation": cfg.reservation,
         }
       )
     else:  # TpuConfig
       entry = pool.name.apply(
-        lambda pn, a=accel, mn=min_nodes: {
+        lambda pn, a=accel, cfg=np: {
           "type": "TPU",
           "name": a.name,
           "chips": a.chips,
           "topology": a.topology,
           "machine_type": a.machine_type,
           "node_pool": pn,
           "node_count": a.num_nodes,
-          "min_nodes": mn,
+          "min_nodes": cfg.min_nodes,
+          "spot": a.spot,
+          "reservation": cfg.reservation,
         }
       )
     export_outputs.append(entry)

kinetic/cli/infra/program_test.py

@@ -5,6 +5,7 @@
 from absl.testing import absltest, parameterized
 
 from kinetic.cli.config import NodePoolConfig
+from kinetic.cli.infra import stack_manager
 from kinetic.core.accelerators import GpuConfig, TpuConfig
 
 # Patch pulumi provider modules before importing program, so the module-level
@@ -182,6 +183,150 @@ def test_force_destroy_is_exported(self):
     self.assertFalse(exported["force_destroy"])
 
 
+_SPOT_GPU = GpuConfig(
+  "a100", 1, "nvidia-tesla-a100", "a2-highgpu-1g", spot=True
+)
+_SPOT_TPU = TpuConfig(
+  "v6e", 8, "2x4", "tpu-v6e-slice", "ct6e-standard-4t", 2, spot=True
+)
+_ON_DEMAND_GPU = GpuConfig("l4", 1, "nvidia-l4", "g2-standard-4")
+_ON_DEMAND_TPU = TpuConfig(
+  "v5p", 8, "2x2x2", "tpu-v5p-slice", "ct5p-hightpu-4t", 2
+)
+
+
+class _ResolvedOutputs(list):
+  """Stands in for ``pulumi.Output.all`` — the resolved values, as a list.
+
+  ``apply`` stays lazy (it returns a mock without running the callback),
+  matching how the fully mocked ``pulumi`` behaves elsewhere in this
+  file. ``_build_kubeconfig`` calls it with mock cluster attributes that
+  its ``json.dumps`` cannot serialize.
+  """
+
+  def apply(self, fn):
+    del fn
+    return mock.MagicMock()
+
+
+def _run_program_capturing_exports(config):
+  """Run the Pulumi program, returning (stack exports, gcp mock).
+
+  ``NodePool.name.apply`` is stubbed to evaluate eagerly, and
+  ``Output.all`` to return the resolved list, so the exported
+  accelerator entries are plain dicts rather than Pulumi Outputs.
+  """
+  exports = {}
+
+  with (
+    mock.patch.object(program, "pulumi") as pulumi_mock,
+    mock.patch.object(program, "command"),
+    mock.patch.object(program, "gcp") as gcp_mock,
+    mock.patch.object(program, "k8s"),
+  ):
+    pulumi_mock.export.side_effect = exports.__setitem__
+    pulumi_mock.Output.all.side_effect = lambda *entries: _ResolvedOutputs(
+      entries
+    )
+
+    def make_node_pool(_resource_name, **kwargs):
+      pool = mock.MagicMock()
+      pool.name.apply.side_effect = lambda fn: fn(kwargs["name"])
+      return pool
+
+    gcp_mock.container.NodePool.side_effect = make_node_pool
+    program.create_program(config)()
+
+  return exports, gcp_mock
+
+
+class TestAcceleratorExportRoundTrip(parameterized.TestCase):
+  """Pool settings must survive export → load_state → re-apply.
+
+  ``pool add``/``pool remove`` rebuild the whole pool list from the
+  ``accelerators`` stack export and re-declare every existing pool. A
+  setting the export omits comes back as its default, and because GKE
+  node_config changes force replacement, the pool is silently rebuilt
+  without it.
+  """
+
+  @parameterized.named_parameters(
+    dict(testcase_name="gpu_spot", accel=_SPOT_GPU, min_nodes=1),
+    dict(testcase_name="tpu_spot", accel=_SPOT_TPU, min_nodes=2),
+    dict(testcase_name="gpu_on_demand", accel=_ON_DEMAND_GPU, min_nodes=0),
+    dict(testcase_name="tpu_on_demand", accel=_ON_DEMAND_TPU, min_nodes=0),
+  )
+  def test_spot_survives_round_trip(self, accel, min_nodes):
+    pool = NodePoolConfig("pool-abcd", accel, min_nodes=min_nodes)
+    exports, _ = _run_program_capturing_exports(_make_config([pool]))
+
+    (entry,) = exports["accelerators"]
+    self.assertEqual(entry["spot"], accel.spot)
+    self.assertEqual(stack_manager._export_to_node_pool(entry), pool)
+
+  @parameterized.named_parameters(
+    dict(testcase_name="gpu", accel=_ON_DEMAND_GPU),
+    dict(testcase_name="tpu", accel=_ON_DEMAND_TPU),
+  )
+  def test_reservation_survives_round_trip(self, accel):
+    pool = NodePoolConfig("pool-abcd", accel, reservation="my-reservation")
+    exports, _ = _run_program_capturing_exports(_make_config([pool]))
+
+    (entry,) = exports["accelerators"]
+    self.assertEqual(entry["reservation"], "my-reservation")
+    self.assertEqual(stack_manager._export_to_node_pool(entry), pool)
+
+  def test_multiple_pools_keep_their_own_settings(self):
+    pools = [
+      NodePoolConfig("gpu-a100-abcd", _SPOT_GPU),
+      NodePoolConfig("gpu-l4-ef01", _ON_DEMAND_GPU, reservation="l4-res"),
+    ]
+    exports, _ = _run_program_capturing_exports(_make_config(pools))
+
+    restored = [
+      stack_manager._export_to_node_pool(e) for e in exports["accelerators"]
+    ]
+    self.assertEqual(restored, pools)
+
+  @parameterized.named_parameters(
+    dict(testcase_name="spot", accel=_SPOT_GPU, reservation=None),
+    dict(
+      testcase_name="reservation", accel=_ON_DEMAND_GPU, reservation="my-res"
+    ),
+  )
+  def test_reapplying_exported_state_keeps_node_config(
+    self, accel, reservation
+  ):
+    """The regression: a second `pool add` must not reset the first pool.
+
+    Simulates `pool add --spot` followed by another pool command —
+    export, read back via load_state, re-declare — and checks the node
+    config Pulumi would apply the second time still carries the
+    setting.
+    """
+    original = NodePoolConfig("pool-abcd", accel, reservation=reservation)
+    exports, _ = _run_program_capturing_exports(_make_config([original]))
+
+    restored = [
+      stack_manager._export_to_node_pool(e) for e in exports["accelerators"]
+    ]
+    new_pool = NodePoolConfig("gpu-l4-new1", _ON_DEMAND_GPU)
+    _, gcp_mock = _run_program_capturing_exports(
+      _make_config(restored + [new_pool])
+    )
+
+    node_config = gcp_mock.container.NodePoolNodeConfigArgs.call_args_list[0]
+    self.assertEqual(node_config.kwargs["spot"], accel.spot)
+    if reservation is None:
+      self.assertIsNone(node_config.kwargs["reservation_affinity"])
+    else:
+      gcp_mock.container.NodePoolNodeConfigReservationAffinityArgs.assert_called_once_with(
+        consume_reservation_type="SPECIFIC_RESERVATION",
+        key="compute.googleapis.com/reservation-name",
+        values=[reservation],
+      )
+
+
 class TestClusterResourceLabels(absltest.TestCase):
   """The GKE cluster must carry a kinetic resource label.
 

kinetic/cli/infra/stack_manager.py

@@ -123,17 +123,32 @@ def get_current_node_pools(stack) -> list[NodePoolConfig]:
 
 
 def _export_to_node_pool(entry: dict) -> NodePoolConfig:
-  """Convert a stack export dict back to a NodePoolConfig."""
+  """Convert a stack export dict back to a NodePoolConfig.
+
+  Every field must survive this round trip: the result is fed straight
+  back into the Pulumi program, which re-declares the pool. A dropped
+  field is applied as its default, and because GKE node_config changes
+  force replacement, that silently rebuilds the pool without it.
+
+  ``spot`` and ``reservation`` are absent from stacks last updated
+  before they were exported, and default to False / None there.
+  """
   pool_name = entry["node_pool"]
+  spot = bool(entry.get("spot", False))
   accelerator: accelerators.GpuConfig | accelerators.TpuConfig
   if entry["type"] == "GPU":
-    accelerator = accelerators.make_gpu(entry["name"], entry["count"])
+    accelerator = accelerators.make_gpu(
+      entry["name"], entry["count"], spot=spot
+    )
   elif entry["type"] == "TPU":
-    accelerator = accelerators.make_tpu(entry["name"], entry["chips"])
+    accelerator = accelerators.make_tpu(
+      entry["name"], entry["chips"], spot=spot
+    )
   else:
     raise ValueError(f"Unknown accelerator type in node pool export: {entry}")
   return NodePoolConfig(
     name=pool_name,
     accelerator=accelerator,
     min_nodes=entry.get("min_nodes", 0),
+    reservation=entry.get("reservation") or None,
   )