def testSingleNodeCustomPython(self):
expected_custom_job_spec = {
'name': 'my-custom-job',
'jobSpec': {
'workerPoolSpecs': [{
'machineSpec': {
'machineType': 'n1-standard-8'
},
'replicaCount': '1',
'pythonPackageSpec': {
'executorImageUri':
'my_image:latest',
'packageUris':
['gs://my-bucket/my-training-prgram.tar.gz'],
'pythonModule':
'my_trainer',
'args': [
'--input_path',
"{{$.inputs.artifacts['examples'].path}}",
'--output_path',
"{{$.outputs.artifacts['model'].path}}",
'--optimizer',
"{{$.inputs.parameters['optimizer']}}",
'--output_param_path',
"{{$.outputs.parameters['out_param'].output_file}}"
]
}
}]
}
}
task = aiplatform.custom_job(
name='my-custom-job',
input_artifacts={
'examples':
dsl.PipelineParam(name='output',
op_name='ingestor',
param_type='Dataset')
},
input_parameters={'optimizer': 'sgd'},
output_artifacts={'model': ontology_artifacts.Model},
output_parameters={'out_param': str},
executor_image_uri='my_image:latest',
package_uris=['gs://my-bucket/my-training-prgram.tar.gz'],
python_module='my_trainer',
args=[
'--input_path',
structures.InputPathPlaceholder('examples'), '--output_path',
structures.OutputPathPlaceholder('model'), '--optimizer',
structures.InputValuePlaceholder('optimizer'),
'--output_param_path',
structures.OutputPathPlaceholder('out_param')
],
machine_type='n1-standard-8',
)
self.assertDictEqual(expected_custom_job_spec, task.custom_job_spec)
self.assertDictEqual(_EXPECTED_COMPONENT_SPEC,
json_format.MessageToDict(task.component_spec))
self.assertDictEqual(_EXPECTED_TASK_SPEC,
json_format.MessageToDict(task.task_spec))
def _transform_arg(arg: Union[str, BasePlaceholder]) -> Any:
if isinstance(arg, str):
return arg
elif isinstance(arg, InputValuePlaceholder):
return v1_components.InputValuePlaceholder(arg.name)
elif isinstance(arg, InputPathPlaceholder):
return v1_components.InputPathPlaceholder(arg.name)
elif isinstance(arg, InputUriPlaceholder):
return v1_components.InputUriPlaceholder(arg.name)
elif isinstance(arg, OutputPathPlaceholder):
return v1_components.OutputPathPlaceholder(arg.name)
elif isinstance(arg, OutputUriPlaceholder):
return v1_components.OutputUriPlaceholder(arg.name)
else:
# TODO(chensun): transform additional placeholders: if, concat, etc.?
raise ValueError(
f'Unexpected command/argument type: "{arg}" of type "{type(arg)}".'
)
def testAdvancedTrainingSpec(self):
expected_custom_job_spec = {
'name': 'my-custom-job',
'jobSpec': {
"workerPoolSpecs": [{
"replicaCount": 1,
"containerSpec": {
"imageUri":
"my-master-image:latest",
"command": ["python3", "master_entrypoint.py"],
"args": [
"--input_path",
"{{$.inputs.artifacts['examples'].path}}",
"--output_path",
"{{$.outputs.artifacts['model'].path}}",
"--optimizer",
"{{$.inputs.parameters['optimizer']}}",
"--output_param_path",
"{{$.outputs.parameters['out_param'].output_file}}"
]
},
"machineSpec": {
"machineType": "n1-standard-4"
}
}, {
"replicaCount": 4,
"containerSpec": {
"imageUri":
"gcr.io/my-project/my-worker-image:latest",
"command": ["python3", "worker_entrypoint.py"],
"args": [
"--input_path",
"{{$.inputs.artifacts['examples'].path}}",
"--output_path",
"{{$.outputs.artifacts['model'].path}}",
"--optimizer",
"{{$.inputs.parameters['optimizer']}}"
]
},
"machineSpec": {
"machineType": "n1-standard-4",
"acceleratorType": "NVIDIA_TESLA_K80",
"acceleratorCount": 1
}
}]
}
}
task = aiplatform.custom_job(
name='my-custom-job',
input_artifacts={
'examples':
dsl.PipelineParam(name='output',
op_name='ingestor',
param_type='Dataset')
},
input_parameters={'optimizer': 'sgd'},
output_artifacts={'model': ontology_artifacts.Model},
output_parameters={'out_param': str},
additional_job_spec={
'workerPoolSpecs': [
{
'replicaCount': 1,
'containerSpec': {
'imageUri':
'my-master-image:latest',
'command': ['python3', 'master_entrypoint.py'],
'args': [
'--input_path',
structures.InputPathPlaceholder('examples'),
'--output_path',
structures.OutputPathPlaceholder('model'),
'--optimizer',
structures.InputValuePlaceholder('optimizer'),
'--output_param_path',
structures.OutputPathPlaceholder('out_param')
],
},
'machineSpec': {
'machineType': 'n1-standard-4'
}
},
{
'replicaCount': 4,
'containerSpec': {
'imageUri':
'gcr.io/my-project/my-worker-image:latest',
'command': ['python3', 'worker_entrypoint.py'],
'args': [
'--input_path',
structures.InputPathPlaceholder('examples'),
'--output_path',
structures.OutputPathPlaceholder('model'),
'--optimizer',
structures.InputValuePlaceholder('optimizer')
]
},
# Optionally one can also attach accelerators.
'machineSpec': {
'machineType': 'n1-standard-4',
'acceleratorType': 'NVIDIA_TESLA_K80',
'acceleratorCount': 1
}
}
]
})
self.assertDictEqual(expected_custom_job_spec, task.custom_job_spec)
self.assertDictEqual(_EXPECTED_COMPONENT_SPEC,
json_format.MessageToDict(task.component_spec))
self.assertDictEqual(_EXPECTED_TASK_SPEC,
json_format.MessageToDict(task.task_spec))