def _create_launcher_context(self, component_config=None):
test_dir = self.get_temp_dir()
connection_config = metadata_store_pb2.ConnectionConfig()
connection_config.sqlite.SetInParent()
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_config=connection_config,
beam_pipeline_args=[],
additional_pipeline_args={},
component_config=component_config)
return {'launcher': launcher, 'input_artifact': input_artifact}
def setUp(self):
super(StubComponentLauncherTest, self).setUp()
test_dir = os.path.join(
os.environ.get('TEST_UNDECLARED_OUTPUTS_DIR', self.get_temp_dir()),
self._testMethodName)
connection_config = metadata_store_pb2.ConnectionConfig()
connection_config.sqlite.SetInParent()
self.metadata_connection = metadata.Metadata(connection_config)
self.pipeline_root = os.path.join(test_dir, 'Test')
self.input_dir = os.path.join(test_dir, 'input')
self.output_dir = os.path.join(test_dir, 'output')
self.record_dir = os.path.join(test_dir, 'record')
tf.io.gfile.makedirs(self.input_dir)
tf.io.gfile.makedirs(self.output_dir)
tf.io.gfile.makedirs(self.record_dir)
input_artifact = test_utils._InputArtifact() # pylint: disable=protected-access
input_artifact.uri = os.path.join(self.input_dir, 'result.txt')
output_artifact = test_utils._OutputArtifact() # pylint: disable=protected-access
output_artifact.uri = os.path.join(self.output_dir, 'result.txt')
self.component = test_utils._FakeComponent( # pylint: disable=protected-access
name='FakeComponent',
input_channel=channel_utils.as_channel([input_artifact]),
output_channel=channel_utils.as_channel([output_artifact]))
self.driver_args = data_types.DriverArgs(enable_cache=True)
self.pipeline_info = data_types.PipelineInfo(
pipeline_name='Test',
pipeline_root=self.pipeline_root,
run_id='123')
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 testFindComponentLaunchInfoFailWithNoLauncherClassFound(self):
input_artifact = test_utils._InputArtifact()
component = test_utils._FakeComponent(
name='FakeComponent',
input_channel=channel_utils.as_channel([input_artifact]))
p_config = pipeline_config.PipelineConfig(supported_launcher_classes=[
docker_component_launcher.DockerComponentLauncher
])
with self.assertRaises(RuntimeError):
# DockerComponentLauncher cannot launch class executor.
config_utils.find_component_launch_info(p_config, component)
def testFindComponentLaunchInfoReturnDefaultLaunchInfo(self):
input_artifact = types.Artifact(type_name='InputPath')
component = test_utils._FakeComponent(
name='FakeComponent',
input_channel=channel_utils.as_channel([input_artifact]))
p_config = pipeline_config.PipelineConfig()
(launcher_class,
c_config) = config_utils.find_component_launch_info(p_config, component)
self.assertEqual(in_process_component_launcher.InProcessComponentLauncher,
launcher_class)
self.assertIsNone(c_config)
def testRun(self, mock_publisher):
mock_publisher.return_value.publish_execution.return_value = {}
test_dir = os.path.join(
os.environ.get('TEST_UNDECLARED_OUTPUTS_DIR', self.get_temp_dir()),
self._testMethodName)
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_path = os.path.join(test_dir, 'input')
fileio.makedirs(os.path.dirname(input_path))
file_io.write_string_to_file(input_path, 'test')
input_artifact = test_utils._InputArtifact()
input_artifact.uri = input_path
component = test_utils._FakeComponent(
name='FakeComponent',
input_channel=channel_utils.as_channel([input_artifact]))
pipeline_info = data_types.PipelineInfo(pipeline_name='Test',
pipeline_root=pipeline_root,
run_id='123')
driver_args = data_types.DriverArgs(enable_cache=True)
# We use InProcessComponentLauncher to test BaseComponentLauncher logics.
launcher = in_process_component_launcher.InProcessComponentLauncher.create(
component=component,
pipeline_info=pipeline_info,
driver_args=driver_args,
metadata_connection=metadata_connection,
beam_pipeline_args=[],
additional_pipeline_args={})
self.assertEqual(
launcher._component_info.component_type, '.'.join([
test_utils._FakeComponent.__module__,
test_utils._FakeComponent.__name__
]))
launcher.launch()
output_path = component.outputs['output'].get()[0].uri
self.assertTrue(fileio.exists(output_path))
contents = file_io.read_file_to_string(output_path)
self.assertEqual('test', contents)
