Jet amplitude setter: failed state on device errors, commit after upload

FixedPatchStepper.set_jet_amplitude validates index/amplitude and builds
the new rows before touching the device; invalid input leaves the
instance usable. A copy/sync error marks it failed (as initialize/step/
download do) and re-raises; host state (_jets, blocked velocity table)
commits only after a successful upload. CurvedLinkSoA.set_jet_amplitude
likewise uploads a new array first and commits jets/surface_velocity
afterwards. Focused mocked-driver tests for dense and blocked copy
failure, invalid input and IndexError.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

src/CelerisLab/lbm/curved_links.py

@@ -102,16 +102,20 @@ class CurvedLinkSoA:
             self.jet_shape, self.jet_id = jet_link_shapes(self.jets, hits)
 
     def set_jet_amplitude(self, index, amplitude, stream=None):
-        """Rewrite jet rows from cached shapes; re-upload only the velocity column."""
+        """Rewrite jet rows from cached shapes; re-upload only the velocity column.
+
+        Host state is committed only after the upload returns; on a device error
+        the host keeps the previous amplitude (device contents are then unknown).
+        """
         from ..body.jets import jet_link_velocity
         jets = list(self.jets)
         jets[index] = jets[index].with_amplitude(amplitude)  # validates before any change
-        self.jets = tuple(jets)
+        velocity = self.surface_velocity.copy()
         driven = self.jet_id >= 0
-        self.surface_velocity[driven] = jet_link_velocity(
-            self.jets, self.jet_shape, self.jet_id)[driven]
+        velocity[driven] = jet_link_velocity(jets, self.jet_shape, self.jet_id)[driven]
         if self.surface_velocity_gpu is not None:
-            _upload_columns([(self.surface_velocity_gpu, self.surface_velocity)], stream)
+            _upload_columns([(self.surface_velocity_gpu, velocity)], stream)
+        self.jets, self.surface_velocity = tuple(jets), velocity
 
     def surface_exchange(self, body_id):
         """Latest boundary step's (Fx,Fy,Tz), split by aperture; not nozzle thrust.

src/CelerisLab/lbm/refinement_stepper.py

@@ -355,20 +355,32 @@ class FixedPatchStepper:
                      grid=((cl.count + self._block[0] - 1) // self._block[0], 1), **kw)
 
     def set_jet_amplitude(self, index, amplitude):
-        """Replace one jet's amplitude between steps; no geometry/table rebuild."""
+        """Replace one jet's amplitude between steps; no geometry/table rebuild.
+
+        Invalid input raises before any change. A device/sync error marks the
+        instance failed (as step/initialize do); host state commits only after
+        a successful upload.
+        """
         with self._lock, self.ctx.activate():
             self._require("new", "ready")
             jets = list(self._jets)
             jets[index] = jets[index].with_amplitude(amplitude)  # validates first
-            if self._fine_curved is not None:
-                self._fine_curved.set_jet_amplitude(index, amplitude, self._stream)
-            elif self._jet_cache is not None:
+            velocity = None
+            if self._fine_curved is None and self._jet_cache is not None:
                 shape, jet_id = self._jet_cache
                 velocity = self._fine_tables["surface_velocity"].copy()
                 driven = jet_id >= 0
                 velocity[driven] = jet_link_velocity(jets, shape, jet_id)[driven]
-                self._stream.synchronize()  # no in-flight step reads the column
-                cuda.memcpy_htod(self._link_tables["surface_velocity"], velocity)
+            try:
+                if self._fine_curved is not None:
+                    self._fine_curved.set_jet_amplitude(index, amplitude, self._stream)
+                elif velocity is not None:
+                    self._stream.synchronize()  # no in-flight step reads the column
+                    cuda.memcpy_htod(self._link_tables["surface_velocity"], velocity)
+            except BaseException:
+                self._status = "failed"
+                raise
+            if velocity is not None:
                 self._fine_tables["surface_velocity"] = velocity
             self._jets = tuple(jets)
 

tests/validation/multigrid/test_refinement_jets.py

@@ -75,6 +75,38 @@ class JetStepperCPU(unittest.TestCase):
         self.assertEqual(backend._jets[0].amplitude, -SPEED)
         backend.close()
 
+    def test_copy_failure_marks_failed_and_keeps_host_amplitude(self):
+        for fine_block in (None, FINE_BLOCKS[1]):  # dense SoA and blocked tables
+            backend, _ = mocked([circle()], [jet()], fine_block)
+            status = backend._status
+            for bad, error in ((1.5, ValueError), (float("nan"), ValueError)):
+                with self.assertRaises(error):
+                    backend.set_jet_amplitude(0, bad)
+                self.assertEqual(backend._status, status)  # invalid input stays usable
+            with self.assertRaises(IndexError):
+                backend.set_jet_amplitude(1, -SPEED)
+            self.assertEqual(backend._status, status)
+            dense = backend._fine_curved
+            before = (dense.surface_velocity if dense is not None
+                      else backend._fine_tables["surface_velocity"]).copy()
+            boom = Mock(side_effect=RuntimeError("copy failed"))
+            with patch.object(cuda, "memcpy_htod", boom), patch.object(cuda, "memcpy_htod_async", boom), \
+                    patch.object(cuda, "pagelocked_empty", side_effect=lambda shape, dtype: np.empty(shape, dtype)):
+                with self.assertRaisesRegex(RuntimeError, "copy failed"):
+                    backend.set_jet_amplitude(0, -SPEED)
+            self.assertEqual(backend._status, "failed", fine_block)
+            self.assertEqual(backend._jets[0].amplitude, SPEED)
+            if dense is not None:
+                self.assertEqual(dense.jets[0].amplitude, SPEED)
+                np.testing.assert_array_equal(dense.surface_velocity, before)
+            else:
+                np.testing.assert_array_equal(backend._fine_tables["surface_velocity"], before)
+            with self.assertRaisesRegex(RuntimeError, "failed"):
+                backend.step(1)
+            with self.assertRaisesRegex(RuntimeError, "failed"):
+                backend.set_jet_amplitude(0, -SPEED)
+            backend.close()
+
     def test_jet_overlapping_legacy_aperture_rejected_before_cuda(self):
         with patch.object(cuda, "Stream", side_effect=AssertionError("CUDA touched")):
             with self.assertRaisesRegex(ValueError, "SurfaceActuation"):