perf(sampler): keep vLLM's compile cache on the shared volume

src/server/scheduler_worker_manager.py

@@ -123,6 +123,12 @@ def pod_env(worker: Worker) -> list[dict[str, Any]]:
   for name in ("VLLM_GPU_MEMORY_UTILIZATION", "VLLM_MAX_MODEL_LEN", "OPEN_RL_TRAIN_TOKEN_BUDGET", "MAX_JOBS"):
     if os.getenv(name):
       values[name] = os.environ[name]
+  # vLLM keys its torch.compile artifacts by model, GPU type and engine config.
+  # Inside the container they die with the pod and every sampler compiles
+  # again; on the shared volume the first sampler of a key compiles and the
+  # rest load the result.
+  if worker.role == "sampler":
+    values["VLLM_CACHE_ROOT"] = f"{tmp_dir}/vllm-cache"
   # No other worker shares an exclusive worker's GPUs, so it never parks.
   if worker.exclusive:
     values["OPEN_RL_TIME_SLICING"] = "off"

tests/test_scheduler_worker_manager.py

@@ -142,6 +142,18 @@ def test_workers_get_the_deployment_settings(self) -> None:
       env = {e["name"]: e.get("value") for e in pod["spec"]["template"]["spec"]["containers"][0]["env"]}
       self.assertEqual({k: env.get(k) for k in settings}, settings)
 
+  def test_samplers_keep_the_vllm_compile_cache_on_the_shared_volume(self) -> None:
+    s = self.store_with("job-lora-1", {"base_model": "Qwen/Qwen2.5-0.5B", "fine_tuning_type": "lora"})
+    with patch.dict(os.environ, {"OPEN_RL_TMP_DIR": "/mnt/shared/open-rl"}), patch("server.worker_manager.get_state_store", return_value=s):
+      self.manager.ensure("job-lora-1", "sampler")
+      self.manager.ensure("job-lora-1", "trainer")
+
+    envs = {}
+    for pod in self.api.created:
+      envs[pod["spec"]["role"]] = {e["name"]: e.get("value") for e in pod["spec"]["template"]["spec"]["containers"][0]["env"]}
+    self.assertEqual(envs["sampler"]["VLLM_CACHE_ROOT"], "/mnt/shared/open-rl/vllm-cache")
+    self.assertNotIn("VLLM_CACHE_ROOT", envs["trainer"])
+
   def test_launch_is_idempotent(self) -> None:
     s = self.store_with("job-lora-1", {"base_model": "Qwen/Qwen2.5-0.5B", "fine_tuning_type": "lora"})
     with patch("server.worker_manager.get_state_store", return_value=s):