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 testInitSucceed(self): # Init with default parameters pipeline_config.PipelineConfig() # Init with custom parameters pipeline_config.PipelineConfig( supported_launcher_classes=[ in_process_component_launcher.InProcessComponentLauncher ], default_component_configs=[ docker_component_config.DockerComponentConfig() ], component_config_overrides={ 'comp-1', docker_component_config.DockerComponentConfig() })
def testLaunchSucceedsWithConfig(self, mock_docker_client, mock_publisher): mock_publisher.return_value.publish_execution.return_value = {} mock_run = mock_docker_client.return_value.containers.run mock_run.return_value.logs.return_value = [] mock_run.return_value.wait.return_value = {'StatusCode': 0} docker_config = docker_component_config.DockerComponentConfig( docker_server_url='http://mock.docker.server', environment={'name': 'value'}, privileged=True, volumes=['/local/etc:/local/etc'], ports={'2222/tcp': 3333}) context = self._create_launcher_context(docker_config) context['launcher'].launch() mock_run.assert_called_once() _, mock_kwargs = mock_run.call_args self.assertEqual('gcr://test', mock_kwargs['image']) self.assertListEqual([context['input_artifact'].uri], mock_kwargs['command']) mock_docker_client.assert_called_with( base_url='http://mock.docker.server') self.assertDictEqual({'name': 'value'}, mock_kwargs['environment']) self.assertTrue(mock_kwargs['privileged']) self.assertListEqual(['/local/etc:/local/etc'], mock_kwargs['volumes']) self.assertDictEqual({'2222/tcp': 3333}, mock_kwargs['ports'])
def run_executor( self, execution_info: data_types.ExecutionInfo ) -> execution_result_pb2.ExecutorOutput: """Execute underlying component implementation.""" context = placeholder_utils.ResolutionContext( exec_info=execution_info, executor_spec=self._executor_spec, platform_config=self._platform_config) component_executor_spec = ( executor_specs.TemplatedExecutorContainerSpec( image=self._container_executor_spec.image, command=[ placeholder_utils.resolve_placeholder_expression( cmd, context) for cmd in self._container_executor_spec.commands ])) docker_config = docker_component_config.DockerComponentConfig() logging.info('Container spec: %s', vars(component_executor_spec)) logging.info('Docker config: %s', vars(docker_config)) # Call client.containers.run and wait for completion. # ExecutorContainerSpec follows k8s container spec which has different # names to Docker's container spec. It's intended to set command to docker's # entrypoint and args to docker's command. if docker_config.docker_server_url: client = docker.DockerClient( base_url=docker_config.docker_server_url) else: client = docker.from_env() run_args = docker_config.to_run_args() container = client.containers.run( image=component_executor_spec.image, command=component_executor_spec.command, detach=True, **run_args) # Streaming logs for log in container.logs(stream=True): logging.info('Docker: %s', log.decode('utf-8')) exit_code = container.wait()['StatusCode'] if exit_code != 0: raise RuntimeError( 'Container exited with error code "{}"'.format(exit_code)) # TODO(b/141192583): Report data to publisher # - report container digest # - report replaced command line entrypoints # - report docker run args return execution_result_pb2.ExecutorOutput()
def testToRunArgs(self): docker_config = docker_component_config.DockerComponentConfig( docker_server_url='http://mock.docker.server', environment={'name': 'value'}, privileged=True, volumes=['/local/etc:/local/etc'], ports={'2222/tcp': 3333}) run_args = docker_config.to_run_args() self.assertEqual('http://mock.docker.server', docker_config.docker_server_url) self.assertDictEqual({'name': 'value'}, run_args['environment']) self.assertTrue(run_args['privileged']) self.assertListEqual(['/local/etc:/local/etc'], run_args['volumes']) self.assertDictEqual({'2222/tcp': 3333}, run_args['ports'])
def _run_executor(self, execution_id: int, input_dict: Dict[Text, List[types.Artifact]], output_dict: Dict[Text, List[types.Artifact]], exec_properties: Dict[Text, Any]) -> None: """Execute underlying component implementation.""" executor_container_spec = cast(executor_spec.ExecutorContainerSpec, self._component_executor_spec) if self._component_config: docker_config = cast(docker_component_config.DockerComponentConfig, self._component_config) else: docker_config = docker_component_config.DockerComponentConfig() # Replace container spec with jinja2 template. executor_container_spec = container_common.resolve_container_template( executor_container_spec, input_dict, output_dict, exec_properties) absl.logging.info('Container spec: %s' % vars(executor_container_spec)) absl.logging.info('Docker config: %s' % vars(docker_config)) # Call client.containers.run and wait for completion. # ExecutorContainerSpec follows k8s container spec which has different # names to Docker's container spec. It's intended to set command to docker's # entrypoint and args to docker's command. if docker_config.docker_server_url: client = docker.DockerClient( base_url=docker_config.docker_server_url) else: client = docker.from_env() run_args = docker_config.to_run_args() container = client.containers.run( image=executor_container_spec.image, entrypoint=executor_container_spec.command, command=executor_container_spec.args, detach=True, **run_args) # Streaming logs for log in container.logs(stream=True): absl.logging.info('Docker: ' + log.decode('utf-8')) exit_code = container.wait()['StatusCode'] if exit_code != 0: raise RuntimeError( 'Container exited with error code "{}"'.format(exit_code))
def testDockerComponentLauncherInBeam(self): beam_dag_runner.BeamDagRunner(config=pipeline_config.PipelineConfig( supported_launcher_classes=[ docker_component_launcher.DockerComponentLauncher ], default_component_configs=[ docker_component_config.DockerComponentConfig() ])).run( _create_pipeline(pipeline_name=self._pipeline_name, pipeline_root=self._pipeline_root, metadata_path=self._metadata_path, name='docker_e2e_test_in_beam')) metadata_config = metadata.sqlite_metadata_connection_config( self._metadata_path) with metadata.Metadata(metadata_config) as m: self.assertEqual(1, len(m.store.get_executions()))
def testInitFailWithDupDefaultComponentConfigClasses(self): with self.assertRaises(ValueError): pipeline_config.PipelineConfig(default_component_configs=[ docker_component_config.DockerComponentConfig(), docker_component_config.DockerComponentConfig(), ])