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()