# SPDX-License-Identifier: Apache-2.0
# tt-model-manager container manifest (schema 5.1) for pi-0.5 (lerobot/pi05_base) on 2x Blackhole p300 (4 chips).
#
# The model is the multi-config pi0.5 megakernel of changh95/pi05-base-p150 (code/models/experimental/pi0: VISION,
# PREFIX and EXPERT as persistent generic_op programs on 110 cores of a chip, presets built from cameras 1-4 x prompt
# buckets 32/64/128/224 (+ extra buckets c1 L192, c2 L160, c4 L96) x action-row buckets 32/64) plus, in this package:
#   - the single-robot mesh (tt/ttnn_pi05_sp.py): ONE model on all 4 chips of a 1x4 FABRIC_1D_RING mesh --
#     camera-parallel VISION, sequence-parallel PREFIX with an in-kernel K / V exchange over the fabric, EXPERT on one
#     chip -- with a static thermal-aware chip placement (tt/sp_thermal.py); actions bit-identical to one chip;
#   - multi-robot serving with an SLO (tt/mr_sched.py, tt/ttnn_pi05_multichip.py dp): every chip runs the whole model,
#     robots are bound to chips, the robot limit R_max comes from a start-up calibration against the p50 / p99 SLO and
#     a new robot beyond it gets HTTP 503;
#   - the batched EXPERT program (tt/megakernel/program_b.py, kernels/mkb_*), bit-identical per request;
#   - prefix / expert disaggregation (chips 0-2 VISION + PREFIX, chip 3 the batched expert; K / V through the host);
#   - faster / lower-energy programs (fp32 residual in L1, K / V-only last QKV, mask skip, VLM down K groups), all
#     bit-identical to the previous release;
#   - stability on QuietBox 2 (this revision): the PREFIX K / V feed runs its 4 full-grid multicasts one at a time with
#     an acknowledged handoff (PE_KV_SERIAL + PE_KV_HANDOFF, bit-identical; ~11x fewer early-run chip hangs), and the
#     server has a hang watchdog (PI05_HANG_S / PI05_HANG_ACTION / PI05_RESET_ON_START; see SERVING.md). The hang
#     marker lives on the /cache volume so that the next container start resets the chips after a hang.
#
# Three serve profiles of one image (camera count / H / N fixed at server start; every request the same config):
#   multi-robot (default)  4 chips, each the whole model; robots bound to chips, R_max from PI05_SLO_MS / _P99_MS
#   single-robot           all 4 chips as ONE model (SP mesh): the lowest latency for one robot
#   disaggregated          chips 0-2 VISION + PREFIX, chip 3 the batched expert (an experiment; below multi-robot)
#
# Run every tt-model command FROM THIS DIRECTORY (`root: code` resolves against the CWD):
#   tt-model package --container tt-model.yaml --out <build dir>
#   tt-model serve   <build dir>/pi05-base-p300x2/tt_kernel_manifest.json [--profile single-robot]
#   tt-model push    <build dir>/pi05-base-p300x2 --publish
schema: "5.1"

repo: changh95/pi05-base-p300x2
name: pi05-base-p300x2

weights:
  repo: lerobot/pi05_base
  revision: b211f3d44c36b6acfcf7ae94a64e8e96f75a64ba
  allow_patterns:
    - model.safetensors
    - config.json
    - policy_preprocessor.json
    - policy_postprocessor.json

kind: tt-dit-server
arch: blackhole

source:
  # tenstorrent/tt-metal main @ f856a38a361 (v0.80.0-dev20261001-17): the tree the p150 megakernel was validated on and
  # this package was measured on (2x p300c). The port needs no tt-metal change: its kernels compile at run time.
  tt_metal:
    repo: https://github.com/tenstorrent/tt-metal
    ref: f856a38a361939888f92d88f9e69b2f8a83fb713
  code:
    - models/common/lightweightmodule.py
  extra_code:
    - root: code
      paths:
        - models
  ubuntu: "22.04"
  python: "3.12"

runtime:
  app: models.experimental.pi0_5.server.app:app
  mesh_shape_env: TT_MESH_SHAPE
  packages:
    - "numpy>=1.24.4,<2"
    - safetensors
    - huggingface_hub
    - transformers==5.12.1
    - sentencepiece
    - protobuf

serve:
  port: 20000
  hardware: p300x2
  mesh_device: P300x2
  env:
    TT_WEIGHTS_REVISION: "b211f3d44c36b6acfcf7ae94a64e8e96f75a64ba"
    PI05_MEGAKERNEL: "mc"
    # Fixed at server start: cameras 1..4, action horizon 1..64, flow-matching steps 1..10.
    PI05_NUM_IMAGES: "2"
    PI05_ACTION_HORIZON: "50"
    PI05_NUM_STEPS: "10"
    PI05_TOKEN_LEN: "224"
    PI05_TOKENIZER_REQUIRED: "1"
    PI05_BATCH_WINDOW_MS: "0"
    PI05_MODEL_NAME: "pi05-base-p300x2"
    # Hang watchdog (server defaults, spelled out): a chip stuck 10 s is quarantined (multi-robot keeps serving on the
    # others) or the server exits 75 (single-robot / disaggregated); a start after a hang resets the chips (tt-smi -r).
    # The marker must outlive the container, so it goes on the persistent /cache volume.
    PI05_HANG_S: "10"
    PI05_HANG_ACTION: "degrade"
    PI05_RESET_ON_START: "auto"
    PI05_HANG_MARKER: "/cache/pi05_hang_marker"

default_profile: multi-robot
serve_profiles:
  - name: multi-robot
    description: "All 4 chips, each the whole model; each robot (request field robot_id) is bound to one chip. At start the server times every chip's cycle and derives the robots per chip and R_max from the SLO (p50 100 ms, p99 110 ms, clock guard 1.69 for chips clamped to 800 MHz when hot); a new robot beyond R_max gets HTTP 503. Hang watchdog: a stuck chip is quarantined, its robots move to the healthy chips (PI05_HANG_ACTION=exit: exit 75 for a supervisor restart). Measured in-process, 4 robots at 2 cameras / H 50 / N 10: p50 53 ms cool, p50 83.0 / p99 84.1 ms at thermal steady state, 52 req/s, 6.65 J per request."
    env:
      PI05_PROFILE: multi-robot
      TT_MESH_SHAPE: "1x4"
      PI05_LAYOUT: dp
      PI05_SLO_MS: "100"
      PI05_SLO_P99_MS: "110"
      PI05_SLO_CLOCK_GUARD: "1.69"
  - name: single-robot
    description: "All 4 chips as ONE model (1x4 FABRIC_1D_RING mesh): camera-parallel VISION, sequence-parallel PREFIX, EXPERT on a cool chip (thermal placement); actions bit-identical to one chip. Host-to-host p50 at 2 cameras / 224 tokens / H 50: 38.5 ms cool (one chip: 52.6 ms). Per call at thermal steady state (trace replay, one chip at 800 MHz): 42.1 ms vs 38.0 ms cool; the legacy placement: 61.4 ms. On a hang the server exits 75 for a supervisor restart (the next start resets the chips)."
    env:
      PI05_PROFILE: single-robot
      TT_MESH_SHAPE: "1x4"
      PI05_LAYOUT: sp
  - name: disaggregated
    description: "Prefix / expert disaggregation: chips 0-2 run VISION + PREFIX, chip 3 the batched action expert on up to 4 requests (K / V caches copied through the host). Requests are accepted at any time; same config for every request. No SLO or admission control; per-request heterogeneous configs and LLM-style join/leave serving are future work. Measured below multi-robot on this model (see the card). On a hang the server exits 75 for a supervisor restart."
    env:
      PI05_PROFILE: disaggregated
      TT_MESH_SHAPE: "1x4"
      PI05_LAYOUT: pd
      PI05_MAX_BATCH: "4"

verify:
  - "import models.experimental.pi0.tt.ttnn_pi05_model as m; assert m.PI05MegakernelTTNN and m.PI05_DEVICE_PARAMS['worker_l1_size'] == 1395712, m.PI05_DEVICE_PARAMS"
  - "from models.experimental.pi0.tt.megakernel import presets as PS; assert len(PS.PRESETS) == 38 and PS.CAMERAS == (1, 2, 3, 4) and PS.PROMPT_BUCKETS == (32, 64, 128, 224) and PS.SUFFIX_BUCKETS == (32, 64) and PS.EXTRA_BUCKETS == {1: (192,), 2: (160,), 4: (96,)}, (len(PS.PRESETS), PS.EXTRA_BUCKETS)"
  - "from models.experimental.pi0.tt.megakernel import program_b as pb; assert pb.NB_MAX == 8 and pb.BatchedExpertMegakernel and pb.KVSlab"
  - "import models.experimental.pi0.tt.ttnn_pi05_batched as b, models.experimental.pi0.tt.ttnn_pi05_multichip as mc; assert b.PI05BatchedTTNN and mc.PI05MultiChip"
  - "import models.experimental.pi0.tt.ttnn_pi05_sp as sp, models.experimental.pi0.tt.sp_thermal as T, models.experimental.pi0.tt.mr_sched as R; assert sp.PI05SingleRobotMesh and T.ENV_PLACEMENT == 'PI05_SP_PLACEMENT' and T.ENV_HOT == 'PI05_SP_HOT' and T.ENV_MR_GUARD == 'PI05_MR_GUARD' and R.AdmissionRefused"
  - "from models.experimental.pi0_5.server.mcx_backend import max_batch_for as f; assert (f(2, 50), f(4, 50), f(2, 10)) == (6, 4, 8), (f(2, 50), f(4, 50), f(2, 10))"
  - "import os; [os.environ.pop(k, None) for k in ('PI05_NUM_IMAGES', 'PI05_ACTION_HORIZON', 'PI05_NUM_STEPS', 'PI05_SLO_MS', 'PI05_SLO_P99_MS', 'PI05_PROFILE')]; os.environ.update(PI05_MEGAKERNEL='mc', TT_MESH_SHAPE='1x4', PI05_LAYOUT='dp', PI05_MAX_BATCH='4'); from models.experimental.pi0_5.server import app as a; c = a.load_config(); assert (c.backend, c.layout, c.chips, c.max_batch) == ('mcx', 'dp', 4, 4), c"
  - "import os; [os.environ.pop(k, None) for k in ('PI05_NUM_IMAGES', 'PI05_ACTION_HORIZON', 'PI05_NUM_STEPS', 'PI05_MAX_BATCH', 'PI05_PROFILE')]; os.environ.update(PI05_MEGAKERNEL='mc', TT_MESH_SHAPE='1x4', PI05_LAYOUT='dp', PI05_SLO_MS='100', PI05_SLO_P99_MS='110', PI05_SLO_CLOCK_GUARD='1.69'); from models.experimental.pi0_5.server import app as a; c = a.load_config(); assert (c.backend, c.layout, c.routing, c.slo_ms, c.slo_p99_ms, c.slo_clock_guard, c.max_batch) == ('mcx', 'dp', 'robot', 100.0, 110.0, 1.69, 6), c"
  - "import os; [os.environ.pop(k, None) for k in ('PI05_NUM_IMAGES', 'PI05_ACTION_HORIZON', 'PI05_NUM_STEPS', 'PI05_SLO_MS', 'PI05_SLO_P99_MS', 'PI05_MAX_BATCH', 'PI05_PROFILE')]; os.environ.update(PI05_MEGAKERNEL='mc', TT_MESH_SHAPE='1x4', PI05_LAYOUT='sp'); from models.experimental.pi0_5.server import app as a; c = a.load_config(); assert (c.backend, c.layout, c.chips, c.profile) == ('mc', 'sp', 4, 'single-robot'), c"
  - "import os; [os.environ.pop(k, None) for k in ('PI05_LAYOUT', 'PI05_SLO_MS', 'PI05_SLO_P99_MS', 'PI05_MAX_BATCH')]; os.environ.update(PI05_MEGAKERNEL='mc', TT_MESH_SHAPE='1x1'); from models.experimental.pi0_5.server import app as a; c = a.load_config(); assert c.backend == 'mc' and c.layout == 'mesh', c"
  - "import re; from pathlib import Path; T = Path('/opt/tt-metal'); R = T / 'models/experimental/pi0'; inc = T / 'tt_metal/hw/inc'; ks = sorted(k for k in R.rglob('*') if k.suffix in ('.cpp', '.hpp', '.h')); assert len(ks) == 21, ks; pairs = [(k, i) for k in ks for i in re.findall(r'#include\\s+\"([^\"]+)\"', k.read_text())]; bad = [(str(k), i) for k, i in pairs if not ((k.parent / i).is_file() or (inc / i).is_file() or (T / i).is_file())]; assert not bad, bad"
  - "import os; [os.environ.pop(k, None) for k in ('PI05_HANG_S', 'PI05_HANG_ACTION', 'PI05_RESET_ON_START', 'PI05_HANG_MARKER', 'PI05_LAYOUT', 'PI05_SLO_MS', 'PI05_SLO_P99_MS', 'PI05_MAX_BATCH', 'PI05_PROFILE')]; os.environ.update(PI05_MEGAKERNEL='mc', TT_MESH_SHAPE='1x4', PI05_LAYOUT='dp'); from models.experimental.pi0_5.server import app as a; from models.experimental.pi0.tt import mr_sched as R; c = a.load_config(); assert (c.hang_s, c.hang_action, c.hang_exit_code, c.reset_on_start) == (10.0, 'degrade', 75, 'auto') and R.ChipHung, c"
  - "import ttnn; assert hasattr(ttnn, 'generic_op') and hasattr(ttnn, 'CreateDevices') and hasattr(ttnn, 'set_fabric_config') and hasattr(ttnn.FabricConfig, 'FABRIC_1D_RING') and hasattr(ttnn, 'open_mesh_device')"
  - "import safetensors, huggingface_hub, sentencepiece, google.protobuf"
  - "import transformers; assert transformers.__version__ == '5.12.1', transformers.__version__"
  - "import numpy; assert numpy.__version__.startswith('1.'), numpy.__version__"
  - "from pathlib import Path; assert Path('/opt/tt-metal/models/experimental/pi0_5/server/smoke_test.py').is_file()"