def _check_job_list(self, input_ids: List[str] = []) -> Tuple[List[str], List[str]]:
"""
Deduplicates the input job list, maintaining insertion order
Any jobs not present in self._running_jobs are added to an error list
:param input_ids: a list of putative job IDs
:return results: tuple with items "job_ids", containing valid IDs;
and "error_ids", for jobs that the narrative backend does not know about
"""
if not isinstance(input_ids, list):
raise JobRequestException(f"{JOBS_TYPE_ERR}: {input_ids}")
job_ids = []
error_ids = []
for input_id in input_ids:
if input_id and input_id not in job_ids + error_ids:
if input_id in self._running_jobs:
job_ids.append(input_id)
else:
error_ids.append(input_id)
if not len(job_ids) + len(error_ids):
raise JobRequestException(JOBS_MISSING_ERR, input_ids)
return job_ids, error_ids
def __init__(self, rq: dict):
self.raw_request = copy.deepcopy(rq)
self.msg_id = rq.get("msg_id") # might be useful later?
self.rq_data = rq.get("content", {}).get("data")
if self.rq_data is None:
raise JobRequestException(INVALID_REQUEST_ERR)
self.request_type = self.rq_data.get("request_type")
if self.request_type is None:
raise JobRequestException(MISSING_REQUEST_TYPE_ERR)
input_type_count = 0
for input_type in [PARAM["JOB_ID"], PARAM["JOB_ID_LIST"], PARAM["BATCH_ID"]]:
if input_type in self.rq_data:
input_type_count += 1
if input_type_count > 1:
raise JobRequestException(ONE_INPUT_TYPE_ONLY_ERR)
def _modify_job_updates(self, req: JobRequest) -> dict:
"""
Modifies how many things want to listen to a job update.
If this is a request to start a job update, then this starts the update loop that
returns update messages across the job channel.
If this is a request to stop a job update, then this sends that request to the
JobManager, which might have the side effect of shutting down the update loop if there's
no longer anything requesting job status.
If the given job_id in the request doesn't exist in the current Narrative, or is None,
this raises a JobRequestException.
"""
job_id_list = self._get_job_ids(req)
update_type = req.request_type
if update_type == MESSAGE_TYPE["START_UPDATE"]:
update_refresh = True
elif update_type == MESSAGE_TYPE["STOP_UPDATE"]:
update_refresh = False
else:
# this should be impossible
raise JobRequestException("Unknown request")
self._jm.modify_job_refresh(job_id_list, update_refresh)
if update_refresh:
self.start_job_status_loop()
output_states = self._jm.get_job_states(job_id_list)
self.send_comm_message(MESSAGE_TYPE["STATUS"], output_states)
return output_states
def update_batch_job(self, batch_id: str) -> List[str]:
"""
Update a batch job and create child jobs if necessary
"""
batch_job = self.get_job(batch_id)
if not batch_job.batch_job:
raise JobRequestException(JOB_NOT_BATCH_ERR, batch_id)
child_ids = batch_job.child_jobs
reg_child_jobs = []
unreg_child_ids = []
for job_id in child_ids:
if job_id in self._running_jobs:
reg_child_jobs.append(self.get_job(job_id))
else:
unreg_child_ids.append(job_id)
unreg_child_jobs = []
if unreg_child_ids:
unreg_child_jobs = Job.from_job_ids(unreg_child_ids)
for job in unreg_child_jobs:
self.register_new_job(
job=job,
refresh=not job.was_terminal(),
)
batch_job.update_children(reg_child_jobs + unreg_child_jobs)
return [batch_id] + child_ids
def _get_job_ids(self, req: JobRequest = None):
if req.has_batch_id():
return self._jm.update_batch_job(req.batch_id)
try:
return req.job_id_list
except Exception as ex:
raise JobRequestException(ONE_INPUT_TYPE_ONLY_ERR) from ex
def get_job(self, job_id):
"""
Returns a Job with the given job_id.
Raises a JobRequestException if not found.
"""
if job_id not in self._running_jobs:
raise JobRequestException(JOB_NOT_REG_ERR, job_id)
return self._running_jobs[job_id]["job"]
def _get_job_ids_by_cell_id(self, cell_id_list: List[str] = None) -> tuple:
"""
Finds jobs with a cell_id in cell_id_list
Mappings of job ID to cell ID are added when new jobs are registered
Returns a list of job IDs and a mapping of cell IDs to the list of
job IDs associated with the cell.
"""
if not cell_id_list:
raise JobRequestException(CELLS_NOT_PROVIDED_ERR)
cell_to_job_mapping = {
id: self._jobs_by_cell_id[id] if id in self._jobs_by_cell_id else set()
for id in cell_id_list
}
# union of all the job_ids in the cell_to_job_mapping
job_id_list = set().union(*cell_to_job_mapping.values())
return (job_id_list, cell_to_job_mapping)
def _handle_comm_message(self, msg: dict) -> dict:
"""
Handles comm messages that come in from the other end of the KBaseJobs channel.
Messages get translated into one or more JobRequest objects, which are then
passed to the right handler, based on the request.
A handler dictionary is created on JobComm creation.
Any unknown request is returned over the channel with message type 'job_error', and a
JobRequestException is raised.
"""
with exc_to_msg(msg):
request = JobRequest(msg)
kblogging.log_event(
self._log, "handle_comm_message", {"msg": request.request_type}
)
if request.request_type not in self._msg_map:
raise JobRequestException(
f"Unknown KBaseJobs message '{request.request_type}'"
)
return self._msg_map[request.request_type](request)
def job_id_list(self):
if PARAM["JOB_ID_LIST"] in self.rq_data:
return self.rq_data[PARAM["JOB_ID_LIST"]]
if PARAM["JOB_ID"] in self.rq_data:
return [self.rq_data[PARAM["JOB_ID"]]]
raise JobRequestException(JOBS_NOT_PROVIDED_ERR)
def job_id(self):
if PARAM["JOB_ID"] in self.rq_data:
return self.rq_data[PARAM["JOB_ID"]]
raise JobRequestException(JOB_NOT_PROVIDED_ERR)
def cell_id_list(self):
if PARAM["CELL_ID_LIST"] in self.rq_data:
return self.rq_data[PARAM["CELL_ID_LIST"]]
raise JobRequestException(CELLS_NOT_PROVIDED_ERR)
def batch_id(self):
if PARAM["BATCH_ID"] in self.rq_data:
return self.rq_data[PARAM["BATCH_ID"]]
raise JobRequestException(BATCH_NOT_PROVIDED_ERR)
def _construct_job_output_state_set(
self, job_ids: List[str], states: dict = None
) -> dict:
"""
Builds a set of job states for the list of job ids.
:param job_ids: list of job IDs (may be empty)
:param states: dict, where each value is a state is from EE2
"""
if not isinstance(job_ids, list):
raise JobRequestException("job_ids must be a list")
if not len(job_ids):
return {}
output_states = dict()
jobs_to_lookup = list()
# Fetch from cache of terminated jobs, where available.
# These are already post-processed and ready to return.
for job_id in job_ids:
job = self.get_job(job_id)
if job.was_terminal():
output_states[job_id] = job.output_state()
elif states and job_id in states:
state = states[job_id]
output_states[job_id] = job.output_state(state)
else:
jobs_to_lookup.append(job_id)
fetched_states = dict()
# Get the rest of states direct from EE2.
if len(jobs_to_lookup):
try:
fetched_states = clients.get("execution_engine2").check_jobs(
{
"job_ids": jobs_to_lookup,
"exclude_fields": EXCLUDED_JOB_STATE_FIELDS,
"return_list": 0,
}
)
except Exception as e:
error_message = str(e)
kblogging.log_event(
self._log,
"_construct_job_output_state_set",
{"exception": error_message},
)
# fill in the output states for the missing jobs
# if the job fetch failed, add an error message to the output
# and return the cached job state
for job_id in jobs_to_lookup:
job = self.get_job(job_id)
if job_id in fetched_states:
output_states[job_id] = job.output_state(fetched_states[job_id])
else:
# fetch the current state without updating it
output_states[job_id] = job.output_state({})
# add an error field with the error message from the failed look up
output_states[job_id]["error"] = error_message
return output_states