Add statuses monitoring for kf jobs
@@ -5,6 +5,7 @@ from typing import Mapping
import conf
import ocular
+from constants.experiment_jobs_uuids import get_experiment_job_uuid
from constants.jobs import JobLifeCycle
from db.redis.containers import RedisJobContainers
from polyaxon.celery_api import celery_app
@@ -60,6 +61,8 @@ def run(k8s_manager: 'K8SManager') -> None:
namespace=conf.get('K8S_NAMESPACE'),
container_names=(
conf.get('CONTAINER_NAME_EXPERIMENT_JOB'),
+ conf.get('CONTAINER_NAME_TF_JOB'),
+ conf.get('CONTAINER_NAME_PYTORCH_JOB'),
conf.get('CONTAINER_NAME_PLUGIN_JOB'),
conf.get('CONTAINER_NAME_JOB'),
conf.get('CONTAINER_NAME_DOCKERIZER_JOB')),
@@ -80,6 +83,16 @@ def run(k8s_manager: 'K8SManager') -> None:
or (status and labels['app'] == conf.get('APP_LABELS_EXPERIMENT'))
)
+ tf_job_condition = (
+ conf.get('CONTAINER_NAME_TF_JOB') in pod_state['details']['container_statuses']
+ or (status and labels['app'] == conf.get('APP_LABELS_EXPERIMENT'))
+ )
+
+ pytorch_job_condition = (
+ conf.get('CONTAINER_NAME_PYTORCH_JOB') in pod_state['details']['container_statuses']
+ or (status and labels['app'] == conf.get('APP_LABELS_EXPERIMENT'))
+ )
+
job_condition = (
conf.get('CONTAINER_NAME_JOB') in pod_state['details']['container_statuses'] or
(status and labels['app'] == conf.get('APP_LABELS_JOB'))
@@ -104,6 +117,35 @@ def run(k8s_manager: 'K8SManager') -> None:
K8SEventsCeleryTasks.K8S_EVENTS_HANDLE_EXPERIMENT_JOB_STATUSES,
kwargs={'payload': pod_state})
+ if tf_job_condition:
+ update_job_containers(event_object, status, conf.get('CONTAINER_NAME_TF_JOB'))
+ logger.debug("Sending state to handler %s, %s", status, labels)
+
+ pod_state['details']['labels']['job_uuid'] = get_experiment_job_uuid(
+ experiment_uuid=labels['experiment_uuid'],
+ task_type=labels['task_type'],
+ task_index=labels['tf-replica-index']
+ )
+ # We augment the payload with standard Polyaxon requirement
+ # Handle experiment job statuses
+ celery_app.send_task(
+ K8SEventsCeleryTasks.K8S_EVENTS_HANDLE_EXPERIMENT_JOB_STATUSES,
+ kwargs={'payload': pod_state})
+
+ if pytorch_job_condition:
+ update_job_containers(event_object, status, conf.get('CONTAINER_NAME_PYTORCH_JOB'))
+ logger.debug("Sending state to handler %s, %s", status, labels)
+ # We augment the payload with standard Polyaxon requirement
+ pod_state['details']['labels']['job_uuid'] = get_experiment_job_uuid(
+ experiment_uuid=labels['experiment_uuid'],
+ task_type=labels['task_type'],
+ task_index=labels['pytorch-replica-index']
+ )
+ # Handle experiment job statuses
+ celery_app.send_task(
+ K8SEventsCeleryTasks.K8S_EVENTS_HANDLE_EXPERIMENT_JOB_STATUSES,
+ kwargs={'payload': pod_state})
+
elif job_condition:
update_job_containers(event_object, status, conf.get('CONTAINER_NAME_JOB'))
logger.debug("Sending state to handler %s, %s", status, labels)
