feat(backend/sdk):Add download_to_workspace option to dsl.importer

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

Author: VaniHaripriya

api/v2alpha1/go/pipelinespec/pipeline_spec.pb.go

@@ -873,6 +873,9 @@ message PipelineDeploymentConfig {
 
     // Whether or not import an artifact regardless it has been imported before.
     bool reimport = 5;
+
+    // If true, download artifact into the pipeline workspace.
+    bool download_to_workspace = 7;
   }
 
   // ResolverSpec resolves artifacts from historical metadata and returns them

api/v2alpha1/pipeline_spec.proto

@@ -91,6 +91,32 @@ func (c *workflowCompiler) addImporterTemplate() string {
 	if value, ok := os.LookupEnv(PublishLogsEnvVar); ok {
 		args = append(args, "--publish_logs", value)
 	}
+	// Add workspace volume only if the workflow defines a workspace PVC
+	hasWorkspacePVC := false
+	for _, pvc := range c.wf.Spec.VolumeClaimTemplates {
+		if pvc.Name == workspaceVolumeName {
+			hasWorkspacePVC = true
+			break
+		}
+	}
+
+	var volumeMounts []k8score.VolumeMount
+	var volumes []k8score.Volume
+	if hasWorkspacePVC {
+		volumeMounts = append(volumeMounts, k8score.VolumeMount{
+			Name:      workspaceVolumeName,
+			MountPath: component.WorkspaceMountPath,
+		})
+		volumes = append(volumes, k8score.Volume{
+			Name: workspaceVolumeName,
+			VolumeSource: k8score.VolumeSource{
+				PersistentVolumeClaim: &k8score.PersistentVolumeClaimVolumeSource{
+					ClaimName: fmt.Sprintf("{{workflow.name}}-%s", workspaceVolumeName),
+				},
+			},
+		})
+	}
+
 	importerTemplate := &wfapi.Template{
 		Name: name,
 		Inputs: wfapi.Inputs{
@@ -102,13 +128,15 @@ func (c *workflowCompiler) addImporterTemplate() string {
 			},
 		},
 		Container: &k8score.Container{
-			Image:     c.launcherImage,
-			Command:   c.launcherCommand,
-			Args:      args,
-			EnvFrom:   []k8score.EnvFromSource{metadataEnvFrom},
-			Env:       commonEnvs,
-			Resources: driverResources,
+			Image:        c.launcherImage,
+			Command:      c.launcherCommand,
+			Args:         args,
+			EnvFrom:      []k8score.EnvFromSource{metadataEnvFrom},
+			Env:          commonEnvs,
+			Resources:    driverResources,
+			VolumeMounts: volumeMounts,
 		},
+		Volumes: volumes,
 	}
 
 	// If the apiserver is TLS-enabled, add the custom CA bundle to the importer template.

backend/src/v2/compiler/argocompiler/importer.go

@@ -4,10 +4,13 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"os"
+	"path/filepath"
 	"strings"
 
 	"github.com/kubeflow/pipelines/backend/src/common/util"
 
+	"github.com/kubeflow/pipelines/backend/src/v2/config"
 	"github.com/kubeflow/pipelines/backend/src/v2/objectstore"
 
 	pb "github.com/kubeflow/pipelines/third_party/ml-metadata/go/ml_metadata"
@@ -126,12 +129,17 @@ func (l *ImportLauncher) Execute(ctx context.Context) (err error) {
 		return err
 	}
 	ecfg := &metadata.ExecutionConfig{
-		TaskName:      l.task.GetTaskInfo().GetName(),
-		PodName:       l.launcherV2Options.PodName,
-		PodUID:        l.launcherV2Options.PodUID,
-		Namespace:     l.launcherV2Options.Namespace,
-		ExecutionType: metadata.ImporterExecutionTypeName,
-		ParentDagID:   l.importerLauncherOptions.ParentDagID,
+		TaskName:  l.task.GetTaskInfo().GetName(),
+		PodName:   l.launcherV2Options.PodName,
+		PodUID:    l.launcherV2Options.PodUID,
+		Namespace: l.launcherV2Options.Namespace,
+		ExecutionType: func() metadata.ExecutionType {
+			if l.importer.GetDownloadToWorkspace() {
+				return metadata.ImporterWorkspaceExecutionTypeName
+			}
+			return metadata.ImporterExecutionTypeName
+		}(),
+		ParentDagID: l.importerLauncherOptions.ParentDagID,
 	}
 	createdExecution, err := l.metadataClient.CreateExecution(ctx, pipeline, ecfg)
 	if err != nil {
@@ -253,15 +261,17 @@ func (l *ImportLauncher) ImportSpecToMLMDArtifact(ctx context.Context) (artifact
 	}
 
 	if strings.HasPrefix(artifactUri, "oci://") {
+		// OCI artifacts are not supported when workspace is used
+		if l.importer.GetDownloadToWorkspace() {
+			return nil, fmt.Errorf("importer workspace download does not support OCI registries")
+		}
 		artifactType, err := metadata.SchemaToArtifactType(schema)
 		if err != nil {
 			return nil, fmt.Errorf("converting schema to artifact type failed: %w", err)
 		}
-
 		if *artifactType.Name != "system.Model" {
 			return nil, fmt.Errorf("the %s artifact type does not support OCI registries", *artifactType.Name)
 		}
-
 		return artifact, nil
 	}
 
@@ -283,6 +293,40 @@ func (l *ImportLauncher) ImportSpecToMLMDArtifact(ctx context.Context) (artifact
 	}
 	storeSessionInfoStr := string(storeSessionInfoJSON)
 	artifact.CustomProperties["store_session_info"] = metadata.StringValue(storeSessionInfoStr)
+
+	// Download the artifact into the workspace
+	if l.importer.GetDownloadToWorkspace() {
+		bucketConfig, err := objectstore.ParseBucketConfigForArtifactURI(artifactUri)
+		if err != nil {
+			return nil, fmt.Errorf("failed to parse bucket config while downloading artifact into workspace with uri %q: %w", artifactUri, err)
+		}
+		// Resolve and attach session info from kfp-launcher config for the artifact provider
+		if cfg, cfgErr := config.FromConfigMap(ctx, l.k8sClient, l.launcherV2Options.Namespace); cfgErr != nil {
+			glog.Warningf("failed to load launcher config for workspace download: %v", cfgErr)
+		} else if cfg != nil {
+			if sess, sessErr := cfg.GetStoreSessionInfo(artifactUri); sessErr != nil {
+				glog.Warningf("failed to resolve store session info for %q: %v", artifactUri, sessErr)
+			} else {
+				bucketConfig.SessionInfo = &sess
+			}
+		}
+		blobKey, err := bucketConfig.KeyFromURI(artifactUri)
+		if err != nil {
+			return nil, fmt.Errorf("failed to derive blob key from uri %q while downloading artifact into workspace: %w", artifactUri, err)
+		}
+		workspaceRoot := filepath.Join(WorkspaceMountPath, ".artifacts")
+		if err := os.MkdirAll(workspaceRoot, 0755); err != nil {
+			return nil, fmt.Errorf("failed to create workspace directory %q: %w", workspaceRoot, err)
+		}
+		bucket, err := objectstore.OpenBucket(ctx, l.k8sClient, l.launcherV2Options.Namespace, bucketConfig)
+		if err != nil {
+			return nil, fmt.Errorf("failed to open bucket for uri %q: %w", artifactUri, err)
+		}
+		defer bucket.Close()
+		if err := objectstore.DownloadBlob(ctx, bucket, workspaceRoot, blobKey); err != nil {
+			return nil, fmt.Errorf("failed to download artifact to workspace: %w", err)
+		}
+	}
 	return artifact, nil
 }
 

backend/src/v2/component/importer_launcher.go

@@ -713,6 +713,12 @@ 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
+			// Skip downloading if the artifact is flagged as already present in the workspace
+			if inputArtifact.GetMetadata() != nil {
+				if v, ok := inputArtifact.GetMetadata().GetFields()["_kfp_workspace"]; ok && v.GetBoolValue() {
+					continue
+				}
+			}
 			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)
@@ -856,6 +862,22 @@ func getPlaceholders(executorInput *pipelinespec.ExecutorInput) (placeholders ma
 		key := fmt.Sprintf(`{{$.inputs.artifacts['%s'].uri}}`, name)
 		placeholders[key] = inputArtifact.Uri
 
+		// If the artifact is marked as already in the workspace, map the workspace path.
+		if inputArtifact.GetMetadata() != nil {
+			if v, ok := inputArtifact.GetMetadata().GetFields()["_kfp_workspace"]; ok && v.GetBoolValue() {
+				bucketConfig, err := objectstore.ParseBucketConfigForArtifactURI(inputArtifact.Uri)
+				if err == nil {
+					blobKey, err := bucketConfig.KeyFromURI(inputArtifact.Uri)
+					if err == nil {
+						localPath := filepath.Join(WorkspaceMountPath, ".artifacts", blobKey)
+						key = fmt.Sprintf(`{{$.inputs.artifacts['%s'].path}}`, name)
+						placeholders[key] = localPath
+						continue
+					}
+				}
+			}
+		}
+
 		localPath, err := LocalPathForURI(inputArtifact.Uri)
 		if err != nil {
 			// Input Artifact does not have a recognized storage URI

backend/src/v2/component/launcher_v2.go

@@ -20,6 +20,7 @@ import (
 	"errors"
 	"io"
 	"os"
+	"path/filepath"
 	"testing"
 
 	"github.com/kubeflow/pipelines/backend/src/v2/cacheutils"
@@ -176,6 +177,29 @@ func Test_executeV2_publishLogs(t *testing.T) {
 	}
 }
 
+func Test_getPlaceholders_WorkspaceArtifactPath(t *testing.T) {
+	execIn := &pipelinespec.ExecutorInput{
+		Inputs: &pipelinespec.ExecutorInput_Inputs{
+			Artifacts: map[string]*pipelinespec.ArtifactList{
+				"data": {
+					Artifacts: []*pipelinespec.RuntimeArtifact{
+						{Uri: "minio://mlpipeline/sample/sample.txt", Metadata: &structpb.Struct{Fields: map[string]*structpb.Value{"_kfp_workspace": structpb.NewBoolValue(true)}}},
+					},
+				},
+			},
+		},
+	}
+	ph, err := getPlaceholders(execIn)
+	if err != nil {
+		t.Fatalf("getPlaceholders error: %v", err)
+	}
+	got := ph["{{$.inputs.artifacts['data'].path}}"]
+	want := filepath.Join(WorkspaceMountPath, ".artifacts", "sample", "sample.txt")
+	if got != want {
+		t.Fatalf("placeholder path mismatch: got=%q want=%q", got, want)
+	}
+}
+
 func Test_executorInput_compileCmdAndArgs(t *testing.T) {
 	executorInputJSON := `{
 		"inputs": {

backend/src/v2/component/launcher_v2_test.go

@@ -1189,6 +1189,68 @@ func TestWorkspaceMount_PassthroughVolumes_ApplyAndCapture(t *testing.T) {
 	}
 }
 
+func TestWorkspaceMount_TriggeredByArtifactMetadata(t *testing.T) {
+	proxy.InitializeConfigWithEmptyForTests()
+	containerSpec := &pipelinespec.PipelineDeploymentConfig_PipelineContainerSpec{Image: "python:3.9"}
+	componentSpec := &pipelinespec.ComponentSpec{
+		TaskConfigPassthroughs: []*pipelinespec.TaskConfigPassthrough{
+			{
+				Field:       pipelinespec.TaskConfigPassthroughType_KUBERNETES_VOLUMES,
+				ApplyToTask: true,
+			},
+		},
+	}
+
+	// Build an ExecutorInput that does NOT reference workspace path in params,
+	// but contains an input artifact marked as already in workspace.
+	execInput := &pipelinespec.ExecutorInput{
+		Inputs: &pipelinespec.ExecutorInput_Inputs{
+			Artifacts: map[string]*pipelinespec.ArtifactList{
+				"data": {
+					Artifacts: []*pipelinespec.RuntimeArtifact{
+						{
+							Uri: "minio://mlpipeline/sample/sample.txt",
+							Metadata: &structpb.Struct{Fields: map[string]*structpb.Value{
+								"_kfp_workspace": structpb.NewBoolValue(true),
+							}},
+						},
+					},
+				},
+			},
+		},
+	}
+
+	taskCfg := &TaskConfig{}
+	podSpec, err := initPodSpecPatch(
+		containerSpec, componentSpec, execInput,
+		27, "test", "run", "my-run-name", "1", "false", "false", taskCfg, false, false, "",
+	)
+	assert.Nil(t, err)
+
+	// Expect workspace volume mounted
+	if assert.Len(t, podSpec.Volumes, 1) {
+		assert.Equal(t, "kfp-workspace", podSpec.Volumes[0].Name)
+		if assert.NotNil(t, podSpec.Volumes[0].PersistentVolumeClaim) {
+			assert.Equal(t, "my-run-name-kfp-workspace", podSpec.Volumes[0].PersistentVolumeClaim.ClaimName)
+		}
+	}
+	if assert.Len(t, podSpec.Containers, 1) {
+		if assert.Len(t, podSpec.Containers[0].VolumeMounts, 1) {
+			assert.Equal(t, "kfp-workspace", podSpec.Containers[0].VolumeMounts[0].Name)
+			assert.Equal(t, "/kfp-workspace", podSpec.Containers[0].VolumeMounts[0].MountPath)
+		}
+	}
+
+	// Also captured to TaskConfig
+	if assert.Len(t, taskCfg.Volumes, 1) {
+		assert.Equal(t, "kfp-workspace", taskCfg.Volumes[0].Name)
+	}
+	if assert.Len(t, taskCfg.VolumeMounts, 1) {
+		assert.Equal(t, "kfp-workspace", taskCfg.VolumeMounts[0].Name)
+		assert.Equal(t, "/kfp-workspace", taskCfg.VolumeMounts[0].MountPath)
+	}
+}
+
 func Test_initPodSpecPatch_TaskConfig_Env_Passthrough_CaptureOnly(t *testing.T) {
 	proxy.InitializeConfigWithEmptyForTests()
 	containerSpec := &pipelinespec.PipelineDeploymentConfig_PipelineContainerSpec{

backend/src/v2/driver/driver_test.go

@@ -33,6 +33,17 @@ import (
 
 var ErrResolvedParameterNull = errors.New("the resolved input parameter is null")
 
+// setWorkspaceFlag sets the _kfp_workspace metadata flag on the provided
+// runtime artifact when the producing execution is an ImporterWorkspace.
+func setWorkspaceFlag(execution *metadata.Execution, runtimeArtifact *pipelinespec.RuntimeArtifact) {
+	if execution.GetExecution().GetType() == string(metadata.ImporterWorkspaceExecutionTypeName) {
+		if runtimeArtifact.Metadata == nil {
+			runtimeArtifact.Metadata = &structpb.Struct{Fields: map[string]*structpb.Value{}}
+		}
+		runtimeArtifact.Metadata.Fields["_kfp_workspace"] = structpb.NewBoolValue(true)
+	}
+}
+
 // resolveUpstreamOutputsConfig is just a config struct used to store the input
 // parameters of the resolveUpstreamParameters and resolveUpstreamArtifacts
 // functions.
@@ -763,6 +774,8 @@ func resolveUpstreamArtifacts(cfg resolveUpstreamOutputsConfig) (*pipelinespec.A
 			if err != nil {
 				cfg.err(err)
 			}
+			// If produced by workspace importer, Set _kfp_workspace=True
+			setWorkspaceFlag(currentTask, runtimeArtifact)
 			// Base case
 			return &pipelinespec.ArtifactList{
 				Artifacts: []*pipelinespec.RuntimeArtifact{runtimeArtifact},
@@ -980,6 +993,8 @@ func collectContainerOutput(
 		if err != nil {
 			return nil, nil, cfg.err(err)
 		}
+		// If produced by workspace importer, Set _kfp_workspace=True
+		setWorkspaceFlag(currentTask, artifact)
 		glog.V(4).Infof("runtimeArtifact: %v", artifact)
 	} else {
 		_, outputParameters, err := currentTask.GetParameters()