diff --git a/src/server/scheduler_worker_manager.py b/src/server/scheduler_worker_manager.py index 29743bed..753bfbae 100644 --- a/src/server/scheduler_worker_manager.py +++ b/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" diff --git a/tests/test_scheduler_worker_manager.py b/tests/test_scheduler_worker_manager.py index 40becf0e..4ea245b2 100644 --- a/tests/test_scheduler_worker_manager.py +++ b/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):