Revert LocalPathForURI() removal and update retrieveArtifactPath()

Signed-off-by: agoins <[email protected]>

Author: alyssacgoins

backend/src/v2/component/launcher_v2.go

@@ -140,7 +140,7 @@ func stopWaitingArtifacts(artifacts map[string]*pipelinespec.ArtifactList) {
 				continue
 			}
 
-			localPath, err := retrieveArtifactPath(inputArtifact)
+			localPath, err := LocalPathForURI(inputArtifact.Uri)
 			if err != nil {
 				continue
 			}
@@ -472,7 +472,7 @@ func getLogWriter(artifacts map[string]*pipelinespec.ArtifactList) (writer io.Wr
 	}
 
 	logURI := logsArtifactList.Artifacts[0].Uri
-	logFilePath, err := retrieveArtifactPath(logsArtifactList.Artifacts[0])
+	logFilePath, err := LocalPathForURI(logURI)
 	if err != nil {
 		glog.Errorf("Error converting log artifact URI, %s, to file path.", logURI)
 		return os.Stdout
@@ -718,7 +718,7 @@ func downloadArtifacts(ctx context.Context, executorInput *pipelinespec.Executor
 		for _, artifact := range artifactList.Artifacts {
 			// Iterating through the artifact list allows for collected artifacts to be properly consumed.
 			inputArtifact := artifact
-			localPath, err := retrieveArtifactPath(inputArtifact)
+			localPath, err := LocalPathForURI(inputArtifact.Uri)
 			if err != nil {
 				glog.Warningf("Input Artifact %q does not have a recognized storage URI %q. Skipping downloading to local path.", name, inputArtifact.Uri)
 
@@ -861,7 +861,7 @@ func getPlaceholders(executorInput *pipelinespec.ExecutorInput) (placeholders ma
 		key := fmt.Sprintf(`{{$.inputs.artifacts['%s'].uri}}`, name)
 		placeholders[key] = inputArtifact.Uri
 
-		localPath, err := retrieveArtifactPath(inputArtifact)
+		localPath, err := LocalPathForURI(inputArtifact.Uri)
 		if err != nil {
 			// Input Artifact does not have a recognized storage URI
 			continue
@@ -880,7 +880,7 @@ func getPlaceholders(executorInput *pipelinespec.ExecutorInput) (placeholders ma
 		outputArtifact := artifactList.Artifacts[0]
 		placeholders[fmt.Sprintf(`{{$.outputs.artifacts['%s'].uri}}`, name)] = outputArtifact.Uri
 
-		localPath, err := retrieveArtifactPath(outputArtifact)
+		localPath, err := LocalPathForURI(outputArtifact.Uri)
 		if err != nil {
 			return nil, fmt.Errorf("resolve output artifact %q's local path: %w", name, err)
 		}
@@ -983,26 +983,29 @@ func getExecutorOutputFile(path string) (*pipelinespec.ExecutorOutput, error) {
 	return executorOutput, nil
 }
 
+func LocalPathForURI(uri string) (string, error) {
+	if strings.HasPrefix(uri, "gs://") {
+		return "/gcs/" + strings.TrimPrefix(uri, "gs://"), nil
+	}
+	if strings.HasPrefix(uri, "minio://") {
+		return "/minio/" + strings.TrimPrefix(uri, "minio://"), nil
+	}
+	if strings.HasPrefix(uri, "s3://") {
+		return "/s3/" + strings.TrimPrefix(uri, "s3://"), nil
+	}
+	if strings.HasPrefix(uri, "oci://") {
+		return "/oci/" + strings.ReplaceAll(strings.TrimPrefix(uri, "oci://"), "/", "_") + "/models", nil
+	}
+	return "", fmt.Errorf("failed to generate local path for URI %s: unsupported storage scheme", uri)
+}
+
 func retrieveArtifactPath(artifact *pipelinespec.RuntimeArtifact) (string, error) {
 	// If artifact custom path is set, use custom path. Otherwise, use URI.
 	customPath := artifact.CustomPath
 	if customPath != nil {
 		return *customPath, nil
 	} else {
-		uri := artifact.Uri
-		if strings.HasPrefix(uri, "gs://") {
-			return "/gcs/" + strings.TrimPrefix(uri, "gs://"), nil
-		}
-		if strings.HasPrefix(uri, "minio://") {
-			return "/minio/" + strings.TrimPrefix(uri, "minio://"), nil
-		}
-		if strings.HasPrefix(uri, "s3://") {
-			return "/s3/" + strings.TrimPrefix(uri, "s3://"), nil
-		}
-		if strings.HasPrefix(uri, "oci://") {
-			return "/oci/" + strings.ReplaceAll(strings.TrimPrefix(uri, "oci://"), "/", "_") + "/models", nil
-		}
-		return "", fmt.Errorf("failed to generate local path for URI %s: unsupported storage scheme", uri)
+		return LocalPathForURI(artifact.Uri)
 	}
 }
 
@@ -1021,7 +1024,7 @@ func prepareOutputFolders(executorInput *pipelinespec.ExecutorInput) error {
 
 		for _, outputArtifact := range artifactList.Artifacts {
 
-			localPath, err := retrieveArtifactPath(outputArtifact)
+			localPath, err := LocalPathForURI(outputArtifact.Uri)
 			if err != nil {
 				return fmt.Errorf("failed to generate local storage path for output artifact %q: %w", name, err)
 			}

backend/src/v2/driver/driver.go

@@ -545,7 +545,7 @@ func addModelcarsToPodSpec(
 	for i, name := range modelcarArtifactNames {
 		inputArtifact := modelcarArtifacts[name]
 
-		localPath, err := localPathForURI(inputArtifact.Uri)
+		localPath, err := component.LocalPathForURI(inputArtifact.Uri)
 		if err != nil {
 			continue
 		}
@@ -716,7 +716,7 @@ func provisionOutputs(
 
 	// Place the executor output file under localTaskRoot to enable Pythonic artifacts. The SDK's pythonic artifact
 	// runtime derives CONTAINER_TASK_ROOT from the directory of OutputFile to use it in dsl.get_uri.
-	if localTaskRoot, err := localPathForURI(taskRootRemote); err == nil {
+	if localTaskRoot, err := component.LocalPathForURI(taskRootRemote); err == nil {
 		outputs.OutputFile = filepath.Join(localTaskRoot, "output_metadata.json")
 	} else {
 		// Fallback to legacy path if the pipeline root scheme is not recognized.
@@ -742,19 +742,3 @@ func validateVolumeMounts(podSpec *k8score.PodSpec) error {
 
 	return nil
 }
-
-func localPathForURI(uri string) (string, error) {
-	if strings.HasPrefix(uri, "gs://") {
-		return "/gcs/" + strings.TrimPrefix(uri, "gs://"), nil
-	}
-	if strings.HasPrefix(uri, "minio://") {
-		return "/minio/" + strings.TrimPrefix(uri, "minio://"), nil
-	}
-	if strings.HasPrefix(uri, "s3://") {
-		return "/s3/" + strings.TrimPrefix(uri, "s3://"), nil
-	}
-	if strings.HasPrefix(uri, "oci://") {
-		return "/oci/" + strings.ReplaceAll(strings.TrimPrefix(uri, "oci://"), "/", "_") + "/models", nil
-	}
-	return "", fmt.Errorf("failed to generate local path for URI %s: unsupported storage scheme", uri)
-}