v2: tunnel-free proxy via SPACE_HOST+Jobs expose_public, heartbeat watchdog (30min stall -> fail), keep UploadCheckpointsCallback

run_job.sh

@@ -64,9 +64,32 @@ vllm serve "${MODEL:-LiquidAI/LFM2.5-2.6B}" \
 VLLM_PID=$!
 
 echo "=== [start] env server (CPU) ==="
+# Tunnel-free proxy exposure: with SPACE_HOST set, OpenEnv mounts the capture
+# proxy at /capture on the env server itself, reached through the Jobs proxy at
+# https://<JOB_ID>--8000.hf.jobs — no gradio tunnel to die mid-run
+# (that failure stalled the first full run silently for 13h).
+export SPACE_HOST="${JOB_ID}--8000.hf.jobs"
+echo "capture proxy public base: https://${SPACE_HOST}/capture"
 PREPARED=/app/prepared bash serve_env.sh &
 ENV_PID=$!
 
+echo "=== [start] watchdog (fail the job if no training step for 30 min) ==="
+(
+  while true; do
+    sleep 300
+    if [ -f /tmp/rl-heartbeat ]; then
+      age=$(( $(date +%s) - $(stat -c %Y /tmp/rl-heartbeat) ))
+      if [ "$age" -gt 1800 ]; then
+        echo "WATCHDOG: no training step for ${age}s — rollout path is dead, failing the job"
+        pkill -f train_lora.py || true
+        sleep 30
+        exit 1
+      fi
+    fi
+  done
+) &
+WATCHDOG_PID=$!
+
 echo "=== [wait] servers ==="
 python3 - <<'EOF'
 import time, urllib.request
@@ -135,5 +158,5 @@ if repo:
 EOF
 fi
 
-kill $VLLM_PID $ENV_PID 2>/dev/null || true
+kill $VLLM_PID $ENV_PID $WATCHDOG_PID 2>/dev/null || true
 echo "=== done ==="

train_lora.py

@@ -84,6 +84,19 @@ class UploadCheckpointsCallback(TrainerCallback):
         return control
 
 
+class HeartbeatCallback(TrainerCallback):
+    """Touch a file at every optimizer step; run_job.sh's watchdog kills the
+    job if it goes stale. The first full run stalled silently for 13h when the
+    rollout path died — steps just stopped, and nothing noticed."""
+
+    PATH = "/tmp/rl-heartbeat"
+
+    def on_step_end(self, args, state, control, **kwargs):
+        with open(self.PATH, "w") as f:
+            f.write(str(state.global_step))
+        return control
+
+
 def correctness_plus_efficiency(outcome):
     """1/0 correctness, plus <= 0.1 for a correct answer in fewer tool calls.
 
@@ -190,7 +203,7 @@ def main():
         processing_class=tokenizer,
         rollout_worker=worker,
         peft_config=peft_config,
-        callbacks=[UploadCheckpointsCallback(ADAPTER_REPO)],
+        callbacks=[UploadCheckpointsCallback(ADAPTER_REPO), HeartbeatCallback()],
     )
     trainer.train()
     trainer.save_model("./checkpoints/final")   # adapter_model.safetensors + adapter_config.json