Add nodes toleration to Ray pods

kfp/doc/simple_transform_pipeline.md

@@ -111,6 +111,7 @@ The parameters used here are as follows:
   * memory - memory
   * image - image to use
   * image_pull_secret - image pull secret
+  * tolerations - (optional) tolerations for the ray pods
 * ray_worker_options: worker node options (we here are using only 1 worker pool), containing the following:
   * replicas - number of replicas to create
   * max_replicas - max number of replicas
@@ -119,6 +120,7 @@ The parameters used here are as follows:
   * memory - memory
   * image - image to use
   * image_pull_secret - image pull secret
+  * tolerations - (optional) tolerations for the ray pods
 * server_url - server url
 * additional_params: additional (support) parameters, containing the following:
   * wait_interval - wait interval for API server, sec

kfp/kfp_support_lib/shared_workflow_support/src/runtime_utils/remote_jobs_utils.py

@@ -25,9 +25,9 @@
     ClusterSpec,
     HeadNodeSpec,
     RayJobRequest,
-    Template,
     WorkerNodeSpec,
     environment_variables_decoder,
+    template_decoder,
     volume_decoder,
 )
 from ray.job_submission import JobStatus
@@ -121,41 +121,37 @@ def create_ray_cluster(
         """
         # start with templates
         # head_node
-        cpus = head_node.get("cpu", 1)
-        memory = head_node.get("memory", 1)
-        gpus = head_node.get("gpu", 0)
-        accelerator = head_node.get("gpu_accelerator", None)
+        dct = {}
+        dct["cpu"] = head_node.get("cpu", 1)
+        dct["memory"] = head_node.get("memory", 1)
+        dct["gpu"] = head_node.get("gpu", 0)
+        dct["gpu_accelerator"] = head_node.get("gpu_accelerator", None)
         head_node_template_name = f"{name}-head-template"
-        _, _ = self.api_server_client.delete_compute_template(ns=namespace, name=head_node_template_name)
-        head_template = Template(
-            name=head_node_template_name,
-            namespace=namespace,
-            cpu=cpus,
-            memory=memory,
-            gpu=gpus,
-            gpu_accelerator=accelerator,
-        )
+        dct["name"] = head_node_template_name
+        dct["namespace"] = namespace
+        if "tolerations" in head_node:
+            dct["tolerations"] = head_node.get("tolerations")
+        _, _ = self.api_server_client.delete_compute_template(ns=namespace, name=dct["name"])
+        head_template = template_decoder(dct)
         status, error = self.api_server_client.create_compute_template(head_template)
         if status != 200:
             return status, error
         worker_template_names = [""] * len(worker_nodes)
         index = 0
         # For every worker group
         for worker_node in worker_nodes:
-            cpus = worker_node.get("cpu", 1)
-            memory = worker_node.get("memory", 1)
-            gpus = worker_node.get("gpu", 0)
-            accelerator = worker_node.get("gpu_accelerator", None)
+            dct = {}
+            dct["cpu"] = worker_node.get("cpu", 1)
+            dct["memory"] = worker_node.get("memory", 1)
+            dct["gpu"] = worker_node.get("gpu", 0)
+            dct["gpu_accelerator"] = worker_node.get("gpu_accelerator", None)
             worker_node_template_name = f"{name}-worker-template-{index}"
-            _, _ = self.api_server_client.delete_compute_template(ns=namespace, name=worker_node_template_name)
-            worker_template = Template(
-                name=worker_node_template_name,
-                namespace=namespace,
-                cpu=cpus,
-                memory=memory,
-                gpu=gpus,
-                gpu_accelerator=accelerator,
-            )
+            dct["name"] = worker_node_template_name
+            dct["namespace"] = namespace
+            if "tolerations" in worker_node:
+                dct["tolerations"] = worker_node.get("tolerations")
+            _, _ = self.api_server_client.delete_compute_template(ns=namespace, name=dct["name"])
+            worker_template = template_decoder(dct)
             status, error = self.api_server_client.create_compute_template(worker_template)
             if status != 200:
                 return status, error

kfp/pipeline_generator/single-pipeline/templates/simple_pipeline.py

@@ -135,6 +135,7 @@ def {{ pipeline_name }}(
         memory - memory
         image - image to use
         image_pull_secret - image pull secret
+        tolerations - (optional) tolerations for the ray pods
     :param ray_worker_options: worker node options (we here are using only 1 worker pool), containing the following:
         replicas - number of replicas to create
         max_replicas - max number of replicas
@@ -143,6 +144,7 @@ def {{ pipeline_name }}(
         memory - memory
         image - image to use
         image_pull_secret - image pull secret
+        tolerations - (optional) tolerations for the ray pods
     :param server_url - server url
     :param additional_params: additional (support) parameters, containing the following:
         wait_interval - wait interval for API server, sec

transforms/code/code2parquet/kfp_ray/code2parquet_wf.py

@@ -138,6 +138,7 @@ def code2parquet(
         memory - memory
         image - image to use
         image_pull_secret - image pull secret
+        tolerations - (optional) tolerations for the ray pods
     :param ray_worker_options: worker node options (we here are using only 1 worker pool), containing the following:
         replicas - number of replicas to create
         max_replicas - max number of replicas
@@ -146,6 +147,7 @@ def code2parquet(
         memory - memory
         image - image to use
         image_pull_secret - image pull secret
+        tolerations - (optional) tolerations for the ray pods
     :param server_url - server url
     :param additional_params: additional (support) parameters, containing the following:
         wait_interval - wait interval for API server, sec

transforms/code/code_quality/kfp_ray/code_quality_wf.py

@@ -135,6 +135,7 @@ def code_quality(
         memory - memory
         image - image to use
         image_pull_secret - image pull secret
+        tolerations - (optional) tolerations for the ray pods
     :param ray_worker_options: worker node options (we here are using only 1 worker pool), containing the following:
         replicas - number of replicas to create
         max_replicas - max number of replicas
@@ -143,6 +144,7 @@ def code_quality(
         memory - memory
         image - image to use
         image_pull_secret - image pull secret
+        tolerations - (optional) tolerations for the ray pods
     :param server_url - server url
     :param additional_params: additional (support) parameters, containing the following:
         wait_interval - wait interval for API server, sec

transforms/code/header_cleanser/kfp_ray/header_cleanser_wf.py

@@ -132,6 +132,7 @@ def header_cleanser(
         memory - memory
         image - image to use
         image_pull_secret - image pull secret
+        tolerations - (optional) tolerations for the ray pods
     :param ray_worker_options: worker node options (we here are using only 1 worker pool), containing the following:
         replicas - number of replicas to create
         max_replicas - max number of replicas
@@ -140,6 +141,7 @@ def header_cleanser(
         memory - memory
         image - image to use
         image_pull_secret - image pull secret
+        tolerations - (optional) tolerations for the ray pods
     :param server_url - server url
     :param additional_params: additional (support) parameters, containing the following:
         wait_interval - wait interval for API server, sec

transforms/code/malware/kfp_ray/malware_wf.py

@@ -125,6 +125,7 @@ def malware(
         memory - memory
         image - image to use
         image_pull_secret - image pull secret
+        tolerations - (optional) tolerations for the ray pods
     :param ray_worker_options: worker node options (we here are using only 1 worker pool), containing the following:
         replicas - number of replicas to create
         max_replicas - max number of replicas
@@ -133,6 +134,7 @@ def malware(
         memory - memory
         image - image to use
         image_pull_secret - image pull secret
+        tolerations - (optional) tolerations for the ray pods
     :param server_url - server url
     :param additional_params: additional (support) parameters, containing the following:
         wait_interval - wait interval for API server, sec

transforms/code/proglang_select/kfp_ray/proglang_select_wf.py

@@ -128,6 +128,7 @@ def lang_select(
         memory - memory
         image - image to use
         image_pull_secret - image pull secret
+        tolerations - (optional) tolerations for the ray pods
     :param ray_worker_options: worker node options (we here are using only 1 worker pool), containing the following:
         replicas - number of replicas to create
         max_replicas - max number of replicas
@@ -136,6 +137,7 @@ def lang_select(
         memory - memory
         image - image to use
         image_pull_secret - image pull secret
+        tolerations - (optional) tolerations for the ray pods
     :param server_url - server url
     :param additional_params: additional (support) parameters, containing the following:
         wait_interval - wait interval for API server, sec

transforms/code/repo_level_ordering/kfp_ray/repo_level_order_wf.py

@@ -155,6 +155,7 @@ def repo_level_order(
         memory - memory
         image - image to use
         image_pull_secret - image pull secret
+        tolerations - (optional) tolerations for the ray pods
     :param ray_worker_options: worker node options (we here are using only 1 worker pool), containing the following:
         replicas - number of replicas to create
         max_replicas - max number of replicas
@@ -163,6 +164,7 @@ def repo_level_order(
         memory - memory
         image - image to use
         image_pull_secret - image pull secret
+        tolerations - (optional) tolerations for the ray pods
     :param server_url - server url
     :param additional_params: additional (support) parameters, containing the following:
         wait_interval - wait interval for API server, sec

transforms/language/doc_chunk/kfp_ray/doc_chunk_multiple_wf.py

@@ -131,6 +131,7 @@ def doc_chunk(
         memory - memory
         image - image to use
         image_pull_secret - image pull secret
+        tolerations - (optional) tolerations for the ray pods
     :param ray_worker_options: worker node options (we here are using only 1 worker pool), containing the following:
         replicas - number of replicas to create
         max_replicas - max number of replicas
@@ -139,6 +140,7 @@ def doc_chunk(
         memory - memory
         image - image to use
         image_pull_secret - image pull secret
+        tolerations - (optional) tolerations for the ray pods
     :param server_url - server url
     :param additional_params: additional (support) parameters, containing the following:
         wait_interval - wait interval for API server, sec

transforms/language/doc_chunk/kfp_ray/doc_chunk_wf.py

@@ -132,6 +132,7 @@ def doc_chunk(
         memory - memory
         image - image to use
         image_pull_secret - image pull secret
+        tolerations - (optional) tolerations for the ray pods
     :param ray_worker_options: worker node options (we here are using only 1 worker pool), containing the following:
         replicas - number of replicas to create
         max_replicas - max number of replicas
@@ -140,6 +141,7 @@ def doc_chunk(
         memory - memory
         image - image to use
         image_pull_secret - image pull secret
+        tolerations - (optional) tolerations for the ray pods
     :param server_url - server url
     :param additional_params: additional (support) parameters, containing the following:
         wait_interval - wait interval for API server, sec

transforms/language/doc_quality/kfp_ray/doc_quality_multiple_wf.py

@@ -128,6 +128,7 @@ def doc_quality(
         memory - memory
         image - image to use
         image_pull_secret - image pull secret
+        tolerations - (optional) tolerations for the ray pods
     :param ray_worker_options: worker node options (we here are using only 1 worker pool), containing the following:
         replicas - number of replicas to create
         max_replicas - max number of replicas
@@ -136,6 +137,7 @@ def doc_quality(
         memory - memory
         image - image to use
         image_pull_secret - image pull secret
+        tolerations - (optional) tolerations for the ray pods
     :param server_url - server url
     :param additional_params: additional (support) parameters, containing the following:
         wait_interval - wait interval for API server, sec

transforms/language/doc_quality/kfp_ray/doc_quality_wf.py

@@ -128,6 +128,7 @@ def doc_quality(
         memory - memory
         image - image to use
         image_pull_secret - image pull secret
+        tolerations - (optional) tolerations for the ray pods
     :param ray_worker_options: worker node options (we here are using only 1 worker pool), containing the following:
         replicas - number of replicas to create
         max_replicas - max number of replicas
@@ -136,6 +137,7 @@ def doc_quality(
         memory - memory
         image - image to use
         image_pull_secret - image pull secret
+        tolerations - (optional) tolerations for the ray pods
     :param server_url - server url
     :param additional_params: additional (support) parameters, containing the following:
         wait_interval - wait interval for API server, sec

transforms/language/lang_id/kfp_ray/lang_id_multiple_wf.py

@@ -137,6 +137,7 @@ def lang_id(
         memory - memory
         image - image to use
         image_pull_secret - image pull secret
+        tolerations - (optional) tolerations for the ray pods
     :param ray_worker_options: worker node options (we here are using only 1 worker pool), containing the following:
         replicas - number of replicas to create
         max_replicas - max number of replicas
@@ -145,6 +146,7 @@ def lang_id(
         memory - memory
         image - image to use
         image_pull_secret - image pull secret
+        tolerations - (optional) tolerations for the ray pods
     :param server_url - server url
     :param additional_params: additional (support) parameters, containing the following:
         wait_interval - wait interval for API server, sec

transforms/language/lang_id/kfp_ray/lang_id_wf.py

@@ -138,6 +138,7 @@ def lang_id(
         memory - memory
         image - image to use
         image_pull_secret - image pull secret
+        tolerations - (optional) tolerations for the ray pods
     :param ray_worker_options: worker node options (we here are using only 1 worker pool), containing the following:
         replicas - number of replicas to create
         max_replicas - max number of replicas
@@ -146,6 +147,7 @@ def lang_id(
         memory - memory
         image - image to use
         image_pull_secret - image pull secret
+        tolerations - (optional) tolerations for the ray pods
     :param server_url - server url
     :param additional_params: additional (support) parameters, containing the following:
         wait_interval - wait interval for API server, sec

transforms/language/pdf2parquet/kfp_ray/pdf2parquet_multiple_wf.py

@@ -125,6 +125,7 @@ def pdf2parquet(
         memory - memory
         image - image to use
         image_pull_secret - image pull secret
+        tolerations - (optional) tolerations for the ray pods
     :param ray_worker_options: worker node options (we here are using only 1 worker pool), containing the following:
         replicas - number of replicas to create
         max_replicas - max number of replicas
@@ -133,6 +134,7 @@ def pdf2parquet(
         memory - memory
         image - image to use
         image_pull_secret - image pull secret
+        tolerations - (optional) tolerations for the ray pods
     :param server_url - server url
     :param additional_params: additional (support) parameters, containing the following:
         wait_interval - wait interval for API server, sec

transforms/language/pdf2parquet/kfp_ray/pdf2parquet_wf.py

@@ -129,6 +129,7 @@ def pdf2parquet(
         memory - memory
         image - image to use
         image_pull_secret - image pull secret
+        tolerations - (optional) tolerations for the ray pods
     :param ray_worker_options: worker node options (we here are using only 1 worker pool), containing the following:
         replicas - number of replicas to create
         max_replicas - max number of replicas
@@ -137,6 +138,7 @@ def pdf2parquet(
         memory - memory
         image - image to use
         image_pull_secret - image pull secret
+        tolerations - (optional) tolerations for the ray pods
     :param server_url - server url
     :param additional_params: additional (support) parameters, containing the following:
         wait_interval - wait interval for API server, sec

transforms/language/pii_redactor/kfp_ray/pii_redactor_wf.py

@@ -123,6 +123,7 @@ def pii_redactor(
         memory - memory
         image - image to use
         image_pull_secret - image pull secret
+        tolerations - (optional) tolerations for the ray pods
     :param ray_worker_options: worker node options (we here are using only 1 worker pool), containing the following:
         replicas - number of replicas to create
         max_replicas - max number of replicas
@@ -131,6 +132,7 @@ def pii_redactor(
         memory - memory
         image - image to use
         image_pull_secret - image pull secret
+        tolerations - (optional) tolerations for the ray pods
     :param server_url - server url
     :param additional_params: additional (support) parameters, containing the following:
         wait_interval - wait interval for API server, sec

transforms/language/text_encoder/kfp_ray/text_encoder_multiple_wf.py

@@ -128,6 +128,7 @@ def text_encoder(
         memory - memory
         image - image to use
         image_pull_secret - image pull secret
+        tolerations - (optional) tolerations for the ray pods
     :param ray_worker_options: worker node options (we here are using only 1 worker pool), containing the following:
         replicas - number of replicas to create
         max_replicas - max number of replicas
@@ -136,6 +137,7 @@ def text_encoder(
         memory - memory
         image - image to use
         image_pull_secret - image pull secret
+        tolerations - (optional) tolerations for the ray pods
     :param server_url - server url
     :param additional_params: additional (support) parameters, containing the following:
         wait_interval - wait interval for API server, sec

transforms/language/text_encoder/kfp_ray/text_encoder_wf.py

@@ -129,6 +129,7 @@ def text_encoder(
         memory - memory
         image - image to use
         image_pull_secret - image pull secret
+        tolerations - (optional) tolerations for the ray pods
     :param ray_worker_options: worker node options (we here are using only 1 worker pool), containing the following:
         replicas - number of replicas to create
         max_replicas - max number of replicas
@@ -137,6 +138,7 @@ def text_encoder(
         memory - memory
         image - image to use
         image_pull_secret - image pull secret
+        tolerations - (optional) tolerations for the ray pods
     :param server_url - server url
     :param additional_params: additional (support) parameters, containing the following:
         wait_interval - wait interval for API server, sec

transforms/universal/doc_id/kfp_ray/doc_id_wf.py

@@ -139,6 +139,7 @@ def doc_id(
         memory - memory
         image - image to use
         image_pull_secret - image pull secret
+        tolerations - (optional) tolerations for the ray pods
     :param ray_worker_options: worker node options (we here are using only 1 worker pool), containing the following:
         replicas - number of replicas to create
         max_replicas - max number of replicas
@@ -147,6 +148,7 @@ def doc_id(
         memory - memory
         image - image to use
         image_pull_secret - image pull secret
+        tolerations - (optional) tolerations for the ray pods
     :param server_url - server url
     :param additional_params: additional (support) parameters, containing the following:
         wait_interval - wait interval for API server, sec