#!/usr/bin/env bash
# Launch ONE of the pusht training runs with George's lift_new harness
# (port of dish_utensil_train with only the dataset and configs
# swapped). The sbatch wrappers run stage_pusht_local.sh on the training node
# first.
#
#   Usage: run_one.sh <task_name> <main_process_port> [extra hydra overrides...]
#   e.g.:  NUM_GPUS=4 bash run_one.sh pusht_fastwam_scratch 29566
#
# task_name is one of the configs/task/*.yaml here; the output root and wandb
# run name are derived from it. wandb entity/project/mode are baked in (the
# checkout default points at a different entity, offline).
set -uo pipefail

TASK="${1:?task name required (e.g. pusht_fastwam_scratch)}"
PORT="${2:?main_process_port required (distinct per concurrent job)}"
shift 2

FASTWAM_ROOT=/shared_work/physical_intelligence/policies/Fast-WAM/fastwam
CONFIG_DIR=/shared_work/george/real_world_wam/pusht_train/configs
OUTPUT_ROOT=/shared_work/george/real_world_wam/runs/${TASK}
NUM_GPUS="${NUM_GPUS:-2}"
RUN_ID="$(date +%Y-%m-%d_%H-%M-%S)"
OUTPUT_DIR="${OUTPUT_ROOT}/${RUN_ID}"

ENV_PY=/shared_work/environments/miniconda3/envs/fastwam-robotwin-rui/bin
export PATH="${ENV_PY}:${PATH}"

# glibc malloc tuning — two constraints learned the hard way:
#  (a) mmap churn: workers allocating/freeing ~16MB numpy frames above the
#      default 128KB threshold produced tens of thousands of mmap calls/s
#      (kernel page-allocator churn) → keep mid-size buffers pooled in arenas.
#  (b) earlyoom on gpu-h200-103: the old blanket MALLOC_MMAP_THRESHOLD_=1GB
#      let every dataloader worker hoard ~10GB of freed decode buffers
#      (~1TB per 4-GPU job); two concurrent 4-GPU jobs overflowed the node's
#      2TB and earlyoom SIGTERM'd the fattest rank (killed jobs 1677 then
#      1676 on 2026-08-13, ~30 min into each overlap).
# 64MB threshold still pools the 16MB frame buffers (no churn) while rare
# big allocations go to mmap and are returned on free; ARENA_MAX caps
# per-process arena fragmentation; TRIM returns freed arena space to the OS.
export MALLOC_MMAP_THRESHOLD_=67108864
export MALLOC_ARENA_MAX=2
export MALLOC_TRIM_THRESHOLD_=134217728

# NB: do NOT set PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True — amin's
# runs proved it makes DeepSpeed ZeRO-1 bf16 go loss=nan from ~step 10.

mkdir -p "${OUTPUT_DIR}"

# cwd must be the fastwam checkout: its configs and the Wan2.2 base weights
# (checkpoints/...) are referenced by relative paths. We only read from it;
# all outputs go to OUTPUT_DIR.
cd "${FASTWAM_ROOT}"

accelerate launch \
  --config_file scripts/accelerate_configs/accelerate_zero1_ds.yaml \
  --num_processes "${NUM_GPUS}" \
  --main_process_port "${PORT}" \
  scripts/train.py \
  --config-dir "${CONFIG_DIR}" \
  "task=${TASK}" \
  "output_dir=${OUTPUT_DIR}" \
  "wandb.name=${TASK}_${RUN_ID}" \
  wandb.workspace=huaweiwam \
  wandb.project=fastwam-realrobot \
  "wandb.mode=${WANDB_MODE_OVERRIDE:-online}" \
  hydra/job_logging=stdout \
  "$@" 2>&1 | tee "${OUTPUT_DIR}/train_console.log"
# hydra/job_logging=stdout: Hydra's default file handler would try to write
# train.log into the read-only fastwam checkout (cwd) and crash at startup.
exit "${PIPESTATUS[0]}"