feat(sdk): expose new ContainerSpec API (#8144)

* rename ContainerSpec to ContainerSpecImplementation for internal usage

* fix nit
This commit is contained in:
Scott_Xu 2022-08-15 13:33:34 -07:00 committed by GitHub
parent 582aefa56d
commit 35fccb1fae
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
7 changed files with 71 additions and 41 deletions

View File

@ -32,7 +32,7 @@ component_op = TestComponent(
component_spec=structures.ComponentSpec(
name='component_1',
implementation=structures.Implementation(
container=structures.ContainerSpec(
container=structures.ContainerSpecImplementation(
image='alpine',
command=[
'sh',

View File

@ -419,7 +419,7 @@ def create_component_from_func(func: Callable,
component_spec = extract_component_interface(func)
component_spec.implementation = structures.Implementation(
container=structures.ContainerSpec(
container=structures.ContainerSpecImplementation(
image=component_image,
command=packages_to_install_command + command,
args=args,
@ -481,6 +481,9 @@ def create_container_component_from_func(
arg_list.append(placeholders.InputValuePlaceholder(io_name))
container_spec = func(*arg_list)
component_spec.implementation = structures.Implementation(container_spec)
container_spec_implementation = structures.ContainerSpecImplementation.from_container_spec(
container_spec)
component_spec.implementation = structures.Implementation(
container_spec_implementation)
component_spec.validate_placeholders()
return container_component.ContainerComponent(component_spec, func)

View File

@ -43,7 +43,7 @@ def container_component(
return ContainerSpec(
image='gcr.io/my-image',
command=['sh', 'my_component.sh'],
arguments=[
args=[
'--dataset_path', dataset_path,
'--model_path', model.path,
'--output_parameter_path', output_parameter,

View File

@ -207,8 +207,9 @@ class PipelineTask:
self,
component_spec: structures.ComponentSpec,
args: Mapping[str, str],
) -> structures.ContainerSpec:
"""Resolves the command line argument placeholders in a container spec.
) -> structures.ContainerSpecImplementation:
"""Resolves the command line argument placeholders in a
ContainerSpecImplementation.
Args:
component_spec: The component definition.

View File

@ -82,7 +82,7 @@ class PipelineTaskTest(parameterized.TestCase):
expected_component_spec = structures.ComponentSpec(
name='component1',
implementation=structures.Implementation(
container=structures.ContainerSpec(
container=structures.ContainerSpecImplementation(
image='alpine',
command=['sh', '-c', 'echo "$0" >> "$1"'],
args=[
@ -104,7 +104,7 @@ class PipelineTaskTest(parameterized.TestCase):
dependent_tasks=[],
component_ref='component1',
)
expected_container_spec = structures.ContainerSpec(
expected_container_spec = structures.ContainerSpecImplementation(
image='alpine',
command=['sh', '-c', 'echo "$0" >> "$1"'],
args=[

View File

@ -179,7 +179,7 @@ class ResourceSpec(base_model.BaseModel):
class ContainerSpec(base_model.BaseModel):
"""Container implementation definition.
"""Container definition.
This is only used for pipeline authors when constructing a containerized component
using @container_component decorator.
@ -214,6 +214,18 @@ class ContainerSpec(base_model.BaseModel):
args: Optional[List[placeholders.CommandLineElement]] = None
"""Arguments to the container entrypoint."""
class ContainerSpecImplementation(base_model.BaseModel):
"""Container implementation definition."""
image: str
"""Container image."""
command: Optional[List[placeholders.CommandLineElement]] = None
"""Container entrypoint."""
args: Optional[List[placeholders.CommandLineElement]] = None
"""Arguments to the container entrypoint."""
env: Optional[Mapping[str, placeholders.CommandLineElement]] = None
"""Environment variables to be passed to the container."""
@ -233,17 +245,29 @@ class ContainerSpec(base_model.BaseModel):
self.env = None if self.env == {} else self.env
@classmethod
def from_container_dict(cls, container_dict: Dict[str,
Any]) -> 'ContainerSpec':
"""Creates a ContainerSpec from a PipelineContainerSpec message in dict
format (pipeline_spec.deploymentSpec.executors.<executor-
key>.container).
def from_container_spec(
cls,
container_spec: ContainerSpec) -> 'ContainerSpecImplementation':
return ContainerSpecImplementation(
image=container_spec.image,
command=container_spec.command,
args=container_spec.args,
env=None,
resources=None)
@classmethod
def from_container_dict(
cls, container_dict: Dict[str,
Any]) -> 'ContainerSpecImplementation':
"""Creates a ContainerSpecImplementation from a PipelineContainerSpec
message in dict format
(pipeline_spec.deploymentSpec.executors.<executor- key>.container).
Args:
container_dict (Dict[str, Any]): PipelineContainerSpec message in dict format.
Returns:
ContainerSpec: The ContainerSpec instance.
ContainerSpecImplementation: The ContainerSpecImplementation instance.
"""
args = container_dict.get('args')
if args is not None:
@ -257,7 +281,7 @@ class ContainerSpec(base_model.BaseModel):
placeholders.maybe_convert_placeholder_string_to_placeholder(c)
for c in command
]
return ContainerSpec(
return ContainerSpecImplementation(
image=container_dict['image'],
command=command,
args=args,
@ -377,7 +401,7 @@ class Implementation(base_model.BaseModel):
graph: graph implementation details.
importer: importer implementation details.
"""
container: Optional[ContainerSpec] = None
container: Optional[ContainerSpecImplementation] = None
graph: Optional[DagSpec] = None
importer: Optional[ImporterSpec] = None
@ -396,7 +420,8 @@ class Implementation(base_model.BaseModel):
"""
executor_key = utils._EXECUTOR_LABEL_PREFIX + component_name
container = deployment_spec_dict['executors'][executor_key]['container']
container_spec = ContainerSpec.from_container_dict(container)
container_spec = ContainerSpecImplementation.from_container_dict(
container)
return Implementation(container=container_spec)
@ -496,14 +521,14 @@ class ComponentSpec(base_model.BaseModel):
if getattr(implementation, 'container', None) is None:
return
containerSpec: ContainerSpec = implementation.container
containerSpecImplementation: ContainerSpecImplementation = implementation.container
valid_inputs = [] if self.inputs is None else list(self.inputs.keys())
valid_outputs = [] if self.outputs is None else list(
self.outputs.keys())
for arg in itertools.chain((containerSpec.command or []),
(containerSpec.args or [])):
for arg in itertools.chain((containerSpecImplementation.command or []),
(containerSpecImplementation.args or [])):
_check_valid_placeholder_reference(valid_inputs, valid_outputs, arg)
@classmethod
@ -549,7 +574,7 @@ class ComponentSpec(base_model.BaseModel):
command, component_dict=component_dict)
for key, command in container.get('env', {}).items()
}
container_spec = ContainerSpec.from_container_dict({
container_spec = ContainerSpecImplementation.from_container_dict({
'image': container['image'],
'command': container['command'],
'args': container['args'],

View File

@ -49,7 +49,7 @@ V1_YAML_IF_PLACEHOLDER = textwrap.dedent("""\
COMPONENT_SPEC_IF_PLACEHOLDER = structures.ComponentSpec(
name='component_if',
implementation=structures.Implementation(
container=structures.ContainerSpec(
container=structures.ContainerSpecImplementation(
image='alpine',
args=[
placeholders.IfPresentPlaceholder(
@ -83,7 +83,7 @@ V1_YAML_CONCAT_PLACEHOLDER = textwrap.dedent("""\
COMPONENT_SPEC_CONCAT_PLACEHOLDER = structures.ComponentSpec(
name='component_concat',
implementation=structures.Implementation(
container=structures.ContainerSpec(
container=structures.ContainerSpecImplementation(
image='alpine',
args=[
placeholders.ConcatPlaceholder(items=[
@ -122,7 +122,7 @@ V1_YAML_NESTED_PLACEHOLDER = textwrap.dedent("""\
COMPONENT_SPEC_NESTED_PLACEHOLDER = structures.ComponentSpec(
name='component_nested',
implementation=structures.Implementation(
container=structures.ContainerSpec(
container=structures.ContainerSpecImplementation(
image='alpine',
args=[
placeholders.ConcatPlaceholder(items=[
@ -169,7 +169,7 @@ V1_YAML_EXECUTOR_INPUT_PLACEHOLDER = textwrap.dedent("""\
COMPONENT_SPEC_EXECUTOR_INPUT_PLACEHOLDER = structures.ComponentSpec(
name='component_executor_input',
implementation=structures.Implementation(
container=structures.ContainerSpec(
container=structures.ContainerSpecImplementation(
image='alpine',
command=[
'python',
@ -197,7 +197,7 @@ class StructuresTest(parameterized.TestCase):
structures.ComponentSpec(
name='component_1',
implementation=structures.Implementation(
container=structures.ContainerSpec(
container=structures.ContainerSpecImplementation(
image='alpine',
command=[
'sh',
@ -220,7 +220,7 @@ class StructuresTest(parameterized.TestCase):
structures.ComponentSpec(
name='component_1',
implementation=structures.Implementation(
container=structures.ContainerSpec(
container=structures.ContainerSpecImplementation(
image='alpine',
command=[
'sh',
@ -241,7 +241,7 @@ class StructuresTest(parameterized.TestCase):
original_component_spec = structures.ComponentSpec(
name='component_1',
implementation=structures.Implementation(
container=structures.ContainerSpec(
container=structures.ContainerSpecImplementation(
image='alpine',
command=[
'sh',
@ -326,7 +326,7 @@ sdkVersion: kfp-2.0.0-alpha.2
expected_spec = structures.ComponentSpec(
name='component-1',
implementation=structures.Implementation(
container=structures.ContainerSpec(
container=structures.ContainerSpecImplementation(
image='alpine',
command=[
'sh',
@ -397,7 +397,7 @@ sdkVersion: kfp-2.0.0-alpha.2
expected_spec = structures.ComponentSpec(
name='Component with 2 inputs and 2 outputs',
implementation=structures.Implementation(
container=structures.ContainerSpec(
container=structures.ContainerSpecImplementation(
image='busybox',
command=[
'sh',
@ -428,27 +428,28 @@ sdkVersion: kfp-2.0.0-alpha.2
self.assertEqual(generated_spec, expected_spec)
class TestContainerSpec(unittest.TestCase):
class TestContainerSpecImplementation(unittest.TestCase):
def test_command_and_args(self):
obj = structures.ContainerSpec(
obj = structures.ContainerSpecImplementation(
image='image', command=['command'], args=['args'])
self.assertEqual(obj.command, ['command'])
self.assertEqual(obj.args, ['args'])
obj = structures.ContainerSpec(image='image', command=[], args=[])
obj = structures.ContainerSpecImplementation(
image='image', command=[], args=[])
self.assertEqual(obj.command, None)
self.assertEqual(obj.args, None)
def test_env(self):
obj = structures.ContainerSpec(
obj = structures.ContainerSpecImplementation(
image='image',
command=['command'],
args=['args'],
env={'env': 'env'})
self.assertEqual(obj.env, {'env': 'env'})
obj = structures.ContainerSpec(
obj = structures.ContainerSpecImplementation(
image='image', command=[], args=[], env={})
self.assertEqual(obj.env, None)
@ -456,7 +457,7 @@ class TestContainerSpec(unittest.TestCase):
component_spec = structures.ComponentSpec(
name='test',
implementation=structures.Implementation(
container=structures.ContainerSpec(
container=structures.ContainerSpecImplementation(
image='python:3.7',
command=[
'sh', '-c',
@ -494,7 +495,7 @@ class TestContainerSpec(unittest.TestCase):
'image': 'python:3.7'
}
loaded_container_spec = structures.ContainerSpec.from_container_dict(
loaded_container_spec = structures.ContainerSpecImplementation.from_container_dict(
container_dict)
def test_raise_error_if_access_artifact_by_itself(self):
@ -705,7 +706,7 @@ sdkVersion: kfp-2.0.0-alpha.2""")
component_spec = structures.ComponentSpec(
name='component1',
implementation=structures.Implementation(
container=structures.ContainerSpec(
container=structures.ContainerSpecImplementation(
image='alpine',
command=['sh', '-c', 'echo "$0" >> "$1"'],
args=[
@ -772,7 +773,7 @@ sdkVersion: kfp-2.0.0-alpha.2""")
component_spec = structures.ComponentSpec(
name='if',
implementation=structures.Implementation(
container=structures.ContainerSpec(
container=structures.ContainerSpecImplementation(
image='alpine',
command=['sh', '-c', 'echo "$0" "$1"'],
args=[
@ -844,7 +845,7 @@ sdkVersion: kfp-2.0.0-alpha.2""")
component_spec = structures.ComponentSpec(
name='concat',
implementation=structures.Implementation(
container=structures.ContainerSpec(
container=structures.ContainerSpecImplementation(
image='alpine',
command=[
'sh', '-c', 'echo "$0"',