def dstat_get_jobs(statuses=None,
job_ids=None,
task_ids=None,
labels=None,
create_time_min=None,
create_time_max=None):
statuses = statuses or {'*'}
labels = labels or {}
labels['test-token'] = test_setup.TEST_TOKEN
labels['test-name'] = test_setup.TEST_NAME
labels_set = {job_model.LabelParam(k, v) for (k, v) in labels.items()}
return dstat.dstat_job_producer(provider=get_dsub_provider(),
statuses=statuses,
job_ids=job_ids,
task_ids=task_ids,
labels=labels_set,
create_time_min=create_time_min,
create_time_max=create_time_max,
full_output=True).next()
def api_to_dsub(query):
"""Construct the dstat query parameters from an QueryJobsRequest.
Args:
query: An API query request
Returns:
dict: Key value pairs of query parameters, formatted for dstat
"""
dstat_params = {}
dstat_params['statuses'] = {
job_statuses.api_to_dsub(s)
for s in query.status
} if query.status else {'*'}
if query.name:
dstat_params['job_names'] = {query.name}
if query.labels:
if query.labels.get('job-id'):
dstat_params['job_ids'] = {query.labels['job-id']}
if query.labels.get('task-id'):
dstat_params['task_ids'] = {query.labels['task-id']}
if query.labels.get('attempt'):
dstat_params['task_attempts'] = {query.labels['attempt']}
dstat_params['labels'] = {
job_model.LabelParam(k, v)
for (k, v) in query.labels.items()
if k not in ['job-id', 'task-id', 'attempt']
}
if query.submission:
dstat_params['create_time'] = query.submission
if query.extensions:
if query.extensions.user_id:
dstat_params['user_ids'] = {query.extensions.user_id}
return dstat_params
def dstat_get_jobs(statuses=None,
job_ids=None,
task_ids=None,
labels=None,
create_time_min=None,
create_time_max=None):
"""Build up test parameters and call dstat.dstat_job_producer()."""
statuses = statuses or {'*'}
labels = labels or {}
labels['test-token'] = test_setup.TEST_TOKEN
labels['test-name'] = test_setup.TEST_NAME
labels_set = {job_model.LabelParam(k, v) for (k, v) in labels.items()}
return six.advance_iterator(
dstat.dstat_job_producer(provider=get_dsub_provider(),
statuses=statuses,
job_ids=job_ids,
task_ids=task_ids,
labels=labels_set,
create_time_min=create_time_min,
create_time_max=create_time_max,
full_output=True))
def dsub_start_job(command,
job_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
labels['test-name'] = test_setup.TEST_NAME
logging = param_util.build_logging_param(test.LOGGING)
job_resources = job_model.Resources(image='ubuntu',
logging=logging,
zones=['us-central1-*'])
env_data = {job_model.EnvParam(k, v) for (k, v) in envs.items()}
label_data = {job_model.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_params = {
'envs': env_data,
'inputs': input_data,
'outputs': output_data,
'labels': label_data,
}
task_descriptors = [
job_model.TaskDescriptor({'task-id': None}, {
'envs': set(),
'labels': set(),
'inputs': set(),
'outputs': set(),
}, job_model.Resources())
]
return dsub.run(get_dsub_provider(),
job_resources,
job_params,
task_descriptors,
name=job_name,
command=command,
wait=wait,
disable_warning=True)
def test_bad_labels(self, unused_name, name, value):
del unused_name
with self.assertRaises(ValueError):
job_model.LabelParam(name, value)
def test_good_labels(self, unused_name, name, value):
del unused_name
label_param = job_model.LabelParam(name, value)
self.assertEqual(name, label_param.name)
self.assertEqual(value, label_param.value)
""")
_LABELS_JOB_DESCRIPTOR = job_model.JobDescriptor(
job_metadata={
'dsub-version': '0.1.4',
'job-id': 'echo--dsubuser--171201-142229-050417',
'job-name': 'echo',
'task-ids': '1-2',
'user-id': 'dsubuser',
'create-time': CREATE_TIME
},
job_resources=job_model.Resources(
logging='gs://b/dsub/sh/local/labels/logging/labels.log'),
job_params={
'envs': set(),
'labels': {job_model.LabelParam('batch', 'hello-world')},
'inputs': set(),
'outputs': set(),
'mounts': set(),
},
task_descriptors=[
job_model.TaskDescriptor(
task_metadata={'task-id': '2'},
task_resources=job_model.Resources(
logging_path='gs://b/dsub/sh/local/labels/logging/labels.2.log'
),
task_params={
'envs': set(),
'labels': {job_model.LabelParam('item-number', '2')},
'inputs': set(),
'outputs': set(),
def start_job(self,
command,
name=None,
envs={},
labels={},
inputs={},
inputs_recursive={},
outputs={},
outputs_recursive={},
task_count=1,
wait=False):
logging = param_util.build_logging_param(self.log_path)
resources = job_model.Resources(image=DOCKER_IMAGE,
logging=logging,
zones=['us-central1*'])
env_data = {param_util.EnvParam(k, v) for (k, v) in envs.items()}
label_data = {
job_model.LabelParam(k, v)
for (k, v) in labels.items()
}
# This is mostly an extraction dsubs argument parsing here:
# https://github.com/googlegenomics/dsub/blob/master/dsub/lib/param_util.py#L720
# Reworked it to handle dictionaries rather than a list of items
# of the form 'key=val'
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_params = {
'envs': env_data,
'inputs': input_data,
'outputs': output_data,
'labels': label_data,
}
if task_count > 1:
task_descriptors = [
job_model.TaskDescriptor({'task-id': i + 1}, {
'envs': env_data,
'inputs': input_data,
'outputs': output_data,
'labels': label_data,
}, job_model.Resources()) for i in xrange(task_count)
]
all_task_data = [{
'task-id': i + 1
} for i in xrange(task_count)]
else:
task_descriptors = [
job_model.TaskDescriptor({'task-id': None}, {
'labels': set(),
'envs': set(),
'inputs': set(),
'outputs': set()
}, job_model.Resources())
]
return execute_redirect_stdout(
lambda: dsub.run(self.provider,
resources,
job_params,
task_descriptors,
name=name,
command=command,
wait=wait,
disable_warning=True))
