Add support for Kueue.

This commit adds support for using Kueue to submit jobs/pods into Kubernetes.
There are two config options:
- KUEUE_ENABLED: set to True/False
- KUEUE_LOCALQUEUE_NAME: set to the name of the localqueue configured with Kueue.
  See [this](https://kueue.sigs.k8s.io/docs/concepts/local_queue) for details

The config options can be set in the main metaflow config or via the
@kubernetes decorator.

Testing Done:

- Verified that specifying kueue config options in Metaflow config (~/.metaflowconfig/json) works as expected.
- Verified that specifying kueue config options in @kubernetes works as expected
- Verified that @kubernetes options take precedence over the global config
   - If the global KUEUE_ENABLED config is True, but locally set to False for
     a particular step, the step does not run with Kueue.
- Verified that the kueue labels and annotations are set correctly and kueue actually
  runs the jobs.
- Verified that if kueue is configured to manage "pod", Metaflow create argo-workflow
  pods are scheduled by kueue.

- Verified that the default behavior is to not use Kueue and everything works
  correctly as before (jobs and argo-workflows)

metaflow/metaflow_config.py

@@ -313,6 +313,10 @@
 ARGO_WORKFLOWS_KUBERNETES_SECRETS = from_conf("ARGO_WORKFLOWS_KUBERNETES_SECRETS", "")
 ARGO_WORKFLOWS_ENV_VARS_TO_SKIP = from_conf("ARGO_WORKFLOWS_ENV_VARS_TO_SKIP", "")
 
+## Kueue Support
+KUEUE_ENABLED = from_conf("KUEUE_ENABLED", False)
+KUEUE_LOCALQUEUE_NAME = from_conf("KUEUE_LOCALQUEUE_NAME", "")
+
 ##
 # Argo Events Configuration
 ##

metaflow/plugins/argo/argo_workflows.py

@@ -1368,15 +1368,30 @@ def _container_templates(self):
  tmpfs_size = resources["tmpfs_size"]
  tmpfs_path = resources["tmpfs_path"]
  tmpfs_tempdir = resources["tmpfs_tempdir"]
- # Set shared_memory to 0 if it isn't specified. This results
- # in Kubernetes using it's default value when the pod is created.
- shared_memory = resources.get("shared_memory", 0)
 
  tmpfs_enabled = use_tmpfs or (tmpfs_size and not use_tmpfs)
 
  if tmpfs_enabled and tmpfs_tempdir:
  env["METAFLOW_TEMPDIR"] = tmpfs_path
 
+ # Set shared_memory to 0 if it isn't specified. This results
+ # in Kubernetes using it's default value when the pod is created.
+ shared_memory = resources.get("shared_memory", 0)
+
+ kueue_enabled = resources["kueue_enabled"]
+ kueue_localqueue_name = resources["kueue_localqueue_name"]
+ kueue_annotations = {}
+ kueue_labels = {}
+ if kueue_enabled:
+ kueue_annotations["kueue.x-k8s.io/retriable-in-group"] = "false"
+ kueue_annotations["kueue.x-k8s.io/pod-group-total-count"] = str(
+ 1
+ ) # For now, might change with @parallel support
+ kueue_labels["kueue.x-k8s.io/queue-name"] = kueue_localqueue_name
+ kueue_labels["kueue.x-k8s.io/managed"] = "true"
+ kueue_labels["kueue.x-k8s.io/pod-group-name"] = (
+ "{{workflow.name}}-" + node.name + "{{inputs.parameters.split-index}}"
+ )
  # Create a ContainerTemplate for this node. Ideally, we would have
  # liked to inline this ContainerTemplate and avoid scanning the workflow
  # twice, but due to issues with variable substitution, we will have to
@@ -1399,13 +1414,16 @@ def _container_templates(self):
  minutes_between_retries=minutes_between_retries,
  )
  .metadata(
- ObjectMeta().annotation("metaflow/step_name", node.name)
+ ObjectMeta()
+ .annotation("metaflow/step_name", node.name)
  # Unfortunately, we can't set the task_id since it is generated
  # inside the pod. However, it can be inferred from the annotation
  # set by argo-workflows - `workflows.argoproj.io/outputs` - refer
  # the field 'task-id' in 'parameters'
  # .annotation("metaflow/task_id", ...)
  .annotation("metaflow/attempt", retry_count)
+ .annotations(kueue_annotations)
+ .labels(kueue_labels)
  )
  # Set emptyDir volume for state management
  .empty_dir_volume("out")

metaflow/plugins/kubernetes/kubernetes.py

@@ -175,6 +175,8 @@ def create_job(
  tolerations=None,
  labels=None,
  shared_memory=None,
+ kueue_enabled=None,
+ kueue_localqueue_name=None,
  ):
  if env is None:
  env = {}
@@ -183,6 +185,8 @@ def create_job(
  KubernetesClient()
  .job(
  generate_name="t-{uid}-".format(uid=str(uuid4())[:8]),
+ run_id=run_id,
+ task_id=task_id,
  namespace=namespace,
  service_account=service_account,
  secrets=secrets,
@@ -215,6 +219,8 @@ def create_job(
  tmpfs_path=tmpfs_path,
  persistent_volume_claims=persistent_volume_claims,
  shared_memory=shared_memory,
+ kueue_enabled=kueue_enabled,
+ kueue_localqueue_name=kueue_localqueue_name,
  )
  .environment_variable("METAFLOW_CODE_SHA", code_package_sha)
  .environment_variable("METAFLOW_CODE_URL", code_package_url)

metaflow/plugins/kubernetes/kubernetes_cli.py

@@ -1,4 +1,5 @@
 import os
+
 import sys
 import time
 import traceback
@@ -108,6 +109,15 @@ def kubernetes():
  multiple=False,
 )
 @click.option("--shared-memory", default=None, help="Size of shared memory in MiB")
+@click.option(
+ "--kueue-enabled",
+ is_flag=True,
+ default=None,
+ help="Whether to use Kueue for scheduling Kubernetes jobs/pods",
+)
+@click.option(
+ "--kueue-localqueue-name", help="Name of the LocalQueue configured with kueue"
+)
 @click.pass_context
 def step(
  ctx,
@@ -134,6 +144,8 @@ def step(
  persistent_volume_claims=None,
  tolerations=None,
  shared_memory=None,
+ kueue_enabled=None,
+ kueue_localqueue_name=None,
  **kwargs
 ):
  def echo(msg, stream="stderr", job_id=None, **kwargs):
@@ -248,6 +260,8 @@ def _sync_metadata():
  persistent_volume_claims=persistent_volume_claims,
  tolerations=tolerations,
  shared_memory=shared_memory,
+ kueue_enabled=kueue_enabled,
+ kueue_localqueue_name=kueue_localqueue_name,
  )
  except Exception as e:
  traceback.print_exc(chain=False)

metaflow/plugins/kubernetes/kubernetes_decorator.py

@@ -21,6 +21,8 @@
  KUBERNETES_TOLERATIONS,
  KUBERNETES_SERVICE_ACCOUNT,
  KUBERNETES_SHARED_MEMORY,
+ KUEUE_ENABLED,
+ KUEUE_LOCALQUEUE_NAME,
 )
 from metaflow.plugins.resources_decorator import ResourcesDecorator
 from metaflow.plugins.timeout_decorator import get_run_time_limit_for_task
@@ -90,6 +92,10 @@ class KubernetesDecorator(StepDecorator):
  volumes to the path to which the volume is to be mounted, e.g., `{'pvc-name': '/path/to/mount/on'}`.
  shared_memory: int, optional
  Shared memory size (in MiB) required for this step
+ kueue_enabled: bool, optional
+ Whether Kubernetes job/Argo workflow pod should submitted using Kueue
+ kueue_localqueue_name: str, optional
+ The name of the localqueue object configured in Kueue to use for submitting jobs/pods
  """
 
  name = "kubernetes"
@@ -113,6 +119,8 @@ class KubernetesDecorator(StepDecorator):
  "tmpfs_path": "/metaflow_temp",
  "persistent_volume_claims": None, # e.g., {"pvc-name": "/mnt/vol", "another-pvc": "/mnt/vol2"}
  "shared_memory": None,
+ "kueue_enabled": None,
+ "kueue_localqueue_name": None,
  }
  package_url = None
  package_sha = None
@@ -201,6 +209,15 @@ def __init__(self, attributes=None, statically_defined=False):
  if not self.attributes["shared_memory"]:
  self.attributes["shared_memory"] = KUBERNETES_SHARED_MEMORY
 
+ # Process config options related to KUEUE
+ if self.attributes["kueue_enabled"] is None:
+ self.attributes["kueue_enabled"] = KUEUE_ENABLED
+ if (
+ "kueue_localqueue_name" not in self.attributes
+ or self.attributes["kueue_localqueue_name"] is None
+ ):
+ self.attributes["kueue_localqueue_name"] = KUEUE_LOCALQUEUE_NAME
+
  # Refer https://github.com/Netflix/metaflow/blob/master/docs/lifecycle.png
  def step_init(self, flow, graph, step, decos, environment, flow_datastore, logger):
  # Executing Kubernetes jobs requires a non-local datastore.

metaflow/plugins/kubernetes/kubernetes_job.py

@@ -83,14 +83,31 @@ def create(self):
  else None
  )
 
+ annotations = self._kwargs.get("annotations", {})
+ labels = self._kwargs.get("labels", {})
+
+ kueue_enabled = bool(self._kwargs["kueue_enabled"])
+ localqueue_name = self._kwargs["kueue_localqueue_name"]
+ if kueue_enabled:
+ labels["kueue.x-k8s.io/queue-name"] = localqueue_name
+ labels["kueue.x-k8s.io/pod-group-name"] = (
+ self._kwargs["run_id"]
+ + "-"
+ + self._kwargs["step_name"]
+ + "-"
+ + self._kwargs["task_id"]
+ )
+ annotations["kueue.x-k8s.io/retriable-in-group"] = "false"
+ annotations["kueue.x-k8s.io/pod-group-total-count"] = str(1)
+
  self._job = client.V1Job(
  api_version="batch/v1",
  kind="Job",
  metadata=client.V1ObjectMeta(
  # Annotations are for humans
- annotations=self._kwargs.get("annotations", {}),
+ annotations=annotations,
  # While labels are for Kubernetes
- labels=self._kwargs.get("labels", {}),
+ labels=labels,
  generate_name=self._kwargs["generate_name"],
  namespace=self._kwargs["namespace"], # Defaults to `default`
  ),