Add Get Message in Agent

Signed-off-by: Future Outlier <eric901201@gmai.com>
flyteorg · Sep 14, 2023 · 6b69092 · 6b69092
1 parent d0c1602
commit 6b69092
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 1 deletion.
diff --git a/go/tasks/plugins/webapi/agent/plugin.go b/go/tasks/plugins/webapi/agent/plugin.go
@@ -38,6 +38,7 @@ type Plugin struct {
 type ResourceWrapper struct {
 	State   admin.State
 	Outputs *flyteIdl.LiteralMap
+	Message string
 }
 
 type ResourceMetaWrapper struct {
@@ -133,6 +134,7 @@ func (p Plugin) Get(ctx context.Context, taskCtx webapi.GetContext) (latest weba
 	return &ResourceWrapper{
 		State:   res.Resource.State,
 		Outputs: res.Resource.Outputs,
+		Message: res.Resource.Message,
 	}, nil
 }
 
@@ -164,7 +166,7 @@ func (p Plugin) Status(ctx context.Context, taskCtx webapi.StatusContext) (phase
 
 	switch resource.State {
 	case admin.State_PENDING:
-		return core.PhaseInfoInitializing(time.Now(), core.DefaultPhaseVersion, "job submitted", taskInfo), nil
+		return core.PhaseInfoInitializing(time.Now(), core.DefaultPhaseVersion, resource.Message, taskInfo), nil
 	case admin.State_RUNNING:
 		return core.PhaseInfoRunning(core.DefaultPhaseVersion, taskInfo), nil
 	case admin.State_PERMANENT_FAILURE:

diff --git a/go/tasks/plugins/webapi/agent/plugin_test.go b/go/tasks/plugins/webapi/agent/plugin_test.go
@@ -107,6 +107,7 @@ func TestPlugin(t *testing.T) {
 		taskContext.On("Resource").Return(&ResourceWrapper{
 			State:   admin.State_PENDING,
 			Outputs: nil,
+			Message: "Waiting for cluster",
 		})
 
 		phase, err := plugin.Status(context.Background(), taskContext)
@@ -119,6 +120,7 @@ func TestPlugin(t *testing.T) {
 		taskContext.On("Resource").Return(&ResourceWrapper{
 			State:   admin.State_RUNNING,
 			Outputs: nil,
+			Message: "",
 		})
 
 		phase, err := plugin.Status(context.Background(), taskContext)
@@ -131,6 +133,7 @@ func TestPlugin(t *testing.T) {
 		taskContext.On("Resource").Return(&ResourceWrapper{
 			State:   admin.State_PERMANENT_FAILURE,
 			Outputs: nil,
+			Message: "",
 		})
 
 		phase, err := plugin.Status(context.Background(), taskContext)
@@ -143,6 +146,7 @@ func TestPlugin(t *testing.T) {
 		taskContext.On("Resource").Return(&ResourceWrapper{
 			State:   admin.State_RETRYABLE_FAILURE,
 			Outputs: nil,
+			Message: "",
 		})
 
 		phase, err := plugin.Status(context.Background(), taskContext)
@@ -155,6 +159,7 @@ func TestPlugin(t *testing.T) {
 		taskContext.On("Resource").Return(&ResourceWrapper{
 			State:   admin.State_SUCCEEDED,
 			Outputs: nil,
+			Message: "",
 		})
 
 		phase, err := plugin.Status(context.Background(), taskContext)
@@ -167,6 +172,7 @@ func TestPlugin(t *testing.T) {
 		taskContext.On("Resource").Return(&ResourceWrapper{
 			State:   5,
 			Outputs: nil,
+			Message: "",
 		})
 
 		phase, err := plugin.Status(context.Background(), taskContext)