From a63b90e41a416979c23dcb56882f105b44f4fe61 Mon Sep 17 00:00:00 2001 From: Mohamad Rostami Date: Sun, 6 Sep 2026 22:25:18 +0200 Subject: [PATCH 1/2] runner: gpusnap warm/snapshot pause hook MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add Runner.gpusnap_warm_pause(), called in lifespan right after setup()+warmup and BEFORE the queue worker starts / health goes green. When GPUSNAP_WARM= is set (only while gpusnap builds a snapshot), the worker signals WARM_READY and blocks until restored + RESUMEd — so a snapshot is taken at a clean point (model warm, no job claimed) and a restored worker resumes straight into serving. Unset on a normal boot => no-op, behaviour unchanged. Works for queue mode, HTTP mode, and TorchRun multi-GPU workers. Co-Authored-By: Claude Opus 4.8 --- sprocket/sprocket.py | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/sprocket/sprocket.py b/sprocket/sprocket.py index 27f4769..bc5c0d6 100644 --- a/sprocket/sprocket.py +++ b/sprocket/sprocket.py @@ -355,6 +355,26 @@ async def maybe_run_warmup(self) -> None: SHUTDOWN_REQUESTED.touch() sys.exit(0) + async def gpusnap_warm_pause(self) -> None: + """gpusnap snapshot hook. When this worker is being warmed to build a snapshot (GPUSNAP_WARM set to + a directory), pause here — after setup()+warmup, so the model is loaded and JIT-warm, but BEFORE + the queue worker claims any real job and before /health goes green — signal the snapshotter that we + are warm (WARM_READY), and block until it restores us and writes RESUME. This guarantees a snapshot + never captures (and then loses) a claimed job, and that a restored worker resumes cleanly into + serving. On a normal boot GPUSNAP_WARM is unset and this is a no-op, so non-gpusnap behaviour is + unchanged. Works for both queue and HTTP mode and for TorchRun (multi-GPU) workers.""" + warm_dir = os.getenv("GPUSNAP_WARM") + if not warm_dir: + return + d = Path(warm_dir) + d.mkdir(parents=True, exist_ok=True) + (d / "WARM_READY").touch() + logger.info(f"gpusnap: warm (setup+warmup done); paused pre-queue, signalled WARM_READY in {d}") + resume = d / "RESUME" + while not resume.exists(): + await asyncio.sleep(0.1) + logger.info("gpusnap: RESUME received; proceeding to serve") + @contextlib.asynccontextmanager async def lifespan(self, _: Starlette) -> AsyncIterator[None]: if isinstance(self.sprocket, AsyncSprocket): @@ -362,6 +382,8 @@ async def lifespan(self, _: Starlette) -> AsyncIterator[None]: else: self.sprocket.setup() await self.maybe_run_warmup() + # gpusnap: pause here (warm, pre-queue) while a snapshot is taken; no-op on a normal boot. + await self.gpusnap_warm_pause() if self.queue_mode: asyncio.create_task(self.run_queue_worker()) self.healthy = True From 8c96c89d29a73576518feb91c77b1112d238696f Mon Sep 17 00:00:00 2001 From: Mohamad Rostami Date: Mon, 7 Sep 2026 09:59:37 +0200 Subject: [PATCH 2/2] runner: expose GPUSNAP_GPUS_CHANGED so multi-GPU workers reinit NCCL only on a real GPU move The gpusnap tool records whether the physical GPUs moved vs dump time and writes GPUS_CHANGED in the warm dir. The warm-pause hook reads it and exports GPUSNAP_GPUS_CHANGED, so a TorchRun/multi-GPU worker reinitialises its process group only on a cross-GPU retarget and keeps its valid communicators for an in-place restore. Co-Authored-By: Claude Opus 4.8 --- sprocket/sprocket.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/sprocket/sprocket.py b/sprocket/sprocket.py index bc5c0d6..c9216aa 100644 --- a/sprocket/sprocket.py +++ b/sprocket/sprocket.py @@ -373,7 +373,13 @@ async def gpusnap_warm_pause(self) -> None: resume = d / "RESUME" while not resume.exists(): await asyncio.sleep(0.1) - logger.info("gpusnap: RESUME received; proceeding to serve") + # The tool records whether the physical GPUs moved vs dump time (a cross-GPU retarget). Expose it + # as GPUSNAP_GPUS_CHANGED so a multi-GPU (TorchRun) worker reinitialises its process group / NCCL + # ONLY when the GPUs actually changed — an in-place restore keeps its still-valid communicators. + gc = d / "GPUS_CHANGED" + changed = gc.read_text().strip() != "0" if gc.exists() else True + os.environ["GPUSNAP_GPUS_CHANGED"] = "1" if changed else "0" + logger.info(f"gpusnap: RESUME received (gpus_changed={changed}); proceeding to serve") @contextlib.asynccontextmanager async def lifespan(self, _: Starlette) -> AsyncIterator[None]: