class _DummyComponent(base_component.BaseComponent):
SPEC_CLASS = _DummySpec
EXECUTOR_SPEC = executor_spec.ExecutorContainerSpec(image='dummy:latest',
command=['ls'])
def __init__(self):
super().__init__(_DummySpec())
def _create_launcher_context(self, component_config=None):
test_dir = self.get_temp_dir()
connection_config = metadata_store_pb2.ConnectionConfig()
connection_config.sqlite.SetInParent()
metadata_connection = metadata.Metadata(connection_config)
pipeline_root = os.path.join(test_dir, 'Test')
input_artifact = test_utils._InputArtifact()
input_artifact.uri = os.path.join(test_dir, 'input')
component = test_utils._FakeComponent(
name='FakeComponent',
input_channel=channel_utils.as_channel([input_artifact]),
custom_executor_spec=executor_spec.ExecutorContainerSpec(
image='gcr://test', args=['{{input_dict["input"][0].uri}}']))
pipeline_info = data_types.PipelineInfo(pipeline_name='Test',
pipeline_root=pipeline_root,
run_id='123')
driver_args = data_types.DriverArgs(enable_cache=True)
launcher = kubernetes_component_launcher.KubernetesComponentLauncher.create(
component=component,
pipeline_info=pipeline_info,
driver_args=driver_args,
metadata_connection=metadata_connection,
beam_pipeline_args=[],
additional_pipeline_args={},
component_config=component_config)
return {'launcher': launcher, 'input_artifact': input_artifact}
def testFindComponentLaunchInfoReturnConfigOverride(self):
input_artifact = test_utils._InputArtifact()
component = test_utils._FakeComponent(
name='FakeComponent',
input_channel=channel_utils.as_channel([input_artifact]),
custom_executor_spec=executor_spec.ExecutorContainerSpec(
image='gcr://test', args=['{{input_dict["input"][0].uri}}']))
default_config = docker_component_config.DockerComponentConfig()
override_config = docker_component_config.DockerComponentConfig(
name='test')
p_config = pipeline_config.PipelineConfig(
supported_launcher_classes=[
docker_component_launcher.DockerComponentLauncher
],
default_component_configs=[default_config],
component_config_overrides={
'_FakeComponent.FakeComponent': override_config
})
(launcher_class, c_config) = config_utils.find_component_launch_info(
p_config, component)
self.assertEqual(docker_component_launcher.DockerComponentLauncher,
launcher_class)
self.assertEqual(override_config, c_config)
def testCanLaunch(self):
self.assertTrue(
kubernetes_component_launcher.KubernetesComponentLauncher.
can_launch(executor_spec.ExecutorContainerSpec(image='test')))
self.assertFalse(
kubernetes_component_launcher.KubernetesComponentLauncher.
can_launch(
executor_spec.ExecutorClassSpec(base_executor.BaseExecutor)))
def testCanLaunch(self):
self.assertTrue(
docker_component_launcher.DockerComponentLauncher.can_launch(
executor_spec.ExecutorContainerSpec(image='test'),
component_config=None))
self.assertFalse(
docker_component_launcher.DockerComponentLauncher.can_launch(
executor_spec.ExecutorClassSpec(base_executor.BaseExecutor),
component_config=None))
def testExecutorContainerSpecCopy(self):
spec = executor_spec.ExecutorContainerSpec(image='path/to:image',
command=['command'],
args=['args'])
spec_copy = spec.copy()
del spec
self.assertEqual(spec_copy.image, 'path/to:image')
self.assertEqual(spec_copy.command, ['command'])
self.assertEqual(spec_copy.args, ['args'])
class ByeWorldComponent(BaseComponent):
"""Consumer component."""
SPEC_CLASS = _ByeWorldSpec
EXECUTOR_SPEC = executor_spec.ExecutorContainerSpec(
image='bash:latest',
command=['echo'],
args=['received {{input_dict["hearing"][0].value}}'])
def __init__(self, hearing):
super(ByeWorldComponent, self).__init__(_ByeWorldSpec(hearing=hearing))
class _HelloWorldComponent(base_component.BaseComponent):
SPEC_CLASS = _HelloWorldSpec
EXECUTOR_SPEC = executor_spec.ExecutorContainerSpec(
# TODO(b/143965964): move the image to private repo if the test is flaky
# due to docker hub.
image='alpine:latest',
command=['echo'],
args=['hello {{exec_properties.name}}'])
def __init__(self, name):
super(_HelloWorldComponent, self).__init__(_HelloWorldSpec(name=name))
def resolve_container_template(
container_spec_tmpl: Union[
executor_spec.ExecutorContainerSpec,
executor_specs.TemplatedExecutorContainerSpec],
input_dict: Dict[Text, List[types.Artifact]],
output_dict: Dict[Text, List[types.Artifact]],
exec_properties: Dict[Text,
Any]) -> executor_spec.ExecutorContainerSpec:
"""Resolves Jinja2 template languages from an executor container spec.
Args:
container_spec_tmpl: the container spec template to be resolved.
input_dict: Dictionary of input artifacts consumed by this component.
output_dict: Dictionary of output artifacts produced by this component.
exec_properties: Dictionary of execution properties.
Returns:
A resolved container spec.
"""
context = {
'input_dict': input_dict,
'output_dict': output_dict,
'exec_properties': exec_properties,
}
if isinstance(container_spec_tmpl,
executor_specs.TemplatedExecutorContainerSpec):
return executor_spec.ExecutorContainerSpec(
image=container_spec_tmpl.image,
command=_resolve_container_command_line(
cmd_args=container_spec_tmpl.command,
input_dict=input_dict,
output_dict=output_dict,
exec_properties=exec_properties,
),
)
return executor_spec.ExecutorContainerSpec(
image=_render_text(container_spec_tmpl.image, context),
command=_render_items(container_spec_tmpl.command, context),
args=_render_items(container_spec_tmpl.args, context))
class HelloWorldComponent(BaseComponent):
"""Producer component."""
SPEC_CLASS = _HelloWorldSpec
EXECUTOR_SPEC = executor_spec.ExecutorContainerSpec(
# TODO(b/143965964): move the image to private repo if the test is flaky
# due to docker hub.
image='google/cloud-sdk:latest',
command=['sh', '-c'],
args=[
'echo "hello {{exec_properties.word}}" | gsutil cp - {{output_dict["greeting"][0].uri}}'
])
def __init__(self, word, greeting=None):
if not greeting:
artifact = standard_artifacts.String()
greeting = channel_utils.as_channel([artifact])
super(HelloWorldComponent,
self).__init__(_HelloWorldSpec(word=word, greeting=greeting))
def testResolveContainerTemplate(self):
container_spec = executor_spec.ExecutorContainerSpec(
image='gcr.io/my/trainer:{{exec_properties.version}}',
command=['{{exec_properties.model}}_trainer'],
args=[
'--steps',
'{{exec_properties.train_args.num_steps}}',
'--examples',
'{{input_dict["examples"]|join(",",attribute="uri")}}',
'--model-path',
'{{output_dict["model"][0].uri}}',
])
examples_artifact_1 = standard_artifacts.Examples()
examples_artifact_1.uri = 'gcs://examples/1'
examples_artifact_2 = standard_artifacts.Examples()
examples_artifact_2.uri = 'gcs://examples/2'
model = standard_artifacts.Model()
model.uri = 'gcs://model'
input_dict = {'examples': [examples_artifact_1, examples_artifact_2]}
output_dict = {'model': [model]}
exec_properties = {
'version': 'v1',
'model': 'cnn',
'train_args': trainer_pb2.TrainArgs(num_steps=10000),
}
actual_spec = container_common.resolve_container_template(
container_spec, input_dict, output_dict, exec_properties)
self.assertEqual('gcr.io/my/trainer:v1', actual_spec.image)
self.assertListEqual(['cnn_trainer'], actual_spec.command)
self.assertListEqual([
'--steps',
'10000',
'--examples',
'gcs://examples/1,gcs://examples/2',
'--model-path',
'gcs://model',
], actual_spec.args)
def run_executor(
self, execution_info: data_types.ExecutionInfo
) -> execution_result_pb2.ExecutorOutput:
"""Execute underlying component implementation.
Runs executor container in a Kubernetes Pod and wait until it goes into
`Succeeded` or `Failed` state.
Args:
execution_info: All the information that the launcher provides.
Raises:
RuntimeError: when the pod is in `Failed` state or unexpected failure from
Kubernetes API.
Returns:
An ExecutorOutput instance
"""
context = placeholder_utils.ResolutionContext(
exec_info=execution_info,
executor_spec=self._executor_spec,
platform_config=self._platform_config)
container_spec = executor_spec_lib.ExecutorContainerSpec(
image=self._container_executor_spec.image,
command=[
placeholder_utils.resolve_placeholder_expression(cmd, context)
for cmd in self._container_executor_spec.commands
] or None,
args=[
placeholder_utils.resolve_placeholder_expression(arg, context)
for arg in self._container_executor_spec.args
] or None,
)
# Replace container spec with jinja2 template.
input_dict = execution_info.input_dict
output_dict = execution_info.output_dict
exec_properties = execution_info.exec_properties
container_spec = container_common.resolve_container_template(
container_spec, input_dict, output_dict, exec_properties)
pod_name = self._build_pod_name(execution_info)
# TODO(hongyes): replace the default value from component config.
try:
namespace = kube_utils.get_kfp_namespace()
except RuntimeError:
namespace = 'kubeflow'
pod_manifest = self._build_pod_manifest(pod_name, container_spec)
core_api = kube_utils.make_core_v1_api()
if kube_utils.is_inside_kfp():
launcher_pod = kube_utils.get_current_kfp_pod(core_api)
pod_manifest['spec'][
'serviceAccount'] = launcher_pod.spec.service_account
pod_manifest['spec'][
'serviceAccountName'] = launcher_pod.spec.service_account_name
pod_manifest['metadata'][
'ownerReferences'] = container_common.to_swagger_dict(
launcher_pod.metadata.owner_references)
else:
pod_manifest['spec'][
'serviceAccount'] = kube_utils.TFX_SERVICE_ACCOUNT
pod_manifest['spec'][
'serviceAccountName'] = kube_utils.TFX_SERVICE_ACCOUNT
logging.info('Looking for pod "%s:%s".', namespace, pod_name)
resp = kube_utils.get_pod(core_api, pod_name, namespace)
if not resp:
logging.info('Pod "%s:%s" does not exist. Creating it...',
namespace, pod_name)
logging.info('Pod manifest: %s', pod_manifest)
try:
resp = core_api.create_namespaced_pod(namespace=namespace,
body=pod_manifest)
except client.rest.ApiException as e:
raise RuntimeError(
'Failed to created container executor pod!\nReason: %s\nBody: %s'
% (e.reason, e.body))
# Wait up to 300 seconds for the pod to move from pending to another status.
logging.info('Waiting for pod "%s:%s" to start.', namespace, pod_name)
kube_utils.wait_pod(
core_api,
pod_name,
namespace,
exit_condition_lambda=kube_utils.pod_is_not_pending,
condition_description='non-pending status',
timeout_sec=300)
logging.info('Start log streaming for pod "%s:%s".', namespace,
pod_name)
try:
logs = core_api.read_namespaced_pod_log(
name=pod_name,
namespace=namespace,
container=kube_utils.ARGO_MAIN_CONTAINER_NAME,
follow=True,
_preload_content=False).stream()
except client.rest.ApiException as e:
raise RuntimeError(
'Failed to stream the logs from the pod!\nReason: %s\nBody: %s'
% (e.reason, e.body))
for log in logs:
logging.info(log.decode().rstrip('\n'))
# Wait indefinitely for the pod to complete.
resp = kube_utils.wait_pod(
core_api,
pod_name,
namespace,
exit_condition_lambda=kube_utils.pod_is_done,
condition_description='done state')
if resp.status.phase == kube_utils.PodPhase.FAILED.value:
raise RuntimeError('Pod "%s:%s" failed with status "%s".' %
(namespace, pod_name, resp.status))
logging.info('Pod "%s:%s" is done.', namespace, pod_name)
return execution_result_pb2.ExecutorOutput()
