Custom schema implementation - add schema file in temporary location

Author: lmizzoni

Diff for: calrissian/custom_schema/schema.yaml

@@ -0,0 +1,48 @@
+$base: https://calrissian-cwl.github.io/schema#
+$namespaces:
+  cwl: "https://w3id.org/cwl/cwl#"
+$graph:
+- $import: https://w3id.org/cwl/CommonWorkflowLanguage.yml
+
+- name: DaskGatewayRequirement
+  type: record
+  extends: cwl:ProcessRequirement
+  inVocab: false
+  doc: "Indicates that a process requires a Dask cluster procured via [Dask Gateway](https://gateway.dask.org/) runtime."
+  fields:
+    class:
+      type: 'string'
+      doc: "Always 'DaskGatewayRequirement'"
+      jsonldPredicate:
+        "_id": "@type"
+        "_type": "@vocab"
+    workerCores:
+      type:
+        - 'int'
+        - 'cwl:Expression'
+      doc: |
+        Number of cpu-cores available for a dask worker.
+    workerCoresLimit:
+      type:
+        - 'int'
+        - 'cwl:Expression'
+      doc: |
+        Maximum number of cpu-cores available for a dask worker.
+    workerMemory:
+      type:
+        - 'string'
+        - 'cwl:Expression'
+      doc: |
+        Maximum number of bytes available for a dask worker.
+    clustermaxCore:
+      type:
+        - 'int'
+        - 'cwl:Expression'
+      doc: |
+        Maximum number of cores to configure the cluster
+    clusterMaxMemory:
+      type:
+        - 'string'
+        - 'cwl:Expression'
+      doc: |
+        Maximum amount of memory to configure the cluster

Diff for: calrissian/job.py

@@ -490,7 +490,7 @@ def init_containers(self):
                 'volumeMounts': self.volume_mounts,
             })
 
-        dask_requirement = next((elem for elem in self.hints if elem['class'] == 'https://calrissian-cwl.github.io/schema#DaskGatewayRequirement'), None)
+        dask_requirement = next((elem for elem in self.requirements if elem['class'] == 'https://calrissian-cwl.github.io/schema#DaskGatewayRequirement'), None)
 
         init_dask_command = [
             'python',
@@ -500,17 +500,17 @@ def init_containers(self):
             '--gateway-url',
             self.gateway_url,
             '--image',
-            str(dask_requirement['dockerPull']),
+            str(self.container_image),
             '--worker-cores',
             str(dask_requirement["workerCores"]),
             '--worker-memory',
             str(dask_requirement["workerMemory"]),
             '--worker-cores-limit',
             str(dask_requirement["workerCoresLimit"]),
             '--max-cores',
-            str(dask_requirement["coresMax"]),
+            str(dask_requirement["clustermaxCore"]),
             '--max-ram',
-            str(dask_requirement["ramMax"])
+            str(dask_requirement["clusterMaxMemory"])
         ]
 
         log.info(init_dask_command)
@@ -693,8 +693,6 @@ def check_requirements(self, runtimeContext):
                     if not field in self.supported_features[feature]:
                         raise UnsupportedRequirement('Error: feature {}.{} is not supported'.format(feature, field))
 
-        
-
     def _get_container_image(self):
         docker_requirement, _ = self.get_requirement('DockerRequirement')
         if docker_requirement:
@@ -706,14 +704,6 @@ def _get_container_image(self):
             raise CalrissianCommandLineJobException('Unable to create Job - Please ensure tool has a DockerRequirement with dockerPull or specify a default_container')
         return container_image
 
-    def _get_worker_image(self):
-        docker_requirement, _ = self.get_requirement('https://calrissian-cwl.github.io/schema#DaskGatewayRequirement')
-        if docker_requirement:
-            container_image = docker_requirement['dockerPull']
-        if not container_image:
-            raise CalrissianCommandLineJobException('Unable to create Job - Please ensure tool has a DaskGatewayRequirement with dockerPull')
-        return container_image
-
     def quoted_command_line(self):
         return quoted_arg_list(self.command_line)
 
@@ -800,7 +790,7 @@ def create_kubernetes_runtime(self, runtimeContext):
 
             k8s_builder = KubernetesDaskPodBuilder(
                 self.name,
-                self._get_worker_image(),
+                self._get_container_image(),
                 self.environment,
                 self.volume_builder.volume_mounts,
                 self.volume_builder.volumes,

Diff for: calrissian/main.py

@@ -1,4 +1,8 @@
 import contextlib
+import re
+
+import cwltool
+from cwltool.process import use_custom_schema, get_schema
 from calrissian.executor import ThreadPoolJobExecutor
 from calrissian.context import CalrissianLoadingContext, CalrissianRuntimeContext
 from calrissian.version import version
@@ -124,14 +128,30 @@ def flush_tees():
     sys.stderr.flush()
 
 
+def add_custom_schema():
+    cwltool.command_line_tool.ACCEPTLIST_EN_RELAXED_RE = re.compile(r".*")
+    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:
+        schema_content = f.read()
+
+    for s in supported_versions:
+        use_custom_schema(s, "https://calrissian-cwl.github.io/schema", schema_content)
+        get_schema(s)
+
+    cwltool.process.supportedProcessRequirements.extend([
+        "https://calrissian-cwl.github.io/schema#DaskGatewayRequirement"
+    ])
+
+
 def main():
     parser = arg_parser()
     add_arguments(parser)
     parsed_args = parse_arguments(parser)
     level = get_log_level(parsed_args)
     activate_logging(level)
     install_tees(parsed_args.stdout, parsed_args.stderr)
-    log.info("CALRISSIAN LOG")
     max_ram_megabytes = MemoryParser.parse_to_megabytes(parsed_args.max_ram)
     max_cores = CPUParser.parse(parsed_args.max_cores)
     max_gpus = int(parsed_args.max_gpus) if parsed_args.max_gpus else 0
@@ -146,7 +166,8 @@ def main():
                          loadingContext=CalrissianLoadingContext(),
                          runtimeContext=runtime_context,
                          versionfunc=version,
-                         )
+                         custom_schema_callback=(add_custom_schema if parsed_args.gateway_url else None)
+                        )
     finally:
         # Always clean up after cwlmain
         delete_pods()