diff --git a/.trae/documents/训练SDK化与注册鉴权实现.md b/.trae/documents/训练SDK化与注册鉴权实现.md new file mode 100644 index 0000000..5ba6df0 --- /dev/null +++ b/.trae/documents/训练SDK化与注册鉴权实现.md @@ -0,0 +1,52 @@ +# 总体架构 +- 服务端统一暴露训练服务与进程管理接口；引入客户端注册与鉴权，训练仅对已注册客户端开放。 +- 训练能力保持原子化：按 `train_type`（pretrain/sft/lora/dpo/ppo/grpo/spo）分派脚本与参数，保留现有脚本接口。 +- 提供 Python SDK，封装注册、发起训练、查询进程与日志等操作，客户端只需少量模型参数即可启动训练。 +- Web 端保留现有操作界面，新增注册流程与鉴权头注入。 + +# 服务端改造（Flask） +- 新增注册接口：`POST /api/register`，入参：`name`、`email`；生成并返回 `client_id` 与 `api_key`，并在服务端保存（`trainer_web/clients.json`）。 +- 新增鉴权中间件：校验请求头 `Authorization: Bearer `；解析出 `client_id` 并附加到请求上下文。 +- 训练相关接口全部鉴权并按 `client_id` 隔离： + - 保留并保护现有接口：`POST /train`（trainer_web/train_web_ui.py:247-266）、`GET /processes`（trainer_web/train_web_ui.py:267-286）、`GET /logs/`（trainer_web/train_web_ui.py:288-374）、`POST /stop/`（trainer_web/train_web_ui.py:466-485）、`POST /delete/`（trainer_web/train_web_ui.py:487-520）、`GET /logfiles`、`GET /logfile-content/`、`DELETE /delete-logfile/`。 + - `training_processes` 结构增加 `client_id` 字段；`training_processes.json` 持久化同步该字段（trainer_web/train_web_ui.py:540-607）。 + - 在读取日志、停止/删除进程时先校验归属（仅允许操作本客户端的进程与日志）。 +- 抽象原子训练分发器：新增 `dispatcher` 模块，统一 `train_type → 脚本路径 + 参数构造`，内部直接复用现有命令拼装逻辑（trainer_web/train_web_ui.py:59-175）。脚本与参数名均保持不变。 +- 安全与配置： + - `api_key` 服务端按 `sha256` 存储，避免明文；可选 `SERVER_SECRET` 用于签发/验证简易令牌。 + - 新增布尔配置 `ALLOW_ANONYMOUS=false`，默认禁止匿名训练。 + +# 客户端 SDK（Python） +- 包名：`minimind_sdk`；类：`MinimindClient(base_url, api_key=None)`。 +- 方法： + - `register(name, email) -> {client_id, api_key}` + - `start_training(train_type, **params) -> process_id`（保留与 `/train` 相同字段约定） + - `get_processes() -> list`、`get_logs(process_id) -> str` + - `stop(process_id)`、`delete(process_id)`、`get_logfiles()`、`get_logfile_content(filename)`、`delete_logfile(filename)` +- 简化封装：按训练类型提供高层方法：`train_sft(...)`、`train_dpo(...)`、`train_ppo(...)` 等，只暴露核心参数，其余按当前脚本默认值填充（例如 SFT 参考 trainer/train_full_sft.py:83-105）。 +- 错误处理：统一抛出带有服务端错误信息的异常；返回类型标准化。 + +# Web 端适配 +- 新增“注册”入口：表单提交到 `POST /api/register`，成功后将 `api_key` 保存到 `localStorage`。 +- `apiClient.js` 在所有请求中自动注入 `Authorization` 头（trainer_web/static/js/services/apiClient.js:1-73）。 +- 进程与日志列表仅展示当前客户端数据；保留现有 Tab 与交互（trainer_web/static/js/app.js:1-24、train/form.js:1-127）。 + +# 接口保留与兼容 +- 现有 REST 路径与请求体字段保持不变；仅增加注册接口与鉴权要求。 +- 训练脚本 CLI 接口完全保留（`argparse` 参数不改动），通过分发器集中调用。 +- 对历史脚本/工具的兼容：如果需要，可保留一个仅本机可用的“匿名模式”开关用于调试（默认关闭）。 + +# 数据与安全 +- 客户注册数据文件：`trainer_web/clients.json`；字段：`client_id`、`name`、`email`、`api_key_hash`、`created_at`。 +- 进程记录增加 `client_id`，日志操作前进行归属校验；多租户隔离。 +- 不在日志中打印敏感信息；不落盘明文 `api_key`。 + +# 交付与变更点 +- 服务端：`train_web_ui.py` 引入鉴权与 `client_id` 归属；新增 `auth.py` 与 `dispatcher.py`。 +- SDK：`minimind_sdk` 源码与示例；示例脚本：注册→发起 SFT 训练→轮询进程→查看日志→停止。 +- Web：新增注册 UI、`apiClient.js` 注入鉴权、列表隔离逻辑。 + +# 验收与测试 +- 未注册请求调用 `/train` 返回 `401`；注册后可正常发起训练。 +- 用 SDK 完成一次 SFT 训练，能看到进程、日志、停止与删除操作均正常并隔离。 +- 多客户端并发训练与日志隔离验证；重启服务后通过 `training_processes.json` 成功恢复状态。 diff --git a/minimind_sdk/__init__.py b/minimind_sdk/__init__.py new file mode 100644 index 0000000..4f9b7ed --- /dev/null +++ b/minimind_sdk/__init__.py @@ -0,0 +1 @@ +from .client import MinimindClient \ No newline at end of file diff --git a/minimind_sdk/client.py b/minimind_sdk/client.py new file mode 100644 index 0000000..d104e1c --- /dev/null +++ b/minimind_sdk/client.py @@ -0,0 +1,65 @@ +import json +import urllib.request +import urllib.error + +class MinimindClient: + def __init__(self, base_url, api_key=None, timeout=10): + self.base_url = base_url.rstrip('/') + self.api_key = api_key or '' + self.timeout = timeout + + def _request(self, method, path, body=None, expect_text=False): + url = f"{self.base_url}{path}" + headers = { + 'Content-Type': 'application/json', + 'Cache-Control': 'no-cache' + } + if self.api_key: + headers['Authorization'] = f"Bearer {self.api_key}" + data = None + if body is not None: + data = json.dumps(body).encode('utf-8') + req = urllib.request.Request(url, data=data, headers=headers, method=method) + try: + with urllib.request.urlopen(req, timeout=self.timeout) as resp: + raw = resp.read() + if expect_text: + return raw.decode('utf-8', errors='replace') + return json.loads(raw.decode('utf-8')) + except urllib.error.HTTPError as e: + msg = e.read().decode('utf-8', errors='replace') + raise RuntimeError(f"HTTP {e.code}: {msg}") + except urllib.error.URLError as e: + raise RuntimeError(str(e)) + + def register(self, name, email): + res = self._request('POST', '/api/register', {'name': name, 'email': email}) + self.api_key = res.get('api_key', self.api_key) + return res + + def start_training(self, train_type, **params): + payload = {'train_type': train_type} + payload.update(params or {}) + res = self._request('POST', '/train', payload) + return res + + def get_processes(self): + return self._request('GET', '/processes', None) + + def get_logs(self, process_id): + return self._request('GET', f"/logs/{process_id}", None, expect_text=True) + + def stop(self, process_id): + return self._request('POST', f"/stop/{process_id}", None) + + def delete(self, process_id): + return self._request('POST', f"/delete/{process_id}", None) + + def get_logfiles(self): + return self._request('GET', '/logfiles', None) + + def get_logfile_content(self, filename): + return self._request('GET', f"/logfile-content/{filename}", None, expect_text=True) + + def delete_logfile(self, filename): + return self._request('DELETE', f"/delete-logfile/{filename}", None) \ No newline at end of file diff --git a/trainer_web/auth.py b/trainer_web/auth.py new file mode 100644 index 0000000..62b3910 --- /dev/null +++ b/trainer_web/auth.py @@ -0,0 +1,49 @@ +import os +import json +import hashlib +import secrets +import time + +CLIENTS_FILE = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'clients.json') + +def _load_clients(): + if os.path.exists(CLIENTS_FILE): + with open(CLIENTS_FILE, 'r', encoding='utf-8') as f: + return json.load(f) + return {} + +def _save_clients(clients): + with open(CLIENTS_FILE, 'w', encoding='utf-8') as f: + json.dump(clients, f, ensure_ascii=False, indent=2) + +def _hash_api_key(api_key): + return hashlib.sha256(api_key.encode('utf-8')).hexdigest() + +def register_client(name, email): + clients = _load_clients() + client_id = secrets.token_hex(8) + str(int(time.time())) + api_key = secrets.token_hex(32) + clients[client_id] = { + 'client_id': client_id, + 'name': name, + 'email': email, + 'api_key_hash': _hash_api_key(api_key), + 'created_at': int(time.time()) + } + _save_clients(clients) + return {'client_id': client_id, 'api_key': api_key} + +def verify_api_key(api_key): + clients = _load_clients() + target = _hash_api_key(api_key) + for _, info in clients.items(): + if info.get('api_key_hash') == target: + return info + return None + +def get_client_from_request(request): + auth = request.headers.get('Authorization') + if not auth or not auth.startswith('Bearer '): + return None + token = auth.split(' ', 1)[1].strip() + return verify_api_key(token) \ No newline at end of file diff --git a/trainer_web/dispatcher.py b/trainer_web/dispatcher.py new file mode 100644 index 0000000..092f7b5 --- /dev/null +++ b/trainer_web/dispatcher.py @@ -0,0 +1,81 @@ +import sys +import os + +def build_command(train_type, params, gpu_num, use_torchrun): + if train_type == 'pretrain': + script_path = '../trainer/train_pretrain.py' + cmd = ['torchrun', '--nproc_per_node', str(gpu_num), script_path] if use_torchrun else [sys.executable, script_path] + if 'save_weight' in params: + cmd.extend(['--save_weight', params['save_weight']]) + elif train_type == 'sft': + script_path = '../trainer/train_full_sft.py' + cmd = ['torchrun', '--nproc_per_node', str(gpu_num), script_path] if use_torchrun else [sys.executable, script_path] + if 'save_weight' in params: + cmd.extend(['--save_weight', params['save_weight']]) + elif train_type == 'lora': + script_path = '../trainer/train_lora.py' + cmd = ['torchrun', '--nproc_per_node', str(gpu_num), script_path] if use_torchrun else [sys.executable, script_path] + if 'lora_name' in params: + cmd.extend(['--lora_name', params['lora_name']]) + elif train_type == 'dpo': + script_path = '../trainer/train_dpo.py' + cmd = ['torchrun', '--nproc_per_node', str(gpu_num), script_path] if use_torchrun else [sys.executable, script_path] + if 'beta' in params and params['beta']: + cmd.extend(['--beta', params['beta']]) + if 'accumulation_steps' in params and params['accumulation_steps']: + cmd.extend(['--accumulation_steps', params['accumulation_steps']]) + if 'grad_clip' in params and params['grad_clip']: + cmd.extend(['--grad_clip', params['grad_clip']]) + elif train_type == 'ppo': + script_path = '../trainer/train_ppo.py' + cmd = ['torchrun', '--nproc_per_node', str(gpu_num), script_path] if use_torchrun else [sys.executable, script_path] + if 'clip_epsilon' in params and params['clip_epsilon']: + cmd.extend(['--clip_epsilon', params['clip_epsilon']]) + if 'vf_coef' in params and params['vf_coef']: + cmd.extend(['--vf_coef', params['vf_coef']]) + if 'kl_coef' in params and params['kl_coef']: + cmd.extend(['--kl_coef', params['kl_coef']]) + if 'reasoning' in params and params['reasoning']: + cmd.extend(['--reasoning', params['reasoning']]) + if 'update_old_actor_freq' in params and params['update_old_actor_freq']: + cmd.extend(['--update_old_actor_freq', params['update_old_actor_freq']]) + if 'reward_model_path' in params and params['reward_model_path']: + cmd.extend(['--reward_model_path', params['reward_model_path']]) + elif train_type == 'grpo': + script_path = '../trainer/train_grpo.py' + cmd = ['torchrun', '--nproc_per_node', str(gpu_num), script_path] if use_torchrun else [sys.executable, script_path] + if 'beta' in params and params['beta']: + cmd.extend(['--beta', params['beta']]) + if 'num_generations' in params and params['num_generations']: + cmd.extend(['--num_generations', params['num_generations']]) + if 'reasoning' in params and params['reasoning']: + cmd.extend(['--reasoning', params['reasoning']]) + if 'reward_model_path' in params and params['reward_model_path']: + cmd.extend(['--reward_model_path', params['reward_model_path']]) + elif train_type == 'spo': + script_path = '../trainer/train_spo.py' + cmd = ['torchrun', '--nproc_per_node', str(gpu_num), script_path] if use_torchrun else [sys.executable, script_path] + if 'beta' in params and params['beta']: + cmd.extend(['--beta', params['beta']]) + if 'reasoning' in params and params['reasoning']: + cmd.extend(['--reasoning', params['reasoning']]) + if 'reward_model_path' in params and params['reward_model_path']: + cmd.extend(['--reward_model_path', params['reward_model_path']]) + else: + return None + + for key, value in params.items(): + if key in ['train_type', 'save_weight', 'lora_name', 'train_monitor', 'beta', 'accumulation_steps', 'grad_clip', 'gpu_num', 'clip_epsilon', 'vf_coef', 'kl_coef', 'reasoning', 'update_old_actor_freq', 'reward_model_path', 'num_generations'] or ((train_type == 'ppo' or train_type == 'grpo' or train_type == 'spo') and key == 'from_weight'): + continue + elif key == 'from_resume': + cmd.extend([f'--{key}', str(value)]) + else: + cmd.extend([f'--{key}', str(value)]) + + if 'train_monitor' in params: + if params['train_monitor'] == 'wandb' or params['train_monitor'] == 'swanlab': + cmd.append('--use_wandb') + if params['train_monitor'] == 'wandb': + cmd.extend(['--wandb_project', 'minimind_training']) + + return cmd \ No newline at end of file diff --git a/trainer_web/static/js/services/apiClient.js b/trainer_web/static/js/services/apiClient.js index 0da4516..c269d60 100644 --- a/trainer_web/static/js/services/apiClient.js +++ b/trainer_web/static/js/services/apiClient.js @@ -1,8 +1,10 @@ +import { getApiKey } from './authClient.js'; const defaultTimeout = 10000; export function fetchWithTimeoutAndRetry(url, options = {}, timeout = defaultTimeout, retries = 3) { const controller = new AbortController(); const timeoutId = setTimeout(() => controller.abort(), timeout); + const apiKey = getApiKey(); const fetchOptions = { ...options, headers: { @@ -10,6 +12,7 @@ export function fetchWithTimeoutAndRetry(url, options = {}, timeout = defaultTim 'Cache-Control': 'no-cache, no-store, must-revalidate', Pragma: 'no-cache', Expires: '0', + ...(apiKey ? { Authorization: `Bearer ${apiKey}` } : {}), }, signal: controller.signal, }; diff --git a/trainer_web/static/js/services/authClient.js b/trainer_web/static/js/services/authClient.js new file mode 100644 index 0000000..18470c5 --- /dev/null +++ b/trainer_web/static/js/services/authClient.js @@ -0,0 +1,29 @@ +const KEY = 'minimind_api_key'; + +export function getApiKey() { + try { + return localStorage.getItem(KEY) || ''; + } catch (_) { + return ''; + } +} + +export function setApiKey(k) { + try { + localStorage.setItem(KEY, k || ''); + } catch (_) {} +} + +export function registerClient(payload) { + return fetch('/api/register', { + method: 'POST', + headers: { 'Content-Type': 'application/json', 'Cache-Control': 'no-cache' }, + body: JSON.stringify(payload || {}), + }).then((r) => { + if (!r.ok) throw new Error('register_failed'); + return r.json(); + }).then((res) => { + if (res && res.api_key) setApiKey(res.api_key); + return res; + }); +} \ No newline at end of file diff --git a/trainer_web/templates/index.html b/trainer_web/templates/index.html index d6eb535..788819b 100644 --- a/trainer_web/templates/index.html +++ b/trainer_web/templates/index.html @@ -11,6 +11,26 @@

MiniMind Training Lab

+
+
+
+

注册客户端

+
+
+ + +
+
+ + +
+
+
+ +
+
+
+
@@ -271,5 +291,26 @@
+ diff --git a/trainer_web/train_web_ui.py b/trainer_web/train_web_ui.py index 146e302..38f560e 100644 --- a/trainer_web/train_web_ui.py +++ b/trainer_web/train_web_ui.py @@ -7,6 +7,7 @@ import socket import atexit import signal from flask import Flask, render_template, request, jsonify, redirect, url_for +from flask import g import time import psutil @@ -55,8 +56,10 @@ PROCESSES_FILE = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'train # PID文件 PID_FILE = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'train_web_ui.pid') +from .auth import register_client, get_client_from_request + # 启动训练进程 -def start_training_process(train_type, params): +def start_training_process(train_type, params, client_id): # 获取脚本所在目录的绝对路径 script_dir = os.path.dirname(os.path.abspath(__file__)) # 使用详细的时间戳作为进程ID和日志文件名 @@ -73,113 +76,11 @@ def start_training_process(train_type, params): gpu_num = int(params.get('gpu_num', 0)) if 'gpu_num' in params else 0 use_torchrun = HAS_TORCH and GPU_COUNT > 0 and gpu_num > 1 - # 构建命令 - if train_type == 'pretrain': - script_path = '../trainer/train_pretrain.py' - if use_torchrun: - cmd = ['torchrun', '--nproc_per_node', str(gpu_num), script_path] - else: - cmd = [sys.executable, script_path] - if 'save_weight' in params: - cmd.extend(['--save_weight', params['save_weight']]) - elif train_type == 'sft': - script_path = '../trainer/train_full_sft.py' - if use_torchrun: - cmd = ['torchrun', '--nproc_per_node', str(gpu_num), script_path] - else: - cmd = [sys.executable, script_path] - if 'save_weight' in params: - cmd.extend(['--save_weight', params['save_weight']]) - elif train_type == 'lora': - script_path = '../trainer/train_lora.py' - if use_torchrun: - cmd = ['torchrun', '--nproc_per_node', str(gpu_num), script_path] - else: - cmd = [sys.executable, script_path] - if 'lora_name' in params: - cmd.extend(['--lora_name', params['lora_name']]) - elif train_type == 'dpo': - script_path = '../trainer/train_dpo.py' - if use_torchrun: - cmd = ['torchrun', '--nproc_per_node', str(gpu_num), script_path] - else: - cmd = [sys.executable, script_path] - # 添加DPO特定参数 - if 'beta' in params and params['beta']: - cmd.extend(['--beta', params['beta']]) - if 'accumulation_steps' in params and params['accumulation_steps']: - cmd.extend(['--accumulation_steps', params['accumulation_steps']]) - if 'grad_clip' in params and params['grad_clip']: - cmd.extend(['--grad_clip', params['grad_clip']]) - elif train_type == 'ppo': - script_path = '../trainer/train_ppo.py' - if use_torchrun: - cmd = ['torchrun', '--nproc_per_node', str(gpu_num), script_path] - else: - cmd = [sys.executable, script_path] - # 添加PPO特定参数 - if 'clip_epsilon' in params and params['clip_epsilon']: - cmd.extend(['--clip_epsilon', params['clip_epsilon']]) - if 'vf_coef' in params and params['vf_coef']: - cmd.extend(['--vf_coef', params['vf_coef']]) - if 'kl_coef' in params and params['kl_coef']: - cmd.extend(['--kl_coef', params['kl_coef']]) - if 'reasoning' in params and params['reasoning']: - cmd.extend(['--reasoning', params['reasoning']]) - if 'update_old_actor_freq' in params and params['update_old_actor_freq']: - cmd.extend(['--update_old_actor_freq', params['update_old_actor_freq']]) - if 'reward_model_path' in params and params['reward_model_path']: - cmd.extend(['--reward_model_path', params['reward_model_path']]) - elif train_type == 'grpo': - script_path = '../trainer/train_grpo.py' - if use_torchrun: - cmd = ['torchrun', '--nproc_per_node', str(gpu_num), script_path] - else: - cmd = [sys.executable, script_path] - if 'beta' in params and params['beta']: - cmd.extend(['--beta', params['beta']]) - if 'num_generations' in params and params['num_generations']: - cmd.extend(['--num_generations', params['num_generations']]) - if 'reasoning' in params and params['reasoning']: - cmd.extend(['--reasoning', params['reasoning']]) - if 'reward_model_path' in params and params['reward_model_path']: - cmd.extend(['--reward_model_path', params['reward_model_path']]) - elif train_type == 'spo': - script_path = '../trainer/train_spo.py' - if use_torchrun: - cmd = ['torchrun', '--nproc_per_node', str(gpu_num), script_path] - else: - cmd = [sys.executable, script_path] - if 'beta' in params and params['beta']: - cmd.extend(['--beta', params['beta']]) - if 'reasoning' in params and params['reasoning']: - cmd.extend(['--reasoning', params['reasoning']]) - if 'reward_model_path' in params and params['reward_model_path']: - cmd.extend(['--reward_model_path', params['reward_model_path']]) - else: + from .dispatcher import build_command + cmd = build_command(train_type, params, gpu_num, use_torchrun) + if cmd is None: return None - # 添加通用参数 - for key, value in params.items(): - # 跳过特殊参数和DPO、PPO特有参数，以及gpu_num参数（因为已经在torchrun命令中使用） - # 对于PPO/GRPO/SPO训练，跳过--from_weight参数 - if key in ['train_type', 'save_weight', 'lora_name', 'train_monitor', 'beta', 'accumulation_steps', 'grad_clip', 'gpu_num', 'clip_epsilon', 'vf_coef', 'kl_coef', 'reasoning', 'update_old_actor_freq', 'reward_model_path', 'num_generations'] or ((train_type == 'ppo' or train_type == 'grpo' or train_type == 'spo') and key == 'from_weight'): - continue - # 对于from_resume参数，需要正确传递参数值 - elif key == 'from_resume': - # 确保传递参数名和参数值 - cmd.extend([f'--{key}', str(value)]) - else: - # 确保log_interval和save_interval参数正确传递 - cmd.extend([f'--{key}', str(value)]) - - # 单独处理训练监控参数，确保它不会被错误地添加值 - if 'train_monitor' in params: - if params['train_monitor'] == 'wandb' or params['train_monitor'] == 'swanlab': - cmd.append('--use_wandb') # 对于wandb和swanlab，只添加标志，不添加值 - if params['train_monitor'] == 'wandb': - cmd.extend(['--wandb_project', 'minimind_training']) - # 创建日志文件 with open(log_file, 'w') as f: f.write(f"开始训练 {train_type} 进程\n") @@ -204,7 +105,8 @@ def start_training_process(train_type, params): 'error': False, 'train_monitor': params.get('train_monitor', 'none'), # 保存训练监控设置 'swanlab_url': None, - 'next_line_is_swanlab_url': False + 'next_line_is_swanlab_url': False, + 'client_id': client_id } # 开始读取输出 @@ -246,6 +148,9 @@ def index(): @app.route('/train', methods=['POST']) def train(): + client = get_client_from_request(request) + if not client: + return jsonify({'success': False, 'error': '未授权'}), 401 data = request.json train_type = data.get('train_type') @@ -257,7 +162,7 @@ def train(): params['from_resume'] = '0' # 启动训练进程 - process_id = start_training_process(train_type, params) + process_id = start_training_process(train_type, params, client['client_id']) if process_id: return jsonify({'success': True, 'process_id': process_id}) @@ -266,8 +171,14 @@ def train(): @app.route('/processes') def processes(): + client = get_client_from_request(request) + if not client: + return jsonify([]), 401 + cid = client['client_id'] result = [] for process_id, info in training_processes.items(): + if info.get('client_id') != cid: + continue # 确定状态 status = '运行中' if info['running'] else \ '手动停止' if 'manually_stopped' in info and info['manually_stopped'] else \ @@ -287,6 +198,12 @@ def processes(): @app.route('/logs/') def logs(process_id): + client = get_client_from_request(request) + if not client: + return '未授权', 401 + cid = client['client_id'] + if process_id not in training_processes or training_processes[process_id].get('client_id') != cid: + return '无权限', 403 # 直接从本地logfile目录读取日志文件 # 获取脚本所在目录的绝对路径 script_dir = os.path.dirname(os.path.abspath(__file__)) @@ -374,6 +291,10 @@ def logs(process_id): @app.route('/logfiles') def get_logfiles(): + client = get_client_from_request(request) + if not client: + return jsonify([]), 401 + cid = client['client_id'] # 获取脚本所在目录的绝对路径 script_dir = os.path.dirname(os.path.abspath(__file__)) # 构建logfile目录的绝对路径 @@ -381,6 +302,7 @@ def get_logfiles(): log_dir = os.path.abspath(log_dir) logfiles = [] + client_pids = {pid for pid, info in training_processes.items() if info.get('client_id') == cid} if os.path.exists(log_dir): for filename in os.listdir(log_dir): if filename.endswith('.log') and filename.startswith('train_'): @@ -388,6 +310,8 @@ def get_logfiles(): try: modified_time = os.path.getmtime(file_path) formatted_time = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(modified_time)) + if not any(filename.endswith(f'{pid}.log') for pid in client_pids): + continue logfiles.append({ 'filename': filename, 'modified_time': formatted_time, @@ -401,6 +325,9 @@ def get_logfiles(): @app.route('/logfile-content/') def get_logfile_content(filename): + client = get_client_from_request(request) + if not client: + return jsonify({'error': '未授权'}), 401 # 安全检查：确保文件名不包含路径遍历字符 if '..' in filename or '/' in filename or '\\' in filename: return jsonify({'error': 'Invalid filename'}), 400 @@ -411,6 +338,9 @@ def get_logfile_content(filename): log_dir = os.path.join(script_dir, '../logfile') log_dir = os.path.abspath(log_dir) log_file = os.path.join(log_dir, filename) + pid = filename.split('.')[-2].split('_')[-1] if filename.endswith('.log') else None + if not pid or pid not in training_processes or training_processes[pid].get('client_id') != client['client_id']: + return jsonify({'error': '无权限'}), 403 try: # 使用二进制模式读取文件，可以更可靠地保留原始换行符 @@ -442,6 +372,9 @@ def get_logfile_content(filename): @app.route('/delete-logfile/', methods=['DELETE']) def delete_logfile(filename): + client = get_client_from_request(request) + if not client: + return jsonify({'success': False, 'message': '未授权'}), 401 # 获取脚本所在目录的绝对路径 script_dir = os.path.dirname(os.path.abspath(__file__)) # 构建logfile目录的绝对路径 @@ -453,6 +386,9 @@ def delete_logfile(filename): return jsonify({'success': False, 'message': '非法的文件名'}) log_file = os.path.join(log_dir, filename) + pid = filename.split('.')[-2].split('_')[-1] if filename.endswith('.log') else None + if not pid or pid not in training_processes or training_processes[pid].get('client_id') != client['client_id']: + return jsonify({'success': False, 'message': '无权限'}), 403 if os.path.exists(log_file) and os.path.isfile(log_file): try: os.remove(log_file) @@ -465,7 +401,10 @@ def delete_logfile(filename): @app.route('/stop/', methods=['POST']) def stop(process_id): - if process_id in training_processes and training_processes[process_id]['running']: + client = get_client_from_request(request) + if not client: + return jsonify({'success': False}), 401 + if process_id in training_processes and training_processes[process_id]['running'] and training_processes[process_id].get('client_id') == client['client_id']: process = training_processes[process_id]['process'] # 在Windows上使用terminate，在Unix上尝试优雅终止 try: @@ -486,7 +425,10 @@ def stop(process_id): @app.route('/delete/', methods=['POST']) def delete(process_id): - if process_id in training_processes: + client = get_client_from_request(request) + if not client: + return jsonify({'success': False}), 401 + if process_id in training_processes and training_processes[process_id].get('client_id') == client['client_id']: # 确保进程已经停止 if training_processes[process_id]['running']: # 如果进程还在运行，先停止它 @@ -552,7 +494,8 @@ def save_processes_info(): 'error': info.get('error', False), 'manually_stopped': info.get('manually_stopped', False), 'train_monitor': info.get('train_monitor', 'none'), # 保存train_monitor - 'swanlab_url': info.get('swanlab_url') # 保存swanlab_url + 'swanlab_url': info.get('swanlab_url'), + 'client_id': info.get('client_id') } with open(PROCESSES_FILE, 'w', encoding='utf-8') as f: @@ -579,6 +522,8 @@ def load_processes_info(): info['error'] = False if 'train_monitor' not in info: info['train_monitor'] = 'none' + if 'client_id' not in info: + info['client_id'] = None if info['running']: try: @@ -650,4 +595,13 @@ if __name__ == '__main__': os.remove(PID_FILE) except: pass - sys.exit(1) \ No newline at end of file + sys.exit(1) +@app.route('/api/register', methods=['POST']) +def api_register(): + data = request.json or {} + name = data.get('name') + email = data.get('email') + if not name or not email: + return jsonify({'error': '缺少参数'}), 400 + creds = register_client(name, email) + return jsonify({'client_id': creds['client_id'], 'api_key': creds['api_key']}) \ No newline at end of file