#!/bin/bash
# Run inside HF Job container. Repo at /w, torch already installed in pytorch image.
set -e
cd /w
export PYTHONPATH="/w:${PYTHONPATH}"

# Install huggingface_hub first so we can download data files
pip install -q 'huggingface_hub>=0.28'

# Fetch large data files from HF (uploaded separately via hf upload)
mkdir -p data
python -c "from huggingface_hub import hf_hub_download; [hf_hub_download(repo_id='Pratham-math/fathom-code', filename=f'data/{f}', local_dir='/w', token='${HF_TOKEN}') for f in ['train.jsonl','eval.jsonl','sft_traces.jsonl']]"
ls -la data/

# Base runtime deps (these don't conflict)
pip install -q openenv-core fastapi 'uvicorn[standard]' pydantic RestrictedPython tiktoken httpx
pip install -q hydra-core omegaconf wandb tyro

# Core ML deps for smoke path.
# NOTE: trl 1.2 imports `is_trackio_available` from transformers, which is not
# present in 4.49.0. Use a newer transformers line in jobs.
pip install -q transformers==4.56.2 accelerate==1.5.2 peft==0.14.0 bitsandbytes==0.45.1
pip install -q datasets==4.7.0
# Keep TRL pinned for OpenEnv path but bypass resolver deadlock with datasets pin.
pip install -q --no-deps trl==1.2.0
pip install -q vllm==0.18.0

# Optional transitive deps often required by quantized loaders / datasets stack
pip install -q safetensors sentencepiece einops scipy xxhash protobuf pyyaml fsspec aiohttp dill multiprocess pyarrow requests filelock packaging tokenizers regex tqdm

# Start env server
nohup uvicorn env.server.app:app --host 0.0.0.0 --port 8001 > env.log 2>&1 &
sleep 10
python -c "import urllib.request, sys; sys.exit(0 if urllib.request.urlopen('http://localhost:8001/healthz', timeout=5).status==200 else 1)" || (echo "env server down"; cat env.log; exit 1)

# Smoke
python -m train.smoke_test --env-url http://localhost:8001
echo "===== SMOKE_RESULT.md ====="
cat outputs/smoke/SMOKE_RESULT.md