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)