Code refactor - add default init and dispose script - new optional configMap with custom init and dispose schema

lmizzoni · lmizzoni · commit a27849cd64d0 · 2025-01-30T17:40:15.000+01:00
diff --git a/calrissian/context.py b/calrissian/context.py
@@ -23,5 +23,6 @@ def __init__(self, kwargs=None):
         self.pod_serviceaccount = None
         self.tool_logs_basepath = None
         self.max_gpus = None
-        self.gateway_url = None
+        self.dask_gateway_url = None
+        self.dask_gateway_extra_config = None
         return super(CalrissianRuntimeContext, self).__init__(kwargs)
diff --git a/calrissian/dask.py b/calrissian/dask.py
diff --git a/calrissian/dask/custom_schema/schema.yaml b/calrissian/dask/custom_schema/schema.yaml
diff --git a/calrissian/dask/dispose-dask.py b/calrissian/dask/dispose-dask.py
@@ -0,0 +1,30 @@
+import time, os, sys
+import logging
+from dask_gateway import Gateway
+
+logger = logging.getLogger(__name__)
+logging.basicConfig(level=logging.INFO)
+
+source = os.environ.get("DASK_CLUSTER_NAME_PATH", None)
+gateway_url = os.environ.get("DASK_GATEWAY_URL", None)
+signal = "/shared/completed"
+
+logger.info(f"Sidecar: Waiting for completion signal ({signal}) from main container...")
+
+# Poll for the existence of the completion file
+while not os.path.exists(signal):
+    logger.info("Sidecar: Waiting for completion signal from main container...")
+    time.sleep(5)
+
+logger.info("Sidecar: Completion signal received. Shutting down Dask cluster...")
+
+# Shut down the Dask cluster
+with open(source, "r") as f:
+    cluster_name = f.read().strip()
+
+gateway = Gateway(gateway_url)
+cluster = gateway.connect(cluster_name)
+logger.info(f"Sidecar: Connected to Dask cluster: {cluster_name}")
+cluster.shutdown()
+logger.info("Sidecar: Dask cluster shut down successfully.")
+sys.exit(0)
diff --git a/calrissian/dask/init-dask.py b/calrissian/dask/init-dask.py
@@ -0,0 +1,84 @@
+# this code is responsible for creating a Dask cluster
+# it's executed by the CWL runner in the context of the Dask Gateway extension
+# this is for the prototyping purposes only
+import os
+import re
+import logging
+from dask_gateway import Gateway
+
+logger = logging.getLogger(__name__)
+logging.basicConfig(level=logging.INFO)
+
+def parse_memory(mem_str):
+    units = {
+        "B": 1,
+        "KB": 1000,
+        "MB": 1000**2,
+        "GB": 1000**3,
+        "TB": 1000**4,
+        "PB": 1000**5,
+        "K": 1024,
+        "M": 1024**2,
+        "G": 1024**3,
+        "T": 1024**4,
+        "P": 1024**5,
+    }
+
+    match = re.match(r"(\d+)([A-Za-z]+)", mem_str)
+    if not match:
+        raise ValueError(f"Invalid memory format: {mem_str}")
+
+    value, unit = match.groups()
+
+    if unit not in units:
+        raise ValueError(f"Unknown unit: {unit}")
+
+    return int(value) * units[unit]
+
+target = os.environ.get("DASK_CLUSTER_NAME_PATH", None)
+gateway_url = os.environ.get("DASK_GATEWAY_URL", None)
+image = os.environ.get("DASK_GATEWAY_IMAGE", None)
+worker_cores = os.environ.get("DASK_GATEWAY_WORKER_CORES", None)
+worker_cores_limit = os.environ.get("DASK_GATEWAY_WORKER_CORES_LIMIT", None)
+worker_memory = os.environ.get("DASK_GATEWAY_WORKER_MEMORY", None)
+max_cores = os.environ.get("DASK_GATEWAY_CLUSTER_MAX_CORES", None)
+max_ram = os.environ.get("DASK_GATEWAY_CLUSTER_MAX_RAM", None)
+
+logger.info(f"Creating Dask cluster and saving the name to {target}")
+
+gateway = Gateway(gateway_url)
+
+cluster_options = gateway.cluster_options()
+
+cluster_options['image'] = image
+cluster_options['worker_cores'] = float(worker_cores)
+cluster_options['worker_cores_limit'] = int(worker_cores_limit)
+
+cluster_options['worker_memory'] = worker_memory
+#cluster_options["worker_extra_pod_labels"] = {"group": "dask"}
+
+logger.info(f"Cluster options: {cluster_options}")
+logger.info(dir(cluster_options))
+cluster = gateway.new_cluster(cluster_options, shutdown_on_close=False)
+
+# resource requirements
+#worker_cores = 0.5
+#worker_cores_limit = 1 # would come from DaskGateway.Requirement.ResourceRequirement.worker_cores_limit (or worker_cores)
+#worker_memory = 2 # would come from DaskGateway.Requirement.ResourceRequirement.worker_memory
+logger.info(f"Resource requirements: {worker_cores} cores, {worker_memory}")
+
+# scale cluster
+# max_cores = 5 # would come from DaskGateway.Requirement.ResourceRequirement.max_cores
+# max_ram = 16  # would come from DaskGateway.Requirement.ResourceRequirement.max_ram
+logger.info(f"Resource limits: {max_cores} cores, {max_ram} GB RAM")
+
+workers = int(min(int(max_cores) // int(worker_cores_limit), parse_memory(max_ram) // parse_memory(worker_memory)))
+
+logger.info(f"Scaling cluster to {workers} workers")
+cluster.scale(workers)
+
+
+# save the cluster name to a file
+with open(target, "w") as f:
+    f.write(cluster.name)
+logger.info(f"Cluster name {cluster.name} saved to {target}")
diff --git a/calrissian/job.py b/calrissian/job.py
@@ -1,3 +1,4 @@
+from copy import deepcopy
 from typing import Dict
 from cwltool.job import ContainerCommandLineJob, needs_shell_quoting_re
 
diff --git a/calrissian/k8s.py b/calrissian/k8s.py
@@ -3,6 +3,7 @@
 from kubernetes.client.models import V1ContainerState, V1Container, V1ContainerStatus
 from kubernetes.client.rest import ApiException
 from kubernetes.config.config_exception import ConfigException
+import yaml
 from calrissian.executor import IncompleteStatusException
 from calrissian.retry import retry_exponential_if_exception_type
 import threading
@@ -156,70 +157,6 @@ def follow_logs(self):
         
         log.info('[{}] follow_logs end'.format(pod_name))
 
-    @retry_exponential_if_exception_type((ApiException, HTTPError,), log)
-    def follow_container_logs(self, status):
-        pod_name = self.pod.metadata.name
-
-        log.info('[{}] follow_logs start'.format(pod_name))
-        for line in self.core_api_instance.read_namespaced_pod_log(self.pod.metadata.name, self.namespace, follow=True,
-                                                                   _preload_content=False, container=status.name).stream():
-            # .stream() is only available if _preload_content=False
-            # .stream() returns a generator, each iteration yields bytes.
-            # kubernetes-client decodes them as utf-8 when _preload_content is True
-            # https://github.com/kubernetes-client/python/blob/fcda6fe96beb21cd05522c17f7f08c5a7c0e3dc3/kubernetes/client/rest.py#L215-L216
-            # So we do the same here
-            if not status.state.running:
-                break
-            line = line.decode('utf-8', errors="ignore").rstrip()
-            log.debug('[{}] {}'.format(pod_name, line))
-            self.tool_log.append(self.format_log_entry(pod_name, line))
-        
-        log.info('[{}] follow_logs end'.format(pod_name))
-
-    @retry_exponential_if_exception_type((ApiException, HTTPError, IncompleteStatusException), log)
-    def wait_for_dask_completion(self) -> 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']
-            # status = self.get_first_or_none(pod.status.container_statuses)
-            last_status = self.get_last_or_none(pod.status.container_statuses)
-            if last_status == None or not self.state_is_terminated(last_status.state):
-                statuses = self.get_list_or_none(pod.status.container_statuses)
-                if statuses == None:
-                    continue
-                for status in statuses:
-                    log.info('pod name {} with id {} has status {}'.format(pod.metadata.name, pod.metadata.uid, status))
-                    if status is None:
-                        continue
-                    if self.state_is_waiting(status.state):
-                        continue
-                    elif self.state_is_running(status.state):
-                        # Can only get logs once container is running
-                        self.follow_container_logs(status) # This will not return until container completes
-                    elif self.state_is_terminated(status.state):
-                        continue
-                    else:
-                        raise CalrissianJobException('Unexpected pod container status', status)
-            elif self.state_is_terminated(last_status.state):
-                log.info('Handling terminated pod name {} with id {}'.format(pod.metadata.name, pod.metadata.uid))
-                container = self.get_last_or_none(pod.spec.containers)
-                self._handle_completion(last_status.state, container)
-                if self.should_delete_pod():
-                    with PodMonitor() as monitor:
-                        self.delete_pod_name(pod.metadata.name)
-                        monitor.remove(pod)
-                self._clear_pod()
-                # stop watching for events, our pod is done. Causes wait loop to exit
-                w.stop()
-            else:
-                raise CalrissianJobException('Unexpected pod container status', last_status)
-        
-        # When the pod is done we should have a completion result
-        # Otherwise it will lead to further exceptions
-        if self.completion_result is None:
-            raise IncompleteStatusException
-        
-        return self.completion_result
         
     @retry_exponential_if_exception_type((ApiException, HTTPError, IncompleteStatusException), log)
     def wait_for_completion(self) -> CompletionResult:
@@ -280,20 +217,6 @@ def state_is_waiting(state):
     def state_is_terminated(state):
         return state.terminated
 
-    @staticmethod
-    def get_list_or_none(container_list: List[Union[V1ContainerStatus, V1Container]]) -> Union[V1ContainerStatus, V1Container]:
-        if not container_list: # None or empty list
-            return None
-        else:
-            return list(container_list)
-
-    @staticmethod
-    def get_last_or_none(container_list: List[Union[V1ContainerStatus, V1Container]]) -> Union[V1ContainerStatus, V1Container]:
-        if not container_list: # None or empty list
-            return None
-        else:
-            return container_list[-1]
-
     @staticmethod
     def get_first_or_none(container_list: List[Union[V1ContainerStatus, V1Container]]) -> Union[V1ContainerStatus, V1Container]:
         """
@@ -350,16 +273,6 @@ def get_current_pod(self):
             raise CalrissianJobException("Missing required environment variable ${}".format(POD_NAME_ENV_VARIABLE))
         return self.get_pod_for_name(pod_name)
 
-    @retry_exponential_if_exception_type((ApiException, HTTPError,), log)
-    def create_dask_gateway_cofig_map(self, gateway_url: str):
-        configmap = client.V1ConfigMap(
-            metadata=client.V1ObjectMeta(name="dask-gateway-cm"),
-            data={
-                "gateway.yaml": f"|\ngateway:\n\taddress:{gateway_url}"
-            }
-        )
-
-        self.core_api_instance.create_namespaced_config_map(namespace=self.namespace, body=configmap)   
 
 class PodMonitor(object):
     """
diff --git a/calrissian/main.py b/calrissian/main.py
@@ -34,7 +34,7 @@ def get_log_level(parsed_args):
 
 
 def activate_logging(level):
-    loggers = ['executor','context','tool','job', 'k8s','main']
+    loggers = ['executor','context','tool','job', 'k8s','main', 'dask']
     for logger in loggers:
         logging.getLogger('calrissian.{}'.format(logger)).setLevel(level)
         logging.getLogger('calrissian.{}'.format(logger)).addHandler(logging.StreamHandler())
@@ -53,7 +53,8 @@ def add_arguments(parser):
     parser.add_argument('--stderr', type=Text, nargs='?', help='Output file name to tee standard error to (includes tool logs)')
     parser.add_argument('--tool-logs-basepath', type=Text, nargs='?', help='Base path for saving the tool logs')
     parser.add_argument('--conf', help='Defines the default values for the CLI arguments', action='append')
-    parser.add_argument('--gateway-url', type=Text, nargs='?', help='Defines the Dask Gateway URL', required=False)
+    parser.add_argument('--dask-gateway-url', type=Text, nargs='?', help='Defines the Dask Gateway URL', required=False)
+    parser.add_argument('--dask-gateway-extra-config', type=Text, nargs='?', help='YAML file of extra k8s config for Dask', required=False)
 
 
 def print_version():
@@ -133,7 +134,7 @@ def add_custom_schema():
     cwltool.command_line_tool.ACCEPTLIST_RE = cwltool.command_line_tool.ACCEPTLIST_EN_RELAXED_RE
     supported_versions = ["v1.0", "v1.1", "v1.2"]
 
-    with open(os.path.join(os.path.dirname(__file__), "custom_schema/schema.yaml")) as f:
+    with open(os.path.join(os.path.dirname(__file__), "dask/custom_schema/schema.yaml")) as f:
         schema_content = f.read()
 
     for s in supported_versions:
@@ -160,13 +161,16 @@ def main():
     runtime_context = CalrissianRuntimeContext(vars(parsed_args))
     runtime_context.select_resources = executor.select_resources
     install_signal_handler()
+
+    parsed_args.enable_ext = True
+
     try:
         result = cwlmain(args=parsed_args,
                          executor=executor,
                          loadingContext=CalrissianLoadingContext(),
                          runtimeContext=runtime_context,
                          versionfunc=version,
-                         custom_schema_callback=(add_custom_schema if parsed_args.gateway_url else None)
+                         custom_schema_callback=(add_custom_schema if parsed_args.dask_gateway_url else None)
                         )
     finally:
         # Always clean up after cwlmain

Original file line number	Diff line number	Diff line change
`@@ -1,3 +1,4 @@`
	`1`	`+from copy import deepcopy`
`1`	`2`	`from typing import Dict`
`2`	`3`	`from cwltool.job import ContainerCommandLineJob, needs_shell_quoting_re`
`3`	`4`