issue-12432 - cleanups

Signed-off-by: Helber Belmiro <[email protected]>

Author: hbelmiro

backend/src/agent/persistence/client/pipeline_client.go

@@ -23,6 +23,7 @@ import (
 	"strings"
 	"time"
 
+	"github.com/golang/glog"
 	"google.golang.org/grpc/metadata"
 
 	api "github.com/kubeflow/pipelines/backend/api/v1beta1/go_client"
@@ -39,7 +40,7 @@ const (
 type PipelineClientInterface interface {
 	ReportWorkflow(workflow util.ExecutionSpec) error
 	ReportScheduledWorkflow(swf *util.ScheduledWorkflow) error
-	ReadArtifactForMetrics(request *util.ArtifactRequest) (*util.ArtifactResponse, error)
+	ReadArtifact(request *util.ReadArtifactRequest) (*util.ReadArtifactResponse, error)
 	ReportRunMetrics(request *api.ReportRunMetricsRequest) (*api.ReportRunMetricsResponse, error)
 }
 
@@ -122,7 +123,10 @@ func (p *PipelineClient) ReportWorkflow(workflow util.ExecutionSpec) error {
 				workflow.ToStringForStore())
 		} else if statusCode.Code() == codes.Unauthenticated && strings.Contains(err.Error(), "service account token has expired") {
 			// If unauthenticated because SA token is expired, re-read/refresh the token and try again
-			p.tokenRefresher.RefreshToken()
+			if refreshErr := p.tokenRefresher.RefreshToken(); refreshErr != nil {
+				return util.NewCustomError(refreshErr, util.CUSTOM_CODE_PERMANENT,
+					"Failed to refresh token: %v", refreshErr.Error())
+			}
 			return util.NewCustomError(err, util.CUSTOM_CODE_TRANSIENT,
 				"Error while reporting workflow resource (code: %v, message: %v): %v, %+v",
 				statusCode.Code(),
@@ -167,7 +171,10 @@ func (p *PipelineClient) ReportScheduledWorkflow(swf *util.ScheduledWorkflow) er
 				swf.ScheduledWorkflow)
 		} else if statusCode.Code() == codes.Unauthenticated && strings.Contains(err.Error(), "service account token has expired") {
 			// If unauthenticated because SA token is expired, re-read/refresh the token and try again
-			p.tokenRefresher.RefreshToken()
+			if refreshErr := p.tokenRefresher.RefreshToken(); refreshErr != nil {
+				return util.NewCustomError(refreshErr, util.CUSTOM_CODE_PERMANENT,
+					"Failed to refresh token: %v", refreshErr.Error())
+			}
 			return util.NewCustomError(err, util.CUSTOM_CODE_TRANSIENT,
 				"Error while reporting workflow resource (code: %v, message: %v): %v, %+v",
 				statusCode.Code(),
@@ -187,15 +194,17 @@ func (p *PipelineClient) ReportScheduledWorkflow(swf *util.ScheduledWorkflow) er
 	return nil
 }
 
-// ReadArtifactForMetrics reads artifact content using the new util.ArtifactRequest/Response types.
-// This method is used by the metrics collection system.
-func (p *PipelineClient) ReadArtifactForMetrics(request *util.ArtifactRequest) (*util.ArtifactResponse, error) {
-	// Construct the HTTP streaming endpoint URL
-	// Format: /apis/v1beta1/runs/{run_id}/nodes/{node_id}/artifacts/{artifact_name}:stream
+// ReadArtifact reads artifact content using HTTP streaming.
+//
+// Error Handling:
+// - Returns nil for artifacts that don't exist (HTTP 404)
+// - Returns CUSTOM_CODE_PERMANENT for client errors (400, 403) and unexpected failures
+// - Returns CUSTOM_CODE_TRANSIENT for retryable errors (401, 500, network issues)
+// - Automatically refreshes tokens on expiry; callers should retry transient errors
+func (p *PipelineClient) ReadArtifact(request *util.ReadArtifactRequest) (*util.ReadArtifactResponse, error) {
 	url := fmt.Sprintf("%s/apis/v1beta1/runs/%s/nodes/%s/artifacts/%s:stream",
 		p.httpBaseURL, request.RunID, request.NodeID, request.ArtifactName)
 
-	// Create HTTP request with timeout
 	ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
 	defer cancel()
 
@@ -205,61 +214,65 @@ func (p *PipelineClient) ReadArtifactForMetrics(request *util.ArtifactRequest) (
 			"Failed to create HTTP request: %v", err.Error())
 	}
 
-	// Add authorization header
 	req.Header.Set("Authorization", "Bearer "+p.tokenRefresher.GetToken())
 
-	// Make the HTTP request
 	resp, err := p.httpClient.Do(req)
 	if err != nil {
 		if strings.Contains(err.Error(), "service account token has expired") {
-			// If unauthenticated because SA token is expired, re-read/refresh the token and try again
-			p.tokenRefresher.RefreshToken()
+			// If unauthenticated because SA token is expired, refresh the token and the caller should retry
+			if refreshErr := p.tokenRefresher.RefreshToken(); refreshErr != nil {
+				return nil, util.NewCustomError(refreshErr, util.CUSTOM_CODE_PERMANENT,
+					"Failed to refresh token: %v", refreshErr.Error())
+			}
 			return nil, util.NewCustomError(err, util.CUSTOM_CODE_TRANSIENT,
 				"Error while reading artifact due to token expiry: %v", err.Error())
 		}
 		return nil, util.NewCustomError(err, util.CUSTOM_CODE_PERMANENT,
 			"Failed to make HTTP request: %v", err.Error())
 	}
-	defer resp.Body.Close()
+	defer func() {
+		if closeErr := resp.Body.Close(); closeErr != nil {
+			glog.Warningf("Failed to close response body: %v", closeErr)
+		}
+	}()
 
-	// Handle HTTP status codes
 	switch resp.StatusCode {
 	case http.StatusOK:
-		// Success case - read the artifact data
 		data, err := io.ReadAll(resp.Body)
 		if err != nil {
 			return nil, util.NewCustomError(err, util.CUSTOM_CODE_PERMANENT,
 				"Failed to read artifact data: %v", err.Error())
 		}
-		return &util.ArtifactResponse{Data: data}, nil
+		return &util.ReadArtifactResponse{Data: data}, nil
 
 	case http.StatusNotFound:
-		// Artifact not found - return nil as per original behavior
 		return nil, nil
 
 	case http.StatusUnauthorized:
 		// Unauthorized - refresh token and return transient error
-		p.tokenRefresher.RefreshToken()
+		if refreshErr := p.tokenRefresher.RefreshToken(); refreshErr != nil {
+			if closeErr := resp.Body.Close(); closeErr != nil {
+				glog.Warningf("Failed to close response body: %v", closeErr)
+			}
+			return nil, util.NewCustomError(refreshErr, util.CUSTOM_CODE_PERMANENT,
+				"Failed to refresh token: %v", refreshErr.Error())
+		}
 		return nil, util.NewCustomError(fmt.Errorf("HTTP 401"), util.CUSTOM_CODE_TRANSIENT,
 			"Failed to read artifact, unauthorized (token may have expired)")
 
 	case http.StatusForbidden:
-		// Forbidden - return permanent error
 		return nil, util.NewCustomError(fmt.Errorf("HTTP 403"), util.CUSTOM_CODE_PERMANENT,
 			"Failed to read artifact, forbidden")
 
 	case http.StatusBadRequest:
-		// Bad request - return permanent error
 		return nil, util.NewCustomError(fmt.Errorf("HTTP 400"), util.CUSTOM_CODE_PERMANENT,
 			"Failed to read artifact, bad request")
 
 	case http.StatusInternalServerError:
-		// Internal server error - return transient error
 		return nil, util.NewCustomError(fmt.Errorf("HTTP 500"), util.CUSTOM_CODE_TRANSIENT,
 			"Failed to read artifact, internal server error")
 
 	default:
-		// Other status codes - return permanent error
 		return nil, util.NewCustomError(fmt.Errorf("HTTP %d", resp.StatusCode), util.CUSTOM_CODE_PERMANENT,
 			"Failed to read artifact, HTTP status: %d", resp.StatusCode)
 	}

backend/src/agent/persistence/client/pipeline_client_fake.go

@@ -24,8 +24,8 @@ type PipelineClientFake struct {
 	workflows                 map[string]util.ExecutionSpec
 	scheduledWorkflows        map[string]*util.ScheduledWorkflow
 	err                       error
-	artifacts                 map[string]*util.ArtifactResponse
-	readArtifactRequest       *util.ArtifactRequest
+	artifacts                 map[string]*util.ReadArtifactResponse
+	readArtifactRequest       *util.ReadArtifactRequest
 	reportedMetricsRequest    *api.ReportRunMetricsRequest
 	reportMetricsResponseStub *api.ReportRunMetricsResponse
 	reportMetricsErrorStub    error
@@ -36,7 +36,7 @@ func NewPipelineClientFake() *PipelineClientFake {
 		workflows:                 make(map[string]util.ExecutionSpec),
 		scheduledWorkflows:        make(map[string]*util.ScheduledWorkflow),
 		err:                       nil,
-		artifacts:                 make(map[string]*util.ArtifactResponse),
+		artifacts:                 make(map[string]*util.ReadArtifactResponse),
 		reportMetricsResponseStub: &api.ReportRunMetricsResponse{},
 	}
 }
@@ -57,7 +57,7 @@ func (p *PipelineClientFake) ReportScheduledWorkflow(swf *util.ScheduledWorkflow
 	return nil
 }
 
-func (p *PipelineClientFake) ReadArtifactForMetrics(request *util.ArtifactRequest) (*util.ArtifactResponse, error) {
+func (p *PipelineClientFake) ReadArtifact(request *util.ReadArtifactRequest) (*util.ReadArtifactResponse, error) {
 	if p.err != nil {
 		return nil, p.err
 	}
@@ -82,11 +82,11 @@ func (p *PipelineClientFake) GetScheduledWorkflow(namespace string, name string)
 	return p.scheduledWorkflows[getKey(namespace, name)]
 }
 
-func (p *PipelineClientFake) StubArtifact(request *util.ArtifactRequest, response *util.ArtifactResponse) {
+func (p *PipelineClientFake) StubArtifact(request *util.ReadArtifactRequest, response *util.ReadArtifactResponse) {
 	p.artifacts[request.String()] = response
 }
 
-func (p *PipelineClientFake) GetReadArtifactRequest() *util.ArtifactRequest {
+func (p *PipelineClientFake) GetReadArtifactRequest() *util.ReadArtifactRequest {
 	return p.readArtifactRequest
 }
 

backend/src/agent/persistence/worker/metrics_reporter.go

@@ -52,7 +52,7 @@ func (r MetricsReporter) ReportMetrics(workflow util.ExecutionSpec) error {
 		// Skip reporting if the workflow doesn't have the run id label
 		return nil
 	}
-	runMetrics, partialFailures := workflow.ExecutionStatus().CollectionMetrics(r.pipelineClient.ReadArtifactForMetrics)
+	runMetrics, partialFailures := workflow.ExecutionStatus().CollectionMetrics(r.pipelineClient.ReadArtifact)
 	if len(runMetrics) == 0 {
 		return aggregateErrors(partialFailures)
 	}

backend/src/agent/persistence/worker/metrics_reporter_test.go

@@ -170,12 +170,12 @@ func TestReportMetrics_Succeed(t *testing.T) {
 	metricsJSON := `{"metrics": [{"name": "accuracy", "numberValue": 0.77}, {"name": "logloss", "numberValue": 1.2}]}`
 	artifactData, _ := util.ArchiveTgz(map[string]string{"file": metricsJSON})
 	pipelineFake.StubArtifact(
-		&util.ArtifactRequest{
+		&util.ReadArtifactRequest{
 			RunID:        "run-1",
 			NodeID:       "node-1",
 			ArtifactName: "mlpipeline-metrics",
 		},
-		&util.ArtifactResponse{
+		&util.ReadArtifactResponse{
 			Data: []byte(artifactData),
 		})
 	pipelineFake.StubReportRunMetrics(&api.ReportRunMetricsResponse{
@@ -233,12 +233,12 @@ func TestReportMetrics_EmptyArchive_Fail(t *testing.T) {
 	})
 	artifactData, _ := util.ArchiveTgz(map[string]string{})
 	pipelineFake.StubArtifact(
-		&util.ArtifactRequest{
+		&util.ReadArtifactRequest{
 			RunID:        "run-1",
 			NodeID:       "node-1",
 			ArtifactName: "mlpipeline-metrics",
 		},
-		&util.ArtifactResponse{
+		&util.ReadArtifactResponse{
 			Data: []byte(artifactData),
 		})
 
@@ -277,12 +277,12 @@ func TestReportMetrics_MultipleFilesInArchive_Fail(t *testing.T) {
 	invalidMetricsJSON := `invalid JSON`
 	artifactData, _ := util.ArchiveTgz(map[string]string{"file1": validMetricsJSON, "file2": invalidMetricsJSON})
 	pipelineFake.StubArtifact(
-		&util.ArtifactRequest{
+		&util.ReadArtifactRequest{
 			RunID:        "run-1",
 			NodeID:       "node-1",
 			ArtifactName: "mlpipeline-metrics",
 		},
-		&util.ArtifactResponse{
+		&util.ReadArtifactResponse{
 			Data: []byte(artifactData),
 		})
 
@@ -320,12 +320,12 @@ func TestReportMetrics_InvalidMetricsJSON_Fail(t *testing.T) {
 	metricsJSON := `invalid JSON`
 	artifactData, _ := util.ArchiveTgz(map[string]string{"file": metricsJSON})
 	pipelineFake.StubArtifact(
-		&util.ArtifactRequest{
+		&util.ReadArtifactRequest{
 			RunID:        "run-1",
 			NodeID:       "node-1",
 			ArtifactName: "mlpipeline-metrics",
 		},
-		&util.ArtifactResponse{
+		&util.ReadArtifactResponse{
 			Data: []byte(artifactData),
 		})
 
@@ -374,21 +374,21 @@ func TestReportMetrics_InvalidMetricsJSON_PartialFail(t *testing.T) {
 	invalidArtifactData, _ := util.ArchiveTgz(map[string]string{"file": invalidMetricsJSON})
 	// Stub two artifacts, node-1 is invalid, node-2 is valid.
 	pipelineFake.StubArtifact(
-		&util.ArtifactRequest{
+		&util.ReadArtifactRequest{
 			RunID:        "run-1",
 			NodeID:       "node-1",
 			ArtifactName: "mlpipeline-metrics",
 		},
-		&util.ArtifactResponse{
+		&util.ReadArtifactResponse{
 			Data: []byte(invalidArtifactData),
 		})
 	pipelineFake.StubArtifact(
-		&util.ArtifactRequest{
+		&util.ReadArtifactRequest{
 			RunID:        "run-1",
 			NodeID:       "node-2",
 			ArtifactName: "mlpipeline-metrics",
 		},
-		&util.ArtifactResponse{
+		&util.ReadArtifactResponse{
 			Data: []byte(validArtifactData),
 		})
 
@@ -444,12 +444,12 @@ func TestReportMetrics_CorruptedArchiveFile_Fail(t *testing.T) {
 		},
 	})
 	pipelineFake.StubArtifact(
-		&util.ArtifactRequest{
+		&util.ReadArtifactRequest{
 			RunID:        "run-1",
 			NodeID:       "node-1",
 			ArtifactName: "mlpipeline-metrics",
 		},
-		&util.ArtifactResponse{
+		&util.ReadArtifactResponse{
 			Data: []byte("invalid tgz content"),
 		})
 
@@ -488,12 +488,12 @@ func TestReportMetrics_MultiplMetricErrors_TransientErrowWin(t *testing.T) {
 		`{"metrics": [{"name": "accuracy", "numberValue": 0.77}, {"name": "log loss", "numberValue": 1.2}, {"name": "accuracy", "numberValue": 1.2}]}`
 	artifactData, _ := util.ArchiveTgz(map[string]string{"file": metricsJSON})
 	pipelineFake.StubArtifact(
-		&util.ArtifactRequest{
+		&util.ReadArtifactRequest{
 			RunID:        "run-1",
 			NodeID:       "node-1",
 			ArtifactName: "mlpipeline-metrics",
 		},
-		&util.ArtifactResponse{
+		&util.ReadArtifactResponse{
 			Data: []byte(artifactData),
 		})
 	pipelineFake.StubReportRunMetrics(&api.ReportRunMetricsResponse{
@@ -551,12 +551,12 @@ func TestReportMetrics_Unauthorized(t *testing.T) {
 	metricsJSON := `{"metrics": [{"name": "accuracy", "numberValue": 0.77}, {"name": "logloss", "numberValue": 1.2}]}`
 	artifactData, _ := util.ArchiveTgz(map[string]string{"file": metricsJSON})
 	pipelineFake.StubArtifact(
-		&util.ArtifactRequest{
+		&util.ReadArtifactRequest{
 			RunID:        "run-1",
 			NodeID:       "node-1",
 			ArtifactName: "mlpipeline-metrics",
 		},
-		&util.ArtifactResponse{
+		&util.ReadArtifactResponse{
 			Data: []byte(artifactData),
 		})
 	pipelineFake.StubReportRunMetrics(&api.ReportRunMetricsResponse{

backend/src/common/util/execution_status.go

@@ -31,24 +31,24 @@ type NodeStatus struct {
 	Children    []string
 }
 
-// ArtifactRequest is a simple artifact request struct to replace the removed protobuf types
-type ArtifactRequest struct {
+// ReadArtifactRequest is a simple artifact request struct to replace the removed protobuf types
+type ReadArtifactRequest struct {
 	RunID        string
 	NodeID       string
 	ArtifactName string
 }
 
 // String returns a string representation for use as a map key
-func (r *ArtifactRequest) String() string {
+func (r *ReadArtifactRequest) String() string {
 	return r.RunID + "/" + r.NodeID + "/" + r.ArtifactName
 }
 
-// ArtifactResponse is a simple artifact response struct to replace the removed protobuf types
-type ArtifactResponse struct {
+// ReadArtifactResponse is a simple artifact response struct to replace the removed protobuf types
+type ReadArtifactResponse struct {
 	Data []byte
 }
 
-type RetrieveArtifact func(request *ArtifactRequest) (*ArtifactResponse, error)
+type RetrieveArtifact func(request *ReadArtifactRequest) (*ReadArtifactResponse, error)
 
 // Abstract interface to encapsulate the resources of the execution runtime specifically
 // for status information. This interface is mainly to access the status related information

backend/src/common/util/workflow.go

@@ -579,7 +579,7 @@ func readNodeMetricsJSONOrEmpty(runID string, nodeStatus *workflowapi.NodeStatus
 		return "", nil // No metrics artifact, skip the reporting
 	}
 
-	artifactRequest := &ArtifactRequest{
+	artifactRequest := &ReadArtifactRequest{
 		RunID:        runID,
 		NodeID:       nodeStatus.ID,
 		ArtifactName: metricsArtifactName,