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()