Random tag for the runtime created configmap - Remove the runtime created configmap when the pod is removed - First basic tests for Dask

Author: lmizzoni

Diff for: calrissian/dask.py

@@ -16,11 +16,13 @@
 from calrissian.retry import retry_exponential_if_exception_type
 from calrissian.job import (
     CalrissianCommandLineJob,
-    KubernetesPodBuilder,
+    KubernetesPodBuilder
 )
 from calrissian.job import (
     quoted_arg_list,
-    read_yaml
+    read_yaml,
+    random_tag,
+    k8s_safe_name,
 )
 from calrissian.job import (
     DEFAULT_INIT_IMAGE,
@@ -246,8 +248,6 @@ class CalrissianCommandLineDaskJob(CalrissianCommandLineJob):
     daskGateway_controller_dir = '/controller'
 
     daskGateway_config_dir = '/etc/dask'
-    daskGateway_cm_name = 'dask-gateway-cm'
-    daskGateway_cm = 'dask-gateway-cm'
 
     daskGateway_controller_cm_name = 'dask-cluster-controller-cm'
 
@@ -256,8 +256,14 @@ def __init__(self, *args, **kwargs):
         super(CalrissianCommandLineDaskJob, self).__init__(*args, **kwargs)
         self.client = KubernetesDaskClient()
 
-    def wait_for_kubernetes_pod(self):
-        return self.client.wait_for_completion()
+        self.dask_cm_name, self.dask_cm_claim_name = self.dask_configmap_name()
+
+    def dask_configmap_name(self):
+        tag = random_tag()
+        return k8s_safe_name('{}-cm-{}'.format('dask', tag)), k8s_safe_name('{}-cm-{}'.format('dask', tag))
+
+    def wait_for_kubernetes_pod(self, cm_name: str):
+        return self.client.wait_for_completion(cm_name = cm_name)
 
     def get_dask_gateway_url(self, runtimeContext):
         return runtimeContext.dask_gateway_url
@@ -301,7 +307,8 @@ def create_kubernetes_runtime(self, runtimeContext):
 
 
         self.client.create_dask_gateway_cofig_map(
-            dask_gateway_url=self.get_dask_gateway_url(runtimeContext))
+            dask_gateway_url=self.get_dask_gateway_url(runtimeContext),
+            cm_name=self.dask_cm_name)
 
         # emptyDir volume at /shared for sharing the Dask cluster name between containers
         self._add_emptydir_volume_and_binding('shared-data', self.container_shared_dir)
@@ -310,8 +317,8 @@ def create_kubernetes_runtime(self, runtimeContext):
         # Need this ConfigMap to simplify configuration by providing defaults, 
         # as explained here: https://gateway.dask.org/configuration-user.html
         self._add_configmap_volume_and_binding(
-            name=self.daskGateway_cm,
-            cm_name=self.daskGateway_cm_name,
+            name=self.dask_cm_name,
+            cm_name=self.dask_cm_claim_name,
             target=self.daskGateway_config_dir)
 
 
@@ -375,7 +382,7 @@ def get_pod_name(pod):
 
         pod = self.create_kubernetes_runtime(runtimeContext) # analogous to create_runtime()
         self.execute_kubernetes_pod(pod) # analogous to _execute()
-        completion_result = self.wait_for_kubernetes_pod()
+        completion_result = self.wait_for_kubernetes_pod(cm_name = self.dask_cm_name)
         if completion_result.exit_code != 0:
             log_main.error(f"ERROR the command below failed in pod {get_pod_name(pod)}:")
             log_main.error("\t" + " ".join(get_pod_command(pod)))
@@ -387,6 +394,15 @@ class KubernetesDaskClient(KubernetesClient):
     def __init__(self):
         super().__init__()
 
+    @retry_exponential_if_exception_type((ApiException, HTTPError,), log)
+    def submit_pod(self, pod_body):
+        with DaskPodMonitor() as monitor:
+            pod = self.core_api_instance.create_namespaced_pod(self.namespace, pod_body)
+            log.info('Created k8s pod name {} with id {}'.format(pod.metadata.name, pod.metadata.uid))
+            monitor.add(pod)
+            self._set_pod(pod)
+
+
     @retry_exponential_if_exception_type((ApiException, HTTPError,), log)
     def follow_logs(self, status):
         pod_name = self.pod.metadata.name
@@ -409,7 +425,7 @@ def follow_logs(self, status):
 
 
     @retry_exponential_if_exception_type((ApiException, HTTPError, IncompleteStatusException), log)
-    def wait_for_completion(self) -> CompletionResult:
+    def wait_for_completion(self, cm_name: str) -> CompletionResult:
         w = watch.Watch()
         for event in w.stream(self.core_api_instance.list_namespaced_pod, self.namespace, field_selector=self._get_pod_field_selector()):
             pod = event['object']
@@ -439,7 +455,7 @@ def wait_for_completion(self) -> CompletionResult:
                 if self.should_delete_pod():
                     with DaskPodMonitor() as monitor:
                         self.delete_pod_name(pod.metadata.name)
-                        self.delete_configmap_name(cm_name="dask-gateway-cm")
+                        self.delete_configmap_name(cm_name=cm_name)
                         monitor.remove(pod)
                 self._clear_pod()
                 # stop watching for events, our pod is done. Causes wait loop to exit
@@ -469,11 +485,11 @@ def get_last_or_none(container_list: List[Union[V1ContainerStatus, V1Container]]
             return container_list[-1]
 
     @retry_exponential_if_exception_type((ApiException, HTTPError,), log)
-    def create_dask_gateway_cofig_map(self, dask_gateway_url: str):
+    def create_dask_gateway_cofig_map(self, dask_gateway_url: str, cm_name: str):
         gateway = {'gateway': {'address': dask_gateway_url}}
 
         configmap = client.V1ConfigMap(
-            metadata=client.V1ObjectMeta(name="dask-gateway-cm"),
+            metadata=client.V1ObjectMeta(name=cm_name),
             data={
                 "gateway.yaml": yaml.dump(gateway)
             }
@@ -518,7 +534,7 @@ def delete_configmap_name(self, cm_name):
 class DaskPodMonitor(PodMonitor):
     def __init__(self):
         super().__init__()
-    
+
     @staticmethod
     def cleanup():
         log.info('Starting Cleanup')
@@ -528,7 +544,6 @@ def cleanup():
                 log.info('PodMonitor deleting pod {}'.format(pod_name))
                 try:
                     k8s_client.delete_pod_name(pod_name)
-                    k8s_client.delete_configmap_name(cm_name="dask-gateway-cm")
                 except Exception:
                     log.error('Error deleting pod named {}, ignoring'.format(pod_name))
             PodMonitor.pod_names = []

Diff for: calrissian/main.py

@@ -8,6 +8,7 @@
 from calrissian.version import version
 from calrissian.k8s import delete_pods
 from calrissian.report import initialize_reporter, write_report, CPUParser, MemoryParser
+from calrissian.dask import DaskPodMonitor
 from cwltool.main import main as cwlmain
 from cwltool.argparser import arg_parser
 from typing_extensions import Text
@@ -101,7 +102,6 @@ def install_signal_handler():
     """
     signal.signal(signal.SIGTERM, handle_sigterm)
 
-
 def install_tees(stdout_path=None, stderr_path=None):
     """
     Reconnects stdout/stderr to `tee` processes via subprocess.PIPE that can write to user-supplied files
@@ -174,7 +174,10 @@ def main():
                         )
     finally:
         # Always clean up after cwlmain
-        delete_pods()
+        if parsed_args.dask_gateway_url:
+            DaskPodMonitor.cleanup()
+        else:
+            delete_pods()
         if parsed_args.usage_report:
             write_report(parsed_args.usage_report)
         flush_tees()

Diff for: pyproject.toml

@@ -122,6 +122,7 @@ RETRY_ATTEMPTS="0"
 [tool.hatch.envs.test.scripts]
 test = "hatch run nose2"
 testv = "hatch run nose2 --verbose"
+testdask = "hatch run nose2 tests.test_dask"
 cov = ["coverage run --source=calrissian -m nose2", "coverage report"]
 
 [[tool.hatch.envs.test.matrix]]

Diff for: tests/test_dask.py

@@ -0,0 +1,106 @@
+import os
+
+from unittest import TestCase
+from unittest.mock import patch, call, Mock, mock_open
+import logging
+
+from cwltool.utils import CWLObjectType
+
+from calrissian.dask import (
+    KubernetesDaskPodBuilder
+)
+
+from calrissian.dask import (
+    dask_req_validate,
+)
+
+
+class ValidateExtensionTestCase(TestCase):
+
+    def setUp(self):
+        self.daskGatewayRequirement: CWLObjectType = {
+            "workerCores": 2,
+            "workerCoresLimit": 2,
+            "workerMemory": "4G",
+            "clustermaxCore": 8,
+            "clusterMaxMemory": "16G",
+            "class": "https://calrissian-cwl.github.io/schema#DaskGatewayRequirement" # From cwl
+        }
+    
+    def tests_validate_extension(self):
+        self.assertTrue(dask_req_validate(self.daskGatewayRequirement))
+
+
+class KubernetesDaskPodBuilderTestCase(TestCase):
+
+    def setUp(self):
+        self.name = 'PodName'
+        self.container_image = 'dockerimage:1.0'
+        self.volume_mounts = [Mock(), Mock()]
+        self.volumes = [Mock()]
+        self.command_line = ['cat']
+        self.stdout = 'stdout.txt'
+        self.stderr = 'stderr.txt'
+        self.stdin = 'stdin.txt'
+        self.resources = {'cores': 1, 'ram': 1024}
+        self.labels = {'key1': 'val1', 'key2': 123}
+        self.nodeselectors = {'disktype': 'ssd', 'cachelevel': 2}
+        self.security_context = { 'runAsUser': os.getuid(),'runAsGroup': os.getgid() }
+        self.pod_serviceaccount = "podmanager"
+        self.dask_gateway_url = "http://dask-gateway-url:80"
+        self.dask_gateway_controller = False
+        self.environment = {
+            'HOME': '/homedir',
+            'PYTHONPATH': '/app',
+        }
+        self.pod_builder = KubernetesDaskPodBuilder(self.name, self.container_image, self.environment, self.volume_mounts,
+                                                self.volumes, self.command_line, self.stdout, self.stderr, self.stdin,
+                                                self.resources, self.labels, self.nodeselectors, self.security_context, self.pod_serviceaccount,
+                                                self.dask_gateway_url, self.dask_gateway_controller)
+        self.pod_builder.dask_requirement = {
+            "workerCores": 2,
+            "workerCoresLimit": 2,
+            "workerMemory": "4G",
+            "clustermaxCore": 8,
+            "clusterMaxMemory": "16G",
+            "class": "https://calrissian-cwl.github.io/schema#DaskGatewayRequirement" # From cwl
+        }
+
+    def test_main_container_args_without_redirects(self):
+        # container_args returns a list with a single item since it is passed to 'sh', '-c'
+        self.pod_builder.stdout = None
+        self.pod_builder.stderr = None
+        self.pod_builder.stdin = None
+        self.assertEqual(['set -e; trap "touch /shared/completed" EXIT;export DASK_CLUSTER=$(cat /shared/dask_cluster_name.txt) ; cat' ], self.pod_builder.container_args())
+
+    
+    def test_container_args_with_redirects(self):
+        self.assertEqual(['set -e; trap "touch /shared/completed" EXIT;export DASK_CLUSTER=$(cat /shared/dask_cluster_name.txt) ; cat > stdout.txt 2> stderr.txt < stdin.txt'], self.pod_builder.container_args())
+
+
+    def test_init_container_command_with_external_script(self):
+        self.pod_builder.dask_gateway_controller = True
+        self.assertEqual(['python', '/controller/init-dask.py'], self.pod_builder.init_container_command())
+
+
+    @patch("builtins.open", new_callable=mock_open, read_data="print('Default script')")
+    @patch("os.path.join", return_value="/mocked/path/init-dask.py")  # Mock path join
+    def test_init_container_command_with_default_script(self, mock_path_join, mock_file):
+        self.pod_builder.dask_gateway_controller = False
+        expected_command = ['python', '-c', "print('Default script')"]
+        self.assertEqual(expected_command, self.pod_builder.init_container_command())
+
+
+    def test_container_environment(self):
+        environment = self.pod_builder.container_environment()
+        self.assertEqual(len(self.environment) + 8, len(environment)) # +8 Because the dask related are added at runtime 
+        self.assertIn({'name': 'HOME', 'value': '/homedir'}, environment)
+        self.assertIn({'name': 'PYTHONPATH', 'value': '/app'}, environment)
+        self.assertIn({'name': 'DASK_GATEWAY_WORKER_CORES', 'value': '2'}, environment)
+        self.assertIn({'name': 'DASK_GATEWAY_WORKER_MEMORY', 'value': '4G'}, environment)
+        self.assertIn({'name': 'DASK_GATEWAY_WORKER_CORES_LIMIT', 'value': '2'}, environment)
+        self.assertIn({'name': 'DASK_GATEWAY_CLUSTER_MAX_CORES', 'value': '8'}, environment)
+        self.assertIn({'name': 'DASK_GATEWAY_CLUSTER_MAX_RAM', 'value': '16G'}, environment)
+        self.assertIn({'name': 'DASK_GATEWAY_URL', 'value': 'http://dask-gateway-url:80'}, environment)
+        self.assertIn({'name': 'DASK_GATEWAY_IMAGE', 'value': 'dockerimage:1.0'}, environment)  # Replace with actual image if needed
+        self.assertIn({'name': 'DASK_CLUSTER_NAME_PATH', 'value': '/shared/dask_cluster_name.txt'}, environment)