Merge pull request #8091 from nuwang/add_k8s_job_cleanup

Added option to delete or preserve k8s pods after completion

lib/galaxy/jobs/runners/kubernetes.py

@@ -57,6 +57,7 @@ class KubernetesJobRunner(AsynchronousJobRunner):
             k8s_default_requests_memory=dict(map=str, default=None),
             k8s_default_limits_cpu=dict(map=str, default=None),
             k8s_default_limits_memory=dict(map=str, default=None),
+            k8s_cleanup_job=dict(map=str, valid=lambda s: s in {"onsuccess", "always", "never"}, default="always"),
             k8s_pod_retries=dict(map=int, valid=lambda x: int >= 0, default=3),
             k8s_pod_retrials=dict(map=int, valid=lambda x: int >= 0, default=3),
             k8s_walltime_limit=dict(map=int, valid=lambda x: int(x) >= 0, default=172800))
@@ -406,7 +407,6 @@ class KubernetesJobRunner(AsynchronousJobRunner):
             active = 0
             failed = 0
 
-            max_pod_retries = 1
             if 'max_pod_retries' in job_destination.params:
                 max_pod_retries = int(job_destination.params['max_pod_retries'])
             elif 'k8s_pod_retries' in self.runner_params:
@@ -417,6 +417,8 @@ class KubernetesJobRunner(AsynchronousJobRunner):
             elif 'k8s_pod_retrials' in self.runner_params:
                 # For backward compatibility
                 max_pod_retries = int(self.runner_params['max_pod_retrials'])
+            else:
+                max_pod_retries = 1
 
             if 'succeeded' in job.obj['status']:
                 succeeded = job.obj['status']['succeeded']
@@ -436,7 +438,8 @@ class KubernetesJobRunner(AsynchronousJobRunner):
                     job_state.job_wrapper.change_state(model.Job.states.RUNNING)
                 return job_state
             elif job_state.job_wrapper.get_job().state == model.Job.states.DELETED:
-                # Job has been deleted via stop_job, cleanup and remove from watched_jobs by returning `None`
+                # Job has been deleted via stop_job and job has not been deleted,
+                # remove from watched_jobs by returning `None`
                 if job_state.job_wrapper.cleanup_job in ("always", "onsuccess"):
                     job_state.job_wrapper.cleanup()
                 return None
@@ -444,6 +447,12 @@ class KubernetesJobRunner(AsynchronousJobRunner):
                 return self._handle_job_failure(job, job_state)
 
         elif len(jobs.response['items']) == 0:
+            if job_state.job_wrapper.get_job().state == model.Job.states.DELETED:
+                # Job has been deleted via stop_job and job has been deleted,
+                # cleanup and remove from watched_jobs by returning `None`
+                if job_state.job_wrapper.cleanup_job in ("always", "onsuccess"):
+                    job_state.job_wrapper.cleanup()
+                return None
             # there is no job responding to this job_id, it is either lost or something happened.
             log.error("No Jobs are available under expected selector app=%s", job_state.job_id)
             with open(job_state.error_file, 'w') as error_file:
@@ -474,9 +483,25 @@ class KubernetesJobRunner(AsynchronousJobRunner):
                 job_state.fail_message = "More pods failed than allowed. See stdout for pods details."
         job_state.running = False
         self.mark_as_failed(job_state)
-        job.scale(replicas=0)
+        self.__cleanup_k8s_job(job)
         return None
 
+    def __cleanup_k8s_job(self, job):
+        k8s_cleanup_job = self.runner_params['k8s_cleanup_job']
+        job_failed = (job.obj['status']['failed'] > 0
+                      if 'failed' in job.obj['status'] else False)
+        # Scale down the job just in case even if cleanup is never
+        job.scale(replicas=0)
+        if (k8s_cleanup_job == "always" or
+                (k8s_cleanup_job == "onsuccess" and not job_failed)):
+            delete_options = {
+                "apiVersion": "v1",
+                "kind": "DeleteOptions",
+                "propagationPolicy": "Background"
+            }
+            r = job.api.delete(json=delete_options, **job.api_kwargs())
+            job.api.raise_for_status(r)
+
     def __job_failed_due_to_walltime_limit(self, job):
         conditions = job.obj['status']['conditions']
         return any(True for c in conditions if c['type'] == 'Failed' and c['reason'] == 'DeadlineExceeded')
@@ -509,7 +534,7 @@ class KubernetesJobRunner(AsynchronousJobRunner):
                 namespace=self.runner_params['k8s_namespace'])
             if len(jobs.response['items']) >= 0:
                 job_to_delete = Job(self._pykube_api, jobs.response['items'][0])
-                job_to_delete.scale(replicas=0)
+                self.__cleanup_k8s_job(job_to_delete)
             # TODO assert whether job parallelism == 0
             # assert not job_to_delete.exists(), "Could not delete job,"+job.job_runner_external_id+" it still exists"
             log.debug("(%s/%s) Terminated at user's request" % (job.id, job.job_runner_external_id))
@@ -541,3 +566,12 @@ class KubernetesJobRunner(AsynchronousJobRunner):
             ajs.old_state = model.Job.states.QUEUED
             ajs.running = False
             self.monitor_queue.put(ajs)
+
+    def finish_job(self, job_state):
+        super(KubernetesJobRunner, self).finish_job(job_state)
+        jobs = Job.objects(self._pykube_api).filter(selector="app=" + job_state.job_id,
+                                                    namespace=self.runner_params['k8s_namespace'])
+        # If more than one job matches selector, leave all jobs intact as it's a configuration error
+        if len(jobs.response['items']) == 1:
+            job = Job(self._pykube_api, jobs.response['items'][0])
+            self.__cleanup_k8s_job(job)

test/integration/test_kubernetes_runner.py

@@ -8,6 +8,8 @@ import subprocess
 import tempfile
 import time
 
+import pytest
+
 from base import integration_util  # noqa: I100,I202
 from base.populators import skip_without_tool
 from .test_containerized_jobs import MulledJobTestCases  # noqa: I201
@@ -96,6 +98,12 @@ def job_config(jobs_directory):
             <param id="k8s_galaxy_instance_id">gx-short-id</param>
             <param id="k8s_walltime_limit">10</param>
         </plugin>
+        <plugin id="k8s_no_cleanup" type="runner" load="galaxy.jobs.runners.kubernetes:KubernetesJobRunner">
+            <param id="k8s_persistent_volume_claims">jobs-directory-claim:$jobs_directory,tool-directory-claim:$tool_directory</param>
+            <param id="k8s_config_path">$k8s_config_path</param>
+            <param id="k8s_galaxy_instance_id">gx-short-id</param>
+            <param id="k8s_cleanup_job">never</param>
+        </plugin>
     </plugins>
     <destinations default="k8s_destination">
         <destination id="k8s_destination" runner="k8s">
@@ -112,12 +120,20 @@ def job_config(jobs_directory):
             <param id="docker_default_container_id">busybox:ubuntu-14.04</param>
             <env id="SOME_ENV_VAR">42</env>
         </destination>
+        <destination id="k8s_destination_no_cleanup" runner="k8s_no_cleanup">
+            <param id="limits_cpu">1.9</param>
+            <param id="limits_memory">10M</param>
+            <param id="docker_enabled">true</param>
+            <param id="docker_default_container_id">busybox:ubuntu-14.04</param>
+            <env id="SOME_ENV_VAR">42</env>
+        </destination>
         <destination id="local_dest" runner="local">
         </destination>
     </destinations>
     <tools>
         <tool id="upload1" destination="local_dest"/>
         <tool id="create_2" destination="k8s_destination_walltime_short"/>
+        <tool id="galaxy_slots_and_memory" destination="k8s_destination_no_cleanup"/>
     </tools>
 </job_conf>
 """)
@@ -222,8 +238,10 @@ class BaseKubernetesIntegrationTestCase(BaseJobEnvironmentIntegrationTestCase, M
 
             # Wait for job to be cancelled in kubernetes
             time.sleep(2)
-            status = json.loads(subprocess.check_output(['kubectl', 'get', 'job', external_id, '-o', 'json']))
-            assert 'active' not in status['status']
+            # The default job config removes jobs, didn't find a better way to check that the job doesn't exist anymore
+            with pytest.raises(subprocess.CalledProcessError) as excinfo:
+                subprocess.check_output(['kubectl', 'get', 'job', external_id, '-o', 'json'], stderr=subprocess.STDOUT)
+            assert "not found" in excinfo.value.output.decode()
 
     @skip_without_tool('job_properties')
     def test_exit_code_127(self):
@@ -258,7 +276,7 @@ class BaseKubernetesIntegrationTestCase(BaseJobEnvironmentIntegrationTestCase, M
 
     @skip_without_tool('galaxy_slots_and_memory')
     def test_slots_and_memory(self):
-        self.dataset_populator.run_tool(
+        running_response = self.dataset_populator.run_tool(
             'galaxy_slots_and_memory',
             {},
             self.history_id,
@@ -270,6 +288,13 @@ class BaseKubernetesIntegrationTestCase(BaseJobEnvironmentIntegrationTestCase, M
         MEM_PER_SLOT = '5'
         assert [CPU, MEM, MEM_PER_SLOT] == dataset_content.split('\n'), dataset_content
 
+        # Tool is mapped to destination without cleanup, make sure job still exists in kubernetes API
+        job_dict = running_response["jobs"][0]
+        job = self.galaxy_interactor.get("jobs/%s" % job_dict['id'], admin=True).json()
+        external_id = job['external_id']
+        status = json.loads(subprocess.check_output(['kubectl', 'get', 'job', external_id, '-o', 'json']))
+        assert 'active' not in status['status']
+
     @skip_without_tool('create_2')
     def test_walltime_limit(self):
         running_response = self.dataset_populator.run_tool(