From 7b47faac5051ca1d6414e08b0a7f265478449e54 Mon Sep 17 00:00:00 2001
From: shuiyu1123 <shuiyu1123@users.noreply.gitee.com>
Date: Thu, 18 Jun 2026 20:18:20 +0000
Subject: [PATCH] fix: add vscode-server dir to Dockerfile, use vscode event
 for SSH

---
 .cnb.yml        | 35 +++++++++++++++++++++--------------
 .ide/Dockerfile | 10 +++++++---
 2 files changed, 28 insertions(+), 17 deletions(-)

diff --git a/.cnb.yml b/.cnb.yml
index 9d5a9c4..9e9807d 100644
--- a/.cnb.yml
+++ b/.cnb.yml
@@ -1,14 +1,21 @@
-version: "3"
-tasks:
-  - name: "训练小狸AI-Epoch100"
-    image:
-      dockerfile: .ide/Dockerfile
-    resources:
-      cpu: 8
-      memory: 32G
-      gpu: 1
-    env:
-      CUDA_VISIBLE_DEVICES: "0"
-      PYTHONUNBUFFERED: "1"
-    commands:
-      - cd /workspace && nvidia-smi && echo "=== GPU就绪 ===" && python3 scripts/train_cnb.py --epochs 100 --batch_size 16
\ No newline at end of file
+$:
+  vscode:
+    - docker:
+        build: .ide/Dockerfile
+      runner:
+        cpus: 8
+        tags: cnb:arch:amd64:gpu
+      services:
+        - vscode
+        - docker
+      env:
+        CUDA_VISIBLE_DEVICES: "0"
+        PYTHONUNBUFFERED: "1"
+      stages:
+        - name: 开始训练
+          script: |
+            cd /workspace
+            nvidia-smi
+            echo "=== GPU就绪，开始训练 ==="
+            python3 scripts/train_cnb.py --epochs 100 --batch_size 16
+            echo "=== 训练完成 ==="
\ No newline at end of file
diff --git a/.ide/Dockerfile b/.ide/Dockerfile
index de9cb1c..b89d6b8 100644
--- a/.ide/Dockerfile
+++ b/.ide/Dockerfile
@@ -1,11 +1,15 @@
 FROM pytorch/pytorch:2.1.0-cuda12.1-cudnn8-runtime
 
-# 基础工具
-RUN apt-get update && apt-get install -y git curl wget nano && rm -rf /var/lib/apt/lists/*
+# 安装必要工具
+RUN apt-get update && apt-get install -y git curl wget nano openssh-server \
+    && rm -rf /var/lib/apt/lists/*
 
-# tqdm用于进度条
+# pip依赖
 RUN pip install --no-cache-dir tqdm numpy requests
 
+# 创建 vscode-server 需要的目录
+RUN mkdir -p /root/.vscode-server /root/.vscode-server/data/Machine
+
 # 设置工作目录
 WORKDIR /workspace