scripts/run_cosmos_transfer.sh

#!/usr/bin/env bash
# Run Cosmos-Transfer2.5 inference on an exported cosmos_input bundle.
#
# Exists because the repo's `uv sync` environment does not, on this host, come
# up in a state where `import transformer_engine` succeeds. Two things have to
# be set first, and both are non-obvious:
#
#   1. CUDA_HOME must point at the venv's own `nvidia/` package tree.
#      transformer_engine's _load_nvrtc() globs "$CUDA_HOME/**/libnvrtc.so*"
#      and, when that misses, falls back to `ldconfig -p | grep libnvrtc`
#      inside subprocess.check_output(). There is no libnvrtc in this host's
#      ldconfig cache, grep therefore exits 1, and check_output raises --
#      before reaching _load_nvrtc's own final LD_LIBRARY_PATH fallback. So the
#      documented fallback chain is unreachable and CUDA_HOME is load-bearing.
#      It points into the venv rather than at the system CUDA on purpose: that
#      NVRTC is the one built against this exact torch/TE pair. (The `cuda128`
#      conda env has nvcc 12.8 but ships no libnvrtc.so at all.)
#
#   2. Every `nvidia/*/lib` directory must be on LD_LIBRARY_PATH.
#      libtransformer_engine.so links libcublas.so.12 and friends, and is
#      dlopen'd directly rather than through torch's loader, so it does not
#      inherit the preloading torch does for itself.
#
# Usage:
#   scripts/run_cosmos_transfer.sh <spec.json> <output_dir> [extra inference.py args...]
#
# Environment:
#   COSMOS_GPUS   comma-separated CUDA devices (default: 1). Check free memory
#                 first -- the model needs 65.4 GB and this is a shared host.
#   HF_TOKEN      required; the checkpoints are gated.

set -euo pipefail

REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
COSMOS_DIR="$REPO_ROOT/third_party/cosmos-transfer2.5"
VENV="$COSMOS_DIR/.venv"
NVIDIA_LIBS="$VENV/lib/python3.10/site-packages/nvidia"

if [[ $# -lt 2 ]]; then
    echo "usage: $0 <spec.json> <output_dir> [extra args...]" >&2
    exit 2
fi
# Both resolved to absolute paths *before* the cd below, or a relative --out-dir
# would silently land under the cosmos repo instead of where the caller meant.
SPEC="$(readlink -f "$1")"; shift
mkdir -p "$1"
OUT="$(readlink -f "$1")"; shift

[[ -x "$VENV/bin/python" ]] || { echo "no venv at $VENV -- run 'uv sync --extra=cu128 --python 3.10' in $COSMOS_DIR" >&2; exit 1; }
[[ -f "$SPEC" ]] || { echo "spec not found: $SPEC" >&2; exit 1; }
[[ -n "${HF_TOKEN:-}" ]] || { echo "HF_TOKEN is unset; the Cosmos checkpoints are gated" >&2; exit 1; }

# Exported explicitly, and HF_TOKEN_PATH pointed away from the default. The
# checkpoint fetcher shells out to `uvx hf download`, a fresh process that
# resolves credentials on its own; if a stale token file exists at
# ~/.cache/huggingface/token belonging to an account without approval for the
# gated Cosmos repos, that download fails with "Access denied. This repository
# requires approval." even though HF_TOKEN here is fine. Pinning both removes
# the ambiguity without touching the user's stored credentials.
export HF_TOKEN
export HF_TOKEN_PATH="${HF_TOKEN_PATH:-/nonexistent-so-HF_TOKEN-wins}"
# Unversioned aliases for the pip-installed CUDA libraries.
#
# The nvidia-* wheels ship only versioned sonames (libcudart.so.12), but
# transformer_engine's fused attention dlopen()s the *unversioned* name at
# runtime -- "Unable to dlopen libcudart.so". dlopen consults LD_LIBRARY_PATH
# but does no soname fuzzy-matching, so the alias has to exist on disk. Built
# once into a side directory; the wheels' own trees are never modified.
CUDA_ALIASES="$COSMOS_DIR/.cuda-so-aliases"
if [[ ! -d "$CUDA_ALIASES" ]]; then
    mkdir -p "$CUDA_ALIASES"
    find "$NVIDIA_LIBS" -name '*.so.*' -type f | while read -r lib; do
        base="$(basename "$lib")"
        stem="${base%%.so.*}.so"
        [[ -e "$CUDA_ALIASES/$stem" ]] || ln -s "$lib" "$CUDA_ALIASES/$stem"
    done
    echo "built $(ls "$CUDA_ALIASES" | wc -l) unversioned CUDA library aliases"
fi

export CUDA_HOME="$NVIDIA_LIBS"

# $REPO_ROOT/.nvshim FIRST, ahead of everything.
#
# This host's NVIDIA kernel module is 570.172.08 while /usr/lib's
# libnvidia-ml.so.1 symlink points at 580.173.02, so nvmlInit_v2() fails. torch
# tolerates that at import time (it only warns), but the CUDA caching allocator
# does not: the first real conv3d dies with
#   RuntimeError: NVML_SUCCESS == DriverAPI::get()->nvmlInit_v2_()
#   INTERNAL ASSERT FAILED at CUDACachingAllocator.cpp:983
# after the model has already loaded -- i.e. ~10 minutes into a run.
#
# .nvshim symlinks the whole NVIDIA userspace set to the matching 570.172.08
# libraries, which are present on this host, just not the ones ldconfig
# resolves. Fixing it needs no root, and it also makes nvidia-smi work again.
export LD_LIBRARY_PATH="$REPO_ROOT/.nvshim:$CUDA_ALIASES:$(ls -d "$NVIDIA_LIBS"/*/lib 2>/dev/null | tr '\n' ':')${LD_LIBRARY_PATH:-}"
export CUDA_VISIBLE_DEVICES="${COSMOS_GPUS:-1}"

# The repo sets this itself when it submits jobs
# (cosmos_transfer2/_src/transfer2/utils/submit_helper.py:66) but examples/inference.py
# run directly does not, so it has to be set here. It matters most on a busy shared
# GPU, where the free memory is fragmented across whatever else is resident and a
# large contiguous reservation is exactly what fails.
export PYTORCH_CUDA_ALLOC_CONF="${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True}"

mkdir -p "$OUT"
cd "$COSMOS_DIR"

echo "spec:    $SPEC"
echo "output:  $OUT"
echo "devices: $CUDA_VISIBLE_DEVICES"
echo "hf_token: ${HF_TOKEN:0:8}... (the gated downloads must use this one)"

# COSMOS_NPROC>1 runs under torchrun, which the repo documents for multi-GPU
# inference. Context parallelism splits the sequence across ranks, so activation
# memory divides roughly by the rank count -- the difference between fitting and
# not fitting on a busy host, since multicontrol loads all four control branches
# and does not fit in a single ~27 GB slice. Weights are replicated per rank, so
# this trades total memory for per-GPU memory, which is exactly the constraint here.
NPROC="${COSMOS_NPROC:-1}"
if [[ "$NPROC" -gt 1 ]]; then
    echo "torchrun with $NPROC processes (context parallel)"
    exec "$VENV/bin/torchrun" --nproc_per_node="$NPROC" --master_port="${COSMOS_PORT:-12341}" \
        examples/inference.py -i "$SPEC" -o "$OUT" "$@"
fi
exec "$VENV/bin/python" examples/inference.py -i "$SPEC" -o "$OUT" "$@"