feat(job-manager): add kubernetes_queue option for MultiKueue (#494)

Also adds monitoring of job status (rather than
just pods) to the job_monitor to support MK since
the status of remote pods cannot be seen.

Closes #493

Author: CameronMcClymont

docs/openapi.json

@@ -85,6 +85,9 @@
         "kubernetes_memory_limit": {
           "type": "string"
         },
+        "kubernetes_queue": {
+          "type": "string"
+        },
         "kubernetes_uid": {
           "format": "int32",
           "type": "integer"

reana_job_controller/config.py

@@ -12,8 +12,6 @@
 import os
 import secrets
 
-from reana_commons.config import REANA_COMPONENT_PREFIX
-
 from werkzeug.utils import import_string
 
 REANA_DB_CLOSE_POOL_CONNECTIONS = bool(
@@ -177,11 +175,11 @@
 SLURM_SSH_AUTH_TIMEOUT = float(os.getenv("SLURM_SSH_AUTH_TIMEOUT", "60"))
 """Seconds to wait for SLURM SSH authentication response."""
 
-USE_KUEUE = bool(strtobool(os.getenv("USE_KUEUE", "False")))
+KUEUE_ENABLED = bool(strtobool(os.getenv("KUEUE_ENABLED", "False")))
 """Whether to use Kueue to manage job execution."""
 
-KUEUE_LOCAL_QUEUE_NAME = "local-queue-job"
-"""Name of the local queue to be used by Kueue."""
+KUEUE_DEFAULT_QUEUE = os.getenv("KUEUE_DEFAULT_QUEUE", "")
+"""Name of the default queue to be used by Kueue."""
 
 REANA_USER_ID = os.getenv("REANA_USER_ID")
 """User UUID of the owner of the workflow."""

reana_job_controller/job_db.py

@@ -115,10 +115,10 @@ def retrieve_job_logs(job_id):
 
 
 def store_job_logs(job_id, logs):
-    """Store job logs.
+    """Store job logs in the database.
 
     :param job_id: Internal REANA job ID.
-    :param logs: Job logs.
+    :param logs: Job logs to store.
     :type job_id: str
     :type logs: str
     """

reana_job_controller/job_monitor.py

@@ -17,8 +17,7 @@
 from kubernetes import client, watch
 from reana_commons.config import REANA_RUNTIME_KUBERNETES_NAMESPACE
 from reana_commons.k8s.api_client import current_k8s_corev1_api_client
-from reana_db.database import Session
-from reana_db.models import Job, JobStatus
+from reana_db.models import JobStatus
 
 from reana_job_controller.config import (
     COMPUTE_BACKENDS,
@@ -32,10 +31,10 @@
     C4P_SSH_TIMEOUT,
     C4P_SSH_BANNER_TIMEOUT,
     C4P_SSH_AUTH_TIMEOUT,
+    KUEUE_ENABLED,
 )
 
 from reana_job_controller.job_db import JOB_DB, store_job_logs, update_job_status
-from reana_job_controller.kubernetes_job_manager import KubernetesJobManager
 from reana_job_controller.utils import (
     SSHClient,
     singleton,
@@ -104,7 +103,8 @@ def get_reana_job_id(self, backend_job_id: str) -> str:
         remaining_jobs = self._get_remaining_jobs()
         return remaining_jobs[backend_job_id]
 
-    def get_backend_job_id(self, job_pod):
+    @staticmethod
+    def get_backend_job_id(job_pod):
         """Get the backend job id for the backend object.
 
         :param job_pod: Compute backend job object (Kubernetes V1Pod
@@ -115,7 +115,7 @@ def get_backend_job_id(self, job_pod):
         """
         return job_pod.metadata.labels["job-name"]
 
-    def should_process_job(self, job_pod) -> bool:
+    def should_process_job_pod(self, job_pod) -> bool:
         """Decide whether the job should be processed or not.
 
         Each job is processed only once, when it reaches a final state (either `failed` or `finished`).
@@ -141,6 +141,27 @@ def should_process_job(self, job_pod) -> bool:
 
         return is_job_in_remaining_jobs and is_job_completed
 
+    def should_process_job(self, job) -> bool:
+        """Decide whether the job should be processed or not.
+
+        Each job is processed only once, when it reaches a final state (either `failed` or `finished`).
+
+        :param job: Compute backend job object (Kubernetes V1Job
+            https://github.com/kubernetes-client/python/blob/master/kubernetes/docs/V1Job.md)
+        """
+        remaining_jobs = self._get_remaining_jobs(
+            statuses_to_skip=[
+                JobStatus.finished.name,
+                JobStatus.failed.name,
+                JobStatus.stopped.name,
+            ]
+        )
+
+        is_job_in_remaining_jobs = job.metadata.name in remaining_jobs
+        is_job_completed = job.status.succeeded and not job.status.active
+
+        return is_job_in_remaining_jobs and is_job_completed
+
     @staticmethod
     def _get_job_container_statuses(job_pod):
         return (job_pod.status.container_statuses or []) + (
@@ -235,48 +256,111 @@ def watch_jobs(self, job_db, app=None):
 
         :param job_db: Dictionary which contains all current jobs.
         """
-        while True:
-            logging.info("Starting a new stream request to watch Jobs")
-            try:
-                w = watch.Watch()
-                for event in w.stream(
-                    current_k8s_corev1_api_client.list_namespaced_pod,
-                    namespace=REANA_RUNTIME_KUBERNETES_NAMESPACE,
-                    label_selector=f"reana-run-job-workflow-uuid={self.workflow_uuid}",
-                ):
-                    logging.info("New Pod event received: {0}".format(event["type"]))
-                    job_pod = event["object"]
-
-                    # Each job is processed once, when reaching a final state
-                    # (either successfully or not)
-                    if self.should_process_job(job_pod):
-                        job_status = self.get_job_status(job_pod)
-                        backend_job_id = self.get_backend_job_id(job_pod)
-                        reana_job_id = self.get_reana_job_id(backend_job_id)
+        # If using MultiKueue, watch jobs instead of pods since worker pods could be
+        # running on a remote cluster that we can't directly monitor
+        if KUEUE_ENABLED:
+            while True:
+                logging.info("Starting a new stream request to watch Jobs")
 
-                        logs = self.job_manager_cls.get_logs(
-                            backend_job_id, job_pod=job_pod
+                try:
+                    w = watch.Watch()
+                    for event in w.stream(
+                        client.BatchV1Api().list_namespaced_job,
+                        namespace=REANA_RUNTIME_KUBERNETES_NAMESPACE,
+                        label_selector=f"reana-run-job-workflow-uuid={self.workflow_uuid}",
+                    ):
+                        logging.info(f"New Job event received: {event["type"]}")
+
+                        job = event["object"]
+                        job_id = job.metadata.name
+                        job_finished = (
+                            job.status.succeeded
+                            and not job.status.active
+                            and not job.status.failed
                         )
+                        job_status = (
+                            JobStatus.finished.name
+                            if job_finished
+                            else (
+                                JobStatus.failed.name
+                                if job.status.failed
+                                else JobStatus.running.name
+                            )
+                        )
+
+                        if self.should_process_job(job):
+                            reana_job_id = self.get_reana_job_id(job_id)
 
-                        if job_status == JobStatus.failed.name:
-                            self.log_disruption(
-                                event["object"].status.conditions, backend_job_id
+                            if job_status == JobStatus.failed.name:
+                                self.log_disruption(
+                                    event["object"].status.conditions, job_id
+                                )
+
+                            # TODO: Fetch logs from the remote job pod on the remote worker when MultiKueue supports this
+                            logs = self.job_manager_cls.get_logs(job_id)
+                            if logs is not None:
+                                store_job_logs(reana_job_id, logs)
+
+                            update_job_status(
+                                reana_job_id,
+                                job_status,
                             )
 
-                        store_job_logs(reana_job_id, logs)
-                        update_job_status(reana_job_id, job_status)
+                            if JobStatus.should_cleanup_job(job_status):
+                                self.clean_job(job_id)
 
-                        if JobStatus.should_cleanup_job(job_status):
-                            self.clean_job(backend_job_id)
-            except client.rest.ApiException as e:
-                logging.exception(
-                    f"Error from Kubernetes API while watching jobs pods: {e}"
-                )
-            except Exception as e:
-                logging.error(traceback.format_exc())
-                logging.error("Unexpected error: {}".format(e))
+                except client.rest.ApiException as e:
+                    logging.exception(
+                        f"Error from Kubernetes API while watching jobs: {e}"
+                    )
+                except Exception as e:
+                    logging.error(traceback.format_exc())
+                    logging.error("Unexpected error: {}".format(e))
+        else:
+            while True:
+                try:
+                    w = watch.Watch()
+                    for event in w.stream(
+                        current_k8s_corev1_api_client.list_namespaced_pod,
+                        namespace=REANA_RUNTIME_KUBERNETES_NAMESPACE,
+                        label_selector=f"reana-run-job-workflow-uuid={self.workflow_uuid}",
+                    ):
+                        logging.info(
+                            "New Pod event received: {0}".format(event["type"])
+                        )
+                        job_pod = event["object"]
+
+                        # Each job is processed once, when reaching a final state
+                        # (either successfully or not)
+                        if self.should_process_job_pod(job_pod):
+                            job_status = self.get_job_status(job_pod)
+                            backend_job_id = self.get_backend_job_id(job_pod)
+                            reana_job_id = self.get_reana_job_id(backend_job_id)
+
+                            logs = self.job_manager_cls.get_logs(
+                                backend_job_id, job_pod=job_pod
+                            )
+
+                            if job_status == JobStatus.failed.name:
+                                self.log_disruption(
+                                    event["object"].status.conditions, backend_job_id
+                                )
 
-    def log_disruption(self, conditions, backend_job_id):
+                            store_job_logs(reana_job_id, logs)
+                            update_job_status(reana_job_id, job_status)
+
+                            if JobStatus.should_cleanup_job(job_status):
+                                self.clean_job(backend_job_id)
+                except client.rest.ApiException as e:
+                    logging.exception(
+                        f"Error from Kubernetes API while watching jobs pods: {e}"
+                    )
+                except Exception as e:
+                    logging.error(traceback.format_exc())
+                    logging.error("Unexpected error: {}".format(e))
+
+    @staticmethod
+    def log_disruption(conditions, backend_job_id):
         """Log disruption message from Kubernetes event conditions.
 
         Usually it is pod eviction but can be any of https://kubernetes.io/docs/concepts/workloads/pods/disruptions/#pod-disruption-conditions.