From da31bd9d999fb71b00ae692d05799212754c3552 Mon Sep 17 00:00:00 2001
From: "gayoun.park" <gayoun.park@tf.dabeeo.com>
Date: Tue, 10 Mar 2026 23:09:39 +0900
Subject: [PATCH] =?UTF-8?q?ai=20=ED=95=99=EC=8A=B5=EC=8B=A4=ED=96=89=20run?=
 =?UTF-8?q?=20command=20=EC=88=98=EC=A0=95?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../train/service/DockerTrainService.java     | 25 ++++++++++++++-----
 1 file changed, 19 insertions(+), 6 deletions(-)

diff --git a/src/main/java/com/kamco/cd/training/train/service/DockerTrainService.java b/src/main/java/com/kamco/cd/training/train/service/DockerTrainService.java
index e52cdf5..4d9bec2 100644
--- a/src/main/java/com/kamco/cd/training/train/service/DockerTrainService.java
+++ b/src/main/java/com/kamco/cd/training/train/service/DockerTrainService.java
@@ -54,6 +54,9 @@ public class DockerTrainService {
   @Value("${train.docker.ipcHost:true}")
   private boolean ipcHost;
 
+  @Value("${spring.profiles.active}")
+  private String profile;
+
   private final ModelTrainJobCoreService modelTrainJobCoreService;
 
   /**
@@ -228,8 +231,11 @@ public class DockerTrainService {
 
     // GPU 전체 사용
     c.add("--gpus");
-    c.add("1"); // 학습서버 GPU 1개인 곳은 1이어야 함
-    //    c.add("all");
+    if ("prod".equals(profile)) {
+      c.add("1"); // 학습서버 GPU 1개인 곳은 1이어야 함
+    } else {
+      c.add("all");
+    }
 
     // IPC host 사용 여부
     if (ipcHost) {
@@ -282,8 +288,12 @@ public class DockerTrainService {
     addArg(c, "--crop-size", req.getCropSize());
     //    addArg(c, "--batch-size", req.getBatchSize());
     //    addArg(c, "--gpu-ids", req.getGpuIds()); // null
-    addArg(c, "--batch-size", 2); // 학습서버 GPU 1개인 곳은 batch-size:2 까지만 가능
-    addArg(c, "--gpu-ids", "0"); // 학습서버 GPU 1개인 곳은 0이어야 함
+    if ("prod".equals(profile)) {
+      addArg(c, "--batch-size", 2); // 학습서버 GPU 1개인 곳은 batch-size:2 까지만 가능
+      addArg(c, "--gpu-ids", "0"); // 학습서버 GPU 1개인 곳은 0이어야 함
+    } else {
+      addArg(c, "--batch-size", req.getBatchSize()); // 학습서버 GPU 1개인 곳은 batch-size:2 까지만 가능
+    }
     addArg(c, "--lr", req.getLearningRate());
     addArg(c, "--backbone", req.getBackbone());
     addArg(c, "--epochs", req.getEpochs());
@@ -448,8 +458,11 @@ public class DockerTrainService {
     c.add("run");
     c.add("--rm");
     c.add("--gpus");
-    c.add("1"); // 학습서버 GPU 1개인 곳은 1이어야 함
-    //    c.add("all");
+    if ("prod".equals(profile)) {
+      c.add("1"); // 학습서버 GPU 1개인 곳은 1이어야 함
+    } else {
+      c.add("all");
+    }
     c.add("--ipc=host");
     c.add("--shm-size=" + shmSize);