From 0451d4bddff6369eca994ecf0a30ce1293fead05 Mon Sep 17 00:00:00 2001 From: Alexey Volkov Date: Tue, 27 Aug 2019 22:02:13 -0700 Subject: [PATCH] SDK - Components - Added output references to TaskSpec Also added TaskSpec.task and ComponentReference.spec attributes --- sdk/python/kfp/components/_components.py | 4 ++- sdk/python/kfp/components/_dsl_bridge.py | 2 +- sdk/python/kfp/components/_structures.py | 23 ++++++++++++++-- .../tests/components/test_components.py | 27 +++++++++++++++++++ 4 files changed, 52 insertions(+), 4 deletions(-) diff --git a/sdk/python/kfp/components/_components.py b/sdk/python/kfp/components/_components.py index 0acd17b5c14..88a0b8bda57 100644 --- a/sdk/python/kfp/components/_components.py +++ b/sdk/python/kfp/components/_components.py @@ -206,7 +206,7 @@ def _create_task_factory_from_component_spec(component_spec:ComponentSpec, compo if component_ref is None: component_ref = ComponentReference(name=component_spec.name or component_filename or _default_component_name) - component_ref._component_spec = component_spec + component_ref.spec = component_spec def create_task_from_component_and_arguments(pythonic_arguments): arguments = {} @@ -238,6 +238,8 @@ def create_task_from_component_and_arguments(pythonic_arguments): component_ref=component_ref, arguments=arguments, ) + task._init_outputs() + if _created_task_transformation_handler: task = _created_task_transformation_handler[-1](task) return task diff --git a/sdk/python/kfp/components/_dsl_bridge.py b/sdk/python/kfp/components/_dsl_bridge.py index 750789f95da..0c4488d8982 100644 --- a/sdk/python/kfp/components/_dsl_bridge.py +++ b/sdk/python/kfp/components/_dsl_bridge.py @@ -20,7 +20,7 @@ def create_container_op_from_task(task_spec: TaskSpec): argument_values = task_spec.arguments - component_spec = task_spec.component_ref._component_spec + component_spec = task_spec.component_ref.spec if not isinstance(component_spec.implementation, ContainerImplementation): raise TypeError('Only container component tasks can be converted to ContainerOp') diff --git a/sdk/python/kfp/components/_structures.py b/sdk/python/kfp/components/_structures.py index b93945399b2..3d3bf1cfe81 100644 --- a/sdk/python/kfp/components/_structures.py +++ b/sdk/python/kfp/components/_structures.py @@ -315,12 +315,13 @@ def __init__(self, digest: Optional[str] = None, tag: Optional[str] = None, url: Optional[str] = None, + spec: Optional[ComponentSpec] = None, ): super().__init__(locals()) self._post_init() def _post_init(self) -> None: - if not any([self.name, self.digest, self.tag, self.url]): + if not any([self.name, self.digest, self.tag, self.url, self.spec]): raise TypeError('Need at least one argument.') @@ -344,10 +345,13 @@ class TaskOutputReference(ModelBase): } def __init__(self, - task_id: str, output_name: str, + task_id: Optional[str] = None, # Used for linking to the upstream task in serialized component file. + task: Optional['TaskSpec'] = None, # Used for linking to the upstream task in runtime since Task does not have an ID until inserted into a graph. ): super().__init__(locals()) + if self.task_id is None and self.task is None: + raise TypeError('task_id and task cannot be None at the same time.') class TaskOutputArgument(ModelBase): #Has additional constructor for convenience @@ -483,6 +487,21 @@ def __init__(self, super().__init__(locals()) #TODO: If component_ref is resolved to component spec, then check that the arguments correspond to the inputs + def _init_outputs(self): + #Adding output references to the task + if self.component_ref.spec is None: + return + task_outputs = OrderedDict() + for output in self.component_ref.spec.outputs or []: + task_output_ref = TaskOutputReference( + output_name=output.name, + task=self, + ) + task_output_arg = TaskOutputArgument(task_output=task_output_ref) + task_outputs[output.name] = task_output_arg + + self.outputs = task_outputs + class GraphSpec(ModelBase): '''Describes the graph component implementation. It represents a graph of component tasks connected to the upstream sources of data using the argument specifications. It also describes the sources of graph output values.''' diff --git a/sdk/python/tests/components/test_components.py b/sdk/python/tests/components/test_components.py index 2c621a000a5..38b15c72de0 100644 --- a/sdk/python/tests/components/test_components.py +++ b/sdk/python/tests/components/test_components.py @@ -15,6 +15,7 @@ import os import sys import unittest +from contextlib import contextmanager from pathlib import Path @@ -23,6 +24,16 @@ from kfp.components._yaml_utils import load_yaml from kfp.dsl.types import InconsistentTypeException + +@contextmanager +def no_task_resolving_context(): + old_handler = kfp.components._components._created_task_transformation_handler + try: + kfp.components._components._created_task_transformation_handler = None + yield None + finally: + kfp.components._components._created_task_transformation_handler = old_handler + class LoadComponentTestCase(unittest.TestCase): def _test_load_component_from_file(self, component_path: str): task_factory1 = comp.load_component_from_file(component_path) @@ -561,6 +572,22 @@ def test_passing_component_metadata_to_container_op(self): self.assertEqual(task1.pod_annotations['key1'], 'value1') self.assertEqual(task1.pod_labels['key1'], 'value1') + def test_check_task_spec_outputs_dictionary(self): + component_text = '''\ +outputs: +- {name: out 1} +- {name: out 2} +implementation: + container: + image: busybox + command: [touch, {outputPath: out 1}, {outputPath: out 2}] +''' + op = comp.load_component_from_text(component_text) + with no_task_resolving_context(): + task = op() + + self.assertEqual(list(task.outputs.keys()), ['out 1', 'out 2']) + def test_type_compatibility_check_for_simple_types(self): component_a = '''\ outputs: