diff --git a/release/benchmarks/distributed/test_many_tasks.py b/release/benchmarks/distributed/test_many_tasks.py index 6f524cf4e3ea5..941795d19b27b 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 d79c65f644186..dbd6b977b261a 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 6ab92abf3fb35..6a03a6ce93c4e 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.