def _maybe_set_run_terminated(active_run, status):
"""
If the passed-in active run is defined and still running (i.e. hasn't already been terminated
within user code), mark it as terminated with the passed-in status.
"""
if active_run is None:
return
run_id = active_run.info.run_id
cur_status = tracking.MlflowClient().get_run(run_id).info.status
if RunStatus.is_terminated(cur_status):
return
tracking.MlflowClient().set_terminated(run_id, status)
def cancel(self):
with self._status_lock:
if not RunStatus.is_terminated(self._status):
_logger.info("Cancelling job.")
kube_api = kubernetes.client.BatchV1Api()
kube_api.delete_namespaced_job(
name=self._job_name,
namespace=self._job_namespace,
body=kubernetes.client.V1DeleteOptions(),
pretty=True)
self._status = RunStatus.KILLED
_logger.info("Job cancelled.")
else:
_logger.info(
"Attempting to cancel a job that is already terminated.")
def _update_status(self, kube_api=kubernetes.client.BatchV1Api()):
api_response = kube_api.read_namespaced_job_status(
name=self._job_name, namespace=self._job_namespace, pretty=True)
status = api_response.status
with self._status_lock:
if RunStatus.is_terminated(self._status):
return self._status
if self._status == RunStatus.SCHEDULED:
if api_response.status.start_time is None:
_logger.info("Waiting for Job to start")
else:
_logger.info("Job started.")
self._status = RunStatus.RUNNING
if status.conditions is not None:
for condition in status.conditions:
if condition.status == "True":
_logger.info(condition.message)
if condition.type == "Failed":
self._status = RunStatus.FAILED
elif condition.type == "Complete":
self._status = RunStatus.FINISHED
return self._status
def test_is_terminated(self):
self.assertTrue(RunStatus.is_terminated(RunStatus.FAILED))
self.assertTrue(RunStatus.is_terminated(RunStatus.FINISHED))
self.assertTrue(RunStatus.is_terminated(RunStatus.KILLED))
self.assertFalse(RunStatus.is_terminated(RunStatus.SCHEDULED))
self.assertFalse(RunStatus.is_terminated(RunStatus.RUNNING))
def get_status(self):
status = self._status
return status if RunStatus.is_terminated(
status) else self._update_status()
def wait(self):
kube_api = kubernetes.client.BatchV1Api()
while not RunStatus.is_terminated(self._update_status(kube_api)):
time.sleep(self.POLL_STATUS_INTERVAL)
return self._status == RunStatus.FINISHED