From 4166389afe29183aa5b2b5ed75246e83d4cced99 Mon Sep 17 00:00:00 2001
From: dentiny <dentinyhao@gmail.com>
Date: Tue, 21 Jan 2025 23:19:50 +0000
Subject: [PATCH] runtime env

Signed-off-by: dentiny <dentinyhao@gmail.com>
---
 release/benchmarks/distributed/test_many_tasks.py | 2 +-
 src/ray/raylet/worker_pool.cc                     | 3 +--
 src/ray/raylet/worker_pool.h                      | 3 ++-
 3 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/release/benchmarks/distributed/test_many_tasks.py b/release/benchmarks/distributed/test_many_tasks.py
index 6f524cf4e3ea..941795d19b27 100644
--- a/release/benchmarks/distributed/test_many_tasks.py
+++ b/release/benchmarks/distributed/test_many_tasks.py
@@ -18,7 +18,7 @@
 def test_max_running_tasks(num_tasks):
     cpus_per_task = 0.25
 
-    @ray.remote(num_cpus=cpus_per_task)
+    @ray.remote(num_cpus=cpus_per_task, runtime_env={"env_vars": {"FOO": "bar"}})
     def task():
         time.sleep(sleep_time)
 
diff --git a/src/ray/raylet/worker_pool.cc b/src/ray/raylet/worker_pool.cc
index d79c65f64418..dbd6b977b261 100644
--- a/src/ray/raylet/worker_pool.cc
+++ b/src/ray/raylet/worker_pool.cc
@@ -1470,8 +1470,7 @@ void WorkerPool::PrestartWorkers(const TaskSpecification &task_spec,
                    << backlog_size << " and available CPUs " << num_available_cpus
                    << " num idle workers " << state.idle.size()
                    << " num registered workers " << state.registered_workers.size();
-    PrestartWorkers(language)
-    PrestartDefaultCpuWorkers(task_spec, num_needed);
+    PrestartWorkersWithRuntimeEnv(task_spec, num_needed);
   }
 }
 
diff --git a/src/ray/raylet/worker_pool.h b/src/ray/raylet/worker_pool.h
index 6ab92abf3fb3..6a03a6ce93c4 100644
--- a/src/ray/raylet/worker_pool.h
+++ b/src/ray/raylet/worker_pool.h
@@ -431,7 +431,8 @@ class WorkerPool : public WorkerPoolInterface, public IOWorkerPoolInterface {
   void PrestartDefaultCpuWorkers(ray::Language language, int64_t num_needed);
 
   /// Prestart worker with the given runtime env.
-  void PrestartWorkers(const TaskSpecification &task_spec, int64_t num_needed);
+  void PrestartWorkersWithRuntimeEnv(const TaskSpecification &task_spec,
+                                     int64_t num_needed);
 
   /// Return the current size of the worker pool for the requested language. Counts only
   /// idle workers.