From 92113b81c4a7717913e7773383d4d325c40d355e Mon Sep 17 00:00:00 2001
From: Juntao Wang <juntwang@redhat.com>
Date: Wed, 6 Sep 2023 11:51:25 -0400
Subject: [PATCH] Add shm to ServingRuntime volumes and volumeMounts

---
 backend/src/types.ts                    |  2 ++
 frontend/src/api/k8s/notebooks.ts       | 12 +---------
 frontend/src/api/k8s/servingRuntimes.ts | 29 +++++++++++++++++++------
 frontend/src/api/k8s/utils.ts           | 12 ++++++++++
 frontend/src/k8sTypes.ts                |  3 +++
 5 files changed, 40 insertions(+), 18 deletions(-)

diff --git a/backend/src/types.ts b/backend/src/types.ts
index 73fef553b3..8dc3a0cac8 100644
--- a/backend/src/types.ts
+++ b/backend/src/types.ts
@@ -901,8 +901,10 @@ export type ServingRuntime = K8sResourceCommon & {
       image: string;
       name: string;
       resources: ContainerResources;
+      volumeMounts?: VolumeMount[];
     }[];
     supportedModelFormats: SupportedModelFormats[];
     replicas: number;
+    volumes?: Volume[];
   };
 };
diff --git a/frontend/src/api/k8s/notebooks.ts b/frontend/src/api/k8s/notebooks.ts
index 8c1ddb8d80..d7a86ba07e 100644
--- a/frontend/src/api/k8s/notebooks.ts
+++ b/frontend/src/api/k8s/notebooks.ts
@@ -26,17 +26,7 @@ import {
 } from '~/concepts/pipelines/elyra/utils';
 import { createRoleBinding } from '~/api';
 import { Volume, VolumeMount } from '~/types';
-import { assemblePodSpecOptions } from './utils';
-
-const getshmVolumeMount = (): VolumeMount => ({
-  name: 'shm',
-  mountPath: '/dev/shm',
-});
-
-const getshmVolume = (): Volume => ({
-  name: 'shm',
-  emptyDir: { medium: 'Memory' },
-});
+import { assemblePodSpecOptions, getshmVolume, getshmVolumeMount } from './utils';
 
 const assembleNotebook = (
   data: StartNotebookData,
diff --git a/frontend/src/api/k8s/servingRuntimes.ts b/frontend/src/api/k8s/servingRuntimes.ts
index 3842828d48..a8b4699cac 100644
--- a/frontend/src/api/k8s/servingRuntimes.ts
+++ b/frontend/src/api/k8s/servingRuntimes.ts
@@ -14,7 +14,7 @@ import { getModelServingRuntimeName } from '~/pages/modelServing/utils';
 import { getDisplayNameFromK8sResource, translateDisplayNameForK8s } from '~/pages/projects/utils';
 import { applyK8sAPIOptions } from '~/api/apiMergeUtils';
 import { getModelServingProjects } from './projects';
-import { assemblePodSpecOptions } from './utils';
+import { assemblePodSpecOptions, getshmVolume, getshmVolumeMount } from './utils';
 
 const assembleServingRuntime = (
   data: CreatingServingRuntimeObject,
@@ -79,12 +79,27 @@ const assembleServingRuntime = (
 
   const { affinity, tolerations, resources } = assemblePodSpecOptions(resourceSettings, gpus);
 
-  updatedServingRuntime.spec.containers = servingRuntime.spec.containers.map((container) => ({
-    ...container,
-    resources,
-    affinity,
-    tolerations,
-  }));
+  const volumes = updatedServingRuntime.spec.volumes || [];
+  if (!volumes.find((volume) => volume.name === 'shm')) {
+    volumes.push(getshmVolume('2Gi'));
+  }
+
+  updatedServingRuntime.spec.volumes = volumes;
+
+  updatedServingRuntime.spec.containers = servingRuntime.spec.containers.map((container) => {
+    const volumeMounts = container.volumeMounts || [];
+    if (!volumeMounts.find((volumeMount) => volumeMount.mountPath === '/dev/shm')) {
+      volumeMounts.push(getshmVolumeMount());
+    }
+
+    return {
+      ...container,
+      resources,
+      affinity,
+      tolerations,
+      volumeMounts,
+    };
+  });
 
   return updatedServingRuntime;
 };
diff --git a/frontend/src/api/k8s/utils.ts b/frontend/src/api/k8s/utils.ts
index 920415d757..ce2867007c 100644
--- a/frontend/src/api/k8s/utils.ts
+++ b/frontend/src/api/k8s/utils.ts
@@ -4,6 +4,8 @@ import {
   PodToleration,
   TolerationSettings,
   ContainerResourceAttributes,
+  VolumeMount,
+  Volume,
 } from '~/types';
 import { determineTolerations } from '~/utilities/tolerations';
 
@@ -54,3 +56,13 @@ export const assemblePodSpecOptions = (
   const tolerations = determineTolerations(gpus > 0, tolerationSettings);
   return { affinity, tolerations, resources };
 };
+
+export const getshmVolumeMount = (): VolumeMount => ({
+  name: 'shm',
+  mountPath: '/dev/shm',
+});
+
+export const getshmVolume = (sizeLimit?: string): Volume => ({
+  name: 'shm',
+  emptyDir: { medium: 'Memory', ...(sizeLimit && { sizeLimit }) },
+});
diff --git a/frontend/src/k8sTypes.ts b/frontend/src/k8sTypes.ts
index 3b4f16524e..2795ef2eed 100644
--- a/frontend/src/k8sTypes.ts
+++ b/frontend/src/k8sTypes.ts
@@ -12,6 +12,7 @@ import {
   TolerationSettings,
   ImageStreamStatusTagItem,
   ImageStreamStatusTagCondition,
+  VolumeMount,
 } from './types';
 import { ServingRuntimeSize } from './pages/modelServing/screens/types';
 
@@ -327,9 +328,11 @@ export type ServingRuntimeKind = K8sResourceCommon & {
       image: string;
       name: string;
       resources: ContainerResources;
+      volumeMounts?: VolumeMount[];
     }[];
     supportedModelFormats: SupportedModelFormats[];
     replicas: number;
+    volumes?: Volume[];
   };
 };