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