Make Stop actually halt the pipeline via a run-scoped sentinel

Gradio's cancels only closes the request generator; the pipeline keeps
running on a worker thread locally and in a forked child on ZeroGPU.
Stop now writes a sentinel file under the run's data dir, and the
motion/denoise hooks — the only code the pipeline re-enters — raise
RunCancelled when they see it.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

fdanyone_app.py

@@ -522,6 +522,28 @@ class Session:
     """Whether the last ``_pump`` call's pipeline thread returned."""
 
 
+class RunCancelled(FourDAnyoneError):
+    """Raised inside a pipeline hook once the run's stop sentinel appears."""
+
+
+def _stop_path(spec: RunSpec) -> Path:
+    return spec.data_dir / "stop-requested"
+
+
+def _check_stop(spec: RunSpec) -> None:
+    """Abort the pipeline from inside one of its hooks after Stop is pressed.
+
+    Gradio's ``cancels`` only closes the request's generator; the pipeline keeps
+    running — on a worker thread here, in a forked ZeroGPU child on the Space.
+    A file under the run's own directory is the one channel that reaches both,
+    and the hooks are the only code of ours the pipeline re-enters, so they are
+    where the run can be unwound.
+    """
+
+    if _stop_path(spec).exists():
+        raise RunCancelled("Stopped.")
+
+
 T = TypeVar("T")
 
 
@@ -570,6 +592,7 @@ def _motion_hook(link: Link, session: Session) -> Callable[[str, dict[str, objec
     recording: rr.RecordingStream = link.recording
 
     def on_stage(stage: str, payload: dict[str, object]) -> None:
+        _check_stop(session.spec)
         if stage == "bboxes":
             boxes: Float[torch.Tensor, "frames 4"] = payload["bbx_xyxy"]  # pyrefly: ignore
             for index, box in enumerate(boxes.numpy()[: INFERENCE.num_frames]):
@@ -604,6 +627,7 @@ def _denoise_hook(link: Link, session: Session) -> Callable[[int, tuple[int, ...
         view_indices: tuple[int, ...],
         x0_hat: Float[torch.Tensor, "views 48 latent_t latent_h latent_w"],
     ) -> None:
+        _check_stop(session.spec)
         if PREVIEW_DECODER is None:
             return
         plan: tuple[tuple[int, int], ...] = preview_slice_plan(x0_hat.shape[2])
@@ -844,6 +868,7 @@ def smoke_motion_phase(link: Link, session: Session) -> Iterator[str]:
 
     rng: np.random.Generator = np.random.default_rng(0)
     for index in range(0, 121, 10):
+        _check_stop(session.spec)
         _set_frame(link.recording, index, session.spec.fps)
         x0: float = 180.0 + 2.0 * index
         link.recording.log(
@@ -864,6 +889,7 @@ def smoke_generate_phase(link: Link, session: Session) -> Iterator[str]:
     link.recording.send_blueprint(diffusion_blueprint(), make_active=True)
     yield "[smoke] diffusion"
     for step in range(4):
+        _check_stop(session.spec)
         link.recording.set_time(DIFFUSION_TIMELINE, sequence=step)
         for view in range(VIEWS):
             noise: UInt8[np.ndarray, "160 88 3"] = rng.integers(
@@ -947,6 +973,16 @@ def run_gpu(spec: RunSpec) -> Iterator[tuple[dict | None, bytes | None, str]]:
     yield session.summary, link.read(), "Publishing the result."
 
 
+def request_stop(spec: RunSpec | None) -> str:
+    """Write the run's stop sentinel; the pipeline unwinds at its next hook."""
+
+    if spec is None:
+        return "Nothing is running."
+    spec.data_dir.mkdir(parents=True, exist_ok=True)
+    _stop_path(spec).touch()
+    return "Stopping — the pipeline halts at its next step."
+
+
 def publish_cpu(spec: RunSpec, summary: dict | None) -> Iterator[tuple[bytes | None, str]]:
     """Attach the finished rig and its six videos, off the GPU allocation."""
 
@@ -1052,7 +1088,12 @@ def build_demo() -> gr.Blocks:
         prepared = started.success(prepare_cpu, spec, [viewer, status])
         generated = prepared.success(run_gpu, spec, [summary, viewer, status])
         published = generated.success(publish_cpu, [spec, summary], [viewer, status])
-        stop_button.click(None, cancels=[started, prepared, generated, published])
+        # `cancels` detaches the UI at once, but the pipeline itself only halts
+        # when it reads the sentinel `request_stop` writes: the worker thread
+        # (and on ZeroGPU, the forked child) never sees a Gradio cancellation.
+        stop_button.click(
+            request_stop, spec, status, cancels=[started, prepared, generated, published]
+        )
     return demo
 
 

tests/test_app_helpers.py

@@ -110,6 +110,44 @@ def test_pump_yields_hook_labels_and_returns_the_work_result() -> None:
     assert session.worker_finished is True
 
 
+def test_stop_sentinel_unwinds_the_pump_worker(tmp_path: Path) -> None:
+    """Stop must reach the pipeline thread itself, not just Gradio's generator.
+
+    ``cancels`` closes the request generator, but the worker thread (or the
+    forked ZeroGPU child) never hears it; the sentinel checked by the hooks is
+    the only thing that actually frees the GPU.
+    """
+
+    session: fdanyone_app.Session = _pump_session()
+    spec: fdanyone_app.RunSpec = fdanyone_app.RunSpec(
+        token=session.spec.token,
+        video_path=session.spec.video_path,
+        start_time=0.0,
+        seed=0,
+        fps=Fraction(25, 1),
+        data_dir=tmp_path,
+    )
+    session = fdanyone_app.Session(spec)
+    assert fdanyone_app.request_stop(None) == "Nothing is running."
+    fdanyone_app.request_stop(spec)
+    assert (tmp_path / "stop-requested").exists()
+
+    ticks: list[int] = []
+
+    def work() -> str:
+        # What every pipeline hook does on entry, once per stage or step.
+        for index in range(1000):
+            fdanyone_app._check_stop(spec)
+            ticks.append(index)
+        return "finished"
+
+    with pytest.raises(fdanyone_app.RunCancelled):
+        list(fdanyone_app._pump(session, work))
+    assert ticks == []
+    # The worker joined before re-raising, so the thread is truly gone.
+    assert session.worker_finished is True
+
+
 def test_pump_reraises_a_worker_failure_on_the_caller_thread() -> None:
     """A pipeline error must surface where Gradio can turn it into a message."""