def test_container_based_custom_job_op_compile(self):
custom_job_op = custom_job.custom_training_job_op(
self._container_component)
@kfp.dsl.pipeline(name="training-test")
def pipeline():
custom_job_task = custom_job_op(
self._test_input_string,
project=self._project,
location=self._location)
compiler.Compiler().compile(
pipeline_func=pipeline, package_path=self._package_path)
with open(self._package_path) as f:
executor_output_json = json.load(f, strict=False)
with open(
os.path.join(
os.path.dirname(__file__),
'../testdata/custom_job_container_component_pipeline.json')
) as ef:
expected_executor_output_json = json.load(ef, strict=False)
# Ignore the kfp SDK & schema version during comparision
del executor_output_json['pipelineSpec']['sdkVersion']
del executor_output_json['pipelineSpec']['schemaVersion']
self.assertEqual(executor_output_json, expected_executor_output_json)
def test_run_as_vertex_ai_custom_with_replica_count_greater_than_1_converts_correctly(
self):
component_factory_function = self._create_a_container_based_component()
expected_sub_results = {
'implementation': {
'container': {
'image':
'test_launcher_image',
'command': [
'python3', '-u', '-m',
'google_cloud_pipeline_components.container.experimental.gcp_launcher.launcher'
],
'args': [
'--type', 'CustomJob', '--payload',
'{"display_name": "ContainerComponent", "job_spec": '
'{"worker_pool_specs": [{"machine_spec": {"machine_type": '
'"n1-standard-4"}, "replica_count": 1, "container_spec": '
'{"image_uri": "google/cloud-sdk:latest", "command": '
'["sh", "-c", "set -e -x\\necho \\"$0, this is an output '
'parameter\\"\\n", '
'"{{$.inputs.parameters[\'input_text\']}}", '
'"{{$.outputs.parameters[\'output_value\'].output_file}}"]},'
' "disk_spec": {"boot_disk_type": "pd-ssd", '
'"boot_disk_size_gb": 100}}, {"machine_spec": '
'{"machine_type": "n1-standard-4"}, "replica_count": "1", '
'"container_spec": {"image_uri": '
'"google/cloud-sdk:latest", "command": ["sh", "-c", "set '
'-e -x\\necho \\"$0, this is an output parameter\\"\\n", '
'"{{$.inputs.parameters[\'input_text\']}}", '
'"{{$.outputs.parameters[\'output_value\'].output_file}}"]},'
' "disk_spec": {"boot_disk_type": "pd-ssd", '
'"boot_disk_size_gb": 100}}], "service_account": '
'"{{$.inputs.parameters[\'service_account\']}}", '
'"network": "{{$.inputs.parameters[\'network\']}}", '
'"tensorboard": '
'"{{$.inputs.parameters[\'tensorboard\']}}", '
'"base_output_directory": {"output_uri_prefix": '
'"{{$.inputs.parameters[\'base_output_directory\']}}"}}}',
'--project', {
'inputValue': 'project'
}, '--location', {
'inputValue': 'location'
}, '--gcp_resources', {
'outputPath': 'gcp_resources'
}
]
}
}
}
custom_job_spec = custom_job.custom_training_job_op(
component_factory_function, replica_count=2)
self.assertDictContainsSubset(
subset=expected_sub_results,
dictionary=custom_job_spec.component_spec.to_dict())
def test_run_as_vertex_ai_custom_with_worker_poolspec_container_spec_converts_correctly(
self):
component_factory_function = self._create_a_container_based_component()
worker_pool_spec = [{
'machine_spec': {
'machine_type': 'test_machine_type'
},
'replica_count': 2,
'container_spec': {
'image_uri': 'test_image_uri',
'command': ['test_command'],
'args': ['test_args']
}
}]
expected_sub_results = {
'implementation': {
'container': {
'image':
'test_launcher_image',
'command': [
'python3', '-u', '-m',
'google_cloud_pipeline_components.container.experimental.gcp_launcher.launcher'
],
'args': [
'--type', 'CustomJob', '--payload',
'{"display_name": "ContainerComponent", "job_spec": '
'{"worker_pool_specs": [{"machine_spec": {"machine_type": '
'"test_machine_type"}, "replica_count": 2, '
'"container_spec": {"image_uri": "test_image_uri", '
'"command": ["test_command"], "args": ["test_args"]}}], '
'"service_account": '
'"{{$.inputs.parameters[\'service_account\']}}", '
'"network": "{{$.inputs.parameters[\'network\']}}", '
'"tensorboard": '
'"{{$.inputs.parameters[\'tensorboard\']}}", '
'"base_output_directory": {"output_uri_prefix": '
'"{{$.inputs.parameters[\'base_output_directory\']}}"}}}',
'--project', {
'inputValue': 'project'
}, '--location', {
'inputValue': 'location'
}, '--gcp_resources', {
'outputPath': 'gcp_resources'
}
]
}
}
}
custom_job_spec = custom_job.custom_training_job_op(
component_factory_function, worker_pool_specs=worker_pool_spec)
self.assertDictContainsSubset(
subset=expected_sub_results,
dictionary=custom_job_spec.component_spec.to_dict())
def test_run_as_vertex_ai_custom_without_container_spec_or_python_package_spec_correctly(
self):
component_factory_function = self._create_a_container_based_component()
worker_pool_spec = [{
'machine_spec': {
'machine_type': 'test_machine_type'
},
'replica_count': 2
}]
with self.assertRaises(ValueError):
custom_job_spec = custom_job.custom_training_job_op(
component_factory_function, worker_pool_specs=worker_pool_spec)
