Add support for shared_memory for Kubernetes workloads. (#1709)

* Add support for shared_memory for Kubernetes workloads.

This includes support for Kubernetes jobs as well as Argo workflows.

Author: shrinandj

metaflow/metaflow_config.py

@@ -299,6 +299,8 @@
 )
 # Toggle for trying to fetch EC2 instance metadata
 KUBERNETES_FETCH_EC2_METADATA = from_conf("KUBERNETES_FETCH_EC2_METADATA", False)
+# Shared memory in MB to use for this step
+KUBERNETES_SHARED_MEMORY = from_conf("KUBERNETES_SHARED_MEMORY", None)
 
 ARGO_WORKFLOWS_KUBERNETES_SECRETS = from_conf("ARGO_WORKFLOWS_KUBERNETES_SECRETS", "")
 ARGO_WORKFLOWS_ENV_VARS_TO_SKIP = from_conf("ARGO_WORKFLOWS_ENV_VARS_TO_SKIP", "")

metaflow/plugins/argo/argo_workflows.py

@@ -1368,6 +1368,9 @@ 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)
 
@@ -1412,6 +1415,7 @@ def _container_templates(self):
                     medium="Memory",
                     size_limit=tmpfs_size if tmpfs_enabled else 0,
                 )
+                .empty_dir_volume("dhsm", medium="Memory", size_limit=shared_memory)
                 .pvc_volumes(resources.get("persistent_volume_claims"))
                 # Set node selectors
                 .node_selectors(resources.get("node_selector"))
@@ -1505,6 +1509,17 @@ def _container_templates(self):
                                 if tmpfs_enabled
                                 else []
                             )
+                            # Support shared_memory
+                            + (
+                                [
+                                    kubernetes_sdk.V1VolumeMount(
+                                        name="dhsm",
+                                        mount_path="/dev/shm",
+                                    )
+                                ]
+                                if shared_memory
+                                else []
+                            )
                             # Support persistent volume claims.
                             + (
                                 [

metaflow/plugins/kubernetes/kubernetes.py

@@ -174,6 +174,7 @@ def create_job(
         persistent_volume_claims=None,
         tolerations=None,
         labels=None,
+        shared_memory=None,
     ):
         if env is None:
             env = {}
@@ -213,6 +214,7 @@ def create_job(
                 tmpfs_size=tmpfs_size,
                 tmpfs_path=tmpfs_path,
                 persistent_volume_claims=persistent_volume_claims,
+                shared_memory=shared_memory,
             )
             .environment_variable("METAFLOW_CODE_SHA", code_package_sha)
             .environment_variable("METAFLOW_CODE_URL", code_package_url)

metaflow/plugins/kubernetes/kubernetes_cli.py

@@ -107,6 +107,7 @@ def kubernetes():
     type=JSONTypeClass(),
     multiple=False,
 )
+@click.option("--shared-memory", default=None, help="Size of shared memory in MiB")
 @click.pass_context
 def step(
     ctx,
@@ -132,6 +133,7 @@ def step(
     run_time_limit=None,
     persistent_volume_claims=None,
     tolerations=None,
+    shared_memory=None,
     **kwargs
 ):
     def echo(msg, stream="stderr", job_id=None, **kwargs):
@@ -245,6 +247,7 @@ def _sync_metadata():
                 env=env,
                 persistent_volume_claims=persistent_volume_claims,
                 tolerations=tolerations,
+                shared_memory=shared_memory,
             )
     except Exception as e:
         traceback.print_exc(chain=False)

metaflow/plugins/kubernetes/kubernetes_decorator.py

@@ -20,6 +20,7 @@
     KUBERNETES_PERSISTENT_VOLUME_CLAIMS,
     KUBERNETES_TOLERATIONS,
     KUBERNETES_SERVICE_ACCOUNT,
+    KUBERNETES_SHARED_MEMORY,
 )
 from metaflow.plugins.resources_decorator import ResourcesDecorator
 from metaflow.plugins.timeout_decorator import get_run_time_limit_for_task
@@ -87,6 +88,8 @@ class KubernetesDecorator(StepDecorator):
     persistent_volume_claims : Dict[str, str], optional, default None
         A map (dictionary) of persistent volumes to be mounted to the pod for this step. The map is from persistent
         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
     """
 
     name = "kubernetes"
@@ -109,6 +112,7 @@ class KubernetesDecorator(StepDecorator):
         "tmpfs_size": None,
         "tmpfs_path": "/metaflow_temp",
         "persistent_volume_claims": None,  # e.g., {"pvc-name": "/mnt/vol", "another-pvc": "/mnt/vol2"}
+        "shared_memory": None,
     }
     package_url = None
     package_sha = None
@@ -194,6 +198,8 @@ def __init__(self, attributes=None, statically_defined=False):
             if not self.attributes["tmpfs_size"]:
                 # default tmpfs behavior - https://man7.org/linux/man-pages/man5/tmpfs.5.html
                 self.attributes["tmpfs_size"] = int(self.attributes["memory"]) // 2
+        if not self.attributes["shared_memory"]:
+            self.attributes["shared_memory"] = KUBERNETES_SHARED_MEMORY
 
     # Refer https://github.com/Netflix/metaflow/blob/master/docs/lifecycle.png
     def step_init(self, flow, graph, step, decos, environment, flow_datastore, logger):
@@ -289,6 +295,17 @@ def step_init(self, flow, graph, step, decos, environment, flow_datastore, logge
                     )
                 )
 
+        if self.attributes["shared_memory"]:
+            if not (
+                isinstance(self.attributes["shared_memory"], int)
+                and int(self.attributes["shared_memory"]) > 0
+            ):
+                raise KubernetesException(
+                    "Invalid shared_memory value: *{size}* for step *{step}* (should be an integer greater than 0)".format(
+                        size=self.attributes["shared_memory"], step=step
+                    )
+                )
+
     def package_init(self, flow, step_name, environment):
         try:
             # Kubernetes is a soft dependency.

metaflow/plugins/kubernetes/kubernetes_job.py

@@ -77,6 +77,11 @@ def create(self):
         use_tmpfs = self._kwargs["use_tmpfs"]
         tmpfs_size = self._kwargs["tmpfs_size"]
         tmpfs_enabled = use_tmpfs or (tmpfs_size and not use_tmpfs)
+        shared_memory = (
+            int(self._kwargs["shared_memory"])
+            if self._kwargs["shared_memory"]
+            else None
+        )
 
         self._job = client.V1Job(
             api_version="batch/v1",
@@ -184,6 +189,15 @@ def create(self):
                                     if tmpfs_enabled
                                     else []
                                 )
+                                + (
+                                    [
+                                        client.V1VolumeMount(
+                                            mount_path="/dev/shm", name="dhsm"
+                                        )
+                                    ]
+                                    if shared_memory
+                                    else []
+                                )
                                 + (
                                     [
                                         client.V1VolumeMount(
@@ -232,6 +246,19 @@ def create(self):
                             if tmpfs_enabled
                             else []
                         )
+                        + (
+                            [
+                                client.V1Volume(
+                                    name="dhsm",
+                                    empty_dir=client.V1EmptyDirVolumeSource(
+                                        medium="Memory",
+                                        size_limit="{}Mi".format(shared_memory),
+                                    ),
+                                )
+                            ]
+                            if shared_memory
+                            else []
+                        )
                         + (
                             [
                                 client.V1Volume(

stubs/test/test_stubs.yml

@@ -501,18 +501,14 @@
         ...
   out: |
     main:3: error: .*incompatible type.*\[arg-type\]
+    main:4: error: .*Too many positional arguments for "kubernetes"\s+\[misc\]
     main:7: error: .*incompatible type.*\[arg-type\]
-    main:12: error: No overload variant of "kubernetes" matches argument.*\[call-overload\]
-    main:12: note: Possible overload variants
-    main:12: note: .*def.*kubernetes.*
-    main:12: note: .*def.*kubernetes.*
-    main:12: note: .*def.*kubernetes.*
-    main:17: error: No overload variant of "kubernetes" matches argument.*\[call-overload\]
-    main:17: note: Possible overload variants
-    main:17: note: .*def.*kubernetes.*
-    main:17: note: .*def.*kubernetes.*
-    main:17: note: .*def.*kubernetes.*
+    main:12: error: .*Too many positional arguments for "kubernetes"\s+\[misc\]
+    main:17: error: .*Unexpected keyword argument "foo" for "kubernetes"\s+\[call-arg\]
+    main:22: error: .*Too many positional arguments for "kubernetes"\s+\[misc\]
     main:22: error: .*incompatible type.*\[arg-type\]
+    main:26: error: .*Too many positional arguments for "kubernetes"\s+\[misc\]
+    main:26: error: .*incompatible type.*\[arg-type\]
 
 - case: environment_decorator_validity
   regex: yes