def testJobResourceCreation(self):
# Default job resources.
resource = job_util.JobResources()
self.check_resources(
{
'min_cores': 1,
'min_ram': 1,
'disk_size': 10,
'boot_disk_size': 10,
'preemptible': False,
'image': None,
'logging': None,
'zones': None,
'scopes': None
}, resource)
# User defined job resources.
resource = job_util.JobResources(zones='us-east1-a,us-central1-a',
min_ram=16,
image='my-docker-img',
min_cores=8)
self.check_resources(
{
'min_cores': 8,
'min_ram': 16,
'disk_size': 10,
'boot_disk_size': 10,
'preemptible': False,
'image': 'my-docker-img',
'logging': None,
'zones': 'us-east1-a,us-central1-a',
'scopes': None
}, resource)
class JobUtilTest(unittest.TestCase):
DEFAULT_RESOURCES = {
'min_cores': 1,
'min_ram': 3.75,
'disk_size': 200,
'boot_disk_size': 10,
'preemptible': False,
'image': None,
'logging': None,
'zones': None,
'scopes': [
'https://www.googleapis.com/auth/bigquery',
],
'accelerator_type': None,
'accelerator_count': 0,
}
USER_RESOURCES = job_util.JobResources(zones='us-east1-a,us-central1-a',
min_ram=16,
image='my-docker-img',
min_cores=8,
accelerator_type='nvidia-tesla-k80',
accelerator_count=2)
EXPECTED_USER_RESOURCES = {
'min_cores': 8,
'min_ram': 16,
'disk_size': 200,
'boot_disk_size': 10,
'preemptible': False,
'image': 'my-docker-img',
'logging': None,
'zones': 'us-east1-a,us-central1-a',
'scopes': [
'https://www.googleapis.com/auth/bigquery',
],
'accelerator_type': 'nvidia-tesla-k80',
'accelerator_count': 2,
}
@parameterized.parameterized.expand([
('test default resource creation', job_util.JobResources(),
DEFAULT_RESOURCES),
('test user defined job resources', USER_RESOURCES,
EXPECTED_USER_RESOURCES)
])
def test_job_resource_creation(self, unused_name, resource, expected):
del unused_name
for name, value in expected.iteritems():
self.assertEqual(value, getattr(resource, name))
def testScriptCreation(self):
script = job_util.Script('the-name', 'the-value')
self.assertEqual('the-name', script.name)
self.assertEqual('the-value', script.value)
def test_submit_validator_passes(self):
resources = job_util.JobResources(
logging=param_util.LoggingParam('gs://buck/logs', PG))
param_util.validate_submit_args_or_fail(job_resources=resources,
all_task_data=self.task_data,
provider_name='MYPROVIDER',
input_providers=[PG],
output_providers=[PG],
logging_providers=[PG])
def test_submit_validator_fails(self, name, path, inwl, outwl, logwl):
resources = job_util.JobResources(
logging=param_util.LoggingParam('gs://buck/logs', PG))
err_expected = 'Unsupported %s path (%s) for provider' % (name, path)
with self.assertRaisesRegexp(ValueError, re.escape(err_expected)):
param_util.validate_submit_args_or_fail(
job_resources=resources,
all_task_data=self.task_data,
provider_name='MYPROVIDER',
input_providers=inwl,
output_providers=outwl,
logging_providers=logwl)
def test_submit_validator_passes(self):
resources = job_util.JobResources(logging=param_util.LoggingParam(
'gs://buck/logs', param_util.P_GCS))
param_util.validate_submit_args_or_fail(
job_resources=resources,
job_data={
'inputs': [],
'outputs': []
},
all_task_data=TASK_DATA,
provider_name='MYPROVIDER',
input_providers=[param_util.P_GCS],
output_providers=[param_util.P_GCS],
logging_providers=[param_util.P_GCS])
def dsub_start_job(command,
name=None,
envs=None,
labels=None,
inputs=None,
inputs_recursive=None,
outputs=None,
outputs_recursive=None,
wait=False):
envs = envs or {}
labels = labels or {}
inputs = inputs or {}
inputs_recursive = inputs_recursive or {}
outputs = outputs or {}
outputs_recursive = outputs_recursive or {}
labels['test-token'] = test_setup.TEST_TOKEN
logging = param_util.build_logging_param(test.LOGGING)
job_resources = job_util.JobResources(image='ubuntu',
logging=logging,
zones=['us-central1-*'])
env_data = {param_util.EnvParam(k, v) for (k, v) in envs.items()}
label_data = {param_util.LabelParam(k, v) for (k, v) in labels.items()}
input_file_param_util = param_util.InputFileParamUtil('input')
input_data = set()
for (recursive, items) in ((False, inputs.items()),
(True, inputs_recursive.items())):
for (name, value) in items:
name = input_file_param_util.get_variable_name(name)
input_data.add(
input_file_param_util.make_param(name, value, recursive))
output_file_param_util = param_util.OutputFileParamUtil('output')
output_data = set()
for (recursive, items) in ((False, outputs.items()),
(True, outputs_recursive.items())):
for (name, value) in items:
name = output_file_param_util.get_variable_name(name)
output_data.add(
output_file_param_util.make_param(name, value, recursive))
job_data = {
'envs': env_data,
'inputs': input_data,
'outputs': output_data,
'labels': label_data,
}
all_task_data = [{
'envs': env_data,
'labels': label_data,
'inputs': input_data,
'outputs': output_data,
}]
return dsub.run(get_dsub_provider(),
job_resources,
job_data,
all_task_data,
name=name,
command=command,
wait=wait,
disable_warning=True)