From bcaac30b14b59caeb0ea82dd70780f36a8bc4957 Mon Sep 17 00:00:00 2001
From: Dmitry Meyer <me@undef.im>
Date: Mon, 29 Sep 2025 12:35:09 +0000
Subject: [PATCH] Kubernetes: configure /dev/shm if requested

Part-of: https://github.com/dstackai/dstack/issues/3126
---
 .../core/backends/kubernetes/compute.py       | 39 ++++++++++++++++---
 1 file changed, 34 insertions(+), 5 deletions(-)

diff --git a/src/dstack/_internal/core/backends/kubernetes/compute.py b/src/dstack/_internal/core/backends/kubernetes/compute.py
index f2bc714232..96dbbc50b4 100644
--- a/src/dstack/_internal/core/backends/kubernetes/compute.py
+++ b/src/dstack/_internal/core/backends/kubernetes/compute.py
@@ -47,7 +47,7 @@
     Resources,
     SSHConnectionParams,
 )
-from dstack._internal.core.models.resources import CPUSpec
+from dstack._internal.core.models.resources import CPUSpec, Memory
 from dstack._internal.core.models.runs import Job, JobProvisioningData, Requirements, Run
 from dstack._internal.core.models.volumes import Volume
 from dstack._internal.utils.common import parse_memory
@@ -171,11 +171,15 @@ def run_job(
                     "jump_pod_port": jump_pod_port,
                 },
             ).start()
-        resources_spec = job.job_spec.requirements.resources
-        assert isinstance(resources_spec.cpu, CPUSpec)
+
         resources_requests: dict[str, str] = {}
         resources_limits: dict[str, str] = {}
         node_affinity: Optional[client.V1NodeAffinity] = None
+        volumes_: list[client.V1Volume] = []
+        volume_mounts: list[client.V1VolumeMount] = []
+
+        resources_spec = job.job_spec.requirements.resources
+        assert isinstance(resources_spec.cpu, CPUSpec)
         if (cpu_min := resources_spec.cpu.count.min) is not None:
             resources_requests["cpu"] = str(cpu_min)
         if (gpu_spec := resources_spec.gpu) is not None:
@@ -231,13 +235,32 @@ def run_job(
                         ),
                     ],
                 )
+
         if (memory_min := resources_spec.memory.min) is not None:
-            resources_requests["memory"] = f"{float(memory_min)}Gi"
+            resources_requests["memory"] = _render_memory(memory_min)
         if (
             resources_spec.disk is not None
             and (disk_min := resources_spec.disk.size.min) is not None
         ):
-            resources_requests["ephemeral-storage"] = f"{float(disk_min)}Gi"
+            resources_requests["ephemeral-storage"] = _render_memory(disk_min)
+        if (shm_size := resources_spec.shm_size) is not None:
+            shm_volume_name = "dev-shm"
+            volumes_.append(
+                client.V1Volume(
+                    name=shm_volume_name,
+                    empty_dir=client.V1EmptyDirVolumeSource(
+                        medium="Memory",
+                        size_limit=_render_memory(shm_size),
+                    ),
+                )
+            )
+            volume_mounts.append(
+                client.V1VolumeMount(
+                    name=shm_volume_name,
+                    mount_path="/dev/shm",
+                )
+            )
+
         pod = client.V1Pod(
             metadata=client.V1ObjectMeta(
                 name=instance_name,
@@ -264,9 +287,11 @@ def run_job(
                             requests=resources_requests,
                             limits=resources_limits,
                         ),
+                        volume_mounts=volume_mounts,
                     )
                 ],
                 affinity=node_affinity,
+                volumes=volumes_,
             ),
         )
         call_api_method(
@@ -452,6 +477,10 @@ def _parse_memory(memory: str) -> int:
     return int(parse_memory(memory, as_untis="M"))
 
 
+def _render_memory(memory: Memory) -> str:
+    return f"{float(memory)}Gi"
+
+
 def _get_gpus_from_node_labels(labels: dict[str, str]) -> tuple[list[Gpu], Optional[str]]:
     # We rely on https://github.com/NVIDIA/k8s-device-plugin/tree/main/docs/gpu-feature-discovery
     # to detect gpus. Note that "nvidia.com/gpu.product" is not a short gpu name like "T4" or