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 _create_launcher_context(self, component_config=None): test_dir = self.get_temp_dir() input_artifact = test_utils._InputArtifact() input_artifact.uri = os.path.join(test_dir, 'input') operator = docker_executor_operator.DockerExecutorOperator( _EXECUTOR_SEPC, _PLATFORM_CONFIG) return {'operator': operator, 'input_artifact': input_artifact}
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 _create_launcher_context(self, component_config=None): test_dir = self.get_temp_dir() connection_config = metadata_store_pb2.ConnectionConfig() connection_config.sqlite.SetInParent() input_artifact = test_utils._InputArtifact() input_artifact.uri = os.path.join(test_dir, 'input') operator = kubernetes_executor_operator.KubernetesExecutorOperator( _SEPC) return {'operator': operator, 'input_artifact': input_artifact}
def testFindComponentLaunchInfoReturnDefaultLaunchInfo(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() (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)