From 2e0379a7e0a0000da21de4b845947c3df4218dd0 Mon Sep 17 00:00:00 2001 From: 1fanwang <1fannnw@gmail.com> Date: Sat, 26 Sep 2026 12:04:39 -0700 Subject: [PATCH 1/2] Honor user error verdicts for Ray tasks Signed-off-by: 1fanwang <1fannnw@gmail.com> --- flyteplugins/go/tasks/plugins/k8s/ray/ray.go | 20 ++++ .../go/tasks/plugins/k8s/ray/ray_test.go | 107 +++++++++++++++++- 2 files changed, 125 insertions(+), 2 deletions(-) diff --git a/flyteplugins/go/tasks/plugins/k8s/ray/ray.go b/flyteplugins/go/tasks/plugins/k8s/ray/ray.go index 7dfbb5850d..196f1bceba 100644 --- a/flyteplugins/go/tasks/plugins/k8s/ray/ray.go +++ b/flyteplugins/go/tasks/plugins/k8s/ray/ray.go @@ -27,9 +27,11 @@ import ( pluginsCore "github.com/flyteorg/flyte/flyteplugins/go/tasks/pluginmachinery/core" "github.com/flyteorg/flyte/flyteplugins/go/tasks/pluginmachinery/flytek8s" "github.com/flyteorg/flyte/flyteplugins/go/tasks/pluginmachinery/flytek8s/config" + "github.com/flyteorg/flyte/flyteplugins/go/tasks/pluginmachinery/ioutils" "github.com/flyteorg/flyte/flyteplugins/go/tasks/pluginmachinery/k8s" "github.com/flyteorg/flyte/flyteplugins/go/tasks/pluginmachinery/tasklog" pluginsUtils "github.com/flyteorg/flyte/flyteplugins/go/tasks/pluginmachinery/utils" + "github.com/flyteorg/flyte/flytestdlib/logger" "github.com/flyteorg/flyte/flytestdlib/utils" ) @@ -712,6 +714,24 @@ func (plugin rayJobResourceHandler) GetTaskPhase(ctx context.Context, pluginCont case rayv1.JobDeploymentStatusFailed: failInfo := fmt.Sprintf("Failed to run Ray job %s with error: [%s] %s", rayJob.Name, rayJob.Status.Reason, rayJob.Status.Message) phaseInfo, err = pluginsCore.PhaseInfoSystemRetryableFailureWithCleanup(flyteerr.TaskFailedWithError, failInfo, info), nil + if writer := pluginContext.OutputWriter(); writer != nil { + reader := ioutils.NewRemoteFileOutputReader(ctx, pluginContext.DataStore(), writer, 0) + hasError, readErr := reader.IsError(ctx) + if readErr != nil { + logger.Warnf(ctx, "Failed to check Ray task error file; retaining system retry: %v", readErr) + } else if hasError { + taskError, readErr := reader.ReadError(ctx) + if readErr != nil { + logger.Warnf(ctx, "Failed to read Ray task error file; retaining system retry: %v", readErr) + } else if taskError.Kind == core.ExecutionError_USER { + if taskError.IsRecoverable { + phaseInfo = pluginsCore.PhaseInfoRetryableFailureWithCleanup(flyteerr.TaskFailedWithError, failInfo, info) + } else { + phaseInfo = pluginsCore.PhaseInfoFailureWithCleanup(flyteerr.TaskFailedWithError, failInfo, info) + } + } + } + } default: // We already handle all known deployment status, so this should never happen unless a future version of ray // introduced a new job status. diff --git a/flyteplugins/go/tasks/plugins/k8s/ray/ray_test.go b/flyteplugins/go/tasks/plugins/k8s/ray/ray_test.go index c8ab6e3815..752563aacc 100644 --- a/flyteplugins/go/tasks/plugins/k8s/ray/ray_test.go +++ b/flyteplugins/go/tasks/plugins/k8s/ray/ray_test.go @@ -4,6 +4,8 @@ import ( "context" "encoding/json" "fmt" + "os" + "path/filepath" "reflect" "testing" "time" @@ -26,10 +28,16 @@ import ( "github.com/flyteorg/flyte/flyteplugins/go/tasks/pluginmachinery/flytek8s" "github.com/flyteorg/flyte/flyteplugins/go/tasks/pluginmachinery/flytek8s/config" pluginIOMocks "github.com/flyteorg/flyte/flyteplugins/go/tasks/pluginmachinery/io/mocks" + "github.com/flyteorg/flyte/flyteplugins/go/tasks/pluginmachinery/ioutils" "github.com/flyteorg/flyte/flyteplugins/go/tasks/pluginmachinery/k8s" mocks2 "github.com/flyteorg/flyte/flyteplugins/go/tasks/pluginmachinery/k8s/mocks" "github.com/flyteorg/flyte/flyteplugins/go/tasks/pluginmachinery/tasklog" + "github.com/flyteorg/flyte/flytestdlib/contextutils" + "github.com/flyteorg/flyte/flytestdlib/promutils" + "github.com/flyteorg/flyte/flytestdlib/promutils/labeled" + "github.com/flyteorg/flyte/flytestdlib/storage" "github.com/flyteorg/flyte/flytestdlib/utils" + "github.com/flyteorg/stow/local" ) const ( @@ -1172,7 +1180,8 @@ func TestGetTaskPhaseFailedRetryable(t *testing.T) { // reason and message from the RayJob status. ctx := context.Background() rayJobResourceHandler := rayJobResourceHandler{} - pluginCtx := newPluginContext(k8s.PluginState{}) + pluginCtx := newPluginContext(k8s.PluginState{}).(*mocks2.PluginContext) + pluginCtx.EXPECT().OutputWriter().Return(nil).Maybe() rayObject := &rayv1.RayJob{ ObjectMeta: metav1.ObjectMeta{ @@ -1191,6 +1200,99 @@ func TestGetTaskPhaseFailedRetryable(t *testing.T) { assert.Contains(t, phaseInfo.Err().GetMessage(), "head node ran out of memory") } +func TestGetTaskPhaseTaskError(t *testing.T) { + labeled.SetMetricKeys(contextutils.ExecIDKey) + cases := []struct { + name string + writeError bool + corrupt bool + oversized bool + origin core.ExecutionError_ErrorKind + recoverable bool + expectedPhase pluginsCore.Phase + expectedKind core.ExecutionError_ErrorKind + }{ + { + name: "permanent user error", writeError: true, origin: core.ExecutionError_USER, + expectedPhase: pluginsCore.PhasePermanentFailure, expectedKind: core.ExecutionError_USER, + }, + { + name: "recoverable user error", writeError: true, origin: core.ExecutionError_USER, recoverable: true, + expectedPhase: pluginsCore.PhaseRetryableFailure, expectedKind: core.ExecutionError_USER, + }, + { + name: "system error keeps system retry", writeError: true, origin: core.ExecutionError_SYSTEM, + expectedPhase: pluginsCore.PhaseRetryableFailure, expectedKind: core.ExecutionError_SYSTEM, + }, + { + name: "missing error keeps system retry", + expectedPhase: pluginsCore.PhaseRetryableFailure, expectedKind: core.ExecutionError_SYSTEM, + }, + { + name: "corrupt error keeps system retry", corrupt: true, + expectedPhase: pluginsCore.PhaseRetryableFailure, expectedKind: core.ExecutionError_SYSTEM, + }, + { + name: "oversized error keeps system retry", oversized: true, + expectedPhase: pluginsCore.PhaseRetryableFailure, expectedKind: core.ExecutionError_SYSTEM, + }, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + ctx := context.Background() + directory := t.TempDir() + store, err := storage.NewDataStore(&storage.Config{ + Type: storage.TypeLocal, + InitContainer: directory, + MultiContainerEnabled: true, + Stow: storage.StowConfig{ + Kind: local.Kind, + Config: map[string]string{local.ConfigKeyPath: "/"}, + }, + Limits: storage.LimitsConfig{GetLimitMegabytes: 2}, + }, promutils.NewTestScope()) + require.NoError(t, err) + paths := ioutils.NewReadOnlyOutputFilePaths(ctx, store, storage.DataReference("file://"+directory)) + writer := ioutils.NewRemoteFileOutputWriter(ctx, store, paths) + if tc.writeError { + kind := core.ContainerError_NON_RECOVERABLE + if tc.recoverable { + kind = core.ContainerError_RECOVERABLE + } + require.NoError(t, store.WriteProtobuf(ctx, paths.GetErrorPath(), storage.Options{}, &core.ErrorDocument{ + Error: &core.ContainerError{Code: "task-error", Message: "task failed", Kind: kind, Origin: tc.origin}, + })) + } + if tc.corrupt { + require.NoError(t, os.WriteFile(filepath.Join(directory, ioutils.ErrorsSuffix), []byte{0xff}, 0o600)) + } + if tc.oversized { + file, err := os.Create(filepath.Join(directory, ioutils.ErrorsSuffix)) + require.NoError(t, err) + require.NoError(t, file.Truncate(storage.GetConfig().Limits.GetLimitMegabytes*storage.MiB+1)) + require.NoError(t, file.Close()) + } + pluginContext := newPluginContext(k8s.PluginState{}).(*mocks2.PluginContext) + pluginContext.EXPECT().OutputWriter().Return(writer) + pluginContext.EXPECT().DataStore().Return(store) + job := &rayv1.RayJob{ + ObjectMeta: metav1.ObjectMeta{Name: "failed-ray-job"}, + Status: rayv1.RayJobStatus{ + JobDeploymentStatus: rayv1.JobDeploymentStatusFailed, + Reason: rayv1.AppFailed, + Message: "driver exited", + }, + } + phase, err := (rayJobResourceHandler{}).GetTaskPhase(ctx, pluginContext, job) + require.NoError(t, err) + assert.Equal(t, tc.expectedPhase, phase.Phase()) + require.NotNil(t, phase.Err()) + assert.Equal(t, tc.expectedKind, phase.Err().Kind) + assert.True(t, phase.CleanupOnFailure()) + }) + } +} + func newPluginContext(pluginState k8s.PluginState) k8s.PluginContext { plg := &mocks2.PluginContext{} @@ -1252,7 +1354,8 @@ func init() { func TestGetTaskPhase(t *testing.T) { ctx := context.Background() rayJobResourceHandler := rayJobResourceHandler{} - pluginCtx := newPluginContext(k8s.PluginState{}) + pluginCtx := newPluginContext(k8s.PluginState{}).(*mocks2.PluginContext) + pluginCtx.EXPECT().OutputWriter().Return(nil).Maybe() testCases := []struct { rayJobPhase rayv1.JobDeploymentStatus From bb545593cb5df6fb7c30b06888b43d708c1a8cb3 Mon Sep 17 00:00:00 2001 From: 1fanwang <1fannnw@gmail.com> Date: Sat, 26 Sep 2026 13:02:24 -0700 Subject: [PATCH 2/2] Use the protobuf getter in Ray verdict assertions Signed-off-by: 1fanwang <1fannnw@gmail.com> --- flyteplugins/go/tasks/plugins/k8s/ray/ray_test.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/flyteplugins/go/tasks/plugins/k8s/ray/ray_test.go b/flyteplugins/go/tasks/plugins/k8s/ray/ray_test.go index 752563aacc..c6184fb1f6 100644 --- a/flyteplugins/go/tasks/plugins/k8s/ray/ray_test.go +++ b/flyteplugins/go/tasks/plugins/k8s/ray/ray_test.go @@ -1287,7 +1287,7 @@ func TestGetTaskPhaseTaskError(t *testing.T) { require.NoError(t, err) assert.Equal(t, tc.expectedPhase, phase.Phase()) require.NotNil(t, phase.Err()) - assert.Equal(t, tc.expectedKind, phase.Err().Kind) + assert.Equal(t, tc.expectedKind, phase.Err().GetKind()) assert.True(t, phase.CleanupOnFailure()) }) }