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,