def test_func(output_artifact_two: OutputArtifact(Metrics)):
output_artifact_two.get().metadata['key_1'] = 'value_1'
output_artifact_two.get().metadata['key_2'] = 2
output_artifact_two.uri = 'new-uri'
# log_metric works here since the schema is specified as Metrics.
output_artifact_two.get().log_metric('metric', 0.9)
def train(
# Use InputPath to get a locally accessible path for the input artifact
# of type `Dataset`.
dataset_one_path: InputPath('Dataset'),
# Use InputArtifact to get a metadata-rich handle to the input artifact
# of type `Dataset`.
dataset_two: InputArtifact(Dataset),
# An input parameter of type string.
message: str,
# Use OutputArtifact to get a metadata-rich handle to the output artifact
# of type `Dataset`.
model: OutputArtifact(Model),
# An input parameter of type int with a default value.
num_steps: int = 100):
'''Dummy Training step'''
with open(dataset_one_path, 'r') as input_file:
dataset_one_contents = input_file.read()
with open(dataset_two.path, 'r') as input_file:
dataset_two_contents = input_file.read()
line = "dataset_one_contents: {} || dataset_two_contents: {} || message: {}\n".format(
dataset_one_contents, dataset_two_contents, message)
with open(model.path, 'w') as output_file:
for i in range(num_steps):
output_file.write("Step {}\n{}\n=====\n".format(i, line))
# Use model.get() to get a Model artifact, which has a .metadata dictionary
# to store arbitrary metadata for the output artifact.
model.get().metadata['accuracy'] = 0.9
def preprocess(
# An input parameter of type string.
message: str,
# Use OutputArtifact to get a metadata-rich handle to the output artifact
# of type `Dataset`.
output_dataset_one: OutputArtifact(Dataset),
# A locally accessible filepath for another output artifact of type
# `Dataset`.
output_dataset_two_path: OutputPath('Dataset'),
# A locally accessible filepath for an output parameter of type string.
output_parameter_path: OutputPath(str)):
'''Dummy preprocessing step'''
# Use OutputArtifact.path to access a local file path for writing.
# One can also use OutputArtifact.uri to access the actual URI file path.
with open(output_dataset_one.path, 'w') as f:
f.write(message)
# OutputPath is used to just pass the local file path of the output artifact
# to the function.
with open(output_dataset_two_path, 'w') as f:
f.write(message)
with open(output_parameter_path, 'w') as f:
f.write(message)
def test_func(output_artifact_one: OutputArtifact(Model)):
# Test that output artifacts always have filename 'data' added.
self.assertEqual(output_artifact_one.uri,
'gs://some-bucket/output_artifact_one/data')
self.assertEqual(
output_artifact_one.path,
os.path.join(self._test_dir, 'some-bucket/output_artifact_one',
'data'))
self.assertEqual(output_artifact_one.get().name, 'output_artifact_one')
def _make_output_artifact(cls, runtime_artifact: Dict):
artifact = create_runtime_artifact(runtime_artifact)
return OutputArtifact(artifact_type=type(artifact), artifact=artifact)