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 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))
def testImageUriValidation(self): training_input = { 'scaleTier': 'BASIC_GPU', 'region': 'us-central1', } with self.assertRaisesRegex(ValueError, 'image_uri is required'): _ = ai_platform_training_component.create_ai_platform_training( name='my_training_step', project_id='my-project', training_input=training_input)
def testRegionValidation(self): training_input = { 'scaleTier': 'BASIC_GPU', 'masterConfig': { 'imageUri': 'gcr.io/my-project/caip-training-test:latest' }, 'args': ['--my-flag'] } with self.assertRaisesRegex(ValueError, 'region is required'): _ = ai_platform_training_component.create_ai_platform_training( name='my_training_step', project_id='my-project', training_input=training_input)
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))
_PIPELINE_NAME = 'aip_training_component_pipeline' _EXAMPLE_IMPORTER = importer_node.ImporterNode( 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,