prestart worker with runtime env

Signed-off-by: dentiny <[email protected]>
ray-project · Jan 21, 2025 · 92e6695 · 92e6695
1 parent 7a372c7
commit 92e6695
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 1 deletion.
diff --git a/src/ray/raylet/worker_pool.cc b/src/ray/raylet/worker_pool.cc
@@ -1470,7 +1470,23 @@ 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();
-    PrestartDefaultCpuWorkers(task_spec.GetLanguage(), num_needed);
+    PrestartWorkers(language)
+    PrestartDefaultCpuWorkers(task_spec, num_needed);
+  }
+}
+
+void WorkerPool::PrestartWorkers(const TaskSpecification &task_spec, int64_t num_needed) {
+  RAY_LOG(DEBUG) << "PrestartWorkers " << num_needed;
+  for (int ii = 0; ii < num_needed; ++ii) {
+    PopWorkerStatus status;
+    StartWorkerProcess(task_spec.GetLanguage(),
+                       rpc::WorkerType::WORKER,
+                       JobID::Nil(),
+                       &status,
+                       /*dynamic_options=*/{},
+                       task_spec.GetRuntimeEnvHash(),
+                       task_spec.SerializedRuntimeEnv(),
+                       task_spec.RuntimeEnvInfo());
   }
 }
 

diff --git a/src/ray/raylet/worker_pool.h b/src/ray/raylet/worker_pool.h
@@ -430,6 +430,9 @@ 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);
+
   /// Return the current size of the worker pool for the requested language. Counts only
   /// idle workers.
   ///