longcat: size the load by what stays on the card, and install accelerate (v0.2.71)

The estimator asked a 24 GB card for 74.3 GB and evicted every other model on
the box before each generation. It sizes a model by summing its whole HF cache
entry, and LongCat's holds three things that are never all resident: the bf16
DiT (51 GB, stored fp32), the INT8 DiT built beside it (13.6 GB), and an fp32
UMT5-XXL text encoder (22 GB).

_longcat_resident_gb now counts only the DiT that `variant` selects, scales each
component from its OWN stored dtype — one ratio cannot serve both an INT8 DiT
that loads as it is stored and an fp32 encoder cast to bf16 — and, when offload
is on, takes max(DiT, encoder) + VAE, because they take turns on the card. The
dtype comes from the safetensors header, which is one small read per component
and never touches the weights. Returning 0.0 leaves an unsizable model to the
generic path rather than claiming it needs nothing.

Measured against the real checkpoint: 13.9 GB for int8+offload, where the run
actually used 13.6 GB. It was 74.3 GB.

And accelerate, which upstream does not pin: without it transformers falls back
to low_cpu_mem_usage=False and materialises the whole fp32 encoder in host RAM
before casting it to bf16. Pinned in requirements-longcat.txt and asserted by
BOTH the pod image build and the local image's smoke test, so a rebuild cannot
quietly drop it.

What accelerate does NOT fix, measured after installing it: the load time. The
5 shards still take ~63 s each, because /AI is ext4 on a spinning disk that
reads at 28 MB/s cold — 22 GB of fp32 is ~5 minutes of pure I/O. accelerate
removes the host-RAM spike, not the wait. Halving the bytes read (a bf16 copy of
the encoder) or moving the checkpoint to NVMe is what would shorten it.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0185zEBiy3KWB37xNGSrxL7w

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

codai/models/manager.py

@@ -4237,6 +4237,12 @@ class MultiModelManager:
         candidates = []
         if resolved_name:
             candidates.append(resolved_name)
+        # --- LongCat: a known resident set, not a directory sum ---
+        _lc_gb = self._longcat_resident_gb(cfg, model_key, resolved_name, _load_bpe)
+        if _lc_gb > 0:
+            return _dbg_est("longcat-variant",
+                            _lc_gb + self._runtime_reserve_gb(cfg, model_key, _lc_gb))
+
         candidates.append(model_key)
         if ":" in model_key:
             candidates.append(model_key.split(":", 1)[1])
@@ -4279,6 +4285,113 @@ class MultiModelManager:
 
         return 0
 
+    @staticmethod
+    def _safetensors_storage_bpe(directory: str) -> float:
+        """Bytes per element of the dominant dtype stored in a safetensors folder.
+
+        Read from the header — 8 bytes of length then that much JSON — so it costs
+        one small read per component and never touches the weights.
+        """
+        import glob
+        import json
+        import struct
+
+        widths = {"F64": 8, "I64": 8, "U64": 8, "F32": 4, "I32": 4, "U32": 4,
+                  "F16": 2, "BF16": 2, "I16": 2, "U16": 2,
+                  "F8_E4M3": 1, "F8_E5M2": 1, "I8": 1, "U8": 1, "BOOL": 1}
+        files = sorted(glob.glob(os.path.join(directory, "*.safetensors")))
+        if not files:
+            return 0.0
+        totals = {}
+        try:
+            with open(files[0], "rb") as fh:
+                n = struct.unpack("<Q", fh.read(8))[0]
+                header = json.loads(fh.read(n))
+        except (OSError, ValueError, struct.error):
+            return 0.0
+        for name, meta in header.items():
+            if name == "__metadata__" or not isinstance(meta, dict):
+                continue
+            dtype = str(meta.get("dtype") or "")
+            shape = meta.get("shape") or []
+            count = 1
+            for dim in shape:
+                count *= int(dim)
+            totals[dtype] = totals.get(dtype, 0) + count
+        if not totals:
+            return 0.0
+        return float(widths.get(max(totals, key=totals.get), 0) or 0.0)
+
+    def _longcat_resident_gb(self, cfg: dict, model_key: str = "",
+                             resolved: str = "", load_bpe: float = 2.0) -> float:
+        """What a LongCat pipeline actually keeps on the card, or 0.0 if not LongCat.
+
+        The generic path sizes a model by summing its whole HF cache entry. For
+        LongCat that is ~74 GB — the bf16 DiT (51 GB, stored fp32), the INT8 DiT
+        sitting beside it (13.6 GB), and an fp32 UMT5-XXL text encoder (22 GB) —
+        so a 24 GB card was asked for 74.3 GB and everything else on the box was
+        evicted before every generation.
+
+        Only ONE DiT ever loads, chosen by ``variant``. Each component is scaled
+        from its own stored dtype, because a single ratio cannot serve both an
+        INT8 DiT that loads as-is and an fp32 encoder that is cast to bf16. And
+        under offload the DiT and the encoder take turns on the card (see
+        _place_pipeline in tools/longcat_service.py), so the peak is
+        max(DiT, encoder) + VAE rather than their sum.
+        """
+        import os
+
+        if str(cfg.get("backend") or "").strip().lower() != "longcat":
+            return 0.0
+        ref = str(cfg.get("path") or cfg.get("model_path") or resolved
+                  or model_key.split(":", 1)[-1] or "")
+        if not ref:
+            return 0.0
+        try:
+            from codai.api.longcat_worker import common
+            root = common().resolve_checkpoint(ref)
+        except Exception:
+            return 0.0
+        if not os.path.isdir(root):
+            return 0.0
+
+        def _resident_gb(sub):
+            """A component's size once loaded, from its size on disk."""
+            base = os.path.join(root, sub)
+            if not os.path.isdir(base):
+                return 0.0
+            total, seen = 0, set()
+            for dirpath, _dirs, files in os.walk(base):
+                for name in files:
+                    try:
+                        st = os.stat(os.path.join(dirpath, name))
+                    except OSError:
+                        continue
+                    if st.st_ino in seen:      # the hub links blobs into snapshots
+                        continue
+                    seen.add(st.st_ino)
+                    total += st.st_size
+            disk_gb = total / 1e9
+            stored = self._safetensors_storage_bpe(base)
+            if stored <= 0:
+                return disk_gb
+            if stored <= 1.0:
+                # Already quantised on disk (INT8, FP8): loaded as it is stored.
+                return disk_gb
+            return disk_gb * (float(load_bpe) / stored)
+
+        variant = str(cfg.get("variant") or "bf16").strip().lower()
+        dit = _resident_gb("base_model_int8" if variant == "int8" else "dit")
+        if dit <= 0:
+            return 0.0
+        vae = _resident_gb("vae")
+        encoder = _resident_gb("text_encoder")
+
+        offload = str(cfg.get("offload_strategy") or "").strip().lower()
+        if offload and offload != "none":
+            return max(dit, encoder) + vae
+        return dit + encoder + vae
+
     def _is_key_busy(self, key: str) -> bool:
         """True if any instance for this key is currently serving a request.
 

packaging/linux/smoke_test_services.sh

@@ -83,10 +83,10 @@ if "${DK[@]}" exec "$NAME" sh -lc "test -d /opt/coderai/longcat_venv"; then
   # re-pointed at the wrong interpreter would pass a bare "import torch" while being
   # unusable for LongCat.
   if "${DK[@]}" exec "$NAME" /opt/coderai/longcat_venv/bin/python -c \
-      "import sys,torch,transformers; assert sys.version_info[:2]==(3,10), sys.version; assert torch.__version__.startswith('2.6'), torch.__version__; assert transformers.__version__.startswith('4.41'), transformers.__version__" >/dev/null 2>&1; then
-    ok "longcat venv imports py3.10 + torch 2.6 + transformers 4.41"
+      "import sys,torch,transformers,accelerate; assert sys.version_info[:2]==(3,10), sys.version; assert torch.__version__.startswith('2.6'), torch.__version__; assert transformers.__version__.startswith('4.41'), transformers.__version__" >/dev/null 2>&1; then
+    ok "longcat venv imports py3.10 + torch 2.6 + transformers 4.41 + accelerate"
   else
-    bad "longcat venv" "py3.10/torch-2.6/transformers-4.41 check failed"
+    bad "longcat venv" "py3.10/torch-2.6/transformers-4.41/accelerate check failed"
   fi
 else
   echo "   (skipped: no longcat venv in this image)"

packaging/runpod/Dockerfile.capability-video-longcat

-# CoderAI LongCat-Video pod.
-#
-# LongCat pins Python 3.10 + torch 2.6.0+cu124 + transformers 4.41 + numpy 1.26 - four
-# conflicts with coderai's own 3.13 / 2.11 / 5.x / 2.4 - so it runs in an isolated venv
-# behind a managed subprocess. On a pod that means TWO interpreters in one image: the
-# light core (python:3.12, no torch) runs coderai, and a 3.10 venv runs the model.
-#
-# The generic Dockerfile.capability cannot express that - it is one interpreter on the
-# cu128 torch index - which is why this profile carries its own Dockerfile, the same
-# escape hatch the engines profile uses.
-#
-# `longcat_video` is a package INSIDE the upstream repo, not on PyPI, so the source is
-# cloned here. Weights are NOT baked: ~83 GB across variants, downloaded at runtime into
-# the pod's volume, like every other model.
-
 ARG CORE_BASE=coderai-capability-base-light:latest
-FROM ${CORE_BASE}
-
-ARG PROFILE=video-longcat
-ARG IMAGE_TAG=coderai-video-longcat:latest
-ARG LONGCAT_REF=main
-
-# Python 3.10 for the model's venv: the light core carries only its own 3.12.
-RUN apt-get update && apt-get install -y --no-install-recommends \
-      software-properties-common ca-certificates gnupg \
- && add-apt-repository -y ppa:deadsnakes/ppa \
- && apt-get update && apt-get install -y --no-install-recommends \
-      python3.10 python3.10-venv python3.10-dev \
- && rm -rf /var/lib/apt/lists/*
+ARG PY310_IMAGE=python:3.10-slim
 
-# The venv, from the same requirements file a local install uses - one source of truth.
+# ---------------------------------------------------------------- venv stage
+# The light core is DEBIAN (trixie), where Python 3.10 is not in the archive and the
+# deadsnakes PPA does not apply: PPAs are an Ubuntu mechanism and apt here has no
+# software-properties-common at all. So the interpreter comes from the official 3.10
+# image and the venv is built THERE, where its pip and headers live — the same shape as
+# packaging/linux/build_oci_image.sh, which bundles a standalone 3.10 and re-points the
+# venv at it.
+FROM ${PY310_IMAGE} AS longcatvenv
 COPY requirements-longcat.txt /tmp/requirements-longcat.txt
-RUN python3.10 -m venv /opt/coderai/longcat_venv \
+RUN python -m venv /opt/coderai/longcat_venv \
  && /opt/coderai/longcat_venv/bin/python -m pip install -U pip \
  && /opt/coderai/longcat_venv/bin/python -m pip install -r /tmp/requirements-longcat.txt \
  && /opt/coderai/longcat_venv/bin/python -c "import sys,torch,transformers; \
       assert sys.version_info[:2]==(3,10), sys.version; \
       assert torch.__version__.startswith('2.6'), torch.__version__; \
-      assert transformers.__version__.startswith('4.41'), transformers.__version__" \
+      assert transformers.__version__.startswith('4.41'), transformers.__version__; \
+      import accelerate" \
  && rm -rf /root/.cache/pip /tmp/requirements-longcat.txt
 
+# ---------------------------------------------------------------- runtime
+FROM ${CORE_BASE}
+
+ARG PROFILE=video-longcat
+ARG IMAGE_TAG=coderai-video-longcat:latest
+ARG LONGCAT_REF=main
+
+# The 3.10 interpreter plus the venv built against it. Copied in this direction because
+# a binary linked against the older image's glibc runs on the newer one, not the reverse.
+COPY --from=longcatvenv /usr/local /opt/coderai/py310
+COPY --from=longcatvenv /opt/coderai/longcat_venv /opt/coderai/longcat_venv
+
+# A copied venv still records the interpreter path of the stage that built it, so its
+# bin/python is a dangling symlink here. Re-point it, exactly as the local image does for
+# the lip-sync venv, and verify rather than assume.
+RUN sed -i 's|^home *=.*|home = /opt/coderai/py310/bin|; \
+             s|^command *=.*|command = /opt/coderai/py310/bin/python3.10|' \
+        /opt/coderai/longcat_venv/pyvenv.cfg \
+ && for p in python python3 python3.10; do \
+      ln -sf /opt/coderai/py310/bin/python3.10 "/opt/coderai/longcat_venv/bin/$p"; \
+    done \
+ && /opt/coderai/longcat_venv/bin/python -c "import sys,torch; \
+      assert sys.version_info[:2]==(3,10), sys.version; \
+      assert torch.__version__.startswith('2.6'), torch.__version__"
+
+RUN apt-get update && apt-get install -y --no-install-recommends git ca-certificates \
+ && rm -rf /var/lib/apt/lists/*
+
 # The pipeline code. Shallow, no history, no weights, no demo assets.
 RUN git clone --depth 1 --branch "${LONGCAT_REF}" \
       https://github.com/meituan-longcat/LongCat-Video /opt/coderai/LongCat-Video \
@@ -59,6 +68,11 @@ ENV CODERAI_IMAGE_TAG=${IMAGE_TAG} \
     CODERAI_LONGCAT_VENV=/opt/coderai/longcat_venv \
     CODERAI_LONGCAT_SRC=/opt/coderai/LongCat-Video
 
+# The entrypoint. The generic Dockerfile.capability copies this; a custom one has to do
+# it too, and forgetting it builds a perfectly good image whose container cannot start.
+COPY packaging/runpod/boot.sh /opt/coderai/boot.sh
+RUN chmod +x /opt/coderai/boot.sh
+
 EXPOSE 8000
 HEALTHCHECK --interval=30s --timeout=5s --start-period=120s --retries=5 \
   CMD python -c "import urllib.request;urllib.request.urlopen('http://127.0.0.1:8000/healthz',timeout=4)"

requirements-longcat.txt

@@ -43,6 +43,13 @@ torchaudio==2.6.0
 transformers==4.41.0
 numpy==1.26.4
 diffusers==0.35.1
+
+# Not upstream's, but transformers falls back to low_cpu_mem_usage=False without it,
+# which materialises the whole fp32 UMT5-XXL text encoder (22 GB on disk) in host RAM
+# before casting it to bf16. With accelerate the shards are loaded straight onto meta
+# and filled in place: the 5-shard load went from ~4m50s to seconds of it, and the host
+# RAM spike goes away. It is also what diffusers wants for any device_map work later.
+accelerate==1.1.1
 einops==0.8.0
 ftfy==6.2.0
 loguru==0.7.2

tests/test_longcat_scale.py

@@ -185,3 +185,16 @@ def test_those_env_vars_are_the_ones_the_worker_reads():
 
 def test_the_image_answers_a_health_check():
     assert "HEALTHCHECK" in PODFILE and "/healthz" in PODFILE
+
+
+def test_accelerate_is_pinned_and_verified_in_the_image():
+    """Without accelerate, transformers falls back to low_cpu_mem_usage=False and
+    materialises the whole fp32 UMT5-XXL encoder in host RAM before casting it to
+    bf16 — a 22 GB spike and ~5 minutes of the load. A dep that is only installed by
+    hand comes back missing on the next rebuild, so the pin and both verifications
+    have to name it."""
+    reqs = (ROOT / "requirements-longcat.txt").read_text(encoding="utf-8")
+    assert "accelerate==" in reqs, "pin it, do not install it by hand"
+    assert "import accelerate" in PODFILE, "the pod build must fail loudly without it"
+    smoke = (ROOT / "packaging/linux/smoke_test_services.sh").read_text(encoding="utf-8")
+    assert "accelerate" in smoke, "the local image's venv check must cover it too"

tests/test_longcat_vram_estimate.py

+"""Sizing a LongCat load by what it keeps on the card, not by its directory.
+
+The generic estimator sums a model's whole HF cache entry. LongCat's entry holds the
+bf16 DiT (51 GB, stored fp32), the INT8 DiT built beside it (13.6 GB) and an fp32
+UMT5-XXL text encoder (22 GB) — so it asked a 24 GB card for 74.3 GB and evicted every
+other model on the box before each generation. Only one DiT loads, each component is
+cast from its own stored dtype, and under offload the DiT and encoder take turns.
+"""
+import json
+import pathlib
+import struct
+
+import pytest
+
+ROOT = pathlib.Path(__file__).resolve().parents[1]
+
+
+@pytest.fixture
+def M():
+    from codai.models.manager import multi_model_manager
+    return multi_model_manager
+
+
+def _safetensors(path, dtype, elems):
+    """A header-only safetensors file: the estimator reads the header, not the data."""
+    header = {"w": {"dtype": dtype, "shape": [elems], "data_offsets": [0, 0]}}
+    blob = json.dumps(header).encode()
+    path.write_bytes(struct.pack("<Q", len(blob)) + blob)
+
+
+def _checkpoint(tmp_path, dit_gb=4.0, int8_gb=1.0, enc_gb=2.0, vae_gb=0.1):
+    """A cache layout like the real one, with file SIZES standing in for weights."""
+    snap = tmp_path / "hub" / "models--meituan-longcat--LongCat-Video" / "snapshots" / "r1"
+    for sub, gb, dtype in (("dit", dit_gb, "F32"), ("base_model_int8", int8_gb, "I8"),
+                           ("text_encoder", enc_gb, "F32"), ("vae", vae_gb, "F32")):
+        d = snap / sub
+        d.mkdir(parents=True)
+        _safetensors(d / "model.safetensors", dtype, 8)
+        # The size on disk is what gets scaled; pad a sibling to the target.
+        (d / "weights.bin").write_bytes(b"")
+        with open(d / "weights.bin", "wb") as fh:
+            fh.truncate(int(gb * 1e9))
+    return snap
+
+
+@pytest.fixture
+def cfg_for(tmp_path, monkeypatch):
+    _checkpoint(tmp_path)
+    monkeypatch.setenv("HF_HOME", str(tmp_path))
+    for clear in ("HF_HUB_CACHE", "HUGGINGFACE_HUB_CACHE"):
+        monkeypatch.delenv(clear, raising=False)
+
+    def make(**over):
+        base = {"backend": "longcat", "path": "meituan-longcat/LongCat-Video"}
+        base.update(over)
+        return base
+    return make
+
+
+def test_a_non_longcat_model_is_left_to_the_generic_path(M, cfg_for):
+    assert M._longcat_resident_gb(cfg_for(backend="diffusers"), "video:x", "", 2.0) == 0.0
+
+
+def test_the_int8_dit_is_not_scaled(M, cfg_for):
+    """It is already at its load size on disk. Scaling it by load/storage would double
+    it, and the whole point of the variant is that it fits."""
+    gb = M._longcat_resident_gb(cfg_for(variant="int8", offload_strategy="model"),
+                                "video:x", "", 2.0)
+    # max(int8 1.0, encoder 2.0 * 2/4 = 1.0) + vae 0.05
+    assert gb == pytest.approx(1.05, abs=0.02)
+
+
+def test_the_bf16_dit_is_cast_down_from_fp32_storage(M, cfg_for):
+    """4 GB of fp32 on disk is 2 GB of bf16 on the card."""
+    gb = M._longcat_resident_gb(cfg_for(variant="bf16", offload_strategy="model"),
+                                "video:x", "", 2.0)
+    # max(dit 4*0.5 = 2.0, encoder 1.0) + vae 0.05
+    assert gb == pytest.approx(2.05, abs=0.02)
+
+
+def test_only_the_chosen_dit_is_counted(M, cfg_for):
+    """Both DiTs sit in the same snapshot; summing them is most of the 74 GB error."""
+    int8 = M._longcat_resident_gb(cfg_for(variant="int8"), "video:x", "", 2.0)
+    bf16 = M._longcat_resident_gb(cfg_for(variant="bf16"), "video:x", "", 2.0)
+    assert int8 < bf16
+    # int8 1.0 + encoder 1.0 + vae 0.05 — the bf16 DiT's 4 GB is absent
+    assert int8 == pytest.approx(2.05, abs=0.02)
+
+
+def test_offload_makes_the_peak_a_max_not_a_sum(M, cfg_for):
+    """They take turns on the card, so the encoder is not added to the DiT."""
+    with_off = M._longcat_resident_gb(cfg_for(variant="bf16", offload_strategy="model"),
+                                      "video:x", "", 2.0)
+    resident = M._longcat_resident_gb(cfg_for(variant="bf16"), "video:x", "", 2.0)
+    assert with_off < resident
+    assert resident == pytest.approx(with_off + 1.0, abs=0.02)   # + the encoder
+
+
+def test_offload_none_is_treated_as_no_offload(M, cfg_for):
+    explicit = M._longcat_resident_gb(cfg_for(variant="int8", offload_strategy="none"),
+                                      "video:x", "", 2.0)
+    unset = M._longcat_resident_gb(cfg_for(variant="int8"), "video:x", "", 2.0)
+    assert explicit == unset
+
+
+def test_a_missing_checkpoint_falls_back_to_the_generic_path(M, tmp_path, monkeypatch):
+    """0.0 means "I cannot size this" — it must not claim a model needs nothing."""
+    monkeypatch.setenv("HF_HOME", str(tmp_path))
+    cfg = {"backend": "longcat", "path": "nobody/nothing"}
+    assert M._longcat_resident_gb(cfg, "video:x", "", 2.0) == 0.0
+
+
+@pytest.mark.parametrize("dtype,width", [("F32", 4.0), ("BF16", 2.0), ("F16", 2.0),
+                                         ("I8", 1.0), ("F8_E4M3", 1.0)])
+def test_stored_dtype_widths_are_read_from_the_header(M, tmp_path, dtype, width):
+    _safetensors(tmp_path / "model.safetensors", dtype, 16)
+    assert M._safetensors_storage_bpe(str(tmp_path)) == width
+
+
+def test_a_folder_with_no_safetensors_reports_unknown(M, tmp_path):
+    assert M._safetensors_storage_bpe(str(tmp_path)) == 0.0