From 4d9f831cc98f215814b779bd4084b2bfc5c3bbfc Mon Sep 17 00:00:00 2001
From: chenzhihang <709011834@qq.com>
Date: Thu, 29 Aug 2024 17:37:58 +0800
Subject: [PATCH] =?UTF-8?q?1=E3=80=81=E9=85=8D=E7=BD=AEpod=E8=AF=B7?=
 =?UTF-8?q?=E6=B1=82=E9=99=90=E5=88=B6=E8=B5=84=E6=BA=90?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../service/impl/JupyterServiceImpl.java      |  2 +-
 .../ruoyi/platform/utils/K8sClientUtil.java   | 27 ++++++++++++++++---
 2 files changed, 25 insertions(+), 4 deletions(-)

diff --git a/ruoyi-modules/management-platform/src/main/java/com/ruoyi/platform/service/impl/JupyterServiceImpl.java b/ruoyi-modules/management-platform/src/main/java/com/ruoyi/platform/service/impl/JupyterServiceImpl.java
index b99bdb29..e6380f3a 100644
--- a/ruoyi-modules/management-platform/src/main/java/com/ruoyi/platform/service/impl/JupyterServiceImpl.java
+++ b/ruoyi-modules/management-platform/src/main/java/com/ruoyi/platform/service/impl/JupyterServiceImpl.java
@@ -105,7 +105,7 @@ public class JupyterServiceImpl implements JupyterService {
         //TODO 设置镜像可配置，这里先用默认镜像启动pod
 
         // 调用修改后的 createPod 方法，传入额外的参数
-        Integer podPort = k8sClientUtil.createConfiguredPod(podName, namespace, port, mountPath, pvc, devEnvironment.getImage(), devEnvironment.getComputingResource(), minioPvcName, datasetPath, modelPath);
+        Integer podPort = k8sClientUtil.createConfiguredPod(podName, namespace, port, mountPath, pvc, devEnvironment, minioPvcName, datasetPath, modelPath);
         String url = masterIp + ":" + podPort;
         redisService.setCacheObject(podName, masterIp + ":" + podPort);
         devEnvironment.setStatus("Pending");
diff --git a/ruoyi-modules/management-platform/src/main/java/com/ruoyi/platform/utils/K8sClientUtil.java b/ruoyi-modules/management-platform/src/main/java/com/ruoyi/platform/utils/K8sClientUtil.java
index d5b5ab21..06e59083 100644
--- a/ruoyi-modules/management-platform/src/main/java/com/ruoyi/platform/utils/K8sClientUtil.java
+++ b/ruoyi-modules/management-platform/src/main/java/com/ruoyi/platform/utils/K8sClientUtil.java
@@ -1,6 +1,7 @@
 package com.ruoyi.platform.utils;
 
 import com.ruoyi.platform.constant.Constant;
+import com.ruoyi.platform.domain.DevEnvironment;
 import com.ruoyi.platform.mapper.ComputingResourceDao;
 import io.kubernetes.client.Exec;
 import io.kubernetes.client.custom.IntOrString;
@@ -13,6 +14,7 @@ import io.kubernetes.client.util.ClientBuilder;
 import io.kubernetes.client.util.credentials.AccessTokenAuthentication;
 import lombok.extern.slf4j.Slf4j;
 import org.apache.commons.lang.StringUtils;
+import org.json.JSONObject;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.beans.factory.annotation.Value;
 import org.springframework.stereotype.Component;
@@ -379,7 +381,7 @@ public class K8sClientUtil {
     }
 
     // 创建配置好的Pod
-    public Integer createConfiguredPod(String podName, String namespace, Integer port, String mountPath, V1PersistentVolumeClaim pvc, String image, String computingResource, String dataPvcName, String datasetPath, String modelPath) {
+    public Integer createConfiguredPod(String podName, String namespace, Integer port, String mountPath, V1PersistentVolumeClaim pvc, DevEnvironment devEnvironment, String dataPvcName, String datasetPath, String modelPath) {
 
         //设置选择节点，pod反亲和性
         Map<String, String> selector = new LinkedHashMap<>();
@@ -450,6 +452,24 @@ public class K8sClientUtil {
         volumes.add(new V1Volume().name("minio-pvc").persistentVolumeClaim(new V1PersistentVolumeClaimVolumeSource().claimName(dataPvcName)));
 
 
+        //配置资源
+        JSONObject standardJson = new JSONObject(devEnvironment.getStandard());
+        JSONObject valueJson = (JSONObject)standardJson.get("value");
+        int cpu = (int)valueJson.get("cpu");
+        String memory = (String)valueJson.get("memory");
+
+        HashMap<String, Quantity> limitMap = new HashMap<>();
+        if (Constant.Computing_Resource_GPU.equals(devEnvironment.getComputingResource())) {
+            limitMap.put("nvidia.com/gpu", new Quantity("1"));
+        }
+        limitMap.put("cpu", new Quantity(String.valueOf(cpu)));
+        limitMap.put("memory", new Quantity(memory));
+        limitMap.put("ephemeral-storage", new Quantity("100Gi"));
+
+        V1ResourceRequirements v1ResourceRequirements = new V1ResourceRequirements();
+        v1ResourceRequirements.setRequests(limitMap);
+        v1ResourceRequirements.setLimits(limitMap);
+
         V1Pod pod = new V1PodBuilder()
                 .withNewMetadata()
                 .withName(podName)
@@ -458,9 +478,10 @@ public class K8sClientUtil {
                 .withNewSpec()
                 .addNewContainer()
                 .withName(podName)
-                .withImage(image)
+                .withImage(devEnvironment.getImage())
                 .withPorts(new V1ContainerPort().containerPort(port).protocol("TCP"))
                 .withVolumeMounts(volumeMounts)
+                .withResources(v1ResourceRequirements)
                 .endContainer()
                 .withVolumes(volumes)
                 .withNodeSelector(nodeSelector)
@@ -471,7 +492,7 @@ public class K8sClientUtil {
         try {
             pod = api.createNamespacedPod(namespace, pod, null, null, null);
             String nodeName = getNodeName(podName, namespace);
-            if (Constant.Computing_Resource_GPU.equals(computingResource)) {
+            if (Constant.Computing_Resource_GPU.equals(devEnvironment.getComputingResource())) {
                 computingResourceDao.updateUsedStateByNode(nodeName, Constant.Used_State_used);
             }
         } catch (ApiException e) {