From befe83f8774f634c26a0e4ef34f93d4ae07f80a1 Mon Sep 17 00:00:00 2001 From: 1fanwang <1fannnw@gmail.com> Date: Thu, 13 Aug 2026 11:50:31 -0700 Subject: [PATCH] fix(plugins/ray): stop RayJobs sharing the plugin's start-params map When a task leaves RayStartParams unset, BuildResource assigns the plugin's configured defaults map directly to the local variable and then writes include-dashboard, node-ip-address and dashboard-host into it. That map belongs to the process-wide plugin config, and it is also handed to the RayJob CR, so every RayJob built from the defaults shares one map with the config and with each other. Copy the resolved parameters into the map that was already allocated, for both the head group and the worker groups. Signed-off-by: 1fanwang <1fannnw@gmail.com> --- flyteplugins/go/tasks/plugins/k8s/ray/ray.go | 12 +- .../k8s/ray/sharedconfig_repro_test.go | 130 ++++++++++++++++++ 2 files changed, 138 insertions(+), 4 deletions(-) create mode 100644 flyteplugins/go/tasks/plugins/k8s/ray/sharedconfig_repro_test.go diff --git a/flyteplugins/go/tasks/plugins/k8s/ray/ray.go b/flyteplugins/go/tasks/plugins/k8s/ray/ray.go index aa32b847ba9..bc4d22bd02a 100644 --- a/flyteplugins/go/tasks/plugins/k8s/ray/ray.go +++ b/flyteplugins/go/tasks/plugins/k8s/ray/ray.go @@ -5,6 +5,7 @@ import ( "encoding/base64" "encoding/json" "fmt" + "maps" "regexp" "strconv" "strings" @@ -94,11 +95,14 @@ func (rayJobResourceHandler) BuildResource(ctx context.Context, taskCtx pluginsC cfg := GetConfig() + // Copy rather than alias: the resolved map is filled in below and handed to the RayJob CR, so + // aliasing would let one task's start params edit the shared plugin config and every other CR + // built from those defaults. headNodeRayStartParams := make(map[string]string) if rayJob.GetRayCluster().GetHeadGroupSpec() != nil && rayJob.RayCluster.HeadGroupSpec.RayStartParams != nil { - headNodeRayStartParams = rayJob.GetRayCluster().GetHeadGroupSpec().GetRayStartParams() + maps.Copy(headNodeRayStartParams, rayJob.GetRayCluster().GetHeadGroupSpec().GetRayStartParams()) } else if headNode := cfg.Defaults.HeadNode; len(headNode.StartParameters) > 0 { - headNodeRayStartParams = headNode.StartParameters + maps.Copy(headNodeRayStartParams, headNode.StartParameters) } if _, exist := headNodeRayStartParams[IncludeDashboard]; !exist { @@ -203,9 +207,9 @@ func constructRayJob(taskCtx pluginsCore.TaskExecutionContext, rayJob *plugins.R workerNodeRayStartParams := make(map[string]string) if spec.RayStartParams != nil { - workerNodeRayStartParams = spec.GetRayStartParams() + maps.Copy(workerNodeRayStartParams, spec.GetRayStartParams()) } else if workerNode := cfg.Defaults.WorkerNode; len(workerNode.StartParameters) > 0 { - workerNodeRayStartParams = workerNode.StartParameters + maps.Copy(workerNodeRayStartParams, workerNode.StartParameters) } if _, exist := workerNodeRayStartParams[NodeIPAddress]; !exist { diff --git a/flyteplugins/go/tasks/plugins/k8s/ray/sharedconfig_repro_test.go b/flyteplugins/go/tasks/plugins/k8s/ray/sharedconfig_repro_test.go new file mode 100644 index 00000000000..4566d7c42ae --- /dev/null +++ b/flyteplugins/go/tasks/plugins/k8s/ray/sharedconfig_repro_test.go @@ -0,0 +1,130 @@ +package ray + +import ( + "context" + "maps" + "testing" + + rayv1 "github.com/ray-project/kuberay/ray-operator/apis/ray/v1" + "github.com/stretchr/testify/assert" + + "github.com/flyteorg/flyte/flyteidl/gen/pb-go/flyteidl/plugins" +) + +// A RayJob that leaves RayStartParams unset — the common case, and the one that falls through to +// the plugin's configured defaults. +func rayJobWithoutStartParams() *plugins.RayJob { + return &plugins.RayJob{ + RayCluster: &plugins.RayCluster{ + HeadGroupSpec: &plugins.HeadGroupSpec{}, + WorkerGroupSpec: []*plugins.WorkerGroupSpec{{GroupName: workerGroupName, Replicas: 1, MinReplicas: 1, MaxReplicas: 1}}, + }, + ShutdownAfterJobFinishes: true, + TtlSecondsAfterFinished: 120, + } +} + +// Other tests in this package replace the global config without restoring it, which can leave the +// default start parameters empty and skip the branch under test. Pin a config with non-empty +// defaults, and put the original back afterwards. +func withDefaultStartParams(t *testing.T) { + origConfig := *GetConfig() + t.Cleanup(func() { assert.NoError(t, SetConfig(&origConfig)) }) + + assert.NoError(t, SetConfig(&Config{ + Defaults: DefaultConfig{ + HeadNode: NodeConfig{ + StartParameters: map[string]string{DisableUsageStatsStartParameter: "true"}, + IPAddress: "$MY_POD_IP", + }, + WorkerNode: NodeConfig{ + StartParameters: map[string]string{DisableUsageStatsStartParameter: "true"}, + IPAddress: "$MY_POD_IP", + }, + }, + })) +} + +// BuildResource must not write into the plugin's shared configuration. +// +// When a task leaves RayStartParams unset, the head/worker start params fall back to +// cfg.Defaults.{Head,Worker}Node.StartParameters. Those maps belong to the process-wide plugin +// config, so filling in include-dashboard / node-ip-address / dashboard-host mutates config that +// every later task reads. +func TestBuildResource_DoesNotMutateSharedConfig(t *testing.T) { + withDefaultStartParams(t) + + cfg := GetConfig() + headDefaults := cfg.Defaults.HeadNode.StartParameters + workerDefaults := cfg.Defaults.WorkerNode.StartParameters + + headBefore := maps.Clone(headDefaults) + workerBefore := maps.Clone(workerDefaults) + + handler := rayJobResourceHandler{} + taskTemplate := dummyRayTaskTemplate("ray-id", rayJobWithoutStartParams()) + taskCtx := dummyRayTaskContext(taskTemplate, resourceRequirements, nil, "", serviceAccount) + + _, err := handler.BuildResource(context.TODO(), taskCtx) + assert.NoError(t, err) + + assert.Equal(t, headBefore, cfg.Defaults.HeadNode.StartParameters, + "BuildResource mutated the shared head-node start parameters") + assert.Equal(t, workerBefore, cfg.Defaults.WorkerNode.StartParameters, + "BuildResource mutated the shared worker-node start parameters") +} + +// Two RayJobs built independently must not share one RayStartParams map. +// +// constructRayJob assigns the resolved start params straight into the CR +// (RayStartParams: headNodeRayStartParams). When they came from the shared defaults, every CR +// built that way — and the plugin config itself — is the same map instance. +func TestBuildResource_RayJobsDoNotShareStartParamsMap(t *testing.T) { + withDefaultStartParams(t) + + handler := rayJobResourceHandler{} + + build := func() *rayv1.RayJob { + taskTemplate := dummyRayTaskTemplate("ray-id", rayJobWithoutStartParams()) + taskCtx := dummyRayTaskContext(taskTemplate, resourceRequirements, nil, "", serviceAccount) + obj, err := handler.BuildResource(context.TODO(), taskCtx) + assert.NoError(t, err) + return obj.(*rayv1.RayJob) + } + + first := build() + second := build() + + first.Spec.RayClusterSpec.HeadGroupSpec.RayStartParams["num-cpus"] = "tainted-by-first-rayjob" + + assert.NotContains(t, second.Spec.RayClusterSpec.HeadGroupSpec.RayStartParams, "num-cpus", + "editing one RayJob's start params changed another RayJob's") + assert.NotContains(t, GetConfig().Defaults.HeadNode.StartParameters, "num-cpus", + "editing a RayJob's start params changed the shared plugin config") +} + +// Start params a task supplies for itself reach the CR alongside the defaults the plugin injects. +// Covers the head and worker branches that copy the task's own map rather than the config's. +func TestBuildResource_KeepsTaskSuppliedStartParams(t *testing.T) { + withDefaultStartParams(t) + + rayJob := rayJobWithoutStartParams() + rayJob.RayCluster.HeadGroupSpec.RayStartParams = map[string]string{"num-cpus": "1"} + rayJob.RayCluster.WorkerGroupSpec[0].RayStartParams = map[string]string{"num-cpus": "2"} + + handler := rayJobResourceHandler{} + taskTemplate := dummyRayTaskTemplate("ray-id", rayJob) + taskCtx := dummyRayTaskContext(taskTemplate, resourceRequirements, nil, "", serviceAccount) + + obj, err := handler.BuildResource(context.TODO(), taskCtx) + assert.NoError(t, err) + built := obj.(*rayv1.RayJob) + + head := built.Spec.RayClusterSpec.HeadGroupSpec.RayStartParams + assert.Equal(t, "1", head["num-cpus"]) + assert.Contains(t, head, NodeIPAddress) + + worker := built.Spec.RayClusterSpec.WorkerGroupSpecs[0].RayStartParams + assert.Equal(t, "2", worker["num-cpus"]) + assert.Contains(t, worker, NodeIPAddress) +}