Decouple ray submitter, worker, and head resources

Signed-off-by: Jason Parraga <[email protected]>
flyteorg · Oct 29, 2024 · 75c0ceb · 75c0ceb
1 parent 13b3d82
commit 75c0ceb
Show file tree

Hide file tree

Showing 8 changed files with 314 additions and 120 deletions.
diff --git a/flyteidl/gen/pb-es/flyteidl/plugins/ray_pb.ts b/flyteidl/gen/pb-es/flyteidl/plugins/ray_pb.ts
diff --git a/flyteidl/gen/pb-go/flyteidl/plugins/ray.pb.go b/flyteidl/gen/pb-go/flyteidl/plugins/ray.pb.go
diff --git a/flyteidl/gen/pb_python/flyteidl/plugins/ray_pb2.py b/flyteidl/gen/pb_python/flyteidl/plugins/ray_pb2.py
diff --git a/flyteidl/gen/pb_python/flyteidl/plugins/ray_pb2.pyi b/flyteidl/gen/pb_python/flyteidl/plugins/ray_pb2.pyi
diff --git a/flyteidl/gen/pb_rust/flyteidl.plugins.rs b/flyteidl/gen/pb_rust/flyteidl.plugins.rs
diff --git a/flyteidl/protos/flyteidl/plugins/ray.proto b/flyteidl/protos/flyteidl/plugins/ray.proto
@@ -4,6 +4,8 @@ package flyteidl.plugins;
 
 option go_package = "github.com/flyteorg/flyte/flyteidl/gen/pb-go/flyteidl/plugins";
 
+import "flyteidl/core/tasks.proto";
+
 // RayJobSpec defines the desired state of RayJob
 message RayJob {
   // RayClusterSpec is the cluster template to run the job
@@ -35,6 +37,8 @@ message HeadGroupSpec {
   // Optional. RayStartParams are the params of the start command: address, object-store-memory.
   // Refer to https://docs.ray.io/en/latest/ray-core/package-ref.html#ray-start
   map<string, string> ray_start_params = 1;
+  // Resource specification for ray head pod
+  core.Resources resources = 2;
 }
 
 // WorkerGroupSpec are the specs for the worker pods
@@ -50,4 +54,6 @@ message WorkerGroupSpec {
   // Optional. RayStartParams are the params of the start command: address, object-store-memory.
   // Refer to https://docs.ray.io/en/latest/ray-core/package-ref.html#ray-start
   map<string, string> ray_start_params = 5;
+  // Resource specification for ray worker pods
+  core.Resources resources = 6;
 }