Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
POC Hack to get s3 tags working for aws backend (breaks gcp, and mini…
…o I expect)
- Loading branch information
1 parent
a80b65a
commit 0b6a982
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,69 @@ | ||
| package driver | ||
|
|
||
| import ( | ||
| "context" | ||
| "encoding/json" | ||
| "fmt" | ||
|
|
||
| "github.com/golang/glog" | ||
| "github.com/kubeflow/pipelines/api/v2alpha1/go/pipelinespec" | ||
| "github.com/kubeflow/pipelines/backend/src/v2/metadata" | ||
| ) | ||
|
|
||
| type ArtifactReader interface { | ||
| GetOutputArtifactsByExecutionId(ctx context.Context, executionId int64) (map[string]*metadata.OutputArtifact, error) | ||
| } | ||
|
|
||
| func resolveUpstreamArtifacts(ctx context.Context, tasks map[string]*metadata.Execution, taskName string, outputArtifactKey string, mlmd ArtifactReader) (runtimeArtifact *pipelinespec.RuntimeArtifact, err error) { | ||
| glog.V(4).Info("taskName: ", taskName) | ||
| glog.V(4).Info("outputArtifactKey: ", outputArtifactKey) | ||
| upstreamTask := tasks[taskName] | ||
| if *upstreamTask.GetExecution().Type == "system.DAGExecution" { | ||
| // recurse | ||
| outputArtifactsCustomProperty, ok := upstreamTask.GetExecution().GetCustomProperties()["output_artifacts"] | ||
| if !ok { | ||
| return nil, fmt.Errorf("cannot find output_artifacts") | ||
| } | ||
| var outputArtifacts map[string]*pipelinespec.DagOutputsSpec_DagOutputArtifactSpec | ||
| glog.V(4).Infof("outputArtifactsCustomProperty: %#v", outputArtifactsCustomProperty) | ||
| glog.V(4).Info("outputArtifactsCustomProperty String: ", outputArtifactsCustomProperty.GetStringValue()) | ||
| err = json.Unmarshal([]byte(outputArtifactsCustomProperty.GetStringValue()), &outputArtifacts) | ||
| if err != nil { | ||
| return nil, err | ||
| } | ||
| glog.V(4).Info("Deserialized outputArtifactsMap: ", outputArtifacts) | ||
| var subTaskName string | ||
| outputArtifactSelectors := outputArtifacts[outputArtifactKey].GetArtifactSelectors() | ||
| for _, outputArtifactSelector := range outputArtifactSelectors { | ||
| subTaskName = outputArtifactSelector.ProducerSubtask | ||
| outputArtifactKey = outputArtifactSelector.OutputArtifactKey | ||
| glog.V(4).Infof("ProducerSubtask: %v", outputArtifactSelector.ProducerSubtask) | ||
| glog.V(4).Infof("OutputArtifactKey: %v", outputArtifactSelector.OutputArtifactKey) | ||
| } | ||
| downstreamParameterMapping, err := resolveUpstreamArtifacts(ctx, tasks, subTaskName, outputArtifactKey, mlmd) | ||
| glog.V(4).Infof("downstreamParameterMapping: %#v", downstreamParameterMapping) | ||
| if err != nil { | ||
| return nil, err | ||
| } | ||
| return downstreamParameterMapping, nil | ||
| } else { | ||
| // base case | ||
| outputs, err := mlmd.GetOutputArtifactsByExecutionId(ctx, upstreamTask.GetID()) | ||
| if err != nil { | ||
| return nil, err | ||
| } | ||
| artifact, ok := outputs[outputArtifactKey] | ||
| if !ok { | ||
| return nil, fmt.Errorf( | ||
| "cannot find output artifact key %q in producer task %q", | ||
| outputArtifactKey, | ||
| taskName, | ||
| ) | ||
| } | ||
| runtimeArtifact, err := artifact.ToRuntimeArtifact() | ||
| if err != nil { | ||
| return nil, err | ||
| } | ||
| return runtimeArtifact, nil | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,75 @@ | ||
| package objectstore | ||
|
|
||
| import ( | ||
| "context" | ||
| "fmt" | ||
|
|
||
| "github.com/aws/aws-sdk-go/aws" | ||
| "github.com/aws/aws-sdk-go/aws/credentials" | ||
| "github.com/aws/aws-sdk-go/aws/session" | ||
| "github.com/aws/aws-sdk-go/service/s3" | ||
| "github.com/golang/glog" | ||
| ) | ||
|
|
||
| func createS3Client(region, accessKey, secretKey string) (*s3.S3, error) { | ||
| sess, err := session.NewSession(&aws.Config{ | ||
| Region: aws.String(region), | ||
| Credentials: credentials.NewStaticCredentials( | ||
| accessKey, secretKey, "", | ||
| ), | ||
| }) | ||
| if err != nil { | ||
| return nil, fmt.Errorf("failed to create AWS session: %w", err) | ||
| } | ||
|
|
||
| s3Client := s3.New(sess) | ||
| return s3Client, nil | ||
| } | ||
|
|
||
| type S3Bucket struct { | ||
| client *s3.S3 | ||
| bucket string | ||
| key string | ||
| } | ||
|
|
||
| // AddTags is an additional method for S3Bucket. | ||
| func (b *S3Bucket) AddTags(ctx context.Context, tags []*s3.Tag) (err error) { | ||
| defer func() { | ||
| if err != nil { | ||
| // wrap error before returning | ||
| err = fmt.Errorf("failed to add tags to S3 bucket: %w", err) | ||
| } | ||
| }() | ||
| glog.Info("Bucket String: ", b.bucket) | ||
| bucket := aws.String(b.bucket) | ||
| glog.Info("Key String: ", b.key) | ||
| key := aws.String(b.key) | ||
| glog.Info("Tags: ", tags) | ||
| s3Tags := s3.Tagging{TagSet: tags} | ||
|
|
||
| if b.client == nil { | ||
| return fmt.Errorf("S3 client is not initialized") | ||
| } | ||
| if b.bucket == "" { | ||
| return fmt.Errorf("bucket name is not set") | ||
| } | ||
| if b.key == "" { | ||
| return fmt.Errorf("object key is not set") | ||
| } | ||
|
|
||
| glog.Info("Calling PutObjectTagging") | ||
| objectTaggingInput := &s3.PutObjectTaggingInput{ | ||
| Bucket: bucket, | ||
| Key: key, | ||
| Tagging: &s3Tags, | ||
| } | ||
|
|
||
| glog.Info("Calling PutObjectTaggingWithContext") | ||
| // sess, err1 := createS3BucketSession(ctx, namespace, config.SessionInfo, k8sClient) | ||
| _, err = b.client.PutObjectTaggingWithContext(ctx, objectTaggingInput) | ||
| if err != nil { | ||
| glog.Errorf("failed to add tags to S3 bucket: %v", err) | ||
| return fmt.Errorf("failed to add tags to S3 bucket: %w", err) | ||
| } | ||
| return err | ||
| } |