runpod: cancel a boot, release a pod, and stop blocking the server (v0.3.80)

Three reports from the same run, all of them about giving pods back.

SCALE BLOCKED THE WHOLE SERVER. The endpoint called into the pool straight
from the coroutine, so while it walked dozens of (card, region) candidates
-- a network round trip each -- nothing else was served. /v1/runpod/status
stopped answering during exactly the operation someone would poll it
about, which makes a slow rental look identical to an outage. It now runs
off the event loop and the reply is bounded: after `wait_s` (default 20)
it answers `pending` with a request id, and the SCALING CONTINUES --
shielded, because cancelling it would leave a half-made lease and possibly
a pod. GET /v1/runpod/scale/{id} has the result for a caller that gave up.

A BOOT COULD NOT BE CANCELLED. `_booting` is an integer, not pod objects --
a pod being created does not exist here yet -- so `shed_now` could not
see one, count one or stop one. A client that asked for a pod, got four on
the way and released them was told `kept:0, needed:1, floor:1` and paid
every boot out in full. A shed now bumps a cancel epoch; an ANTICIPATORY
boot (warm floor, background growth) checks it during the model load and
again just before it joins the fleet, and hands its pod straight back. A
boot a request is blocked on is never cancellable, whatever arrives
meanwhile: the request is what the pod is for. Booting pods also count as
capacity in the surplus arithmetic, and the shed reports how many it
stopped.

A POD COULD NOT BE RELEASED unless some pool floor could be lowered to
lose it. `'pool:digesta-qwen38' has no runpod block` came from code whose
own comment says the caller cannot look that key up -- so a live pool is
now addressable by the key /v1/runpod/status reports. And nothing anywhere
could terminate ONE pod by id: not an adopted pod, not one left by a
crash, not one whose pool was reconfigured underneath it. POST
/v1/runpod/release does, admin-scoped (leasing is bounded and expires;
terminating a named pod is neither), attributing the spend before it drops
the pod so the money does not land nowhere.

And the over-provisioning: `ensure_ready` incremented `_booting` by ONE
and then started `want` boots, so with min_pods 3 three cards were on the
way while the 15s scaler, `_should_grow` and the parallel-boot cap all
believed there was one -- each topping up against a deficit already being
filled.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01HXRMe9ZPUQCCCtBai9uqt3

codai/__init__.py

@@ -16,7 +16,7 @@
 
 # Canonical product version for CoderAI — single source of truth. Both the API
 # metadata and the admin web UI read from here.
-__version__ = "0.3.79"
+__version__ = "0.3.80"
 __tagline__ = "Complete orchestration, distribution and escalation of remotizable advanced inference."
 
 # Configure the CUDA caching allocator BEFORE torch is imported anywhere.

codai/api/app.py

@@ -19,8 +19,10 @@ FastAPI application module for codai API.
 Contains the FastAPI app initialization, lifespan, and core endpoints.
 """
 
+import asyncio
 import logging
 import os
+import threading
 import time
 from contextlib import asynccontextmanager
 from typing import List, Optional, Union
@@ -921,9 +923,31 @@ class RunpodScaleRequest(BaseModel):
                     "OFF by a request: asking for true where the operator set "
                     "false is ignored and reported. Omit to use the model's own "
                     "setting.")
+    wait_s: Optional[float] = Field(
+        None, ge=0, le=600,
+        description="How long to wait for the pool to answer before returning "
+                    "`pending`. The work CONTINUES either way -- this bounds "
+                    "the reply, not the scaling. 0 returns immediately with a "
+                    "request id. Default 20s.")
     model_config = ConfigDict(extra="allow")
 
 
+#: Background scale requests, by id, so a caller that stopped waiting can find
+#: out how it ended. Bounded: a scale is a handful of fields and old ones are
+#: dropped, because this must never become a place memory accumulates.
+_scale_jobs: "dict" = {}
+_scale_jobs_lock = threading.Lock()
+SCALE_JOBS_MAX = 64
+SCALE_WAIT_DEFAULT_S = 20.0
+
+
+def _scale_job_put(job_id: str, payload: dict) -> None:
+    with _scale_jobs_lock:
+        _scale_jobs[job_id] = payload
+        while len(_scale_jobs) > SCALE_JOBS_MAX:
+            _scale_jobs.pop(next(iter(_scale_jobs)))
+
+
 @app.post("/v1/runpod/scale", summary="Ask for warm pods on a model", tags=["Core"])
 async def runpod_scale(body: RunpodScaleRequest, request: Request):
     """Let a client hold N pods warm for a model, up to its configured max_pods.
@@ -956,10 +980,45 @@ async def runpod_scale(body: RunpodScaleRequest, request: Request):
     spend caps still refuse to provision regardless of what is asked here.
     """
     from codai.api import runpod_worker
+    import uuid
+
+    # OFF THE EVENT LOOP, and this is the important half. This endpoint used to
+    # call straight into the pool from the coroutine, so while it hunted for
+    # capacity -- dozens of (card, region) candidates, one network round trip
+    # each -- the whole server was blocked. /v1/runpod/status stopped answering
+    # during exactly the operation someone would poll it about, which makes a
+    # slow rental look identical to an outage.
+    job_id = uuid.uuid4().hex[:12]
+    task = asyncio.create_task(asyncio.to_thread(
+        runpod_worker.set_client_pods,
+        body.model, body.pods, teardown=bool(body.teardown),
+        gpu_types=body.gpu_types, allow_other_gpus=body.allow_other_gpus))
+
+    def _done(t):
+        try:
+            _scale_job_put(job_id, {"id": job_id, "state": "done",
+                                    "result": t.result()})
+        except BaseException as exc:            # noqa: BLE001 - recorded, not raised
+            _scale_job_put(job_id, {"id": job_id, "state": "failed",
+                                    "error": str(exc)})
+    task.add_done_callback(_done)
+
+    wait_s = (SCALE_WAIT_DEFAULT_S if body.wait_s is None
+              else max(0.0, float(body.wait_s)))
     try:
-        return runpod_worker.set_client_pods(
-            body.model, body.pods, teardown=bool(body.teardown),
-            gpu_types=body.gpu_types, allow_other_gpus=body.allow_other_gpus)
+        if wait_s <= 0:
+            raise asyncio.TimeoutError
+        return await asyncio.wait_for(asyncio.shield(task), timeout=wait_s)
+    except asyncio.TimeoutError:
+        # The scaling CONTINUES. Answering is what stopped, so a caller with a
+        # 60s client timeout gets an answer instead of a dead socket, and a
+        # rental that is merely slow stops looking like an outage.
+        _scale_job_put(job_id, {"id": job_id, "state": "running"})
+        return {"accepted": True, "pending": True, "request_id": job_id,
+                "model": body.model, "requested": body.pods,
+                "detail": (f"still finding capacity after {wait_s:.0f}s — the "
+                           f"request is running; poll /v1/runpod/scale/"
+                           f"{job_id}, or /v1/runpod/status for the fleet")}
     except LookupError as exc:
         raise HTTPException(status_code=404, detail=str(exc))
     except ValueError as exc:
@@ -972,6 +1031,58 @@ async def runpod_scale(body: RunpodScaleRequest, request: Request):
         raise HTTPException(status_code=500, detail=f"scale failed: {exc}")
 
 
+@app.get("/v1/runpod/scale/{request_id}", summary="How a scale request ended",
+         tags=["Core"])
+async def runpod_scale_result(request_id: str):
+    """The result of a scale that was still running when its caller gave up."""
+    with _scale_jobs_lock:
+        got = _scale_jobs.get(str(request_id or "").strip())
+    if got is None:
+        raise HTTPException(
+            status_code=404,
+            detail=f"no scale request {request_id!r} — it finished long enough "
+                   f"ago to have been dropped, or it was never made. "
+                   f"/v1/runpod/status has the fleet either way.")
+    return got
+
+
+class RunpodReleaseRequest(BaseModel):
+    pod_id: str = Field(..., description="The pod to terminate, as reported by "
+                                         "/v1/runpod/status.")
+    reason: str = Field("released by client",
+                        description="Recorded in the log and the spend ledger.")
+    model_config = ConfigDict(extra="allow")
+
+
+@app.post("/v1/runpod/release", summary="Terminate one pod by id", tags=["Core"])
+async def runpod_release(body: RunpodReleaseRequest, request: Request):
+    """Stop ONE pod, whatever pool it is in and whether or not it is in one.
+
+    Every other path scales a POOL, which means a pod with no floor to lower
+    could not be stopped from here at all: one adopted from another process,
+    one left behind by a crash, one whose pool was reconfigured underneath it.
+    The only remaining way to stop those billing was the RunPod console, and a
+    pod this process cannot see is precisely the one that needs a button.
+
+    ADMIN-SCOPED, unlike /v1/runpod/scale. Leasing warm pods is bounded by the
+    model's own max_pods and expires by itself; terminating a named pod is
+    neither bounded nor reversible, and the pod may be serving someone else's
+    request. The reply says how many were in flight when it went.
+    """
+    if not _admin_scoped_token(request):
+        raise HTTPException(status_code=403,
+                            detail="terminating a named pod needs an "
+                                   "admin-scoped API key")
+    from codai.api import runpod_worker
+    try:
+        return await asyncio.to_thread(runpod_worker.release_pod,
+                                       body.pod_id, body.reason)
+    except LookupError as exc:
+        raise HTTPException(status_code=404, detail=str(exc))
+    except Exception as exc:
+        raise HTTPException(status_code=500, detail=f"release failed: {exc}")
+
+
 @app.get("/coderai/capabilities", summary="Server capability document", tags=["Core"])
 async def get_broker_capabilities():
     """Return broker capability metadata."""

tests/test_pool_growth_adopts.py

@@ -40,7 +40,7 @@ class _Pool:
         self.calls.append("adopt")
         return self._adopts
 
-    def _provision_one(self):
+    def _provision_one(self, cancellable=False):
         self.calls.append("rent")
 
 
@@ -81,7 +81,12 @@ def test_every_provisioning_path_tries_adoption_first():
     """
     src = (ROOT / "codai/api/runpod_worker.py").read_text()
     # 5 call sites: ensure_ready, acquire, growth, warm floor, seed_region.
-    assert src.count("self._provision_one()") == 5
+    # Two of the five now pass `cancellable=True` -- the anticipatory
+    # boots, which hand their pod back if the fleet is released while
+    # they are on the way. The inline one a request is waiting on never
+    # is. Counted together: every call site still adopts first.
+    assert src.count("self._provision_one(") == 5
+    assert src.count("self._provision_one(cancellable=True)") == 2
 
     # Per PATH rather than a global count of one spelling. The warm floor stopped
     # matching `if not self._adopt_shared_pod():` in 0.3.69 when it became a

tests/test_runpod_client_scale.py

@@ -328,12 +328,19 @@ def test_ensure_ready_marks_the_pool_as_provisioning():
     """A pod is not in self.pods until it answers its health probe, minutes
     later. Without this flag the 15 s scaler reads "0 healthy, floor 1" and
     rents a second one — which is exactly what happened on the first live
-    boot: two A40s where the configuration asked for one."""
+    boot: two A40s where the configuration asked for one.
+
+    And it counts ONE PER POD. This said `+= 1` while booting `want` of them,
+    so with min_pods 3 three cards were on the way and every other reader of
+    `_booting` believed there was one — each then topping up against a deficit
+    already being filled. Asking for N and getting more than N started here."""
     src = __import__("pathlib").Path(rw.__file__).read_text()
     blk = src.split("def ensure_ready")[1].split("def _pick")[0]
-    assert "self._booting += 1" in blk
-    assert "finally:" in blk and "self._booting = max(0, self._booting - 1)" in blk
-    assert blk.index("if self._booting:") < blk.index("self._booting += 1")
+    assert "self._booting += want" in blk
+    assert "self._booting - want" in blk, "what goes up must all come down"
+    assert "self._booting += 1" not in blk
+    assert "finally:" in blk and "self._booting = max(0, self._booting - want)" in blk
+    assert blk.index("if self._booting:") < blk.index("self._booting += want")
 
 
 def test_the_warm_top_up_does_not_stack_on_a_booting_pod():
@@ -394,7 +401,7 @@ def test_the_warm_path_adopts_before_it_rents():
     src = __import__("pathlib").Path(rw.__file__).read_text()
     ready = src.split("def ensure_ready")[1].split("def _pick")[0]
     assert "_adopt_shared_pod()" in ready
-    assert ready.index("_adopt_shared_pod()") < ready.index("self._provision_one()")
+    assert ready.index("_adopt_shared_pod()") < ready.index("self._provision_one(")
     warm = src.split("def maintain")[1].split("def close")[0]
     # Spelling changed in 0.3.68: adoption moved ABOVE the sibling-boot check
     # and now returns early, because a restart's own orphans were sitting behind
@@ -402,7 +409,7 @@ def test_the_warm_path_adopts_before_it_rents():
     # guards -- adopt before renting -- is stronger than before, so assert the
     # intent rather than the old line. See tests/test_runpod_restart_orphans.py.
     assert "_adopt_shared_pod()" in warm
-    assert warm.index("_adopt_shared_pod()") < warm.index("self._provision_one()")
+    assert warm.index("_adopt_shared_pod()") < warm.index("self._provision_one(")
 
 
 def test_adoption_is_one_implementation_for_both_callers():
@@ -454,7 +461,7 @@ def test_the_warm_path_waits_for_a_pod_that_is_still_booting():
     src = __import__("pathlib").Path(rw.__file__).read_text()
     ready = src.split("def ensure_ready")[1].split("def _pick")[0]
     assert "sibling_is_provisioning" in ready
-    assert ready.index("sibling_is_provisioning") < ready.index("self._provision_one()")
+    assert ready.index("sibling_is_provisioning") < ready.index("self._provision_one(")
 
 
 def test_a_booting_pod_from_a_previous_process_counts():

tests/test_runpod_release_and_boot_cancel.py

+"""Giving pods back: the three holes a live run fell into.
+
+1. A boot could not be cancelled. `_booting` is an INTEGER, not a list of pod
+   objects -- a pod being created does not exist here yet -- so `shed_now`
+   could not see one, could not count one and could not stop one. A client that
+   asked for a pod, got four on the way and released them was told
+   `kept:0, needed:1, floor:1` and paid all four boots out in full.
+
+2. A pod could not be released at all unless a POOL floor could be lowered to
+   lose it. `'pool:<name>' has no runpod block` -- from code whose own comment
+   says the caller cannot look that key up -- and nothing anywhere could
+   terminate one pod by id. An adopted pod, a pod left by a crash, a pod whose
+   pool was reconfigured: RunPod console only.
+
+3. Asking for N could produce more than N, because `ensure_ready` counted one
+   boot while starting `want` of them.
+"""
+
+import pathlib
+import sys
+import threading
+
+ROOT = pathlib.Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT))
+
+from codai.api import runpod_worker as rw                         # noqa: E402
+
+
+def _pod(pid, inflight=0, last_used=0.0, draining=False, hourly_usd=0.6,
+         started_at=0.0):
+    return rw.PodHandle(pod_id=pid, url=f"http://{pid}", hourly_usd=hourly_usd,
+                        started_at=started_at, inflight=inflight,
+                        last_used=last_used, draining=draining)
+
+
+def _pool(max_pods=8, min_pods=0, threshold=4, pods=(), booting=0,
+          allow_client_scale=True):
+    pool = rw.RunpodPodPool.__new__(rw.RunpodPodPool)
+    pool._cv = threading.Condition()
+    pool.model_key = "pool:digesta-qwen38"
+    pool._closed = False
+    pool._provisioning = False
+    pool._surplus_since = 0.0
+    pool._booting = booting
+    pool._cancel_epoch = 0
+    pool.pods = list(pods)
+    cfg = rw.RunpodModelConfig()
+    cfg.max_pods = max_pods
+    cfg.min_pods = min_pods
+    cfg.scale_up_inflight_per_pod = threshold
+    cfg.scale_down_after_s = 180
+    cfg.schedule_enabled = False
+    cfg.allow_client_scale = allow_client_scale
+    pool.mcfg = cfg
+    pool.effective_min_pods = lambda: min_pods
+    pool._killed = []
+    pool._terminate = lambda pod, why: pool._killed.append((pod.pod_id, why))
+    return pool
+
+
+# ----------------------------------------------- a boot that can be stopped #
+def test_shedding_bumps_the_cancel_epoch():
+    pool = _pool(booting=4)
+    before = pool._cancel_epoch
+    pool.shed_now()
+    assert pool._cancel_epoch == before + 1
+
+
+def test_a_boot_started_before_the_shed_is_cancelled():
+    pool = _pool(booting=1)
+    epoch = pool._cancel_epoch          # what the boot captured at entry
+    assert pool._boot_cancelled(epoch) is False
+    pool.shed_now()
+    assert pool._boot_cancelled(epoch) is True
+
+
+def test_a_boot_started_after_the_shed_is_not():
+    pool = _pool()
+    pool.shed_now()
+    epoch = pool._cancel_epoch          # a boot starting now captures this
+    assert pool._boot_cancelled(epoch) is False
+
+
+def test_the_boots_in_flight_are_counted_and_reported():
+    """`kept:0, needed:1, floor:1` while four cards were booting said nothing
+    had happened. The count is the difference between that and 'four stopped'."""
+    pool = _pool(booting=4)
+    out = pool.shed_now()
+    assert out["cancelled_boots"] == 4
+
+
+def test_booting_pods_count_as_capacity_in_the_surplus():
+    """A pod on the way is capacity arriving. Measuring surplus only against
+    READY pods is why a client holding four boots was told it had none."""
+    pool = _pool(min_pods=1, pods=[_pod("a"), _pod("b")], booting=2)
+    out = pool.shed_now()
+    # 2 ready + 2 booting against a floor of 1: the ready surplus goes, the
+    # boots are cancelled.
+    assert out["cancelled_boots"] == 2
+    assert len(out["terminated"]) >= 1
+
+
+def test_a_shed_with_nothing_booting_reports_zero_not_nothing():
+    pool = _pool(min_pods=0, pods=[_pod("a")])
+    assert pool.shed_now()["cancelled_boots"] == 0
+
+
+def test_only_anticipatory_boots_are_cancellable():
+    """A boot a request is BLOCKED on is the one thing a shed must not take:
+    the request is what the pod is for. Only the warm floor and background
+    growth pass cancellable."""
+    src = (ROOT / "codai/api/runpod_worker.py").read_text()
+    warm = src.split("def _warm_provision_one")[1].split("\n    def ")[0]
+    grow = src.split("def _grow_in_background")[1].split("\n    def ")[0]
+    acquire = src.split("def acquire")[1].split("\n    def ")[0]
+    ready = src.split("def ensure_ready")[1].split("\n    def ")[0]
+    assert "cancellable=True" in warm
+    assert "cancellable=True" in grow
+    assert "cancellable" not in acquire
+    assert "cancellable" not in ready
+
+
+def test_a_cancelled_boot_is_not_retried_on_another_machine():
+    """A boot FAILURE retries elsewhere; a cancellation must not, or the shed
+    rents exactly the card it just gave back."""
+    src = (ROOT / "codai/api/runpod_worker.py").read_text()
+    blk = src.split("def _provision_one")[1].split("\n    def ")[0]
+    assert "except BootCancelled:" in blk
+    # The BOOT's handler, not the one inside `_give_back`: the retry-on-another
+    # -machine path is the one that must not see a cancellation.
+    boot = blk.split("self._dump_pod_logs")[0]
+    assert boot.rindex("except BootCancelled:") < boot.rindex("except Exception as exc:")
+    assert "retrying on the next machine" in blk
+
+
+def test_a_cancelled_boot_is_benign_at_both_call_sites():
+    src = (ROOT / "codai/api/runpod_worker.py").read_text()
+    for name in ("_warm_provision_one", "_grow_in_background"):
+        blk = src.split(f"def {name}")[1].split("\n    def ")[0]
+        assert "except BootCancelled:" in blk, name
+
+
+# ------------------------------------------------ a pod that can be dropped #
+def test_a_live_pool_can_be_addressed_by_the_key_status_reports():
+    """`pool:digesta-qwen38` is what /v1/runpod/status calls it, and it was the
+    one name the scale endpoint refused."""
+    pool = _pool()
+    rw._pools["pool:digesta-qwen38"] = pool
+    try:
+        assert rw.get_live_pool("pool:digesta-qwen38") is pool
+        out = rw._scale_live_pool(pool, "pool:digesta-qwen38", 0, teardown=True)
+        assert out["model"] == "pool:digesta-qwen38"
+        assert "teardown" in out
+    finally:
+        rw._pools.pop("pool:digesta-qwen38", None)
+
+
+def test_a_pool_that_did_not_opt_in_is_still_refused():
+    pool = _pool(allow_client_scale=False)
+    try:
+        rw._scale_live_pool(pool, "pool:x", 1)
+    except PermissionError as exc:
+        assert "allow_client_scale" in str(exc)
+    else:
+        raise AssertionError("a pool that did not opt in was scaled")
+
+
+def test_an_unknown_name_says_what_WOULD_work():
+    """The old message named a condition the caller could not check."""
+    import pytest
+    rw._pools.pop("nope", None)
+    with pytest.raises(LookupError) as exc:
+        rw.set_client_pods("nope", 1)
+    msg = str(exc.value)
+    assert "no pod pool is running" in msg and "/v1/runpod/status" in msg
+
+
+def test_get_live_pool_is_exact_not_fuzzy():
+    pool = _pool()
+    rw._pools["pool:digesta-qwen38"] = pool
+    try:
+        assert rw.get_live_pool("digesta-qwen38") is None
+        assert rw.get_live_pool("") is None
+        assert rw.get_live_pool(None) is None
+    finally:
+        rw._pools.pop("pool:digesta-qwen38", None)
+
+
+def test_release_pod_needs_a_pod_id():
+    import pytest
+    with pytest.raises(LookupError):
+        rw.release_pod("")
+
+
+def test_the_release_route_exists_and_is_admin_scoped():
+    """Leasing warm pods is bounded and expires; terminating a named pod is
+    neither, and it may be serving someone else's request."""
+    src = (ROOT / "codai/api/app.py").read_text()
+    assert '@app.post("/v1/runpod/release"' in src
+    blk = src.split('@app.post("/v1/runpod/release"')[1].split("\n@app.")[0]
+    assert "_admin_scoped_token(request)" in blk
+
+
+# ------------------------------------------- scale stops blocking the server #
+def test_scale_runs_off_the_event_loop():
+    """It called into the pool from the coroutine, so while it walked dozens of
+    (card, region) candidates the whole server was blocked -- and
+    /v1/runpod/status stopped answering during exactly the operation someone
+    would poll it about. A slow rental looked identical to an outage."""
+    src = (ROOT / "codai/api/app.py").read_text()
+    blk = src.split("async def runpod_scale(")[1].split("\n@app.")[0]
+    assert "asyncio.to_thread" in blk
+    assert "runpod_worker.set_client_pods(\n" not in blk
+
+
+def test_a_slow_scale_answers_pending_rather_than_hanging():
+    src = (ROOT / "codai/api/app.py").read_text()
+    blk = src.split("async def runpod_scale(")[1].split("\n@app.")[0]
+    assert "asyncio.wait_for" in blk and "asyncio.shield" in blk
+    assert '"pending": True' in blk and '"request_id"' in blk
+
+
+def test_the_work_continues_after_the_caller_stops_waiting():
+    """`shield` is the whole point: the timeout ends the WAIT, not the scaling.
+    Cancelling the task would leave a half-made lease and possibly a pod."""
+    src = (ROOT / "codai/api/app.py").read_text()
+    blk = src.split("async def runpod_scale(")[1].split("\n@app.")[0]
+    assert "asyncio.shield(task)" in blk
+
+
+def test_a_caller_that_gave_up_can_still_find_out_how_it_ended():
+    src = (ROOT / "codai/api/app.py").read_text()
+    assert '@app.get("/v1/runpod/scale/{request_id}"' in src
+
+
+def test_the_job_table_cannot_grow_without_bound():
+    from codai.api import app as A
+    A._scale_jobs.clear()
+    for i in range(A.SCALE_JOBS_MAX + 20):
+        A._scale_job_put(f"j{i}", {"id": f"j{i}", "state": "done"})
+    assert len(A._scale_jobs) <= A.SCALE_JOBS_MAX
+    # The newest survive; the oldest are the ones nobody is coming back for.
+    assert f"j{A.SCALE_JOBS_MAX + 19}" in A._scale_jobs
+    A._scale_jobs.clear()

tests/test_runpod_scale_out.py

@@ -181,7 +181,7 @@ def test_the_single_flight_paths_ask_how_many_not_whether():
     and `_booting` is counted so a wave cannot overshoot the floor."""
     src = (ROOT / "codai/api/runpod_worker.py").read_text()
     ready = src.split("def ensure_ready")[1].split("\n    def ")[0]
-    assert "if self._booting:" in ready and "self._booting += 1" in ready
+    assert "if self._booting:" in ready and "self._booting += want" in ready
     maint = src.split("def maintain")[1].split("\n    def ")[0]
     assert "scale_up_max_parallel_boots" in maint
     assert "parallel - self._booting" in maint
@@ -195,7 +195,7 @@ def test_the_counter_returns_to_zero_after_a_scale_up():
     exits, each of which must decrement."""
     pool = _pool(pods=[_pod("a", inflight=16)], waiting=48)
     provisioned = []
-    pool._provision_one = lambda: provisioned.append(1)
+    pool._provision_one = lambda cancellable=False: provisioned.append(1)
     pool._pick = lambda healthy, affinity: healthy[0]
     pool._remember_affinity = lambda a, p: None
     got, url = pool.acquire(timeout=5)
@@ -274,7 +274,7 @@ def test_growth_prefers_a_pod_already_running_over_renting():
     nothing until the reaper's grace expired. Only the cold paths adopted."""
     pool = _pool(pods=[_pod("a", inflight=16)], waiting=48)
     provisioned, adopted = [], []
-    pool._provision_one = lambda: provisioned.append(1)
+    pool._provision_one = lambda cancellable=False: provisioned.append(1)
     pool._adopt_shared_pod = lambda: (adopted.append(1), True)[1]
     pool._pick = lambda healthy, affinity: healthy[0]
     pool._remember_affinity = lambda a, p: None

tests/test_runpod_warm_floor_rerent.py

@@ -272,7 +272,7 @@ def test_the_floor_is_re_established_once_the_ghost_is_gone(monkeypatch):
 
     rented = []
     pool._adopt_shared_pod = lambda: False
-    pool._provision_one = lambda: rented.append("rented")
+    pool._provision_one = lambda cancellable=False: rented.append("rented")
     pool._drain_pod_log = lambda p: None
     pool._scale_down = lambda now: None
     monkeypatch.setattr(rw, "sibling_is_provisioning", lambda key: False)
@@ -300,7 +300,7 @@ def test_a_sibling_already_renting_is_one_pod_not_none(monkeypatch):
     _install(monkeypatch, pool)
     rented = []
     pool._adopt_shared_pod = lambda: False
-    pool._provision_one = lambda: rented.append("rented")
+    pool._provision_one = lambda cancellable=False: rented.append("rented")
     pool._drain_pod_log = lambda p: None
     pool._scale_down = lambda now: None
     monkeypatch.setattr(rw, "sibling_is_provisioning", lambda key: True)
@@ -314,7 +314,7 @@ def test_an_existing_pod_elsewhere_is_adopted_before_renting(monkeypatch):
     _install(monkeypatch, pool)
     rented = []
     pool._adopt_shared_pod = lambda: True
-    pool._provision_one = lambda: rented.append("rented")
+    pool._provision_one = lambda cancellable=False: rented.append("rented")
     pool._drain_pod_log = lambda p: None
     pool._scale_down = lambda now: None
     monkeypatch.setattr(rw, "sibling_is_provisioning", lambda key: False)
@@ -330,7 +330,7 @@ def test_no_floor_means_no_re_rent(monkeypatch):
     rw.reconcile_pools({"ghost": "TERMINATED"})
     rented = []
     pool._adopt_shared_pod = lambda: False
-    pool._provision_one = lambda: rented.append("rented")
+    pool._provision_one = lambda cancellable=False: rented.append("rented")
     pool._drain_pod_log = lambda p: None
     pool._scale_down = lambda now: None
     monkeypatch.setattr(rw, "sibling_is_provisioning", lambda key: False)