Skip to content

Commit e7382fa

Browse files
alyssacgoinsalyssacgoins
committed
Implement custom upload path on runtime artifact.
Signed-off-by: agoins <[email protected]>
1 parent 5c0c47e commit e7382fa

File tree

7 files changed

+130
-18
lines changed

7 files changed

+130
-18
lines changed

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

Lines changed: 17 additions & 4 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎api/v2alpha1/pipeline_spec.proto‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -967,6 +967,9 @@ message RuntimeArtifact {
967967

968968
// Properties of the Artifact.
969969
google.protobuf.Struct metadata = 6;
970+
971+
// Custom path for output artifact.
972+
optional string custom_path = 7;
970973
}
971974

972975
// Message that represents a list of artifacts.

‎sdk/python/kfp/dsl/executor.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -287,6 +287,7 @@ def write_executor_output(self,
287287
'name': artifact.name,
288288
'uri': artifact.uri,
289289
'metadata': artifact.metadata,
290+
'custom_path': artifact.custom_path
290291
}
291292
artifacts_list = {'artifacts': [runtime_artifact]}
292293

‎sdk/python/kfp/dsl/executor_test.py‎

Lines changed: 39 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -318,6 +318,8 @@ def test_func(output_artifact_two: Output[Metrics]):
318318
'projects/123/locations/us-central1/metadataStores/default/artifacts/123',
319319
'uri':
320320
'new-uri',
321+
'custom_path':
322+
'',
321323
'metadata': {
322324
'key_1': 'value_1',
323325
'key_2': 2,
@@ -613,7 +615,8 @@ def test_artifact_output1(self):
613615
"type": {
614616
"schemaTitle": "system.Artifact"
615617
},
616-
"uri": "gs://some-bucket/output"
618+
"uri": "gs://some-bucket/output",
619+
"custom_path": ""
617620
}
618621
]
619622
}
@@ -645,7 +648,9 @@ def test_func(first: str, second: str, output: Output[Artifact]) -> str:
645648
'name':
646649
'projects/123/locations/us-central1/metadataStores/default/artifacts/123',
647650
'uri':
648-
'gs://some-bucket/output'
651+
'gs://some-bucket/output',
652+
'custom_path':
653+
''
649654
}]
650655
}
651656
},
@@ -702,7 +707,9 @@ def test_func(first: str, second: str) -> Artifact:
702707
'name':
703708
'projects/123/locations/us-central1/metadataStores/default/artifacts/123',
704709
'uri':
705-
'gs://some-bucket/output'
710+
'gs://some-bucket/output',
711+
'custom_path':
712+
''
706713
}]
707714
}
708715
},
@@ -815,7 +822,9 @@ def func_returning_plain_tuple() -> NamedTuple('Outputs', [
815822
'name':
816823
'projects/123/locations/us-central1/metadataStores/default/artifacts/123',
817824
'uri':
818-
'gs://some-bucket/output_dataset'
825+
'gs://some-bucket/output_dataset',
826+
'custom_path':
827+
''
819828
}]
820829
}
821830
},
@@ -1054,6 +1063,8 @@ def test_func(
10541063
'',
10551064
'uri':
10561065
'gs://mlpipeline/v2/artifacts/my-test-pipeline-beta/b2b0cdee-b15c-48ff-b8bc-a394ae46c854/train/model',
1066+
'custom_path':
1067+
'',
10571068
'metadata': {
10581069
'accuracy': 0.9
10591070
}
@@ -1288,6 +1299,8 @@ def test_func() -> Artifact:
12881299
'projects/123/locations/us-central1/metadataStores/default/artifacts/123',
12891300
'uri':
12901301
'gs://manually_specified_bucket/foo',
1302+
'custom_path':
1303+
'',
12911304
'metadata': {
12921305
'data': 123
12931306
}
@@ -1339,6 +1352,8 @@ def test_func() -> Artifact:
13391352
'projects/123/locations/us-central1/metadataStores/default/artifacts/123',
13401353
'uri':
13411354
'gs://another_bucket/my_artifact',
1355+
'custom_path':
1356+
'',
13421357
'metadata': {
13431358
'data': 123
13441359
}
@@ -1408,6 +1423,8 @@ def test_func() -> NamedTuple('outputs', a=Artifact, d=Dataset):
14081423
'projects/123/locations/us-central1/metadataStores/default/artifacts/123',
14091424
'uri':
14101425
'gs://another_bucket/artifact',
1426+
'custom_path':
1427+
'',
14111428
'metadata': {
14121429
'data': 123
14131430
}
@@ -1419,6 +1436,8 @@ def test_func() -> NamedTuple('outputs', a=Artifact, d=Dataset):
14191436
'projects/123/locations/us-central1/metadataStores/default/artifacts/321',
14201437
'uri':
14211438
'gs://another_bucket/dataset',
1439+
'custom_path':
1440+
'',
14221441
'metadata': {}
14231442
}]
14241443
}
@@ -1621,7 +1640,8 @@ class TestDictToArtifact(parameterized.TestCase):
16211640
'type': {
16221641
'schemaTitle': 'system.Artifact'
16231642
},
1624-
'uri': 'gs://some-bucket/input_artifact_one'
1643+
'uri': 'gs://some-bucket/input_artifact_one',
1644+
'custom_path': '',
16251645
},
16261646
'artifact_cls': artifact_types.Artifact,
16271647
'expected_type': artifact_types.Artifact,
@@ -1633,7 +1653,8 @@ class TestDictToArtifact(parameterized.TestCase):
16331653
'type': {
16341654
'schemaTitle': 'system.Model'
16351655
},
1636-
'uri': 'gs://some-bucket/input_artifact_one'
1656+
'uri': 'gs://some-bucket/input_artifact_one',
1657+
'custom_path': '',
16371658
},
16381659
'artifact_cls': artifact_types.Model,
16391660
'expected_type': artifact_types.Model,
@@ -1645,7 +1666,8 @@ class TestDictToArtifact(parameterized.TestCase):
16451666
'type': {
16461667
'schemaTitle': 'system.Dataset'
16471668
},
1648-
'uri': 'gs://some-bucket/input_artifact_one'
1669+
'uri': 'gs://some-bucket/input_artifact_one',
1670+
'custom_path': '',
16491671
},
16501672
'artifact_cls': artifact_types.Dataset,
16511673
'expected_type': artifact_types.Dataset,
@@ -1657,7 +1679,8 @@ class TestDictToArtifact(parameterized.TestCase):
16571679
'type': {
16581680
'schemaTitle': 'system.Metrics'
16591681
},
1660-
'uri': 'gs://some-bucket/input_artifact_one'
1682+
'uri': 'gs://some-bucket/input_artifact_one',
1683+
'custom_path': '',
16611684
},
16621685
'artifact_cls': artifact_types.Metrics,
16631686
'expected_type': artifact_types.Metrics,
@@ -1669,7 +1692,8 @@ class TestDictToArtifact(parameterized.TestCase):
16691692
'type': {
16701693
'schemaTitle': 'system.ClassificationMetrics'
16711694
},
1672-
'uri': 'gs://some-bucket/input_artifact_one'
1695+
'uri': 'gs://some-bucket/input_artifact_one',
1696+
'custom_path': '',
16731697
},
16741698
'artifact_cls': artifact_types.ClassificationMetrics,
16751699
'expected_type': artifact_types.ClassificationMetrics,
@@ -1681,7 +1705,8 @@ class TestDictToArtifact(parameterized.TestCase):
16811705
'type': {
16821706
'schemaTitle': 'system.SlicedClassificationMetrics'
16831707
},
1684-
'uri': 'gs://some-bucket/input_artifact_one'
1708+
'uri': 'gs://some-bucket/input_artifact_one',
1709+
'custom_path': '',
16851710
},
16861711
'artifact_cls': artifact_types.SlicedClassificationMetrics,
16871712
'expected_type': artifact_types.SlicedClassificationMetrics,
@@ -1693,7 +1718,8 @@ class TestDictToArtifact(parameterized.TestCase):
16931718
'type': {
16941719
'schemaTitle': 'system.HTML'
16951720
},
1696-
'uri': 'gs://some-bucket/input_artifact_one'
1721+
'uri': 'gs://some-bucket/input_artifact_one',
1722+
'custom_path': '',
16971723
},
16981724
'artifact_cls': None,
16991725
'expected_type': artifact_types.HTML,
@@ -1705,7 +1731,8 @@ class TestDictToArtifact(parameterized.TestCase):
17051731
'type': {
17061732
'schemaTitle': 'system.Markdown'
17071733
},
1708-
'uri': 'gs://some-bucket/input_artifact_one'
1734+
'uri': 'gs://some-bucket/input_artifact_one',
1735+
'custom_path': '',
17091736
},
17101737
'artifact_cls': None,
17111738
'expected_type': artifact_types.Markdown,

‎sdk/python/kfp/dsl/types/artifact_types.py‎

Lines changed: 19 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -78,20 +78,23 @@ def __init__(self,
7878
name: Optional[str] = None,
7979
uri: Optional[str] = None,
8080
metadata: Optional[Dict] = None) -> None:
81-
"""Initializes the Artifact with the given name, URI and metadata."""
81+
"""Initializes the Artifact with the given name, URI, metadata, and blank custom path."""
8282
self.uri = uri or ''
8383
self.name = name or ''
8484
self.metadata = metadata or {}
85+
self._custom_path = ''
8586

8687
@property
8788
def path(self) -> str:
8889
return self._get_path()
8990

9091
@path.setter
9192
def path(self, path: str) -> None:
92-
self._set_path(path)
93+
self._set_custom_path(path)
9394

9495
def _get_path(self) -> Optional[str]:
96+
if self.custom_path is not '':
97+
return self._get_custom_path()
9598
if self.uri.startswith(RemotePrefix.GCS.value):
9699
return _GCS_LOCAL_MOUNT_PREFIX + self.uri[len(RemotePrefix.GCS.value
97100
):]
@@ -108,9 +111,23 @@ def _get_path(self) -> Optional[str]:
108111
# uri == path for local execution
109112
return self.uri
110113

114+
@property
115+
def custom_path(self) -> str:
116+
return self._custom_path
117+
118+
def _get_custom_path(self) -> str:
119+
return self._custom_path
120+
111121
def _set_path(self, path: str) -> None:
112122
self.uri = convert_local_path_to_remote_path(path)
113123

124+
def _set_custom_path(self, value: str) -> None:
125+
self._custom_path = value
126+
127+
@custom_path.setter
128+
def custom_path(self, value: str):
129+
self._custom_path = value
130+
114131

115132
def convert_local_path_to_remote_path(path: str) -> str:
116133
if path.startswith(_GCS_LOCAL_MOUNT_PREFIX):

‎test_data/sdk_compiled_pipelines/valid/critical/component_with_runtime_artifact_custom_path.py‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
1+
from sdk.python.kfp.dsl import Output
2+
from sdk.python.kfp.dsl.types.artifact_types import Artifact
3+
from sdk.python.kfp import dsl
4+
5+
@dsl.component
6+
def create_list() -> list:
7+
return [1, 2, 3, 4]
8+
9+
@dsl.component
10+
def append_to_list(digit: int, input_list: Output[Artifact]) -> list:
11+
input_list.append(digit)
12+
return input_list
13+
14+
@dsl.component
15+
def validate_custom_path(exp_path: str, input_list: Output[Artifact]) -> bool:
16+
#todo: is this the correct comparison? (or should use != instead?)
17+
if input_list.path is not exp_path:
18+
raise ValueError(f"File uri is {input_list.path} but should be {exp_path}.")
19+
20+
def component(artifact: Output[Artifact]) -> bool:
21+
return True
22+
23+
@dsl.pipeline
24+
def pipeline_with_custom_path_artifact():
25+
task1 = create_list()
26+
task1.output.set_custom_path('/etc/test/file/path')
27+
# task2 = validate_custom_path(path='/etc/test/file/path', input_list=task1.output)

‎test_data/sdk_compiled_pipelines/valid/critical/pipeline_with_artifact_custom_path.py‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,24 @@
1+
from kfp.dsl import Output
2+
from kfp.dsl.types.artifact_types import Artifact
3+
from kfp.v2 import dsl
4+
5+
6+
@dsl.component
7+
def generate_artifact() -> list:
8+
return [1, 2, 3, 4]
9+
10+
@dsl.component
11+
def append_to_list(digit: int, input_list: Output[Artifact]) -> list:
12+
input_list.append(digit)
13+
return input_list
14+
15+
@dsl.component
16+
def validate_artifact_custom_path(exp_path: str, input_list: Output[Artifact]) -> bool:
17+
if input_list.path is not exp_path:
18+
raise ValueError(f"File uri is {input_list.path} but should be {exp_path}.")
19+
20+
@dsl.pipeline
21+
def pipeline_with_custom_path_artifact():
22+
output_artifact_task = generate_artifact()
23+
output_artifact_task.output.set_custom_path('/etc/test/file/path')
24+
task2 = validate_artifact_custom_path(path='/etc/test/file/path', input_list=output_artifact_task.output)

0 commit comments

Comments
 (0)