text: # Copyright 2021 The Kubeflow Authors. All Rights Reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  """Test Vertex AI Custom Job Client module."""
  
  from google_cloud_pipeline_components.v1.custom_job import utils
  from kfp import components
  
  import unittest
  
  
  class VertexAICustomJobUtilsTests(unittest.TestCase):
  
    def setUp(self):
      super(VertexAICustomJobUtilsTests, self).setUp()
      utils._DEFAULT_CUSTOM_JOB_CONTAINER_IMAGE = 'test_launcher_image'
  
    def _create_a_container_based_component(self) -> callable:
      """Creates a test container based component factory."""
  
      return components.load_component_from_text("""
  name: ContainerComponent
  inputs:
  - {name: input_text, type: String, description: "Represents an input parameter."}
  outputs:
  - {name: output_value, type: String, description: "Represents an output paramter."}
  implementation:
    container:
      image: google/cloud-sdk:latest
      command:
      - sh
      - -c
      - |
        set -e -x
        echo "$0, this is an output parameter"
      - {inputValue: input_text}
      - {outputPath: output_value}
  """)
  
    def test_run_as_vertex_ai_custom_job_on_container_spec_with_defualts_values_converts_correctly(
        self):
      expected_results = {
          'name': 'ContainerComponent',
          'inputs': [{
              'name': 'input_text',
              'type': 'String',
              'description': 'Represents an input parameter.'
          }, {
              'name': 'base_output_directory',
              'type': 'String',
              'default': '',
              'optional': True
          }, {
              'name': 'tensorboard',
              'type': 'String',
              'default': '',
              'optional': True
          }, {
              'name': 'network',
              'type': 'String',
              'default': '',
              'optional': True
          }, {
              'name': 'service_account',
              'type': 'String',
              'default': '',
              'optional': True
          }, {
              'name': 'project',
              'type': 'String'
          }, {
              'name': 'location',
              'type': 'String'
          }],
          'outputs': [{
              'name': 'output_value',
              'type': 'String',
              'description': 'Represents an output paramter.'
          }, {
              'name': 'gcp_resources',
              'type': 'String'
          }],
          'implementation': {
              'container': {
                  'image':
                      'test_launcher_image',
                  'command': [
                      'python3', '-u', '-m',
                      'google_cloud_pipeline_components.container.v1.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}}], "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'
                      }
                  ]
              }
          }
      }
      component_factory_function = self._create_a_container_based_component()
      custom_job_spec = utils.create_custom_training_job_op_from_component(
          component_factory_function)
      self.assertDictEqual(custom_job_spec.component_spec.to_dict(),
                           expected_results)
  
    def test_run_as_vertex_ai_custom_with_accelerator_type_and_count_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.v1.gcp_launcher.launcher'
                  ],
                  'args': [
                      '--type', 'CustomJob', '--payload',
                      '{"display_name": "ContainerComponent", "job_spec": '
                      '{"worker_pool_specs": [{"machine_spec": {"machine_type": '
                      '"n1-standard-4", "accelerator_type": '
                      '"test_accelerator_type", "accelerator_count": 2}, '
                      '"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 = utils.create_custom_training_job_op_from_component(
          component_factory_function,
          accelerator_type='test_accelerator_type',
          accelerator_count=2)
  
      self.assertDictContainsSubset(
          subset=expected_sub_results,
          dictionary=custom_job_spec.component_spec.to_dict())
  
    def test_run_as_vertex_ai_custom_with_boot_disk_type_and_size_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.v1.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 = utils.create_custom_training_job_op_from_component(
          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_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.v1.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 = utils.create_custom_training_job_op_from_component(
          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_time_out_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.v1.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}}], "scheduling": {"timeout": '
                      '2}, "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 = utils.create_custom_training_job_op_from_component(
          component_factory_function, timeout=2)
  
      self.assertDictContainsSubset(
          subset=expected_sub_results,
          dictionary=custom_job_spec.component_spec.to_dict())
  
    def test_run_as_vertex_ai_custom_with_restart_job_on_worker_restart_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.v1.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}}], "scheduling": '
                      '{"restart_job_on_worker_restart": true}, '
                      '"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 = utils.create_custom_training_job_op_from_component(
          component_factory_function, restart_job_on_worker_restart=True)
  
      self.assertDictContainsSubset(
          subset=expected_sub_results,
          dictionary=custom_job_spec.component_spec.to_dict())
  
    def test_run_as_vertex_ai_custom_with_custom_service_account_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.v1.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}}], "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 = utils.create_custom_training_job_op_from_component(
          component_factory_function, service_account='test_service_account')
  
      self.assertDictContainsSubset(
          subset=expected_sub_results,
          dictionary=custom_job_spec.component_spec.to_dict())
  
    def test_run_as_vertex_ai_custom_with_display_name_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.v1.gcp_launcher.launcher'
                  ],
                  'args': [
                      '--type', 'CustomJob', '--payload',
                      '{"display_name": "test_display_name", "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}}], "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 = utils.create_custom_training_job_op_from_component(
          component_factory_function, display_name='test_display_name')
  
      self.assertDictContainsSubset(
          subset=expected_sub_results,
          dictionary=custom_job_spec.component_spec.to_dict())
  
    def test_run_as_vertex_ai_custom_with_network_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.v1.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}}], "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 = utils.create_custom_training_job_op_from_component(
          component_factory_function, network='test_network')
  
      self.assertDictContainsSubset(
          subset=expected_sub_results,
          dictionary=custom_job_spec.component_spec.to_dict())
  
    def test_run_as_vertex_ai_custom_with_labels_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.v1.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}}], "labels": {"test_key": '
                      '"test_value"}, "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 = utils.create_custom_training_job_op_from_component(
          component_factory_function, labels={'test_key': 'test_value'})
  
      self.assertDictContainsSubset(
          subset=expected_sub_results,
          dictionary=custom_job_spec.component_spec.to_dict())
  
    def test_run_as_vertex_ai_custom_with_reserved_ip_ranges(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.v1.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}}], "reserved_ip_ranges": '
                      '["1.0.0.0", "2.0.0.0"], "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 = utils.create_custom_training_job_op_from_component(
          component_factory_function, reserved_ip_ranges=['1.0.0.0', '2.0.0.0'])
  
      self.assertDictContainsSubset(
          subset=expected_sub_results,
          dictionary=custom_job_spec.component_spec.to_dict())
  