Add kuberenetes resource limits & security context configs

metaflow/metaflow_config.py

@@ -277,6 +277,9 @@
 KUBERNETES_PERSISTENT_VOLUME_CLAIMS = from_conf(
  "KUBERNETES_PERSISTENT_VOLUME_CLAIMS", ""
 )
+KUBERNETES_SECURITY_CONTEXT = from_conf("KUBERNETES_SECURITY_CONTEXT", {})
+KUBERNETES_RESOURCE_LIMITS_MEMORY = from_conf("KUBERNETES_RESOURCE_LIMITS_MEMORY", None)
+KUBERNETES_RESOURCE_LIMITS_CPU = from_conf("KUBERNETES_RESOURCE_LIMITS_CPU", None)
 KUBERNETES_SECRETS = from_conf("KUBERNETES_SECRETS", "")
 # Default labels for kubernetes pods
 KUBERNETES_LABELS = from_conf("KUBERNETES_LABELS", "")

metaflow/plugins/argo/argo_workflows.py

@@ -60,6 +60,7 @@
 )
 
 from .argo_client import ArgoClient
+from ..kubernetes.kuberenetes_utils import make_kubernetes_container
 
 
 class ArgoWorkflowsException(MetaflowException):
@@ -1428,96 +1429,21 @@ def _container_templates(self):
  # resources attributes in particular where the keys maybe user
  # defined.
  to_camelcase(
- kubernetes_sdk.V1Container(
- name=self._sanitize(node.name),
- command=cmds,
- env=[
- kubernetes_sdk.V1EnvVar(name=k, value=str(v))
- for k, v in env.items()
- ]
- # Add environment variables for book-keeping.
- # https://argoproj.github.io/argo-workflows/fields/#fields_155
- + [
- kubernetes_sdk.V1EnvVar(
- name=k,
- value_from=kubernetes_sdk.V1EnvVarSource(
- field_ref=kubernetes_sdk.V1ObjectFieldSelector(
- field_path=str(v)
- )
- ),
- )
- for k, v in {
- "METAFLOW_KUBERNETES_POD_NAMESPACE": "metadata.namespace",
- "METAFLOW_KUBERNETES_POD_NAME": "metadata.name",
- "METAFLOW_KUBERNETES_POD_ID": "metadata.uid",
- "METAFLOW_KUBERNETES_SERVICE_ACCOUNT_NAME": "spec.serviceAccountName",
- "METAFLOW_KUBERNETES_NODE_IP": "status.hostIP",
- }.items()
- ],
- image=resources["image"],
- image_pull_policy=resources["image_pull_policy"],
- resources=kubernetes_sdk.V1ResourceRequirements(
- requests={
- "cpu": str(resources["cpu"]),
- "memory": "%sM" % str(resources["memory"]),
- "ephemeral-storage": "%sM" % str(resources["disk"]),
- },
- limits={
- "%s.com/gpu".lower()
- % resources["gpu_vendor"]: str(resources["gpu"])
- for k in [0]
- if resources["gpu"] is not None
- },
- ),
- # Configure secrets
- env_from=[
- kubernetes_sdk.V1EnvFromSource(
- secret_ref=kubernetes_sdk.V1SecretEnvSource(
- name=str(k),
- # optional=True
- )
- )
- for k in list(
- []
- if not resources.get("secrets")
- else [resources.get("secrets")]
- if isinstance(resources.get("secrets"), str)
- else resources.get("secrets")
- )
- + KUBERNETES_SECRETS.split(",")
- + ARGO_WORKFLOWS_KUBERNETES_SECRETS.split(",")
- if k
- ],
- volume_mounts=[
+ make_kubernetes_container(
+ kubernetes_sdk,
+ self._sanitize(node.name),
+ cmds,
+ {
+ **resources,
  # Assign a volume mount to pass state to the next task.
- kubernetes_sdk.V1VolumeMount(
- name="out", mount_path="/mnt/out"
- )
- ]
- # Support tmpfs.
- + (
- [
- kubernetes_sdk.V1VolumeMount(
- name="tmpfs-ephemeral-volume",
- mount_path=tmpfs_path,
- )
- ]
- if tmpfs_enabled
- else []
- )
- # Support persistent volume claims.
- + (
- [
- kubernetes_sdk.V1VolumeMount(
- name=claim, mount_path=path
- )
- for claim, path in resources.get(
- "persistent_volume_claims"
- ).items()
- ]
- if resources.get("persistent_volume_claims") is not None
- else []
- ),
+ "persistent_volume_claims": {
+ **(resources.get("persistent_volume_claims") or {}),
+ "out": "/mnt/out",
+ },
+ },
+ env,
+ additional_secrets=KUBERNETES_SECRETS.split(",")
+ + ARGO_WORKFLOWS_KUBERNETES_SECRETS.split(","),
  ).to_dict()
  )
  )

metaflow/plugins/kubernetes/kuberenetes_utils.py

@@ -0,0 +1,120 @@
+import json
+from metaflow.tracing import inject_tracing_vars
+
+
+def compute_resource_limits(args):
+ limits_dict = dict()
+ if args.get("resource_limits_memory", None):
+ limits_dict["memory"] = "%sM" % str(args["resource_limits_memory"])
+ if args.get("resource_limits_cpu", None):
+ limits_dict["cpu"] = args["resource_limits_cpu"]
+ if args["gpu"] is not None:
+ limits_dict["%s.com/gpu".lower() % args["gpu_vendor"]] = str(args["gpu"])
+ return limits_dict
+
+
+def get_list_from_untyped(possible_list):
+ return (
+ []
+ if not possible_list
+ else [possible_list]
+ if isinstance(possible_list, str)
+ else possible_list
+ )
+
+
+def compute_tempfs_enabled(args):
+ use_tmpfs = args["use_tmpfs"]
+ tmpfs_size = args["tmpfs_size"]
+ return use_tmpfs or (tmpfs_size and not use_tmpfs)
+
+
+def make_kubernetes_container(
+ client, name, commands, args, envs, additional_secrets=[]
+):
+ from kubernetes.client import V1SecurityContext, ApiClient
+
+ class KubernetesClientDataObj(object):
+ def __init__(self, data_dict, class_name):
+ self.data = json.dumps(data_dict) if data_dict is not None else None
+ self.class_name = class_name
+
+ def get_deserialized_object(self):
+ if self.data is not None:
+ return ApiClient().deserialize(self, self.class_name)
+ else:
+ return None
+
+ security_context = KubernetesClientDataObj(
+ args["security_context"], V1SecurityContext
+ ).get_deserialized_object()
+
+ # tmpfs variables
+ tmpfs_enabled = compute_tempfs_enabled(args)
+
+ return client.V1Container(
+ name=name,
+ command=commands,
+ env=[client.V1EnvVar(name=k, value=str(v)) for k, v in envs.items()]
+ # And some downward API magic. Add (key, value)
+ # pairs below to make pod metadata available
+ # within Kubernetes container.
+ + [
+ client.V1EnvVar(
+ name=k,
+ value_from=client.V1EnvVarSource(
+ field_ref=client.V1ObjectFieldSelector(field_path=str(v))
+ ),
+ )
+ for k, v in {
+ "METAFLOW_KUBERNETES_POD_NAMESPACE": "metadata.namespace",
+ "METAFLOW_KUBERNETES_POD_NAME": "metadata.name",
+ "METAFLOW_KUBERNETES_POD_ID": "metadata.uid",
+ "METAFLOW_KUBERNETES_SERVICE_ACCOUNT_NAME": "spec.serviceAccountName",
+ "METAFLOW_KUBERNETES_NODE_IP": "status.hostIP",
+ }.items()
+ ]
+ + [
+ client.V1EnvVar(name=k, value=str(v))
+ for k, v in inject_tracing_vars({}).items()
+ ],
+ env_from=[
+ client.V1EnvFromSource(
+ secret_ref=client.V1SecretEnvSource(
+ name=str(k),
+ # optional=True
+ )
+ )
+ for k in get_list_from_untyped(args.get("secrets")) + additional_secrets
+ if k
+ ],
+ image=args["image"],
+ security_context=security_context,
+ image_pull_policy=args["image_pull_policy"],
+ resources=client.V1ResourceRequirements(
+ requests={
+ "cpu": str(args["cpu"]),
+ "memory": "%sM" % str(args["memory"]),
+ "ephemeral-storage": "%sM" % str(args["disk"]),
+ },
+ limits=compute_resource_limits(args),
+ ),
+ volume_mounts=(
+ [
+ client.V1VolumeMount(
+ mount_path=args.get("tmpfs_path"),
+ name="tmpfs-ephemeral-volume",
+ )
+ ]
+ if tmpfs_enabled
+ else []
+ )
+ + (
+ [
+ client.V1VolumeMount(mount_path=path, name=claim)
+ for claim, path in args["persistent_volume_claims"].items()
+ ]
+ if args["persistent_volume_claims"] is not None
+ else []
+ ),
+ )

metaflow/plugins/kubernetes/kubernetes.py

@@ -174,6 +174,9 @@ def create_job(
  persistent_volume_claims=None,
  tolerations=None,
  labels=None,
+ security_context=None,
+ resource_limits_memory=None,
+ resource_limits_cpu=None,
  ):
  if env is None:
  env = {}
@@ -213,6 +216,9 @@ def create_job(
  tmpfs_size=tmpfs_size,
  tmpfs_path=tmpfs_path,
  persistent_volume_claims=persistent_volume_claims,
+ security_context=security_context,
+ resource_limits_memory=resource_limits_memory,
+ resource_limits_cpu=resource_limits_cpu,
  )
  .environment_variable("METAFLOW_CODE_SHA", code_package_sha)
  .environment_variable("METAFLOW_CODE_URL", code_package_url)

metaflow/plugins/kubernetes/kubernetes_cli.py

@@ -73,6 +73,15 @@ def kubernetes():
 @click.option("--memory", help="Memory requirement for Kubernetes pod.")
 @click.option("--gpu", help="GPU requirement for Kubernetes pod.")
 @click.option("--gpu-vendor", help="GPU vendor requirement for Kubernetes pod.")
+@click.option("--resource-limits-memory", help="Memory limits for Kubernetes pod.")
+@click.option("--resource-limits-cpu", help="CPU limits for Kubernetes pod.")
+@click.option(
+ "--security-context",
+ default=None,
+ type=JSONTypeClass(),
+ multiple=False,
+ help="Security context Kubernetes pod.",
+)
 @click.option("--run-id", help="Passed to the top-level 'step'.")
 @click.option("--task-id", help="Passed to the top-level 'step'.")
 @click.option("--input-paths", help="Passed to the top-level 'step'.")
@@ -132,6 +141,9 @@ def step(
  run_time_limit=None,
  persistent_volume_claims=None,
  tolerations=None,
+ security_context=None,
+ resource_limits_memory=None,
+ resource_limits_cpu=None,
  **kwargs
 ):
  def echo(msg, stream="stderr", job_id=None, **kwargs):
@@ -245,6 +257,9 @@ def _sync_metadata():
  env=env,
  persistent_volume_claims=persistent_volume_claims,
  tolerations=tolerations,
+ security_context=security_context,
+ resource_limits_memory=resource_limits_memory,
+ resource_limits_cpu=resource_limits_cpu,
  )
  except Exception as e:
  traceback.print_exc(chain=False)

metaflow/plugins/kubernetes/kubernetes_decorator.py

@@ -20,6 +20,9 @@
  KUBERNETES_PERSISTENT_VOLUME_CLAIMS,
  KUBERNETES_TOLERATIONS,
  KUBERNETES_SERVICE_ACCOUNT,
+ KUBERNETES_SECURITY_CONTEXT,
+ KUBERNETES_RESOURCE_LIMITS_MEMORY,
+ KUBERNETES_RESOURCE_LIMITS_CPU,
 )
 from metaflow.plugins.resources_decorator import ResourcesDecorator
 from metaflow.plugins.timeout_decorator import get_run_time_limit_for_task
@@ -108,6 +111,9 @@ class KubernetesDecorator(StepDecorator):
  "tmpfs_size": None,
  "tmpfs_path": "/metaflow_temp",
  "persistent_volume_claims": None, # e.g., {"pvc-name": "/mnt/vol", "another-pvc": "/mnt/vol2"}
+ "security_context": None,
+ "resource_limits_memory": None,
+ "resource_limits_cpu": None,
  }
  package_url = None
  package_sha = None
@@ -122,6 +128,20 @@ def __init__(self, attributes=None, statically_defined=False):
  self.attributes["service_account"] = KUBERNETES_SERVICE_ACCOUNT
  if not self.attributes["gpu_vendor"]:
  self.attributes["gpu_vendor"] = KUBERNETES_GPU_VENDOR
+ if not self.attributes["security_context"] and KUBERNETES_SECURITY_CONTEXT:
+ self.attributes["security_context"] = KUBERNETES_SECURITY_CONTEXT
+ if (
+ not self.attributes["resource_limits_memory"]
+ and KUBERNETES_RESOURCE_LIMITS_MEMORY
+ ):
+ self.attributes[
+ "resource_limits_memory"
+ ] = KUBERNETES_RESOURCE_LIMITS_MEMORY
+ if (
+ not self.attributes["resource_limits_cpu"]
+ and KUBERNETES_RESOURCE_LIMITS_CPU
+ ):
+ self.attributes["resource_limits_cpu"] = KUBERNETES_RESOURCE_LIMITS_CPU
  if not self.attributes["node_selector"] and KUBERNETES_NODE_SELECTOR:
  self.attributes["node_selector"] = KUBERNETES_NODE_SELECTOR
  if not self.attributes["tolerations"] and KUBERNETES_TOLERATIONS:
@@ -339,7 +359,11 @@ def runtime_step_cli(
  "=".join([key, str(val)]) if val else key
  for key, val in v.items()
  ]
- elif k in ["tolerations", "persistent_volume_claims"]:
+ elif k in [
+ "tolerations",
+ "persistent_volume_claims",
+ "security_context",
+ ]:
  cli_args.command_options[k] = json.dumps(v)
  else:
  cli_args.command_options[k] = v