Repository: flyteorg/flyte (language: python)

Issue:
Decouple ray submitter, worker, and head resources
## Tracking issue
Related to https://github.com/flyteorg/flyte/issues/5666

## Why are the changes needed?
If you try to use the Ray plugin for GPU related workloads the submitter pod ends up consuming as many GPUs as the heads/worker nodes. 

## What changes were proposed in this pull request?
This pull request introduces resources fields in the worker and head node specs. It then uses the task resources for the submitter pod. 

## How was this patch tested?
Unit tests

### Setup process

### Screenshots

## Check all the applicable boxes <!-- Follow the above conventions to check the box -->

- [x] I updated the documentation accordingly.
- [x] All new and existing tests passed.
- [x] All commits are signed-off.

## Related PRs

flytekit PR will come

## Docs link

TBA



Discussion / hints:

Fix patch:
diff --git a/flyteidl/gen/pb-es/flyteidl/plugins/ray_pb.ts b/flyteidl/gen/pb-es/flyteidl/plugins/ray_pb.ts
index 10877bcb42..0c0a33e0a1 100644
--- a/flyteidl/gen/pb-es/flyteidl/plugins/ray_pb.ts
+++ b/flyteidl/gen/pb-es/flyteidl/plugins/ray_pb.ts
@@ -5,6 +5,7 @@
 
 import type { BinaryReadOptions, FieldList, JsonReadOptions, JsonValue, PartialMessage, PlainMessage } from "@bufbuild/protobuf";
 import { Message, proto3 } from "@bufbuild/protobuf";
+import { K8sPod } from "../core/tasks_pb.js";
 
 /**
  * RayJobSpec defines the desired state of RayJob
@@ -153,6 +154,13 @@ export class HeadGroupSpec extends Message<HeadGroupSpec> {
    */
   rayStartParams: { [key: string]: string } = {};
 
+  /**
+   * Pod Spec for the ray head pod
+   *
+   * @generated from field: flyteidl.core.K8sPod k8s_pod = 2;
+   */
+  k8sPod?: K8sPod;
+
   constructor(data?: PartialMessage<HeadGroupSpec>) {
     super();
     proto3.util.initPartial(data, this);
@@ -162,6 +170,7 @@ export class HeadGroupSpec extends Message<HeadGroupSpec> {
   static readonly typeName = "flyteidl.plugins.HeadGroupSpec";
   static readonly fields: FieldList = proto3.util.newFieldList(() => [
     { no: 1, name: "ray_start_params", kind: "map", K: 9 /* ScalarType.STRING */, V: {kind: "scalar", T: 9 /* ScalarType.STRING */} },
+    { no: 2, name: "k8s_pod", kind: "message", T: K8sPod },
   ]);
 
   static fromBinary(bytes: Uint8Array, options?: Partial<BinaryReadOptions>): HeadGroupSpec {
@@ -223,6 +232,13 @@ export class WorkerGroupSpec extends Message<WorkerGroupSpec> {
    */
   rayStartParams: { [key: string]: string } = {};
 
+  /**
+   * Pod Spec for ray worker pods
+   *
+   * @generated from field: flyteidl.core.K8sPod k8s_pod = 6;
+   */
+  k8sPod?: K8sPod;
+
   constructor(data?: PartialMessage<WorkerGroupSpec>) {
     super();
     proto3.util.initPartial(data, this);
@@ -236,6 +252,7 @@ export class WorkerGroupSpec extends Message<WorkerGroupSpec> {
     { no: 3, name: "min_replicas", kind: "scalar", T: 5 /* ScalarType.INT32 */ },
     { no: 4, name: "max_replicas", kind: "scalar", T: 5 /* ScalarType.INT32 */ },
     { no: 5, name: "ray_start_params", kind: "map", K: 9 /* ScalarType.STRING */, V: {kind: "scalar", T: 9 /* ScalarType.STRING */} },
+    { no: 6, name: "k8s_pod", kind: "message", T: K8sPod },
   ]);
 
   static fromBinary(bytes: Uint8Array, options?: Partial<BinaryReadOptions>): WorkerGroupSpec {
diff --git a/flyteidl/gen/pb-go/flyteidl/plugins/ray.pb.go b/flyteidl/gen/pb-go/flyteidl/plugins/ray.pb.go
index d0f072db51..18d0c4c1cb 100644
--- a/flyteidl/gen/pb-go/flyteidl/plugins/ray.pb.go
+++ b/flyteidl/gen/pb-go/flyteidl/plugins/ray.pb.go
@@ -7,6 +7,7 @@
 package plugins
 
 import (
+	core "github.com/flyteorg/flyte/flyteidl/gen/pb-go/flyteidl/core"
 	protoreflect "google.golang.org/protobuf/reflect/protoreflect"
 	protoimpl "google.golang.org/protobuf/runtime/protoimpl"
 	reflect "reflect"
@@ -186,6 +187,8 @@ type HeadGroupSpec struct {
 	// 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
 	RayStartParams map[string]string `protobuf:"bytes,1,rep,name=ray_start_params,json=rayStartParams,proto3" json:"ray_start_params,omitempty" protobuf_key:"bytes,1,opt,name=key,proto3" protobuf_val:"bytes,2,opt,name=value,proto3"`
+	// Pod Spec for the ray head pod
+	K8SPod *core.K8SPod `protobuf:"bytes,2,opt,name=k8s_pod,json=k8sPod,proto3" json:"k8s_pod,omitempty"`
 }
 
 func (x *HeadGroupSpec) Reset() {
@@ -227,6 +230,13 @@ func (x *HeadGroupSpec) GetRayStartParams() map[string]string {
 	return nil
 }
 
+func (x *HeadGroupSpec) GetK8SPod() *core.K8SPod {
+	if x != nil {
+		return x.K8SPod
+	}
+	return nil
+}
+
 // WorkerGroupSpec are the specs for the worker pods
 type WorkerGroupSpec struct {
 	state         protoimpl.MessageState
@@ -244,6 +254,8 @@ type WorkerGroupSpec struct {
 	// 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
 	RayStartParams map[string]string `protobuf:"bytes,5,rep,name=ray_start_params,json=rayStartParams,proto3" json:"ray_start_params,omitempty" protobuf_key:"bytes,1,opt,name=key,proto3" protobuf_val:"bytes,2,opt,name=value,proto3"`
+	// Pod Spec for ray worker pods
+	K8SPod *core.K8SPod `protobuf:"bytes,6,opt,name=k8s_pod,json=k8sPod,proto3" json:"k8s_pod,omitempty"`
 }
 
 func (x *WorkerGroupSpec) Reset() {
@@ -313,87 +325,101 @@ func (x *WorkerGroupSpec) GetRayStartParams() map[string]string {
 	return nil
 }
 
+func (x *WorkerGroupSpec) GetK8SPod() *core.K8SPod {
+	if x != nil {
+		return x.K8SPod
+	}
+	return nil
+}
+
 var File_flyteidl_plugins_ray_proto protoreflect.FileDescriptor
 
 var file_flyteidl_plugins_ray_proto_rawDesc = []byte{
 	0x0a, 0x1a, 0x66, 0x6c, 0x79, 0x74, 0x65, 0x69, 0x64, 0x6c, 0x2f, 0x70, 0x6c, 0x75, 0x67, 0x69,
 	0x6e, 0x73, 0x2f, 0x72, 0x61, 0x79, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x12, 0x10, 0x66, 0x6c,
-	0x79, 0x74, 0x65, 0x69, 0x64, 0x6c, 0x2e, 0x70, 0x6c, 0x75, 0x67, 0x69, 0x6e, 0x73, 0x22, 0x92,
-	0x02, 0x0a, 0x06, 0x52, 0x61, 0x79, 0x4a, 0x6f, 0x62, 0x12, 0x3d, 0x0a, 0x0b, 0x72, 0x61, 0x79,
-	0x5f, 0x63, 0x6c, 0x75, 0x73, 0x74, 0x65, 0x72, 0x18, 0x01, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x1c,
-	0x2e, 0x66, 0x6c, 0x79, 0x74, 0x65, 0x69, 0x64, 0x6c, 0x2e, 0x70, 0x6c, 0x75, 0x67, 0x69, 0x6e,
-	0x73, 0x2e, 0x52, 0x61, 0x79, 0x43, 0x6c, 0x75, 0x73, 0x74, 0x65, 0x72, 0x52, 0x0a, 0x72, 0x61,
-	0x79, 0x43, 0x6c, 0x75, 0x73, 0x74, 0x65, 0x72, 0x12, 0x23, 0x0a, 0x0b, 0x72, 0x75, 0x6e, 0x74,
-	0x69, 0x6d, 0x65, 0x5f, 0x65, 0x6e, 0x76, 0x18, 0x02, 0x20, 0x01, 0x28, 0x09, 0x42, 0x02, 0x18,
-	0x01, 0x52, 0x0a, 0x72, 0x75, 0x6e, 0x74, 0x69, 0x6d, 0x65, 0x45, 0x6e, 0x76, 0x12, 0x3d, 0x0a,
-	0x1b, 0x73, 0x68, 0x75, 0x74, 0x64, 0x6f, 0x77, 0x6e, 0x5f, 0x61, 0x66, 0x74, 0x65, 0x72, 0x5f,
-	0x6a, 0x6f, 0x62, 0x5f, 0x66, 0x69, 0x6e, 0x69, 0x73, 0x68, 0x65, 0x73, 0x18, 0x03, 0x20, 0x01,
-	0x28, 0x08, 0x52, 0x18, 0x73, 0x68, 0x75, 0x74, 0x64, 0x6f, 0x77, 0x6e, 0x41, 0x66, 0x74, 0x65,
-	0x72, 0x4a, 0x6f, 0x62, 0x46, 0x69, 0x6e, 0x69, 0x73, 0x68, 0x65, 0x73, 0x12, 0x3b, 0x0a, 0x1a,
-	0x74, 0x74, 0x6c, 0x5f, 0x73, 0x65, 0x63, 0x6f, 0x6e, 0x64, 0x73, 0x5f, 0x61, 0x66, 0x74, 0x65,
-	0x72, 0x5f, 0x66, 0x69, 0x6e, 0x69, 0x73, 0x68, 0x65, 0x64, 0x18, 0x04, 0x20, 0x01, 0x28, 0x05,
-	0x52, 0x17, 0x74, 0x74, 0x6c, 0x53, 0x65, 0x63, 0x6f, 0x6e, 0x64, 0x73, 0x41, 0x66, 0x74, 0x65,
-	0x72, 0x46, 0x69, 0x6e, 0x69, 0x73, 0x68, 0x65, 0x64, 0x12, 0x28, 0x0a, 0x10, 0x72, 0x75, 0x6e,
-	0x74, 0x69, 0x6d, 0x65, 0x5f, 0x65, 0x6e, 0x76, 0x5f, 0x79, 0x61, 0x6d, 0x6c, 0x18, 0x05, 0x20,
-	0x01, 0x28, 0x09, 0x52, 0x0e, 0x72, 0x75, 0x6e, 0x74, 0x69, 0x6d, 0x65, 0x45, 0x6e, 0x76, 0x59,
-	0x61, 0x6d, 0x6c, 0x22, 0xd3, 0x01, 0x0a, 0x0a, 0x52, 0x61, 0x79, 0x43, 0x6c, 0x75, 0x73, 0x74,
-	0x65, 0x72, 0x12, 0x47, 0x0a, 0x0f, 0x68, 0x65, 0x61, 0x64, 0x5f, 0x67, 0x72, 0x6f, 0x75, 0x70,
-	0x5f, 0x73, 0x70, 0x65, 0x63, 0x18, 0x01, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x1f, 0x2e, 0x66, 0x6c,
-	0x79, 0x74, 0x65, 0x69, 0x64, 0x6c, 0x2e, 0x70, 0x6c, 0x75, 0x67, 0x69, 0x6e, 0x73, 0x2e, 0x48,
-	0x65, 0x61, 0x64, 0x47, 0x72, 0x6f, 0x75, 0x70, 0x53, 0x70, 0x65, 0x63, 0x52, 0x0d, 0x68, 0x65,
-	0x61, 0x64, 0x47, 0x72, 0x6f, 0x75, 0x70, 0x53, 0x70, 0x65, 0x63, 0x12, 0x4d, 0x0a, 0x11, 0x77,
-	0x6f, 0x72, 0x6b, 0x65, 0x72, 0x5f, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x73, 0x70, 0x65, 0x63,
-	0x18, 0x02, 0x20, 0x03, 0x28, 0x0b, 0x32, 0x21, 0x2e, 0x66, 0x6c, 0x79, 0x74, 0x65, 0x69, 0x64,
-	0x6c, 0x2e, 0x70, 0x6c, 0x75, 0x67, 0x69, 0x6e, 0x73, 0x2e, 0x57, 0x6f, 0x72, 0x6b, 0x65, 0x72,
-	0x47, 0x72, 0x6f, 0x75, 0x70, 0x53, 0x70, 0x65, 0x63, 0x52, 0x0f, 0x77, 0x6f, 0x72, 0x6b, 0x65,
-	0x72, 0x47, 0x72, 0x6f, 0x75, 0x70, 0x53, 0x70, 0x65, 0x63, 0x12, 0x2d, 0x0a, 0x12, 0x65, 0x6e,
-	0x61, 0x62, 0x6c, 0x65, 0x5f, 0x61, 0x75, 0x74, 0x6f, 0x73, 0x63, 0x61, 0x6c, 0x69, 0x6e, 0x67,
-	0x18, 0x03, 0x20, 0x01, 0x28, 0x08, 0x52, 0x11, 0x65, 0x6e, 0x61, 0x62, 0x6c, 0x65, 0x41, 0x75,
-	0x74, 0x6f, 0x73, 0x63, 0x61, 0x6c, 0x69, 0x6e, 0x67, 0x22, 0xb1, 0x01, 0x0a, 0x0d, 0x48, 0x65,
-	0x61, 0x64, 0x47, 0x72, 0x6f, 0x75, 0x70, 0x53, 0x70, 0x65, 0x63, 0x12, 0x5d, 0x0a, 0x10, 0x72,
-	0x61, 0x79, 0x5f, 0x73, 0x74, 0x61, 0x72, 0x74, 0x5f, 0x70, 0x61, 0x72, 0x61, 0x6d, 0x73, 0x18,
-	0x01, 0x20, 0x03, 0x28, 0x0b, 0x32, 0x33, 0x2e, 0x66, 0x6c, 0x79, 0x74, 0x65, 0x69, 0x64, 0x6c,
-	0x2e, 0x70, 0x6c, 0x75, 0x67, 0x69, 0x6e, 0x73, 0x2e, 0x48, 0x65, 0x61, 0x64, 0x47, 0x72, 0x6f,
-	0x75, 0x70, 0x53, 0x70, 0x65, 0x63, 0x2e, 0x52, 0x61, 0x79, 0x53, 0x74, 0x61, 0x72, 0x74, 0x50,
-	0x61, 0x72, 0x61, 0x6d, 0x73, 0x45, 0x6e, 0x74, 0x72, 0x79, 0x52, 0x0e, 0x72, 0x61, 0x79, 0x53,
-	0x74, 0x61, 0x72, 0x74, 0x50, 0x61, 0x72, 0x61, 0x6d, 0x73, 0x1a, 0x41, 0x0a, 0x13, 0x52, 0x61,
-	0x79, 0x53, 0x74, 0x61, 0x72, 0x74, 0x50, 0x61, 0x72, 0x61, 0x6d, 0x73, 0x45, 0x6e, 0x74, 0x72,
-	0x79, 0x12, 0x10, 0x0a, 0x03, 0x6b, 0x65, 0x79, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x03,
-	0x6b, 0x65, 0x79, 0x12, 0x14, 0x0a, 0x05, 0x76, 0x61, 0x6c, 0x75, 0x65, 0x18, 0x02, 0x20, 0x01,
-	0x28, 0x09, 0x52, 0x05, 0x76, 0x61, 0x6c, 0x75, 0x65, 0x3a, 0x02, 0x38, 0x01, 0x22, 0xb6, 0x02,
-	0x0a, 0x0f, 0x57, 0x6f, 0x72, 0x6b, 0x65, 0x72, 0x47, 0x72, 0x6f, 0x75, 0x70, 0x53, 0x70, 0x65,
-	0x63, 0x12, 0x1d, 0x0a, 0x0a, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x6e, 0x61, 0x6d, 0x65, 0x18,
-	0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x09, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x4e, 0x61, 0x6d, 0x65,
-	0x12, 0x1a, 0x0a, 0x08, 0x72, 0x65, 0x70, 0x6c, 0x69, 0x63, 0x61, 0x73, 0x18, 0x02, 0x20, 0x01,
-	0x28, 0x05, 0x52, 0x08, 0x72, 0x65, 0x70, 0x6c, 0x69, 0x63, 0x61, 0x73, 0x12, 0x21, 0x0a, 0x0c,
-	0x6d, 0x69, 0x6e, 0x5f, 0x72, 0x65, 0x70, 0x6c, 0x69, 0x63, 0x61, 0x73, 0x18, 0x03, 0x20, 0x01,
-	0x28, 0x05, 0x52, 0x0b, 0x6d, 0x69, 0x6e, 0x52, 0x65, 0x70, 0x6c, 0x69, 0x63, 0x61, 0x73, 0x12,
-	0x21, 0x0a, 0x0c, 0x6d, 0x61, 0x78, 0x5f, 0x72, 0x65, 0x70, 0x6c, 0x69, 0x63, 0x61, 0x73, 0x18,
-	0x04, 0x20, 0x01, 0x28, 0x05, 0x52, 0x0b, 0x6d, 0x61, 0x78, 0x52, 0x65, 0x70, 0x6c, 0x69, 0x63,
-	0x61, 0x73, 0x12, 0x5f, 0x0a, 0x10, 0x72, 0x61, 0x79, 0x5f, 0x73, 0x74, 0x61, 0x72, 0x74, 0x5f,
-	0x70, 0x61, 0x72, 0x61, 0x6d, 0x73, 0x18, 0x05, 0x20, 0x03, 0x28, 0x0b, 0x32, 0x35, 0x2e, 0x66,
-	0x6c, 0x79, 0x74, 0x65, 0x69, 0x64, 0x6c, 0x2e, 0x70, 0x6c, 0x75, 0x67, 0x69, 0x6e, 0x73, 0x2e,
-	0x57, 0x6f, 0x72, 0x6b, 0x65, 0x72, 0x47, 0x72, 0x6f, 0x75, 0x70, 0x53, 0x70, 0x65, 0x63, 0x2e,
-	0x52, 0x61, 0x79, 0x53, 0x74, 0x61, 0x72, 0x74, 0x50, 0x61, 0x72, 0x61, 0x6d, 0x73, 0x45, 0x6e,
-	0x74, 0x72, 0x79, 0x52, 0x0e, 0x72, 0x61, 0x79, 0x53, 0x74, 0x61, 0x72, 0x74, 0x50, 0x61, 0x72,
-	0x61, 0x6d, 0x73, 0x1a, 0x41, 0x0a, 0x13, 0x52, 0x61, 0x79, 0x53, 0x74, 0x61, 0x72, 0x74, 0x50,
-	0x61, 0x72, 0x61, 0x6d, 0x73, 0x45, 0x6e, 0x74, 0x72, 0x79, 0x12, 0x10, 0x0a, 0x03, 0x6b, 0x65,
-	0x79, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x03, 0x6b, 0x65, 0x79, 0x12, 0x14, 0x0a, 0x05,
-	0x76, 0x61, 0x6c, 0x75, 0x65, 0x18, 0x02, 0x20, 0x01, 0x28, 0x09, 0x52, 0x05, 0x76, 0x61, 0x6c,
-	0x75, 0x65, 0x3a, 0x02, 0x38, 0x01, 0x42, 0xc0, 0x01, 0x0a, 0x14, 0x63, 0x6f, 0x6d, 0x2e, 0x66,
-	0x6c, 0x79, 0x74, 0x65, 0x69, 0x64, 0x6c, 0x2e, 0x70, 0x6c, 0x75, 0x67, 0x69, 0x6e, 0x73, 0x42,
-	0x08, 0x52, 0x61, 0x79, 0x50, 0x72, 0x6f, 0x74, 0x6f, 0x50, 0x01, 0x5a, 0x3d, 0x67, 0x69, 0x74,
-	0x68, 0x75, 0x62, 0x2e, 0x63, 0x6f, 0x6d, 0x2f, 0x66, 0x6c, 0x79, 0x74, 0x65, 0x6f, 0x72, 0x67,
-	0x2f, 0x66, 0x6c, 0x79, 0x74, 0x65, 0x2f, 0x66, 0x6c, 0x79, 0x74, 0x65, 0x69, 0x64, 0x6c, 0x2f,
-	0x67, 0x65, 0x6e, 0x2f, 0x70, 0x62, 0x2d, 0x67, 0x6f, 0x2f, 0x66, 0x6c, 0x79, 0x74, 0x65, 0x69,
-	0x64, 0x6c, 0x2f, 0x70, 0x6c, 0x75, 0x67, 0x69, 0x6e, 0x73, 0xa2, 0x02, 0x03, 0x46, 0x50, 0x58,
-	0xaa, 0x02, 0x10, 0x46, 0x6c, 0x79, 0x74, 0x65, 0x69, 0x64, 0x6c, 0x2e, 0x50, 0x6c, 0x75, 0x67,
-	0x69, 0x6e, 0x73, 0xca, 0x02, 0x10, 0x46, 0x6c, 0x79, 0x74, 0x65, 0x69, 0x64, 0x6c, 0x5c, 0x50,
-	0x6c, 0x75, 0x67, 0x69, 0x6e, 0x73, 0xe2, 0x02, 0x1c, 0x46, 0x6c, 0x79, 0x74, 0x65, 0x69, 0x64,
-	0x6c, 0x5c, 0x50, 0x6c, 0x75, 0x67, 0x69, 0x6e, 0x73, 0x5c, 0x47, 0x50, 0x42, 0x4d, 0x65, 0x74,
-	0x61, 0x64, 0x61, 0x74, 0x61, 0xea, 0x02, 0x11, 0x46, 0x6c, 0x79, 0x74, 0x65, 0x69, 0x64, 0x6c,
-	0x3a, 0x3a, 0x50, 0x6c, 0x75, 0x67, 0x69, 0x6e, 0x73, 0x62, 0x06, 0x70, 0x72, 0x6f, 0x74, 0x6f,
-	0x33,
+	0x79, 0x74, 0x65, 0x69, 0x64, 0x6c, 0x2e, 0x70, 0x6c, 0x75, 0x67, 0x69, 0x6e, 0x73, 0x1a, 0x19,
+	0x66, 0x6c, 0x79, 0x74, 0x65, 0x69, 0x64, 0x6c, 0x2f, 0x63, 0x6f, 0x72, 0x65, 0x2f, 0x74, 0x61,
+	0x73, 0x6b, 0x73, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x22, 0x92, 0x02, 0x0a, 0x06, 0x52, 0x61,
+	0x79, 0x4a, 0x6f, 0x62, 0x12, 0x3d, 0x0a, 0x0b, 0x72, 0x61, 0x79, 0x5f, 0x63, 0x6c, 0x75, 0x73,
+	0x74, 0x65, 0x72, 0x18, 0x01, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x1c, 0x2e, 0x66, 0x6c, 0x79, 0x74,
+	0x65, 0x69, 0x64, 0x6c, 0x2e, 0x70, 0x6c, 0x75, 0x67, 0x69, 0x6e, 0x73, 0x2e, 0x52, 0x61, 0x79,
+	0x43, 0x6c, 0x75, 0x73, 0x74, 0x65, 0x72, 0x52, 0x0a, 0x72, 0x61, 0x79, 0x43, 0x6c, 0x75, 0x73,
+	0x74, 0x65, 0x72, 0x12, 0x23, 0x0a, 0x0b, 0x72, 0x75, 0x6e, 0x74, 0x69, 0x6d, 0x65, 0x5f, 0x65,
+	0x6e, 0x76, 0x18, 0x02, 0x20, 0x01, 0x28, 0x09, 0x42, 0x02, 0x18, 0x01, 0x52, 0x0a, 0x72, 0x75,
+	0x6e, 0x74, 0x69, 0x6d, 0x65, 0x45, 0x6e, 0x76, 0x12, 0x3d, 0x0a, 0x1b, 0x73, 0x68, 0x75, 0x74,
+	0x64, 0x6f, 0x77, 0x6e, 0x5f, 0x61, 0x66, 0x74, 0x65, 0x72, 0x5f, 0x6a, 0x6f, 0x62, 0x5f, 0x66,
+	0x69, 0x6e, 0x69, 0x73, 0x68, 0x65, 0x73, 0x18, 0x03, 0x20, 0x01, 0x28, 0x08, 0x52, 0x18, 0x73,
+	0x68, 0x75, 0x74, 0x64, 0x6f, 0x77, 0x6e, 0x41, 0x66, 0x74, 0x65, 0x72, 0x4a, 0x6f, 0x62, 0x46,
+	0x69, 0x6e, 0x69, 0x73, 0x68, 0x65, 0x73, 0x12, 0x3b, 0x0a, 0x1a, 0x74, 0x74, 0x6c, 0x5f, 0x73,
+	0x65, 0x63, 0x6f, 0x6e, 0x64, 0x73, 0x5f, 0x61, 0x66, 0x74, 0x65, 0x72, 0x5f, 0x66, 0x69, 0x6e,
+	0x69, 0x73, 0x68, 0x65, 0x64, 0x18, 0x04, 0x20, 0x01, 0x28, 0x05, 0x52, 0x17, 0x74, 0x74, 0x6c,
+	0x53, 0x65, 0x63, 0x6f, 0x6e, 0x64, 0x73, 0x41, 0x66, 0x74, 0x65, 0x72, 0x46, 0x69, 0x6e, 0x69,
+	0x73, 0x68, 0x65, 0x64, 0x12, 0x28, 0x0a, 0x10, 0x72, 0x75, 0x6e, 0x74, 0x69, 0x6d, 0x65, 0x5f,
+	0x65, 0x6e, 0x76, 0x5f, 0x79, 0x61, 0x6d, 0x6c, 0x18, 0x05, 0x20, 0x01, 0x28, 0x09, 0x52, 0x0e,
+	0x72, 0x75, 0x6e, 0x74, 0x69, 0x6d, 0x65, 0x45, 0x6e, 0x76, 0x59, 0x61, 0x6d, 0x6c, 0x22, 0xd3,
+	0x01, 0x0a, 0x0a, 0x52, 0x61, 0x79, 0x43, 0x6c, 0x75, 0x73, 0x74, 0x65, 0x72, 0x12, 0x47, 0x0a,
+	0x0f, 0x68, 0x65, 0x61, 0x64, 0x5f, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x73, 0x70, 0x65, 0x63,
+	0x18, 0x01, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x1f, 0x2e, 0x66, 0x6c, 0x79, 0x74, 0x65, 0x69, 0x64,
+	0x6c, 0x2e, 0x70, 0x6c, 0x75, 0x67, 0x69, 0x6e, 0x73, 0x2e, 0x48, 0x65, 0x61, 0x64, 0x47, 0x72,
+	0x6f, 0x75, 0x70, 0x53, 0x70, 0x65, 0x63, 0x52, 0x0d, 0x68, 0x65, 0x61, 0x64, 0x47, 0x72, 0x6f,
+	0x75, 0x70, 0x53, 0x70, 0x65, 0x63, 0x12, 0x4d, 0x0a, 0x11, 0x77, 0x6f, 0x72, 0x6b, 0x65, 0x72,
+	0x5f, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x73, 0x70, 0x65, 0x63, 0x18, 0x02, 0x20, 0x03, 0x28,
+	0x0b, 0x32, 0x21, 0x2e, 0x66, 0x6c, 0x79, 0x74, 0x65, 0x69, 0x64, 0x6c, 0x2e, 0x70, 0x6c, 0x75,
+	0x67, 0x69, 0x6e, 0x73, 0x2e, 0x57, 0x6f, 0x72, 0x6b, 0x65, 0x72, 0x47, 0x72, 0x6f, 0x75, 0x70,
+	0x53, 0x70, 0x65, 0x63, 0x52, 0x0f, 0x77, 0x6f, 0x72, 0x6b, 0x65, 0x72, 0x47, 0x72, 0x6f, 0x75,
+	0x70, 0x53, 0x70, 0x65, 0x63, 0x12, 0x2d, 0x0a, 0x12, 0x65, 0x6e, 0x61, 0x62, 0x6c, 0x65, 0x5f,
+	0x61, 0x75, 0x74, 0x6f, 0x73, 0x63, 0x61, 0x6c, 0x69, 0x6e, 0x67, 0x18, 0x03, 0x20, 0x01, 0x28,
+	0x08, 0x52, 0x11, 0x65, 0x6e, 0x61, 0x62, 0x6c, 0x65, 0x41, 0x75, 0x74, 0x6f, 0x73, 0x63, 0x61,
+	0x6c, 0x69, 0x6e, 0x67, 0x22, 0xe1, 0x01, 0x0a, 0x0d, 0x48, 0x65, 0x61, 0x64, 0x47, 0x72, 0x6f,
+	0x75, 0x70, 0x53, 0x70, 0x65, 0x63, 0x12, 0x5d, 0x0a, 0x10, 0x72, 0x61, 0x79, 0x5f, 0x73, 0x74,
+	0x61, 0x72, 0x74, 0x5f, 0x70, 0x61, 0x72, 0x61, 0x6d, 0x73, 0x18, 0x01, 0x20, 0x03, 0x28, 0x0b,
+	0x32, 0x33, 0x2e, 0x66, 0x6c, 0x79, 0x74, 0x65, 0x69, 0x64, 0x6c, 0x2e, 0x70, 0x6c, 0x75, 0x67,
+	0x69, 0x6e, 0x73, 0x2e, 0x48, 0x65, 0x61, 0x64, 0x47, 0x72, 0x6f, 0x75, 0x70, 0x53, 0x70, 0x65,
+	0x63, 0x2e, 0x52, 0x61, 0x79, 0x53, 0x74, 0x61, 0x72, 0x74, 0x50, 0x61, 0x72, 0x61, 0x6d, 0x73,
+	0x45, 0x6e, 0x74, 0x72, 0x79, 0x52, 0x0e, 0x72, 0x61, 0x79, 0x53, 0x74, 0x61, 0x72, 0x74, 0x50,
+	0x61, 0x72, 0x61, 0x6d, 0x73, 0x12, 0x2e, 0x0a, 0x07, 0x6b, 0x38, 0x73, 0x5f, 0x70, 0x6f, 0x64,
+	0x18, 0x02, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x15, 0x2e, 0x66, 0x6c, 0x79, 0x74, 0x65, 0x69, 0x64,
+	0x6c, 0x2e, 0x63, 0x6f, 0x72, 0x65, 0x2e, 0x4b, 0x38, 0x73, 0x50, 0x6f, 0x64, 0x52, 0x06, 0x6b,
+	0x38, 0x73, 0x50, 0x6f, 0x64, 0x1a, 0x41, 0x0a, 0x13, 0x52, 0x61, 0x79, 0x53, 0x74, 0x61, 0x72,
+	0x74, 0x50, 0x61, 0x72, 0x61, 0x6d, 0x73, 0x45, 0x6e, 0x74, 0x72, 0x79, 0x12, 0x10, 0x0a, 0x03,
+	0x6b, 0x65, 0x79, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x03, 0x6b, 0x65, 0x79, 0x12, 0x14,
+	0x0a, 0x05, 0x76, 0x61, 0x6c, 0x75, 0x65, 0x18, 0x02, 0x20, 0x01, 0x28, 0x09, 0x52, 0x05, 0x76,
+	0x61, 0x6c, 0x75, 0x65, 0x3a, 0x02, 0x38, 0x01, 0x22, 0xe6, 0x02, 0x0a, 0x0f, 0x57, 0x6f, 0x72,
+	0x6b, 0x65, 0x72, 0x47, 0x72, 0x6f, 0x75, 0x70, 0x53, 0x70, 0x65, 0x63, 0x12, 0x1d, 0x0a, 0x0a,
+	0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x6e, 0x61, 0x6d, 0x65, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09,
+	0x52, 0x09, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x4e, 0x61, 0x6d, 0x65, 0x12, 0x1a, 0x0a, 0x08, 0x72,
+	0x65, 0x70, 0x6c, 0x69, 0x63, 0x61, 0x73, 0x18, 0x02, 0x20, 0x01, 0x28, 0x05, 0x52, 0x08, 0x72,
+	0x65, 0x70, 0x6c, 0x69, 0x63, 0x61, 0x73, 0x12, 0x21, 0x0a, 0x0c, 0x6d, 0x69, 0x6e, 0x5f, 0x72,
+	0x65, 0x70, 0x6c, 0x69, 0x63, 0x61, 0x73, 0x18, 0x03, 0x20, 0x01, 0x28, 0x05, 0x52, 0x0b, 0x6d,
+	0x69, 0x6e, 0x52, 0x65, 0x70, 0x6c, 0x69, 0x63, 0x61, 0x73, 0x12, 0x21, 0x0a, 0x0c, 0x6d, 0x61,
+	0x78, 0x5f, 0x72, 0x65, 0x70, 0x6c, 0x69, 0x63, 0x61, 0x73, 0x18, 0x04, 0x20, 0x01, 0x28, 0x05,
+	0x52, 0x0b, 0x6d, 0x61, 0x78, 0x52, 0x65, 0x70, 0x6c, 0x69, 0x63, 0x61, 0x73, 0x12, 0x5f, 0x0a,
+	0x10, 0x72, 0x61, 0x79, 0x5f, 0x73, 0x74, 0x61, 0x72, 0x74, 0x5f, 0x70, 0x61, 0x72, 0x61, 0x6d,
+	0x73, 0x18, 0x05, 0x20, 0x03, 0x28, 0x0b, 0x32, 0x35, 0x2e, 0x66, 0x6c, 0x79, 0x74, 0x65, 0x69,
+	0x64, 0x6c, 0x2e, 0x70, 0x6c, 0x75, 0x67, 0x69, 0x6e, 0x73, 0x2e, 0x57, 0x6f, 0x72, 0x6b, 0x65,
+	0x72, 0x47, 0x72, 0x6f, 0x75, 0x70, 0x53, 0x70, 0x65, 0x63, 0x2e, 0x52, 0x61, 0x79, 0x53, 0x74,
+	0x61, 0x72, 0x74, 0x50, 0x61, 0x72, 0x61, 0x6d, 0x73, 0x45, 0x6e, 0x74, 0x72, 0x79, 0x52, 0x0e,
+	0x72, 0x61, 0x79, 0x53, 0x74, 0x61, 0x72, 0x74, 0x50, 0x61, 0x72, 0x61, 0x6d, 0x73, 0x12, 0x2e,
+	0x0a, 0x07, 0x6b, 0x38, 0x73, 0x5f, 0x70, 0x6f, 0x64, 0x18, 0x06, 0x20, 0x01, 0x28, 0x0b, 0x32,
+	0x15, 0x2e, 0x66, 0x6c, 0x79, 0x74, 0x65, 0x69, 0x64, 0x6c, 0x2e, 0x63, 0x6f, 0x72, 0x65, 0x2e,
+	0x4b, 0x38, 0x73, 0x50, 0x6f, 0x64, 0x52, 0x06, 0x6b, 0x38, 0x73, 0x50, 0x6f, 0x64, 0x1a, 0x41,
+	0x0a, 0x13, 0x52, 0x61, 0x79, 0x53, 0x74, 0x61, 0x72, 0x74, 0x50, 0x61, 0x72, 0x61, 0x6d, 0x73,
+	0x45, 0x6e, 0x74, 0x72, 0x79, 0x12, 0x10, 0x0a, 0x03, 0x6b, 0x65, 0x79, 0x18, 0x01, 0x20, 0x01,
+	0x28, 0x09, 0x52, 0x03, 0x6b, 0x65, 0x79, 0x12, 0x14, 0x0a, 0x05, 0x76, 0x61, 0x6c, 0x75, 0x65,
+	0x18, 0x02, 0x20, 0x01, 0x28, 0x09, 0x52, 0x05, 0x76, 0x61, 0x6c, 0x75, 0x65, 0x3a, 0x02, 0x38,
+	0x01, 0x42, 0xc0, 0x01, 0x0a, 0x14, 0x63, 0x6f, 0x6d, 0x2e, 0x66, 0x6c, 0x79, 0x74, 0x65, 0x69,
+	0x64, 0x6c, 0x2e, 0x70, 0x6c, 0x75, 0x67, 0x69, 0x6e, 0x73, 0x42, 0x08, 0x52, 0x61, 0x79, 0x50,
+	0x72, 0x6f, 0x74, 0x6f, 0x50, 0x01, 0x5a, 0x3d, 0x67, 0x69, 0x74, 0x68, 0x75, 0x62, 0x2e, 0x63,
+	0x6f, 0x6d, 0x2f, 0x66, 0x6c, 0x79, 0x74, 0x65, 0x6f, 0x72, 0x67, 0x2f, 0x66, 0x6c, 0x79, 0x74,
+	0x65, 0x2f, 0x66, 0x6c, 0x79, 0x74, 0x65, 0x69, 0x64, 0x6c, 0x2f, 0x67, 0x65, 0x6e, 0x2f, 0x70,
+	0x62, 0x2d, 0x67, 0x6f, 0x2f, 0x66, 0x6c, 0x79, 0x74, 0x65, 0x69, 0x64, 0x6c, 0x2f, 0x70, 0x6c,
+	0x75, 0x67, 0x69, 0x6e, 0x73, 0xa2, 0x02, 0x03, 0x46, 0x50, 0x58, 0xaa, 0x02, 0x10, 0x46, 0x6c,
+	0x79, 0x74, 0x65, 0x69, 0x64, 0x6c, 0x2e, 0x50, 0x6c, 0x75, 0x67, 0x69, 0x6e, 0x73, 0xca, 0x02,
+	0x10, 0x46, 0x6c, 0x79, 0x74, 0x65, 0x69, 0x64, 0x6c, 0x5c, 0x50, 0x6c, 0x75, 0x67, 0x69, 0x6e,
+	0x73, 0xe2, 0x02, 0x1c, 0x46, 0x6c, 0x79, 0x74, 0x65, 0x69, 0x64, 0x6c, 0x5c, 0x50, 0x6c, 0x75,
+	0x67, 0x69, 0x6e, 0x73, 0x5c, 0x47, 0x50, 0x42, 0x4d, 0x65, 0x74, 0x61, 0x64, 0x61, 0x74, 0x61,
+	0xea, 0x02, 0x11, 0x46, 0x6c, 0x79, 0x74, 0x65, 0x69, 0x64, 0x6c, 0x3a, 0x3a, 0x50, 0x6c, 0x75,
+	0x67, 0x69, 0x6e, 0x73, 0x62, 0x06, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x33,
 }
 
 var (
@@ -416,18 +442,21 @@ var file_flyteidl_plugins_ray_proto_goTypes = []interface{}{
 	(*WorkerGroupSpec)(nil), // 3: flyteidl.plugins.WorkerGroupSpec
 	nil,                     // 4: flyteidl.plugins.HeadGroupSpec.RayStartParamsEntry
 	nil,                     // 5: flyteidl.plugins.WorkerGroupSpec.RayStartParamsEntry
+	(*core.K8SPod)(nil),     // 6: flyteidl.core.K8sPod
 }
 var file_flyteidl_plugins_ray_proto_depIdxs = []int32{
 	1, // 0: flyteidl.plugins.RayJob.ray_cluster:type_name -> flyteidl.plugins.RayCluster
 	2, // 1: flyteidl.plugins.RayCluster.head_group_spec:type_name -> flyteidl.plugins.HeadGroupSpec
 	3, // 2: flyteidl.plugins.RayCluster.worker_group_spec:type_name -> flyteidl.plugins.WorkerGroupSpec
 	4, // 3: flyteidl.plugins.HeadGroupSpec.ray_start_params:type_name -> flyteidl.plugins.HeadGroupSpec.RayStartParamsEntry
-	5, // 4: flyteidl.plugins.WorkerGroupSpec.ray_start_params:type_name -> flyteidl.plugins.WorkerGroupSpec.RayStartParamsEntry
-	5, // [5:5] is the sub-list for method output_type
-	5, // [5:5] is the sub-list for method input_type
-	5, // [5:5] is the sub-list for extension type_name
-	5, // [5:5] is the sub-list for extension extendee
-	0, // [0:5] is the sub-list for field type_name
+	6, // 4: flyteidl.plugins.HeadGroupSpec.k8s_pod:type_name -> flyteidl.core.K8sPod
+	5, // 5: flyteidl.plugins.WorkerGroupSpec.ray_start_params:type_name -> flyteidl.plugins.WorkerGroupSpec.RayStartParamsEntry
+	6, // 6: flyteidl.plugins.WorkerGroupSpec.k8s_pod:type_name -> flyteidl.core.K8sPod
+	7, // [7:7] is the sub-list for method output_type
+	7, // [7:7] is the sub-list for method input_type
+	7, // [7:7] is the sub-list for extension type_name
+	7, // [7:7] is the sub-list for extension extendee
+	0, // [0:7] is the sub-list for field type_name
 }
 
 func init() { file_flyteidl_plugins_ray_proto_init() }
diff --git a/flyteidl/gen/pb_python/flyteidl/plugins/ray_pb2.py b/flyteidl/gen/pb_python/flyteidl/plugins/ray_pb2.py
index af809dcfcd..c625fd957b 100644
--- a/flyteidl/gen/pb_python/flyteidl/plugins/ray_pb2.py
+++ b/flyteidl/gen/pb_python/flyteidl/plugins/ray_pb2.py
@@ -11,9 +11,10 @@
 _sym_db = _symbol_database.Default()
 
 
+from flyteidl.core import tasks_pb2 as flyteidl_dot_core_dot_tasks__pb2
 
 
-DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x1a\x66lyteidl/plugins/ray.proto\x12\x10\x66lyteidl.plugins\"\x92\x02\n\x06RayJob\x12=\n\x0bray_cluster\x18\x01 \x01(\x0b\x32\x1c.flyteidl.plugins.RayClusterR\nrayCluster\x12#\n\x0bruntime_env\x18\x02 \x01(\tB\x02\x18\x01R\nruntimeEnv\x12=\n\x1bshutdown_after_job_finishes\x18\x03 \x01(\x08R\x18shutdownAfterJobFinishes\x12;\n\x1attl_seconds_after_finished\x18\x04 \x01(\x05R\x17ttlSecondsAfterFinished\x12(\n\x10runtime_env_yaml\x18\x05 \x01(\tR\x0eruntimeEnvYaml\"\xd3\x01\n\nRayCluster\x12G\n\x0fhead_group_spec\x18\x01 \x01(\x0b\x32\x1f.flyteidl.plugins.HeadGroupSpecR\rheadGroupSpec\x12M\n\x11worker_group_spec\x18\x02 \x03(\x0b\x32!.flyteidl.plugins.WorkerGroupSpecR\x0fworkerGroupSpec\x12-\n\x12\x65nable_autoscaling\x18\x03 \x01(\x08R\x11\x65nableAutoscaling\"\xb1\x01\n\rHeadGroupSpec\x12]\n\x10ray_start_params\x18\x01 \x03(\x0b\x32\x33.flyteidl.plugins.HeadGroupSpec.RayStartParamsEntryR\x0erayStartParams\x1a\x41\n\x13RayStartParamsEntry\x12\x10\n\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n\x05value\x18\x02 \x01(\tR\x05value:\x02\x38\x01\"\xb6\x02\n\x0fWorkerGroupSpec\x12\x1d\n\ngroup_name\x18\x01 \x01(\tR\tgroupName\x12\x1a\n\x08replicas\x18\x02 \x01(\x05R\x08replicas\x12!\n\x0cmin_replicas\x18\x03 \x01(\x05R\x0bminReplicas\x12!\n\x0cmax_replicas\x18\x04 \x01(\x05R\x0bmaxReplicas\x12_\n\x10ray_start_params\x18\x05 \x03(\x0b\x32\x35.flyteidl.plugins.WorkerGroupSpec.RayStartParamsEntryR\x0erayStartParams\x1a\x41\n\x13RayStartParamsEntry\x12\x10\n\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n\x05value\x18\x02 \x01(\tR\x05value:\x02\x38\x01\x42\xc0\x01\n\x14\x63om.flyteidl.pluginsB\x08RayProtoP\x01Z=github.com/flyteorg/flyte/flyteidl/gen/pb-go/flyteidl/plugins\xa2\x02\x03\x46PX\xaa\x02\x10\x46lyteidl.Plugins\xca\x02\x10\x46lyteidl\\Plugins\xe2\x02\x1c\x46lyteidl\\Plugins\\GPBMetadata\xea\x02\x11\x46lyteidl::Pluginsb\x06proto3')
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x1a\x66lyteidl/plugins/ray.proto\x12\x10\x66lyteidl.plugins\x1a\x19\x66lyteidl/core/tasks.proto\"\x92\x02\n\x06RayJob\x12=\n\x0bray_cluster\x18\x01 \x01(\x0b\x32\x1c.flyteidl.plugins.RayClusterR\nrayCluster\x12#\n\x0bruntime_env\x18\x02 \x01(\tB\x02\x18\x01R\nruntimeEnv\x12=\n\x1bshutdown_after_job_finishes\x18\x03 \x01(\x08R\x18shutdownAfterJobFinishes\x12;\n\x1attl_seconds_after_finished\x18\x04 \x01(\x05R\x17ttlSecondsAfterFinished\x12(\n\x10runtime_env_yaml\x18\x05 \x01(\tR\x0eruntimeEnvYaml\"\xd3\x01\n\nRayCluster\x12G\n\x0fhead_group_spec\x18\x01 \x01(\x0b\x32\x1f.flyteidl.plugins.HeadGroupSpecR\rheadGroupSpec\x12M\n\x11worker_group_spec\x18\x02 \x03(\x0b\x32!.flyteidl.plugins.WorkerGroupSpecR\x0fworkerGroupSpec\x12-\n\x12\x65nable_autoscaling\x18\x03 \x01(\x08R\x11\x65nableAutoscaling\"\xe1\x01\n\rHeadGroupSpec\x12]\n\x10ray_start_params\x18\x01 \x03(\x0b\x32\x33.flyteidl.plugins.HeadGroupSpec.RayStartParamsEntryR\x0erayStartParams\x12.\n\x07k8s_pod\x18\x02 \x01(\x0b\x32\x15.flyteidl.core.K8sPodR\x06k8sPod\x1a\x41\n\x13RayStartParamsEntry\x12\x10\n\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n\x05value\x18\x02 \x01(\tR\x05value:\x02\x38\x01\"\xe6\x02\n\x0fWorkerGroupSpec\x12\x1d\n\ngroup_name\x18\x01 \x01(\tR\tgroupName\x12\x1a\n\x08replicas\x18\x02 \x01(\x05R\x08replicas\x12!\n\x0cmin_replicas\x18\x03 \x01(\x05R\x0bminReplicas\x12!\n\x0cmax_replicas\x18\x04 \x01(\x05R\x0bmaxReplicas\x12_\n\x10ray_start_params\x18\x05 \x03(\x0b\x32\x35.flyteidl.plugins.WorkerGroupSpec.RayStartParamsEntryR\x0erayStartParams\x12.\n\x07k8s_pod\x18\x06 \x01(\x0b\x32\x15.flyteidl.core.K8sPodR\x06k8sPod\x1a\x41\n\x13RayStartParamsEntry\x12\x10\n\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n\x05value\x18\x02 \x01(\tR\x05value:\x02\x38\x01\x42\xc0\x01\n\x14\x63om.flyteidl.pluginsB\x08RayProtoP\x01Z=github.com/flyteorg/flyte/flyteidl/gen/pb-go/flyteidl/plugins\xa2\x02\x03\x46PX\xaa\x02\x10\x46lyteidl.Plugins\xca\x02\x10\x46lyteidl\\Plugins\xe2\x02\x1c\x46lyteidl\\Plugins\\GPBMetadata\xea\x02\x11\x46lyteidl::Pluginsb\x06proto3')
 
 _globals = globals()
 _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
@@ -28,16 +29,16 @@
   _HEADGROUPSPEC_RAYSTARTPARAMSENTRY._serialized_options = b'8\001'
   _WORKERGROUPSPEC_RAYSTARTPARAMSENTRY._options = None
   _WORKERGROUPSPEC_RAYSTARTPARAMSENTRY._serialized_options = b'8\001'
-  _globals['_RAYJOB']._serialized_start=49
-  _globals['_RAYJOB']._serialized_end=323
-  _globals['_RAYCLUSTER']._serialized_start=326
-  _globals['_RAYCLUSTER']._serialized_end=537
-  _globals['_HEADGROUPSPEC']._serialized_start=540
-  _globals['_HEADGROUPSPEC']._serialized_end=717
-  _globals['_HEADGROUPSPEC_RAYSTARTPARAMSENTRY']._serialized_start=652
-  _globals['_HEADGROUPSPEC_RAYSTARTPARAMSENTRY']._serialized_end=717
-  _globals['_WORKERGROUPSPEC']._serialized_start=720
-  _globals['_WORKERGROUPSPEC']._serialized_end=1030
-  _globals['_WORKERGROUPSPEC_RAYSTARTPARAMSENTRY']._serialized_start=652
-  _globals['_WORKERGROUPSPEC_RAYSTARTPARAMSENTRY']._serialized_end=717
+  _globals['_RAYJOB']._serialized_start=76
+  _globals['_RAYJOB']._serialized_end=350
+  _globals['_RAYCLUSTER']._serialized_start=353
+  _globals['_RAYCLUSTER']._serialized_end=564
+  _globals['_HEADGROUPSPEC']._serialized_start=567
+  _globals['_HEADGROUPSPEC']._serialized_end=792
+  _globals['_HEADGROUPSPEC_RAYSTARTPARAMSENTRY']._serialized_start=727
+  _globals['_HEADGROUPSPEC_RAYSTARTPARAMSENTRY']._serialized_end=792
+  _globals['_WORKERGROUPSPEC']._serialized_start=795
+  _globals['_WORKERGROUPSPEC']._serialized_end=1153
+  _globals['_WORKERGROUPSPEC_RAYSTARTPARAMSENTRY']._serialized_start=727
+  _globals['_WORKERGROUPSPEC_RAYSTARTPARAMSENTRY']._serialized_end=792
 # @@protoc_insertion_point(module_scope)
diff --git a/flyteidl/gen/pb_python/flyteidl/plugins/ray_pb2.pyi b/flyteidl/gen/pb_python/flyteidl/plugins/ray_pb2.pyi
index 81d7382063..239e2fbc1a 100644
--- a/flyteidl/gen/pb_python/flyteidl/plugins/ray_pb2.pyi
+++ b/flyteidl/gen/pb_python/flyteidl/plugins/ray_pb2.pyi
@@ -1,3 +1,4 @@
+from flyteidl.core import tasks_pb2 as _tasks_pb2
 from google.protobuf.internal import containers as _containers
 from google.protobuf import descriptor as _descriptor
 from google.protobuf import message as _message
@@ -30,7 +31,7 @@ class RayCluster(_message.Message):
     def __init__(self, head_group_spec: _Optional[_Union[HeadGroupSpec, _Mapping]] = ..., worker_group_spec: _Optional[_Iterable[_Union[WorkerGroupSpec, _Mapping]]] = ..., enable_autoscaling: bool = ...) -> None: ...
 
 class HeadGroupSpec(_message.Message):
-    __slots__ = ["ray_start_params"]
+    __slots__ = ["ray_start_params", "k8s_pod"]
     class RayStartParamsEntry(_message.Message):
         __slots__ = ["key", "value"]
         KEY_FIELD_NUMBER: _ClassVar[int]
@@ -39,11 +40,13 @@ class HeadGroupSpec(_message.Message):
         value: str
         def __init__(self, key: _Optional[str] = ..., value: _Optional[str] = ...) -> None: ...
     RAY_START_PARAMS_FIELD_NUMBER: _ClassVar[int]
+    K8S_POD_FIELD_NUMBER: _ClassVar[int]
     ray_start_params: _containers.ScalarMap[str, str]
-    def __init__(self, ray_start_params: _Optional[_Mapping[str, str]] = ...) -> None: ...
+    k8s_pod: _tasks_pb2.K8sPod
+    def __init__(self, ray_start_params: _Optional[_Mapping[str, str]] = ..., k8s_pod: _Optional[_Union[_tasks_pb2.K8sPod, _Mapping]] = ...) -> None: ...
 
 class WorkerGroupSpec(_message.Message):
-    __slots__ = ["group_name", "replicas", "min_replicas", "max_replicas", "ray_start_params"]
+    __slots__ = ["group_name", "replicas", "min_replicas", "max_replicas", "ray_start_params", "k8s_pod"]
     class RayStartParamsEntry(_message.Message):
         __slots__ = ["key", "value"]
         KEY_FIELD_NUMBER: _ClassVar[int]
@@ -56,9 +59,11 @@ class WorkerGroupSpec(_message.Message):
     MIN_REPLICAS_FIELD_NUMBER: _ClassVar[int]
     MAX_REPLICAS_FIELD_NUMBER: _ClassVar[int]
     RAY_START_PARAMS_FIELD_NUMBER: _ClassVar[int]
+    K8S_POD_FIELD_NUMBER: _ClassVar[int]
     group_name: str
     replicas: int
     min_replicas: int
     max_replicas: int
     ray_start_params: _containers.ScalarMap[str, str]
-    def __init__(self, group_name: _Optional[str] = ..., replicas: _Optional[int] = ..., min_replicas: _Optional[int] = ..., max_replicas: _Optional[int] = ..., ray_start_params: _Optional[_Mapping[str, str]] = ...) -> None: ...
+    k8s_pod: _tasks_pb2.K8sPod
+    def __init__(self, group_name: _Optional[str] = ..., replicas: _Optional[int] = ..., min_replicas: _Optional[int] = ..., max_replicas: _Optional[int] = ..., ray_start_params: _Optional[_Mapping[str, str]] = ..., k8s_pod: _Optional[_Union[_tasks_pb2.K8sPod, _Mapping]] = ...) -> None: ...
diff --git a/flyteidl/gen/pb_rust/flyteidl.plugins.rs b/flyteidl/gen/pb_rust/flyteidl.plugins.rs
index 0252c9d882..65f187c3e0 100644
--- a/flyteidl/gen/pb_rust/flyteidl.plugins.rs
+++ b/flyteidl/gen/pb_rust/flyteidl.plugins.rs
@@ -255,6 +255,9 @@ pub struct HeadGroupSpec {
     /// Refer to <https://docs.ray.io/en/latest/ray-core/package-ref.html#ray-start>
     #[prost(map="string, string", tag="1")]
     pub ray_start_params: ::std::collections::HashMap<::prost::alloc::string::String, ::prost::alloc::string::String>,
+    /// Pod Spec for the ray head pod
+    #[prost(message, optional, tag="2")]
+    pub k8s_pod: ::core::option::Option<super::core::K8sPod>,
 }
 /// WorkerGroupSpec are the specs for the worker pods
 #[allow(clippy::derive_partial_eq_without_eq)]
@@ -276,6 +279,9 @@ pub struct WorkerGroupSpec {
     /// Refer to <https://docs.ray.io/en/latest/ray-core/package-ref.html#ray-start>
     #[prost(map="string, string", tag="5")]
     pub ray_start_params: ::std::collections::HashMap<::prost::alloc::string::String, ::prost::alloc::string::String>,
+    /// Pod Spec for ray worker pods
+    #[prost(message, optional, tag="6")]
+    pub k8s_pod: ::core::option::Option<super::core::K8sPod>,
 }
 #[allow(clippy::derive_partial_eq_without_eq)]
 #[derive(Clone, Copy, PartialEq, ::prost::Message)]
diff --git a/flyteidl/protos/flyteidl/plugins/ray.proto b/flyteidl/protos/flyteidl/plugins/ray.proto
index c20c6360e7..749444ee04 100644
--- 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;
+  // Pod Spec for the ray head pod
+  core.K8sPod k8s_pod = 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;
+  // Pod Spec for ray worker pods
+  core.K8sPod k8s_pod = 6;
 }
diff --git a/flyteplugins/go/tasks/plugins/k8s/ray/ray.go b/flyteplugins/go/tasks/plugins/k8s/ray/ray.go
index 90388b46a5..675f32c99c 100644
--- a/flyteplugins/go/tasks/plugins/k8s/ray/ray.go
+++ b/flyteplugins/go/tasks/plugins/k8s/ray/ray.go
@@ -118,24 +118,30 @@ func (rayJobResourceHandler) BuildResource(ctx context.Context, taskCtx pluginsC
 
 	podSpec.ServiceAccountName = cfg.ServiceAccount
 
-	headPodSpec := podSpec.DeepCopy()
-
-	rayjob, err := constructRayJob(taskCtx, rayJob, objectMeta, *podSpec, headPodSpec, headNodeRayStartParams, primaryContainerIdx, *primaryContainer)
+	rayjob, err := constructRayJob(taskCtx, &rayJob, objectMeta, *podSpec, headNodeRayStartParams, primaryContainerIdx, *primaryContainer)
 
 	return rayjob, err
 }
 
-func constructRayJob(taskCtx pluginsCore.TaskExecutionContext, rayJob plugins.RayJob, objectMeta *metav1.ObjectMeta, podSpec v1.PodSpec, headPodSpec *v1.PodSpec, headNodeRayStartParams map[string]string, primaryContainerIdx int, primaryContainer v1.Container) (*rayv1.RayJob, error) {
+func constructRayJob(taskCtx pluginsCore.TaskExecutionContext, rayJob *plugins.RayJob, objectMeta *metav1.ObjectMeta, taskPodSpec v1.PodSpec, headNodeRayStartParams map[string]string, primaryContainerIdx int, primaryContainer v1.Container) (*rayv1.RayJob, error) {
 	enableIngress := true
 	cfg := GetConfig()
+
+	headPodSpec := taskPodSpec.DeepCopy()
+	headPodTemplate, err := buildHeadPodTemplate(
+		&headPodSpec.Containers[primaryContainerIdx],
+		headPodSpec,
+		objectMeta,
+		taskCtx,
+		rayJob.RayCluster.HeadGroupSpec,
+	)
+	if err != nil {
+		return nil, err
+	}
+
 	rayClusterSpec := rayv1.RayClusterSpec{
 		HeadGroupSpec: rayv1.HeadGroupSpec{
-			Template: buildHeadPodTemplate(
-				&headPodSpec.Containers[primaryContainerIdx],
-				headPodSpec,
-				objectMeta,
-				taskCtx,
-			),
+			Template:       headPodTemplate,
 			ServiceType:    v1.ServiceType(cfg.ServiceType),
 			EnableIngress:  &enableIngress,
 			RayStartParams: headNodeRayStartParams,
@@ -145,13 +151,17 @@ func constructRayJob(taskCtx pluginsCore.TaskExecutionContext, rayJob plugins.Ra
 	}
 
 	for _, spec := range rayJob.RayCluster.WorkerGroupSpec {
-		workerPodSpec := podSpec.DeepCopy()
-		workerPodTemplate := buildWorkerPodTemplate(
+		workerPodSpec := taskPodSpec.DeepCopy()
+		workerPodTemplate, err := buildWorkerPodTemplate(
 			&workerPodSpec.Containers[primaryContainerIdx],
 			workerPodSpec,
 			objectMeta,
 			taskCtx,
+			spec,
 		)
+		if err != nil {
+			return nil, err
+		}
 
 		workerNodeRayStartParams := make(map[string]string)
 		if spec.RayStartParams != nil {
@@ -206,10 +216,10 @@ func constructRayJob(taskCtx pluginsCore.TaskExecutionContext, rayJob plugins.Ra
 		ttlSecondsAfterFinished = &rayJob.TtlSecondsAfterFinished
 	}
 
-	submitterPodTemplate := buildSubmitterPodTemplate(headPodSpec, objectMeta, taskCtx)
+	submitterPodSpec := taskPodSpec.DeepCopy()
+	submitterPodTemplate := buildSubmitterPodTemplate(submitterPodSpec, objectMeta, taskCtx)
 
 	// TODO: This is for backward compatibility. Remove this block once runtime_env is removed from ray proto.
-	var err error
 	var runtimeEnvYaml string
 	runtimeEnvYaml = rayJob.RuntimeEnvYaml
 	// If runtime_env exists but runtime_env_yaml does not, convert runtime_env to runtime_env_yaml
@@ -319,7 +329,7 @@ func injectLogsSidecar(primaryContainer *v1.Container, podSpec *v1.PodSpec) {
 	podSpec.Containers = append(podSpec.Containers, *sidecar)
 }
 
-func buildHeadPodTemplate(primaryContainer *v1.Container, podSpec *v1.PodSpec, objectMeta *metav1.ObjectMeta, taskCtx pluginsCore.TaskExecutionContext) v1.PodTemplateSpec {
+func buildHeadPodTemplate(primaryContainer *v1.Container, basePodSpec *v1.PodSpec, objectMeta *metav1.ObjectMeta, taskCtx pluginsCore.TaskExecutionContext, spec *plugins.HeadGroupSpec) (v1.PodTemplateSpec, error) {
 	// Some configs are copy from  https://github.com/ray-project/kuberay/blob/b72e6bdcd9b8c77a9dc6b5da8560910f3a0c3ffd/apiserver/pkg/util/cluster.go#L97
 	// They should always be the same, so we could hard code here.
 	primaryContainer.Name = "ray-head"
@@ -357,30 +367,38 @@ func buildHeadPodTemplate(primaryContainer *v1.Container, podSpec *v1.PodSpec, o
 	primaryContainer.Ports = append(primaryContainer.Ports, ports...)
 
 	// Inject a sidecar for capturing and exposing Ray job logs
-	injectLogsSidecar(primaryContainer, podSpec)
+	injectLogsSidecar(primaryContainer, basePodSpec)
+
+	basePodSpec, err := mergeCustomPodSpec(primaryContainer, basePodSpec, spec.K8SPod)
+	if err != nil {
+		return v1.PodTemplateSpec{}, err
+	}
 
 	podTemplateSpec := v1.PodTemplateSpec{
-		Spec:       *podSpec,
+		Spec:       *basePodSpec,
 		ObjectMeta: *objectMeta,
 	}
 	cfg := config.GetK8sPluginConfig()
 	podTemplateSpec.SetLabels(utils.UnionMaps(cfg.DefaultLabels, podTemplateSpec.GetLabels(), utils.CopyMap(taskCtx.TaskExecutionMetadata().GetLabels())))
 	podTemplateSpec.SetAnnotations(utils.UnionMaps(cfg.DefaultAnnotations, podTemplateSpec.GetAnnotations(), utils.CopyMap(taskCtx.TaskExecutionMetadata().GetAnnotations())))
-	return podTemplateSpec
+	return podTemplateSpec, nil
 }
 
 func buildSubmitterPodTemplate(podSpec *v1.PodSpec, objectMeta *metav1.ObjectMeta, taskCtx pluginsCore.TaskExecutionContext) v1.PodTemplateSpec {
+	submitterPodSpec := podSpec.DeepCopy()
+
 	podTemplateSpec := v1.PodTemplateSpec{
-		Spec:       *podSpec,
 		ObjectMeta: *objectMeta,
+		Spec:       *submitterPodSpec,
 	}
+
 	cfg := config.GetK8sPluginConfig()
 	podTemplateSpec.SetLabels(utils.UnionMaps(cfg.DefaultLabels, podTemplateSpec.GetLabels(), utils.CopyMap(taskCtx.TaskExecutionMetadata().GetLabels())))
 	podTemplateSpec.SetAnnotations(utils.UnionMaps(cfg.DefaultAnnotations, podTemplateSpec.GetAnnotations(), utils.CopyMap(taskCtx.TaskExecutionMetadata().GetAnnotations())))
 	return podTemplateSpec
 }
 
-func buildWorkerPodTemplate(primaryContainer *v1.Container, podSpec *v1.PodSpec, objectMetadata *metav1.ObjectMeta, taskCtx pluginsCore.TaskExecutionContext) v1.PodTemplateSpec {
+func buildWorkerPodTemplate(primaryContainer *v1.Container, basePodSpec *v1.PodSpec, objectMetadata *metav1.ObjectMeta, taskCtx pluginsCore.TaskExecutionContext, spec *plugins.WorkerGroupSpec) (v1.PodTemplateSpec, error) {
 	// Some configs are copy from  https://github.com/ray-project/kuberay/blob/b72e6bdcd9b8c77a9dc6b5da8560910f3a0c3ffd/apiserver/pkg/util/cluster.go#L185
 	// They should always be the same, so we could hard code here.
 
@@ -479,13 +497,50 @@ func buildWorkerPodTemplate(primaryContainer *v1.Container, podSpec *v1.PodSpec,
 	}
 	primaryContainer.Ports = append(primaryContainer.Ports, ports...)
 
+	basePodSpec, err := mergeCustomPodSpec(primaryContainer, basePodSpec, spec.K8SPod)
+	if err != nil {
+		return v1.PodTemplateSpec{}, err
+	}
+
 	podTemplateSpec := v1.PodTemplateSpec{
-		Spec:       *podSpec,
+		Spec:       *basePodSpec,
 		ObjectMeta: *objectMetadata,
 	}
 	podTemplateSpec.SetLabels(utils.UnionMaps(podTemplateSpec.GetLabels(), utils.CopyMap(taskCtx.TaskExecutionMetadata().GetLabels())))
 	podTemplateSpec.SetAnnotations(utils.UnionMaps(podTemplateSpec.GetAnnotations(), utils.CopyMap(taskCtx.TaskExecutionMetadata().GetAnnotations())))
-	return podTemplateSpec
+	return podTemplateSpec, nil
+}
+
+// Merges a ray head/worker node custom pod specs onto task's generated pod spec
+func mergeCustomPodSpec(primaryContainer *v1.Container, podSpec *v1.PodSpec, k8sPod *core.K8SPod) (*v1.PodSpec, error) {
+	if k8sPod == nil {
+		return podSpec, nil
+	}
+
+	if k8sPod.PodSpec == nil {
+		return podSpec, nil
+	}
+
+	var customPodSpec *v1.PodSpec
+
+	err := utils.UnmarshalStructToObj(k8sPod.PodSpec, &customPodSpec)
+	if err != nil {
+		return nil, flyteerr.Errorf(flyteerr.BadTaskSpecification,
+			"Unable to unmarshal pod spec [%v], Err: [%v]", k8sPod.PodSpec, err.Error())
+	}
+
+	for _, container := range customPodSpec.Containers {
+		if container.Name != primaryContainer.Name { // Only support the primary container for now
+			continue
+		}
+
+		// Just handle resources for now
+		if len(container.Resources.Requests) > 0 || len(container.Resources.Limits) > 0 {
+			primaryContainer.Resources = container.Resources
+		}
+	}
+
+	return podSpec, nil
 }
 
 func (rayJobResourceHandler) BuildIdentityResource(ctx context.Context, taskCtx pluginsCore.TaskExecutionMetadata) (client.Object, error) {


Test patch:
diff --git a/flyteplugins/go/tasks/plugins/k8s/ray/ray_test.go b/flyteplugins/go/tasks/plugins/k8s/ray/ray_test.go
index 7b555e9f23..19d3ef1b10 100644
--- a/flyteplugins/go/tasks/plugins/k8s/ray/ray_test.go
+++ b/flyteplugins/go/tasks/plugins/k8s/ray/ray_test.go
@@ -2,18 +2,19 @@ package ray
 
 import (
 	"context"
+	"encoding/json"
 	"reflect"
 	"testing"
 	"time"
 
-	structpb "github.com/golang/protobuf/ptypes/struct"
 	rayv1 "github.com/ray-project/kuberay/ray-operator/apis/ray/v1"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/require"
+	"google.golang.org/protobuf/types/known/structpb"
 	corev1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/api/resource"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
-	v1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 
 	"github.com/flyteorg/flyte/flyteidl/gen/pb-go/flyteidl/core"
 	"github.com/flyteorg/flyte/flyteidl/gen/pb-go/flyteidl/plugins"
@@ -148,7 +149,7 @@ func dummyRayTaskContext(taskTemplate *core.TaskTemplate, resources *corev1.Reso
 	taskExecutionMetadata.OnGetNamespace().Return("test-namespace")
 	taskExecutionMetadata.OnGetAnnotations().Return(map[string]string{"annotation-1": "val1"})
 	taskExecutionMetadata.OnGetLabels().Return(map[string]string{"label-1": "val1"})
-	taskExecutionMetadata.OnGetOwnerReference().Return(v1.OwnerReference{
+	taskExecutionMetadata.OnGetOwnerReference().Return(metav1.OwnerReference{
 		Kind: "node",
 		Name: "blah",
 	})
@@ -419,6 +420,123 @@ func TestBuildResourceRayExtendedResources(t *testing.T) {
 	}
 }
 
+func TestBuildResourceRayCustomK8SPod(t *testing.T) {
+	assert.NoError(t, config.SetK8sPluginConfig(&config.K8sPluginConfig{}))
+
+	headResourceEntries := []*core.Resources_ResourceEntry{
+		{Name: core.Resources_CPU, Value: "10"},
+		{Name: core.Resources_MEMORY, Value: "10Gi"},
+		{Name: core.Resources_GPU, Value: "10"},
+	}
+	headResources := &core.Resources{Requests: headResourceEntries, Limits: headResourceEntries}
+
+	expectedHeadResources, err := flytek8s.ToK8sResourceRequirements(headResources)
+	require.NoError(t, err)
+
+	workerResourceEntries := []*core.Resources_ResourceEntry{
+		{Name: core.Resources_CPU, Value: "20"},
+		{Name: core.Resources_MEMORY, Value: "20Gi"},
+		{Name: core.Resources_GPU, Value: "20"},
+	}
+	workerResources := &core.Resources{Requests: workerResourceEntries, Limits: workerResourceEntries}
+
+	expectedWorkerResources, err := flytek8s.ToK8sResourceRequirements(workerResources)
+	require.NoError(t, err)
+
+	headPodSpec := &corev1.PodSpec{
+		Containers: []corev1.Container{
+			{
+				Name:      "ray-head",
+				Resources: *expectedHeadResources,
+			},
+		},
+	}
+	workerPodSpec := &corev1.PodSpec{
+		Containers: []corev1.Container{
+			{
+				Name:      "ray-worker",
+				Resources: *expectedWorkerResources,
+			},
+		},
+	}
+
+	params := []struct {
+		name                       string
+		taskResources              *corev1.ResourceRequirements
+		headK8SPod                 *core.K8SPod
+		workerK8SPod               *core.K8SPod
+		expectedSubmitterResources *corev1.ResourceRequirements
+		expectedHeadResources      *corev1.ResourceRequirements
+		expectedWorkerResources    *corev1.ResourceRequirements
+	}{
+		{
+			name:                       "task resources",
+			taskResources:              resourceRequirements,
+			expectedSubmitterResources: resourceRequirements,
+			expectedHeadResources:      resourceRequirements,
+			expectedWorkerResources:    resourceRequirements,
+		},
+		{
+			name:          "custom worker and head resources",
+			taskResources: resourceRequirements,
+			headK8SPod: &core.K8SPod{
+				PodSpec: transformStructToStructPB(t, headPodSpec),
+			},
+			workerK8SPod: &core.K8SPod{
+				PodSpec: transformStructToStructPB(t, workerPodSpec),
+			},
+			expectedSubmitterResources: resourceRequirements,
+			expectedHeadResources:      expectedHeadResources,
+			expectedWorkerResources:    expectedWorkerResources,
+		},
+	}
+
+	for _, p := range params {
+		t.Run(p.name, func(t *testing.T) {
+			rayJobInput := dummyRayCustomObj()
+
+			if p.headK8SPod != nil {
+				rayJobInput.RayCluster.HeadGroupSpec.K8SPod = p.headK8SPod
+			}
+
+			if p.workerK8SPod != nil {
+				for _, spec := range rayJobInput.RayCluster.WorkerGroupSpec {
+					spec.K8SPod = p.workerK8SPod
+				}
+			}
+
+			taskTemplate := dummyRayTaskTemplate("ray-id", rayJobInput)
+			taskContext := dummyRayTaskContext(taskTemplate, p.taskResources, nil, "", serviceAccount)
+			rayJobResourceHandler := rayJobResourceHandler{}
+			r, err := rayJobResourceHandler.BuildResource(context.TODO(), taskContext)
+			assert.Nil(t, err)
+			assert.NotNil(t, r)
+			rayJob, ok := r.(*rayv1.RayJob)
+			assert.True(t, ok)
+
+			submitterPodResources := rayJob.Spec.SubmitterPodTemplate.Spec.Containers[0].Resources
+			assert.EqualValues(t,
+				p.expectedSubmitterResources,
+				&submitterPodResources,
+			)
+
+			headPodResources := rayJob.Spec.RayClusterSpec.HeadGroupSpec.Template.Spec.Containers[0].Resources
+			assert.EqualValues(t,
+				p.expectedHeadResources,
+				&headPodResources,
+			)
+
+			for _, workerGroupSpec := range rayJob.Spec.RayClusterSpec.WorkerGroupSpecs {
+				workerPodResources := workerGroupSpec.Template.Spec.Containers[0].Resources
+				assert.EqualValues(t,
+					p.expectedWorkerResources,
+					&workerPodResources,
+				)
+			}
+		})
+	}
+}
+
 func TestDefaultStartParameters(t *testing.T) {
 	rayJobResourceHandler := rayJobResourceHandler{}
 	rayJob := &plugins.RayJob{
@@ -1102,3 +1220,14 @@ func TestGetPropertiesRay(t *testing.T) {
 	expected := k8s.PluginProperties{}
 	assert.Equal(t, expected, rayJobResourceHandler.GetProperties())
 }
+
+func transformStructToStructPB(t *testing.T, obj interface{}) *structpb.Struct {
+	data, err := json.Marshal(obj)
+	assert.Nil(t, err)
+	podSpecMap := make(map[string]interface{})
+	err = json.Unmarshal(data, &podSpecMap)
+	assert.Nil(t, err)
+	s, err := structpb.NewStruct(podSpecMap)
+	assert.Nil(t, err)
+	return s
+}