Merge 8607169546 into 3f1a7cc25b

Automated security fix generated by Orbis Security AI

scripts/serve_openai_api.py

@@ -13,8 +13,9 @@ import uvicorn
 
 from threading import Thread
 from queue import Queue
-from fastapi import FastAPI, HTTPException
+from fastapi import FastAPI, HTTPException, Depends
 from fastapi.responses import StreamingResponse
+from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
 from pydantic import BaseModel
 from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
 from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
@@ -24,6 +25,14 @@ warnings.filterwarnings('ignore')
 
 app = FastAPI()
 
+_security = HTTPBearer(auto_error=False)
+
+
+def _verify_api_key(credentials: HTTPAuthorizationCredentials = Depends(_security)):
+    api_key = os.environ.get("MINIMIND_API_KEY")
+    if api_key and (credentials is None or credentials.credentials != api_key):
+        raise HTTPException(status_code=401, detail="Invalid or missing API key")
+
 
 def init_model(args):
     tokenizer = AutoTokenizer.from_pretrained(args.load_from)
@@ -169,7 +178,7 @@ def generate_stream_response(messages, temperature, top_p, max_tokens, tools=Non
 
 
 @app.post("/v1/chat/completions")
-async def chat_completions(request: ChatRequest):
+async def chat_completions(request: ChatRequest, _: None = Depends(_verify_api_key)):
     try:
         if request.stream:
             return StreamingResponse(

trainer/train_agent.py

@@ -487,4 +487,6 @@ if __name__ == "__main__":
         else:
             rl_train_epoch(epoch, loader, len(loader), rollout_engine, ref_model, reward_model, 0, wandb, use_sglang = (args.rollout_engine == "sglang"))
 
-    if dist.is_initialized(): dist.destroy_process_group()
+    if dist.is_initialized():
+        dist.barrier()
+        dist.destroy_process_group()

trainer/train_distillation.py

@@ -243,4 +243,6 @@ if __name__ == "__main__":
             train_epoch(epoch, loader, len(loader), teacher_model, lm_config_student, 0, wandb, args.alpha, args.temperature)
     
     # ========== 9. 清理分布进程 ==========
-    if dist.is_initialized(): dist.destroy_process_group()
+    if dist.is_initialized():
+        dist.barrier()
+        dist.destroy_process_group()

trainer/train_dpo.py

@@ -223,4 +223,6 @@ if __name__ == "__main__":
             train_epoch(epoch, loader, len(loader), ref_model, lm_config, 0, wandb, args.beta)
     
     # ========== 9. 清理分布进程 ==========
-    if dist.is_initialized(): dist.destroy_process_group()
+    if dist.is_initialized():
+        dist.barrier()
+        dist.destroy_process_group()

trainer/train_full_sft.py

@@ -168,4 +168,6 @@ if __name__ == "__main__":
             train_epoch(epoch, loader, len(loader), 0, wandb)
     
     # ========== 9. 清理分布进程 ==========
-    if dist.is_initialized(): dist.destroy_process_group()
+    if dist.is_initialized():
+        dist.barrier()
+        dist.destroy_process_group()

trainer/train_grpo.py

@@ -329,4 +329,6 @@ if __name__ == "__main__":
             grpo_train_epoch(epoch, loader, len(loader), rollout_engine, ref_model, reward_model, 0, wandb, use_sglang = (args.rollout_engine == "sglang"))
     
     # ========== 9. 清理分布进程 ==========
-    if dist.is_initialized(): dist.destroy_process_group()
+    if dist.is_initialized():
+        dist.barrier()
+        dist.destroy_process_group()

trainer/train_lora.py

@@ -181,4 +181,6 @@ if __name__ == "__main__":
             train_epoch(epoch, loader, len(loader), lora_params, 0, wandb)
     
     # ========== 10. 清理分布进程 ==========
-    if dist.is_initialized(): dist.destroy_process_group()
+    if dist.is_initialized():
+        dist.barrier()
+        dist.destroy_process_group()

trainer/train_ppo.py

@@ -432,4 +432,6 @@ if __name__ == "__main__":
             ppo_train_epoch(epoch, loader, len(loader), rollout_engine, ref_model, actor_scheduler, critic_scheduler, reward_model, 0, wandb, use_sglang = (args.rollout_engine == "sglang"))
     
     # ========== 9. 清理分布进程 ==========
-    if dist.is_initialized(): dist.destroy_process_group()
+    if dist.is_initialized():
+        dist.barrier()
+        dist.destroy_process_group()

trainer/train_pretrain.py

@@ -167,4 +167,6 @@ if __name__ == "__main__":
             train_epoch(epoch, loader, len(loader), 0, wandb)
     
     # ========== 9. 清理分布进程 ==========
-    if dist.is_initialized(): dist.destroy_process_group()
+    if dist.is_initialized():
+        dist.barrier()
+        dist.destroy_process_group()