#!/usr/bin/env bash
# Regenerate the hash-locked python requirements for the DS4 from-source build.
#
# The checked-in python locks detached from the GLM-5.2 NF3-Hybrid harness are
# NOT valid here: they pin torch 2.11 / flashinfer 0.6.12. This stack (the
# newest local-inference-lab b12x fork, feat/gg-ds4-0731-launcher) uses
# torch 2.12.0, flashinfer 0.6.15, cutlass-dsl 4.6.0 — so we lock from the
# FORK'S OWN requirements/* at the pinned commit, which are authoritative.
#
# Usage (run on the aarch64 build host, e.g. dgx3.lan, before build.sh):
#     bash gen-locks.sh
# Writes locks/python-build.txt and locks/python-runtime.txt.
#
# Requires: uv (bootstrapped in the Dockerfile's apt-base; or `apt install uv`).
set -euo pipefail
cd "$(dirname "$0")"
source ./versions.env

PY="${PYTHON_VERSION:-3.12}"
PLAT="${PLATFORM:-aarch64-unknown-linux-gnu}"
W="${VLLM_WORKSPACE:-$PWD/.gen/vllm}"   # shallow fork checkout (partial clone)

echo "==> ensuring fork checkout at ${VLLM_COMMIT}"
if [[ ! -d "$W/.git" ]]; then
    mkdir -p "$(dirname "$W")"
    git clone --filter=blob:none --no-checkout "${VLLM_REPO}" "$W"
fi
git -C "$W" fetch --depth 1 origin "${VLLM_COMMIT}"
git -C "$W" checkout --detach "${VLLM_COMMIT}"

common_flags=(--generate-hashes --python-version "$PY" --python-platform "$PLAT"
  --index-url "$PYPI_INDEX_URL" --extra-index-url "$PYTORCH_INDEX_URL"
  --index-strategy unsafe-best-match)

echo "==> python-build.txt (vLLM build deps, torch 2.12.0)"
mkdir -p .gen
# Also install the CUDA-side build companions the source builds need at build
# time (tvm_ffi is required by flashinfer-cubin's backend; torch/compiler deps
# come from the fork's own build/*.txt).
printf -- "-r '${W}/requirements/build/cuda.txt'\n-r '${W}/requirements/build/rust.txt'\napache-tvm-ffi==${TVM_FFI_VERSION}\nnvidia-nvshmem-cu13==${NVSHMEM_VERSION}\nnvidia-ml-py\nnvidia-cudnn-frontend\nnvidia-cudnn-cu13\npillow\nprotobuf\ntqdm\nregex\nrequests\npyproject-hooks\nwheel\nsemantic-version\nvcs-versioning\nsetuptools-scm\n" > .gen/build.in
uv pip compile "${common_flags[@]}" .gen/build.in \
  --output-file locks/python-build.txt

echo "==> python-runtime.txt (vLLM CUDA runtime deps)"
# FlashInfer is built from source and installed as a wheel (see Dockerfile), so
# drop the flashinfer lines (and the flashinfer.ai extra-index directive) from
# the lock — otherwise the runner's install can't resolve flashinfer-cubin and
# the pinned 0.6.14 would also collide with our built 0.6.15 wheel.
sed -e '/flashinfer/I d' -e '/extra-index-url.*flashinfer[^ ]*/I d' \
  "$W/requirements/cuda.txt" \
  | sed "s|^-r common.txt|-r ${W}/requirements/common.txt|" > .gen/runtime.in
# FlashInfer is stripped from the lock above (we build 0.6.15 from source and
# install the wheel with --no-deps in the Dockerfile), so uv never sees its
# install_requires. Re-add its hard transitive deps explicitly so `pip check`
# passes: the runner installs them from this lock. cuda-tile/nccl4py resolve
# to cp312 aarch64 wheels (deep-gemm / vllm-flash-attn / ray are not required
# by the fork's vLLM and have no aarch64 build, so they stay out).
cat >> .gen/runtime.in <<'EOF'

# --- FlashInfer 0.6.15 transitive deps (built from source, installed --no-deps) ---
cuda-tile==1.5.0
nccl4py==0.3.1
EOF
uv pip compile "${common_flags[@]}" .gen/runtime.in \
  --output-file locks/python-runtime.txt

echo "locked. Now: cd docker-images/deepseek-v4-flash-dspark-src && bash build.sh"