def testCreateComponent(self): component = kfp_container_component.load_kfp_yaml_container_component( os.path.join(self._testdata_path, 'kfp_container_component_test.yaml')) ref_component = container_component.create_container_component( name='Test_Kfp_Container_Component', image='image1', command=[ 'command1', 'command2', 'command3', placeholders.InputUriPlaceholder('Directory'), placeholders.InputValuePlaceholder('Subpath'), placeholders.OutputUriPlaceholder('File'), '--arg1', placeholders.InputUriPlaceholder('input1'), '--arg2', placeholders.InputValuePlaceholder('input2'), '--arg3', placeholders.OutputUriPlaceholder('output1'), ], inputs={ 'input1': standard_artifacts.String, 'input2': standard_artifacts.String, }, outputs={ 'output1': standard_artifacts.String, }, parameters={}, ) self.assertEqual(type(component), type(ref_component)) self.assertEqual(ref_component.EXECUTOR_SPEC, component.EXECUTOR_SPEC)
class DummyProducerComponent(base_component.BaseComponent): """Dummy producer component.""" SPEC_CLASS = ProducerSpec EXECUTOR_SPEC = executor_specs.TemplatedExecutorContainerSpec( image='dummy/producer', command=[ 'producer', '--output1', placeholders.OutputUriPlaceholder('output1'), '--param1', placeholders.InputValuePlaceholder('param1'), '--wrapped-param', placeholders.ConcatPlaceholder([ 'prefix-', placeholders.InputValuePlaceholder('param1'), '-suffix', ]), ]) def __init__(self, param1, output1, instance_name=None): spec = ProducerSpec( output1=output1, param1=param1, ) super(DummyProducerComponent, self).__init__(spec=spec) if instance_name: self._id = '{}.{}'.format(self.__class__.__name__, instance_name) else: self._id = self.__class__.__name__
def testConstructSingleVmJob(self): training = ai_platform_training_component.create_ai_platform_training( name='my_training_step', project_id='my-project', region='us-central1', image_uri='gcr.io/my-project/caip-training-test:latest', args=[ '--examples', placeholders.InputUriPlaceholder('examples'), '--n-steps', placeholders.InputValuePlaceholder('n_step'), '--model-dir', placeholders.OutputUriPlaceholder('model') ], scale_tier='BASIC_GPU', inputs={ 'examples': self.examples, }, outputs={'model': standard_artifacts.Model}, parameters={'n_step': 100}) expected_aip_config = { ai_platform_training_executor.PROJECT_CONFIG_KEY: 'my-project', ai_platform_training_executor.TRAINING_JOB_CONFIG_KEY: { 'training_input': { 'scaleTier': 'BASIC_GPU', 'region': 'us-central1', 'masterConfig': { 'imageUri': 'gcr.io/my-project/caip-training-test:latest' }, 'args': [ '--examples', placeholders.InputUriPlaceholder('examples'), '--n-steps', placeholders.InputValuePlaceholder('n_step'), '--model-dir', placeholders.OutputUriPlaceholder('model') ] }, ai_platform_training_executor.LABELS_CONFIG_KEY: None, }, ai_platform_training_executor.JOB_ID_CONFIG_KEY: None, ai_platform_training_executor.LABELS_CONFIG_KEY: None, } # exec_properties has two entries: one is the user-defined 'n_step', another # is the aip_training_config. self.assertLen(training.exec_properties, 2) self.assertEqual(training.outputs['model'].type_name, standard_artifacts.Model.TYPE_NAME) self.assertEqual(training.inputs['examples'].type_name, standard_artifacts.Examples.TYPE_NAME) self.assertEqual(training.exec_properties['n_step'], 100) self.assertEqual( training.exec_properties[ai_platform_training_executor.CONFIG_KEY], json_utils.dumps(expected_aip_config))
class DummyContainerSpecComponent(base_component.BaseComponent): """Dummy ContainerSpec component.""" SPEC_CLASS = TransformerSpec EXECUTOR_SPEC = executor_specs.TemplatedExecutorContainerSpec( image='dummy/transformer', command=[ 'transformer', '--input1', placeholders.InputUriPlaceholder('input1'), '--output1', placeholders.OutputUriPlaceholder('output1'), '--param1', placeholders.InputValuePlaceholder('param1'), ]) def __init__(self, input1, param1, output1, instance_name=None): spec = TransformerSpec( input1=input1, output1=output1, param1=param1, ) super(DummyContainerSpecComponent, self).__init__(spec=spec) if instance_name: self._id = '{}.{}'.format(self.__class__.__name__, instance_name) else: self._id = self.__class__.__name__
def testSuccessfulExecution(self): example_importer = importer.Importer( artifact_type=simple_artifacts.File, reimport=False, source_uri=f'gs://{self._TEST_DATA_BUCKET}/ai-platform-training/mnist' ).with_id('examples') train = ai_platform_training_component.create_ai_platform_training( name='simple_aip_training', project_id=self._GCP_PROJECT_ID, region=self._GCP_REGION, image_uri=self._TRAINING_IMAGE, args=[ '--dataset', placeholders.InputUriPlaceholder('examples'), '--model-dir', placeholders.OutputUriPlaceholder('model'), '--lr', placeholders.InputValuePlaceholder('learning_rate'), ], scale_tier='BASIC', inputs={'examples': example_importer.outputs['result']}, outputs={'model': standard_artifacts.Model}, parameters={'learning_rate': '0.001'}) pipeline_name = _PIPELINE_NAME_PREFIX.format(test_utils.random_id()) aip_training_pipeline = pipeline.Pipeline( pipeline_name=pipeline_name, pipeline_root=self._pipeline_root(pipeline_name), components=[example_importer, train], ) self._run_pipeline(aip_training_pipeline)
def setUp(self): super().setUp() self._text = 'text' self._input_value_placeholder = placeholders.InputValuePlaceholder( 'input_artifact') self._another_input_value_placeholder = placeholders.InputValuePlaceholder( 'input_parameter') self._input_uri_placeholder = placeholders.InputUriPlaceholder( 'input_artifact') self._output_uri_placeholder = placeholders.OutputUriPlaceholder( 'output_artifact') self._concat_placeholder = placeholders.ConcatPlaceholder([ self._text, self._input_value_placeholder, self._input_uri_placeholder, self._output_uri_placeholder, ]) self._text_concat_placeholder = placeholders.ConcatPlaceholder( [self._text, 'text1', placeholders.ConcatPlaceholder(['text2']),])
def tasks_for_pipeline_with_artifact_value_passing(): """A simple pipeline with artifact consumed as value.""" producer_component = container_component.create_container_component( name='Produce', outputs={ 'data': simple_artifacts.File, }, parameters={ 'message': str, }, image='gcr.io/ml-pipeline/mirrors/cloud-sdk', command=[ 'sh', '-exc', """ message="$0" output_data_uri="$1" output_data_path=$(mktemp) # Running the main code echo "Hello $message" >"$output_data_path" # Getting data out of the container gsutil cp -r "$output_data_path" "$output_data_uri" """, placeholders.InputValuePlaceholder('message'), placeholders.OutputUriPlaceholder('data'), ], ) print_value_component = container_component.create_container_component( name='Print', inputs={ 'text': simple_artifacts.File, }, image='gcr.io/ml-pipeline/mirrors/cloud-sdk', command=[ 'echo', placeholders.InputValuePlaceholder('text'), ], ) producer_task = producer_component(message='World!') print_task = print_value_component(text=producer_task.outputs['data'], ) return [producer_task, print_task]
def _get_command_line_argument_type( command: kfp_component_spec_pb2.StringOrPlaceholder ) -> placeholders.CommandlineArgumentType: """Converts a container command to the corresponding type. Args: command: StringOrPlaceholder which encodes a container command. Returns: command to be passed into create_container_component. """ if command.HasField('constantValue'): return command.constantValue if command.HasField('inputValue'): return placeholders.InputValuePlaceholder(command.inputValue) if command.HasField('inputPath'): return placeholders.InputUriPlaceholder(command.inputPath) if command.HasField('outputPath'): return placeholders.OutputUriPlaceholder(command.outputPath) raise ValueError('Unrecognized command %s' % command)
def setUp(self): super(ExecutorSpecsTest, self).setUp() self._text = 'text' self._input_value_placeholder = placeholders.InputValuePlaceholder( 'input_artifact') self._input_uri_placeholder = placeholders.InputUriPlaceholder( 'input_uri') self._output_uri_placeholder = placeholders.OutputUriPlaceholder( 'output_uri') self._concat_placeholder = placeholders.ConcatPlaceholder([ self._text, self._input_value_placeholder, self._input_uri_placeholder, self._output_uri_placeholder, ]) self._text_concat_placeholder = placeholders.ConcatPlaceholder([ self._text, 'text1', placeholders.ConcatPlaceholder(['text2']), ])
def testConstructFullSpec(self): training_input = { 'scaleTier': 'BASIC_GPU', 'region': 'us-central1', 'masterConfig': { 'imageUri': 'gcr.io/my-project/caip-training-test:latest' }, 'args': [ '--examples', placeholders.InputUriPlaceholder('examples'), '--n-steps', placeholders.InputValuePlaceholder('n_step'), '--model-dir', placeholders.OutputUriPlaceholder('model') ] } training = ai_platform_training_component.create_ai_platform_training( name='my_training_step', project_id='my-project', training_input=training_input, inputs={ 'examples': self.examples, }, outputs={'model': standard_artifacts.Model}, parameters={'n_step': 100}) expected_aip_config = { ai_platform_training_executor.PROJECT_CONFIG_KEY: 'my-project', ai_platform_training_executor.TRAINING_JOB_CONFIG_KEY: { 'training_input': training_input, ai_platform_training_executor.LABELS_CONFIG_KEY: None, }, ai_platform_training_executor.JOB_ID_CONFIG_KEY: None, ai_platform_training_executor.LABELS_CONFIG_KEY: None, } self.assertEqual( training.exec_properties[ai_platform_training_executor.CONFIG_KEY], json_utils.dumps(expected_aip_config))
}, outputs={ 'output1': standard_artifacts.Model, }, parameters={ 'param1': str, }, image='dummy/transformer', command=[ 'transformer', '--input1', placeholders.InputUriPlaceholder('input1'), '--output1', placeholders.OutputUriPlaceholder('output1'), '--param1', placeholders.InputValuePlaceholder('param1'), ], ) dummy_producer_component = container_component.create_container_component( name='DummyProducerComponent', outputs={ 'output1': standard_artifacts.Model, }, parameters={ 'param1': str, }, image='dummy/producer', command=[ 'producer', '--output1',
instance_name='examples', artifact_type=simple_artifacts.File, reimport=False, source_uri='gs://tfx-oss-testing-bucket/sample-data/mnist') _TRAIN = ai_platform_training_component.create_ai_platform_training( name='simple_aip_training', project_id='tfx-oss-testing', region='us-central1', image_uri='gcr.io/tfx-oss-testing/caip-training:tfx-test', args=[ '--dataset', placeholders.InputUriPlaceholder('examples'), '--model-dir', placeholders.OutputUriPlaceholder('model'), '--lr', placeholders.InputValuePlaceholder('learning_rate'), ], scale_tier='BASIC', inputs={'examples': _EXAMPLE_IMPORTER.outputs['result']}, outputs={'model': standard_artifacts.Model}, parameters={'learning_rate': '0.001'}) class AiPlatformTrainingComponentIntegrationTest( test_utils.BaseAIPlatformPipelinesTest): def testSuccessfulExecution(self): aip_training_pipeline = pipeline.Pipeline( pipeline_name=_PIPELINE_NAME, pipeline_root=self._pipeline_root(_PIPELINE_NAME), components=[_EXAMPLE_IMPORTER, _TRAIN], )
image='google/cloud-sdk:278.0.0', command=[ 'sh', '-exc', ''' url="$0" output_data_uri="$1"/data # TODO(b/150515270) Remove when fixed. output_data_path=$(mktemp) # Running the main code wget "$0" -O "$output_data_path" || curl "$0" > "$output_data_path" # Getting data out of the container gsutil cp "$output_data_path" "$output_data_uri" ''', placeholders.InputValuePlaceholder('url'), placeholders.OutputUriPlaceholder('data'), ], ) grep_component = container_component.create_container_component( name='FilterWithGrep', inputs={ 'text': standard_artifacts.ExternalArtifact, }, outputs={ 'filtered_text': standard_artifacts.ExternalArtifact, }, parameters={ 'pattern': str, },
def setUp(self): super().setUp() self._project_id = 'my-project' self._job_id = 'my-job-123' self._labels = ['label1', 'label2'] examples_artifact = standard_artifacts.Examples() examples_artifact.split_names = artifact_utils.encode_split_names( ['train', 'eval']) examples_artifact.uri = _EXAMPLE_LOCATION self._inputs = {'examples': [examples_artifact]} model_artifact = standard_artifacts.Model() model_artifact.uri = _MODEL_LOCATION self._outputs = {'model': [model_artifact]} training_job = { 'training_input': { 'scaleTier': 'CUSTOM', 'region': 'us-central1', 'masterType': 'n1-standard-8', 'masterConfig': { 'imageUri': 'gcr.io/my-project/caip-training-test:latest' }, 'workerType': 'n1-standard-8', 'workerCount': 8, 'workerConfig': { 'imageUri': 'gcr.io/my-project/caip-training-test:latest' }, 'args': [ '--examples', placeholders.InputUriPlaceholder('examples'), '--n-steps', placeholders.InputValuePlaceholder('n_step'), '--model-dir', placeholders.OutputUriPlaceholder('model') ] }, 'labels': self._labels, } aip_training_config = { ai_platform_training_executor.PROJECT_CONFIG_KEY: self._project_id, ai_platform_training_executor.TRAINING_JOB_CONFIG_KEY: training_job, ai_platform_training_executor.JOB_ID_CONFIG_KEY: self._job_id, ai_platform_training_executor.LABELS_CONFIG_KEY: self._labels, } self._exec_properties = { ai_platform_training_executor.CONFIG_KEY: json_utils.dumps(aip_training_config), 'n_step': 100 } resolved_training_input = copy.deepcopy(training_job['training_input']) resolved_training_input['args'] = [ '--examples', _EXAMPLE_LOCATION, '--n-steps', '100', '--model-dir', _MODEL_LOCATION ] self._expected_job_spec = { 'job_id': self._job_id, 'training_input': resolved_training_input, 'labels': self._labels, } self._mock_api_client = mock.Mock() mock_discovery = self.enter_context( mock.patch.object(discovery, 'build', autospec=True)) mock_discovery.return_value = self._mock_api_client self._setUpTrainingMocks()