Update launcher_v2 local path helper method.

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 := LocalPathForURI(inputArtifact.Uri)
+			localPath, err := retrieveArtifactPath(inputArtifact)
 			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 := LocalPathForURI(logURI)
+	logFilePath, err := retrieveArtifactPath(logsArtifactList.Artifacts[0])
 	if err != nil {
 		glog.Errorf("Error converting log artifact URI, %s, to file path.", logURI)
 		return os.Stdout
@@ -628,7 +628,12 @@ func uploadOutputArtifacts(ctx context.Context, executorInput *pipelinespec.Exec
 			}
 
 			// Upload artifacts from local path to remote storages.
-			localDir, err := LocalPathForURI(outputArtifact.Uri)
+			localDir, err := retrieveArtifactPath(outputArtifact)
+
+			if outputArtifact.CustomPath != nil {
+				localDir = *outputArtifact.CustomPath
+			}
+
 			if err != nil {
 				glog.Warningf("Output Artifact %q does not have a recognized storage URI %q. Skipping uploading to remote storage.", name, outputArtifact.Uri)
 			} else if !strings.HasPrefix(outputArtifact.Uri, "oci://") {
@@ -713,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 := LocalPathForURI(inputArtifact.Uri)
+			localPath, err := retrieveArtifactPath(inputArtifact)
 			if err != nil {
 				glog.Warningf("Input Artifact %q does not have a recognized storage URI %q. Skipping downloading to local path.", name, inputArtifact.Uri)
 
@@ -856,7 +861,7 @@ func getPlaceholders(executorInput *pipelinespec.ExecutorInput) (placeholders ma
 		key := fmt.Sprintf(`{{$.inputs.artifacts['%s'].uri}}`, name)
 		placeholders[key] = inputArtifact.Uri
 
-		localPath, err := LocalPathForURI(inputArtifact.Uri)
+		localPath, err := retrieveArtifactPath(inputArtifact)
 		if err != nil {
 			// Input Artifact does not have a recognized storage URI
 			continue
@@ -875,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 := LocalPathForURI(outputArtifact.Uri)
+		localPath, err := retrieveArtifactPath(outputArtifact)
 		if err != nil {
 			return nil, fmt.Errorf("resolve output artifact %q's local path: %w", name, err)
 		}
@@ -978,20 +983,27 @@ 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
+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 "", fmt.Errorf("failed to generate local path for URI %s: unsupported storage scheme", uri)
 }
 
 func prepareOutputFolders(executorInput *pipelinespec.ExecutorInput) error {
@@ -1009,7 +1021,7 @@ func prepareOutputFolders(executorInput *pipelinespec.ExecutorInput) error {
 
 		for _, outputArtifact := range artifactList.Artifacts {
 
-			localPath, err := LocalPathForURI(outputArtifact.Uri)
+			localPath, err := retrieveArtifactPath(outputArtifact)
 			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 := component.LocalPathForURI(inputArtifact.Uri)
+		localPath, err := 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 := component.LocalPathForURI(taskRootRemote); err == nil {
+	if localTaskRoot, err := 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,3 +742,19 @@ 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)
+}