Merge pull request #627 from revit13/taint

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