longcat: report the real denoise step, over every segment (v0.2.87)

A 304-frame request sat at "generating 0/16 (0%)" for its whole nine-minute
first segment while the service's log showed Denoising: 15/16. Two defects.

The service set _progress["step"] = 0 at the start of every segment and never
again, so the step count existed only inside the vendored pipeline's
tqdm(total=len(timesteps), desc="Denoising") -- which takes no callback.
Nothing reached GET /progress; the poller in video.py faithfully reported the
0 it was handed. Each pipeline module did `from tqdm import tqdm`, so the bar
is instrumented by rebinding the name on the MODULE (patching tqdm.tqdm would
be too late), and only bars labelled "Denoising" are tracked so a
checkpoint-shard bar cannot move a request's progress.

The totals also counted ONE segment: a four-segment render would have run
0->16 four times against a denominator of 16. `step` now carries the completed
segments' offset so it counts the stage end to end, and the bar spans
steps x segments -- with refinement excluded, being a single pass over the
finished video. A request that asked for a duration rather than a segment
count only learns the real count from the service, so the denominator is
corrected on the first report, through a setter that does not restart the bar
and lose the elapsed time the rate is derived from.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

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.2.86"
+__version__ = "0.2.87"
 __tagline__ = "Complete orchestration, distribution and escalation of remotizable advanced inference."
 
 # Configure the CUDA caching allocator BEFORE torch is imported anywhere.

codai/api/video.py

@@ -86,6 +86,16 @@ def _vid_progress_done():
     _vid_progress["active"] = False
     _vid_progress["phase"] = "idle"
 
+def _vid_progress_total(total: int):
+    """Correct the denominator mid-render, without resetting the bar.
+
+    Some paths only learn how much work there is after the worker has started -- a
+    LongCat request that asked for a duration rather than a segment count is told the
+    segment count by the service -- and restarting the bar there would throw away the
+    elapsed time the rate is derived from."""
+    _vid_progress["total"] = int(total)
+
+
 def _vid_progress_step(step: int):
     _vid_progress["current"] = step
     elapsed = time.monotonic() - _vid_progress["started_at"]
@@ -4240,10 +4250,22 @@ async def _generate_longcat(request: VideoGenerationRequest, model_name: str,
                     else "source" if _k == "longcat_source" else _k] = str(model_cfg[_k])
 
     # Total steps across the stages we are about to run, so the bar spans the whole
-    # request rather than restarting per stage.
-    _total = sum(lc.stage_params(s, request.num_inference_steps,
-                                 request.guidance_scale)["num_inference_steps"]
-                 for s in stages)
+    # request rather than restarting per stage -- and across SEGMENTS, because a long
+    # video runs the generative stages once per segment. Counting a segment's worth
+    # made a four-segment render read 0/16 four times over instead of 0/64.
+    def _grand_total(segments: int) -> int:
+        segments = max(1, int(segments or 1))
+        return sum(
+            lc.stage_params(s, request.num_inference_steps,
+                            request.guidance_scale)["num_inference_steps"]
+            * (1 if s == "refinement" else segments)
+            for s in stages)
+
+    # Refinement is a single pass over the finished video, so it is not multiplied.
+    # The service owns the real segment count (it derives one from total_frames when
+    # the caller asked for a duration); the poll loop corrects this once it reports.
+    _segments_assumed = max(1, int(payload.get("num_segments") or 1))
+    _total = _grand_total(_segments_assumed)
     _vid_progress_reset(_total)
 
     _stop = threading.Event()
@@ -4254,7 +4276,7 @@ async def _generate_longcat(request: VideoGenerationRequest, model_name: str,
     # claimed in its own docstring to "carry cancellation across the process
     # boundary" and never did, so a LongCat render — the longest-running video job
     # there is — was the only one nothing could stop.
-    _lc_steps = sum(lc.STAGE_DEFAULTS.get(_s, (0, 0))[0] for _s in stages) or 1
+    _lc_steps = _total or 1
     _tid = task_registry.register(
         "video", title=(payload.get("prompt") or mode or "")[:80],
         model=model_name or "", total=_lc_steps)
@@ -4270,6 +4292,10 @@ async def _generate_longcat(request: VideoGenerationRequest, model_name: str,
         returns with an error. That is also exactly what eviction does to it, so the
         path is already exercised."""
         _seen_stage = ""
+        # What the stage now running reports as its own total, segments included, so a
+        # finished stage contributes what it actually ran rather than a default.
+        _stage_total = 0
+        _live = {"segments": _segments_assumed, "total": _total}
         while not _stop.wait(2.0):
             try:
                 if not _cancelled["flag"]:
@@ -4292,12 +4318,21 @@ async def _generate_longcat(request: VideoGenerationRequest, model_name: str,
                 st = str(p.get("stage") or "")
                 if st and st != _seen_stage:
                     if _seen_stage:
-                        _done_steps["base"] += lc.STAGE_DEFAULTS.get(
+                        _done_steps["base"] += _stage_total or lc.STAGE_DEFAULTS.get(
                             _seen_stage, (0, 0))[0]
                     _seen_stage = st
+                    _stage_total = 0
+                _stage_total = int(p.get("steps") or 0) or _stage_total
+                _segs_now = int(p.get("segments") or 0)
+                if _segs_now and _segs_now != _live["segments"]:
+                    _live["segments"] = _segs_now
+                    _live["total"] = _grand_total(_segs_now)
+                    _vid_progress_total(_live["total"])
+                # `step` counts the whole stage, every segment of it: the service adds
+                # the completed segments' offset before reporting the bar's position.
                 _step = _done_steps["base"] + int(p.get("step") or 0)
                 _vid_progress_step(_step)
-                task_registry.step(_tid, _step)
+                task_registry.step(_tid, _step, total=_live["total"])
             except Exception:
                 pass
 

tests/test_longcat_progress.py

+"""The denoising progress the caller polls, and the segment count it spans.
+
+A 304-frame request on 2026-10-08 sat at "generating 0/16 (0%)" for its whole
+nine-minute first segment while the service's own log showed
+``Denoising: 94%|...| 15/16``. Two separate defects:
+
+1. ``tools/longcat_service.py`` set ``_progress["step"] = 0`` at the start of every
+   segment and never again. The step count existed only inside the vendored
+   pipeline's ``tqdm(total=len(timesteps), desc="Denoising")``, which takes no
+   callback, so nothing reached ``GET /progress``. The consumer in
+   ``codai/api/video.py`` was already polling it correctly -- it faithfully reported
+   the 0 it was handed.
+
+2. Even with a live step, the totals counted ONE segment. A four-segment render
+   would have run 0->16 four times over against a denominator of 16.
+"""
+import importlib.util
+import pathlib
+import sys
+
+import pytest
+
+ROOT = pathlib.Path(__file__).resolve().parents[1]
+SERVICE = ROOT / "tools" / "longcat_service.py"
+
+
+@pytest.fixture
+def service(monkeypatch):
+    spec = importlib.util.spec_from_file_location("lc_service_progress", SERVICE)
+    mod = importlib.util.module_from_spec(spec)
+    monkeypatch.setitem(sys.modules, spec.name, mod)
+    spec.loader.exec_module(mod)
+    return mod
+
+
+# --------------------------------------------------------------- the instrumented bar
+
+def _bars(service, monkeypatch):
+    """Install the patched tqdm into two stand-in pipeline modules and return it."""
+    import types
+
+    mods = {}
+    for name in ("longcat_video.pipeline_longcat_video",
+                 "longcat_video.pipeline_longcat_video_avatar"):
+        m = types.ModuleType(name)
+        m.tqdm = object()
+        monkeypatch.setitem(sys.modules, name, m)
+        mods[name] = m
+    service._install_progress_bars()
+    return mods
+
+
+def test_the_pipeline_modules_bar_is_replaced(service, monkeypatch):
+    """`from tqdm import tqdm` binds the name on the MODULE, so that is what has to
+    be rebound -- patching tqdm.tqdm after the import would be too late."""
+    mods = _bars(service, monkeypatch)
+    for name, mod in mods.items():
+        assert mod.tqdm.__name__ == "_ProgressTqdm", name
+
+
+def test_installing_twice_does_not_stack_wrappers(service, monkeypatch):
+    mods = _bars(service, monkeypatch)
+    first = mods["longcat_video.pipeline_longcat_video"].tqdm
+    service._install_progress_bars()
+    assert mods["longcat_video.pipeline_longcat_video"].tqdm is first
+
+
+def test_a_denoising_bar_publishes_its_step(service, monkeypatch):
+    """The defect itself: stepping the bar must move what /progress reports."""
+    mods = _bars(service, monkeypatch)
+    bar_cls = mods["longcat_video.pipeline_longcat_video"].tqdm
+    service._progress.update(segments=1)
+    with bar_cls(total=16, desc="Denoising") as bar:
+        assert service._progress["steps"] == 16
+        assert service._progress["step"] == 0
+        for expected in range(1, 17):
+            bar.update()
+            assert service._progress["step"] == expected
+
+
+def test_a_bar_that_is_not_denoising_is_ignored(service, monkeypatch):
+    """Checkpoint-shard bars must not move the request's progress."""
+    mods = _bars(service, monkeypatch)
+    bar_cls = mods["longcat_video.pipeline_longcat_video"].tqdm
+    service._progress.update(step=4, steps=16)
+    with bar_cls(total=5, desc="Loading checkpoint shards") as bar:
+        bar.update()
+    assert (service._progress["step"], service._progress["steps"]) == (4, 16)
+
+
+def test_the_bar_total_sets_the_step_count_for_every_segment(service, monkeypatch):
+    """`steps` is the whole stage: one pass's steps times the segment count."""
+    mods = _bars(service, monkeypatch)
+    bar_cls = mods["longcat_video.pipeline_longcat_video"].tqdm
+    service._progress.update(segments=4)
+    with bar_cls(total=16, desc="Denoising"):
+        pass
+    assert service._progress["steps"] == 64
+
+
+def test_the_avatar_bar_is_instrumented_too(service, monkeypatch):
+    """The avatar pipeline is a different module with its own denoising loops, and
+    its distilled pass is 8 steps rather than 16."""
+    mods = _bars(service, monkeypatch)
+    bar_cls = mods["longcat_video.pipeline_longcat_video_avatar"].tqdm
+    service._progress.update(segments=1)
+    with bar_cls(total=8, desc="Denoising") as bar:
+        bar.update()
+        assert (service._progress["step"], service._progress["steps"]) == (1, 8)
+
+
+# ------------------------------------------------------------------ segment offsets
+
+def test_a_later_segment_continues_the_count(service, monkeypatch):
+    """Segment 2 of 4 starts at 16, not back at 0 -- the bar restarts, the stage
+    does not."""
+    mods = _bars(service, monkeypatch)
+    bar_cls = mods["longcat_video.pipeline_longcat_video"].tqdm
+    service._progress.update(segments=4)
+    service._bar["offset"] = 16
+    with bar_cls(total=16, desc="Denoising") as bar:
+        assert service._progress["step"] == 16
+        bar.update()
+        assert service._progress["step"] == 17
+
+
+def test_the_offset_advances_one_pass_per_segment(service, monkeypatch):
+    """_run_segments is what sets the offset, so drive it with a fake pass."""
+    mods = _bars(service, monkeypatch)
+    bar_cls = mods["longcat_video.pipeline_longcat_video"].tqdm
+    seen = []
+
+    def fake_pass(pipe, task, stage, ctx, cond=None, log=None):
+        with bar_cls(total=16, desc="Denoising") as bar:
+            for _ in range(16):
+                bar.update()
+        seen.append((service._progress["segment"], service._bar["offset"],
+                     service._progress["step"], service._progress["steps"]))
+        return [object()] * 93
+
+    monkeypatch.setattr(service, "_one_pass", fake_pass)
+    monkeypatch.setattr(service, "_frames_from_output", list)
+    monkeypatch.setattr(service, "torch_gc", lambda: None)
+    ctx = {"num_segments": 4, "num_cond_frames": 13, "cond_video": None,
+           "num_inference_steps": None, "guidance_scale": None}
+    service._progress.update(segments=4)
+    service._bar["offset"] = 0
+    service._bar["pass_steps"] = 0
+    acc, yielded = service._run_segments(None, "t2v", "distill", ctx, log=lambda *_: None)
+
+    assert yielded == 0
+    assert [s[0] for s in seen] == [1, 2, 3, 4]
+    assert [s[1] for s in seen] == [0, 16, 32, 48]
+    # The step count rises monotonically across the whole stage and ends at the total.
+    assert [s[2] for s in seen] == [16, 32, 48, 64]
+    assert {s[3] for s in seen} == {64}
+
+
+def test_the_first_segment_falls_back_to_the_stage_default(service, monkeypatch):
+    """Before any bar has existed there is no measured pass length, so segment 1
+    must still report a sane denominator rather than 0."""
+    _bars(service, monkeypatch)
+    calls = []
+
+    def fake_pass(pipe, task, stage, ctx, cond=None, log=None):
+        calls.append((service._progress["step"], service._progress["steps"]))
+        return [object()] * 93
+
+    monkeypatch.setattr(service, "_one_pass", fake_pass)
+    monkeypatch.setattr(service, "_frames_from_output", list)
+    monkeypatch.setattr(service, "torch_gc", lambda: None)
+    service._bar["offset"] = 0
+    service._bar["pass_steps"] = 0
+    ctx = {"num_segments": 2, "num_cond_frames": 13, "cond_video": None,
+           "num_inference_steps": None, "guidance_scale": None}
+    service._run_segments(None, "t2v", "distill", ctx, log=lambda *_: None)
+    # 16 is the distilled stage's default, times two segments.
+    assert calls[0] == (0, 32)
+
+
+def test_the_steps_honour_an_explicit_request_count(service, monkeypatch):
+    _bars(service, monkeypatch)
+
+    def fake_pass(pipe, task, stage, ctx, cond=None, log=None):
+        return [object()] * 93
+
+    monkeypatch.setattr(service, "_one_pass", fake_pass)
+    monkeypatch.setattr(service, "_frames_from_output", list)
+    monkeypatch.setattr(service, "torch_gc", lambda: None)
+    service._bar["offset"] = 0
+    service._bar["pass_steps"] = 0
+    ctx = {"num_segments": 3, "num_cond_frames": 13, "cond_video": None,
+           "num_inference_steps": 20, "guidance_scale": None}
+    service._run_segments(None, "t2v", "distill", ctx, log=lambda *_: None)
+    assert service._progress["steps"] == 60
+
+
+def test_progress_reports_no_private_keys(service):
+    """The offset bookkeeping is ours; /progress returns the documented contract."""
+    assert set(service._progress) == {"active", "stage", "segment", "segments",
+                                      "step", "steps"}
+
+
+# ----------------------------------------------------- the consumer in codai/api/video.py
+# Extracted from source the way tests/test_longcat_cancellable.py does: the poll loop
+# lives inside a long async request handler that cannot be called in a unit test.
+
+def _longcat_body() -> str:
+    import re
+    src = (ROOT / "codai" / "api" / "video.py").read_text(encoding="utf-8")
+    start = src.index("async def _generate_longcat")
+    nxt = re.search(r"\n(async def |def )", src[start + 10:])
+    return src[start:start + 10 + nxt.start()] if nxt else src[start:]
+
+
+def test_the_total_spans_every_segment():
+    """A four-segment render must not be measured against one segment's steps."""
+    body = _longcat_body()
+    assert "_grand_total(" in body
+    assert "num_segments" in body
+
+
+def test_refinement_is_not_multiplied_by_the_segment_count():
+    """It is a single pass over the finished video, not one per segment."""
+    body = _longcat_body()
+    assert '(1 if s == "refinement" else segments)' in body
+
+
+def test_the_segment_count_is_corrected_from_the_service():
+    """A request that asked for a duration rather than a segment count only learns
+    the real count from the service's first progress report."""
+    body = _longcat_body()
+    assert '_vid_progress_total(' in body
+    assert 'p.get("segments")' in body
+
+
+def test_a_finished_stage_contributes_what_it_actually_ran():
+    """lc.STAGE_DEFAULTS is one pass's steps; the stage ran it once per segment, and
+    the service reports that in `steps`."""
+    body = _longcat_body()
+    assert '_stage_total or lc.STAGE_DEFAULTS.get(' in body
+    assert 'p.get("steps")' in body
+
+
+def test_the_task_page_total_is_kept_in_step():
+    body = _longcat_body()
+    assert 'task_registry.step(_tid, _step, total=' in body
+
+
+def test_correcting_the_total_does_not_restart_the_bar():
+    """_vid_progress_reset would zero `current` and the elapsed time the rate is
+    derived from, which is why this is a separate setter."""
+    from codai.api import video
+
+    video._vid_progress_reset(16)
+    video._vid_progress_step(7)
+    started = video._vid_progress["started_at"]
+    rate = video._vid_progress["it_per_s"]
+
+    video._vid_progress_total(64)
+
+    assert video._vid_progress["total"] == 64
+    assert video._vid_progress["current"] == 7
+    assert video._vid_progress["active"] is True
+    assert video._vid_progress["started_at"] == started
+    assert video._vid_progress["it_per_s"] == rate

tools/longcat_service.py

@@ -70,6 +70,13 @@ _state = {
 }
 _progress = {"active": False, "stage": "", "segment": 0, "segments": 0,
              "step": 0, "steps": 0}
+# Denoise-step bookkeeping for /progress. The vendored pipelines render their step
+# count onto a tqdm bar and nothing else, so until this was wired the bar the caller
+# polls reported 0/N for an entire render. ``offset`` is how many steps earlier
+# SEGMENTS of the current stage already contributed, so ``_progress["step"]`` counts
+# the stage end to end instead of restarting at every segment boundary;
+# ``pass_steps`` is one segment's step count, learnt from the first bar's total.
+_bar = {"offset": 0, "pass_steps": 0}
 # Set by POST /yield: stop at the next SEGMENT boundary and return what has been
 # generated, so another model can have the card without the request being lost.
 _yield_flag = threading.Event()
@@ -79,6 +86,58 @@ def log(msg):
     print(f"[longcat] {msg}", flush=True)
 
 
+def _install_progress_bars() -> None:
+    """Teach the pipelines' denoising bars to publish their step count.
+
+    Upstream's generate_* methods drive a bare ``tqdm(total=len(timesteps),
+    desc="Denoising")`` and take no callback, so the only way to observe a step
+    without forking the vendored pipeline is to replace the ``tqdm`` the pipeline
+    module resolved at import time. Each module did ``from tqdm import tqdm``, so the
+    name has to be rebound on the MODULE -- patching ``tqdm.tqdm`` would be too late.
+
+    Only bars labelled "Denoising" are tracked: the pipelines use tqdm for nothing
+    else today, and a checkpoint-shard bar must not move the request's progress."""
+    import tqdm as _tqdm_mod
+
+    base = _tqdm_mod.tqdm
+
+    class _ProgressTqdm(base):
+        # Checked instead of class identity: this function builds a fresh subclass on
+        # every call, so identity would re-wrap an already-wrapped bar.
+        _lc_progress_bar = True
+
+        def __init__(self, *a, **kw):
+            super().__init__(*a, **kw)
+            self._lc_track = str(kw.get("desc") or "") == "Denoising"
+            if not self._lc_track:
+                return
+            with _lock:
+                total = int(self.total or 0)
+                if total:
+                    _bar["pass_steps"] = total
+                segments = max(1, int(_progress.get("segments") or 1))
+                _progress["steps"] = total * segments
+                _progress["step"] = int(_bar["offset"])
+
+        def update(self, n=1):
+            ret = super().update(n)
+            if getattr(self, "_lc_track", False):
+                with _lock:
+                    _progress["step"] = int(_bar["offset"]) + int(self.n or 0)
+            return ret
+
+    patched = []
+    for name in ("longcat_video.pipeline_longcat_video",
+                 "longcat_video.pipeline_longcat_video_avatar"):
+        mod = sys.modules.get(name)
+        if mod is not None and not getattr(
+                getattr(mod, "tqdm", None), "_lc_progress_bar", False):
+            mod.tqdm = _ProgressTqdm
+            patched.append(name.rsplit(".", 1)[-1])
+    if patched:
+        log(f"progress: denoising bars instrumented ({', '.join(patched)})")
+
+
 # ── loading ───────────────────────────────────────────────────────────────────
 
 def load_pipeline(checkpoint_dir: str, source_dir: str, dtype: str = "bfloat16",
@@ -226,6 +285,9 @@ def load_pipeline(checkpoint_dir: str, source_dir: str, dtype: str = "bfloat16",
                 raise RuntimeError(
                     "use_distill needs the DiT's enable_loras() — this build of "
                     "longcat_video does not expose it")
+        # Both pipeline modules are imported by now (the avatar one only for an
+        # avatar family), so this is the first point where the bars can be rebound.
+        _install_progress_bars()
         _place_pipeline(pipe, offload)
         _state.update(pipe=pipe, model=root, source=src, dtype=dtype)
         log("ready")
@@ -1097,10 +1159,15 @@ def _run_segments(pipe, task: str, stage: str, ctx: dict, log=print):
 
     for index in range(segments):
         with _lock:
-            _progress.update(segment=index + 1, segments=segments, step=0,
-                             steps=LC.stage_params(
-                                 stage, ctx.get("num_inference_steps"),
-                                 ctx.get("guidance_scale"))["num_inference_steps"])
+            # One segment's step count comes from the first bar's own total, which is
+            # authoritative (the avatar distilled pass is 8 steps, the base one 16).
+            # Before that bar exists, the stage default is the estimate.
+            per_pass = int(_bar["pass_steps"]) or int(LC.stage_params(
+                stage, ctx.get("num_inference_steps"),
+                ctx.get("guidance_scale"))["num_inference_steps"])
+            _bar["offset"] = index * per_pass
+            _progress.update(segment=index + 1, segments=segments,
+                             step=index * per_pass, steps=per_pass * segments)
         out = _one_pass(pipe, task, stage, ctx,
                         cond=cur if (index or first_is_continuation) else None, log=log)
         new = _frames_from_output(out)
@@ -1354,6 +1421,8 @@ def generate(body: dict) -> dict:
     frames = None
     yielded = 0
     with _lock:
+        _bar["offset"] = 0
+        _bar["pass_steps"] = 0
         _progress.update(active=True, stage="", segment=0, segments=segments,
                          step=0, steps=0)
     try:
@@ -1377,6 +1446,12 @@ def generate(body: dict) -> dict:
                 if image is not None:
                     kw["image"] = image
                     kw["num_cond_frames"] = 1
+                with _lock:
+                    # Refinement is one pass over the whole video, so the segment
+                    # offset from the stage before it must not carry over.
+                    _bar["offset"] = 0
+                    _progress.update(segment=1, segments=1, step=0,
+                                     steps=params["num_inference_steps"])
                 log(f"stage 'refinement': refining {len(frames)} frames")
                 frames = _frames_from_output(pipe.generate_refine(**kw)[0])
             else:
@@ -1386,6 +1461,7 @@ def generate(body: dict) -> dict:
     finally:
         _yield_flag.clear()
         with _lock:
+            _bar["offset"] = 0
             _progress.update(active=False, stage="", step=0)
 
     if stages[-1] == "refinement":