merge: 合并远程 ft_wyt 分支，解决冲突

compute/api/main.py

@@ -221,6 +221,35 @@ def create_app() -> FastAPI:
             return fallback_gpu_resources()
 
         items: list[dict[str, Any]] = []
+        processes_by_uuid: dict[str, list[dict[str, Any]]] = {}
+        try:
+            process_result = subprocess.run(
+                [
+                    "nvidia-smi",
+                    "--query-compute-apps=gpu_uuid,pid,process_name,used_memory",
+                    "--format=csv,noheader,nounits",
+                ],
+                check=True,
+                capture_output=True,
+                text=True,
+                timeout=5,
+            )
+            for process_line in process_result.stdout.splitlines():
+                process_parts = [part.strip() for part in process_line.split(",")]
+                if len(process_parts) < 4:
+                    continue
+                process_uuid, pid, process_name, used_memory = process_parts[:4]
+                processes_by_uuid.setdefault(process_uuid, []).append(
+                    {
+                        "pid": int(_safe_float(pid)),
+                        "name": process_name,
+                        "memory_used_gb": round(_safe_float(used_memory) / 1024, 2),
+                    }
+                )
+        except Exception:
+            # Some driver/runtime combinations do not expose compute-apps;
+            # utilization and memory metrics remain useful without processes.
+            pass
         for line in result.stdout.splitlines():
             parts = [part.strip() for part in line.split(",")]
             if len(parts) < 9:
@@ -244,7 +273,7 @@ def create_app() -> FastAPI:
                     "temperature": int(_safe_float(temp)),
                     "power_w": round(_safe_float(power), 1),
                     "power_limit_w": round(_safe_float(power_limit), 1),
-                    "processes": [],
+                    "processes": processes_by_uuid.get(uuid, []),
                 }
             )
         return items
@@ -747,6 +776,23 @@ def create_app() -> FastAPI:
             infer_backend: str (default: "huggingface")
             infer_dtype: str (default: "auto")
         """
+        requested_gpus = payload.get("gpu_indices")
+        if requested_gpus is None:
+            requested_gpus = payload.get("gpus") or []
+        try:
+            requested_gpus = sorted({int(item) for item in requested_gpus})
+        except (TypeError, ValueError) as exc:
+            raise HTTPException(status_code=400, detail=f"invalid GPU selection: {exc}") from exc
+        if any(item < 0 for item in requested_gpus):
+            raise HTTPException(status_code=400, detail="GPU index must be non-negative")
+        if requested_gpus:
+            known_gpus = {int(item.get("gpu_index", item.get("id", -1))) for item in gpu_resources()}
+            missing = sorted(set(requested_gpus) - known_gpus)
+            if missing:
+                raise HTTPException(status_code=409, detail=f"requested GPU not found: {missing}")
+            conflict = sorted(set(requested_gpus).intersection(process_manager.locked_gpus()))
+            if conflict:
+                raise HTTPException(status_code=409, detail=f"GPU already used by another compute job: {conflict}")
         session = get_inference_session()
         result = session.load(
             model_name_or_path=payload.get("model_name_or_path", ""),
@@ -754,6 +800,7 @@ def create_app() -> FastAPI:
             template=payload.get("template", "qwen"),
             infer_backend=payload.get("infer_backend", "huggingface"),
             infer_dtype=payload.get("infer_dtype", "auto"),
+            gpu_indices=requested_gpus,
         )
         return result
 

compute/engines/llama_factory/inference.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import os
 import threading
 import time
 import uuid
@@ -36,6 +37,7 @@ class InferenceSession:
         self._generating_args: dict[str, Any] = {}
         self._model_name: str = ""
         self._adapter_path: str = ""
+        self._gpu_indices: list[int] = []
         self._loaded_at: float = 0.0
 
     @property
@@ -53,6 +55,7 @@ class InferenceSession:
                 "loaded_at": self._loaded_at,
                 "request_id": self._request_id,
                 "error": self._error,
+                "gpu_indices": list(self._gpu_indices),
             }
 
     def wait_until_loaded(self, timeout: float | None = None) -> dict[str, Any]:
@@ -87,8 +90,12 @@ class InferenceSession:
         template="qwen",
         infer_backend="huggingface",
         infer_dtype="auto",
+        gpu_indices=None,
         **kwargs,
     ) -> dict[str, Any]:
+        requested_gpus = sorted({int(item) for item in (gpu_indices or [])})
+        if any(item < 0 for item in requested_gpus):
+            return {"loaded": False, "status": "error", "error": "GPU index must be non-negative"}
         with self._state_lock:
             if self._status == "loading":
                 # A model is already loading — dedupe, reuse the same request id.
@@ -97,6 +104,7 @@ class InferenceSession:
             self._status = "loading"
             self._error = ""
             self._request_id = uuid.uuid4().hex[:12]
+            self._gpu_indices = requested_gpus
             self._cancel_requested = False
             self._load_args = {
                 "model_name_or_path": model_name_or_path,
@@ -115,13 +123,21 @@ class InferenceSession:
 
     def _load_worker(self) -> None:
         """Build the ChatModel off the state lock so info() never blocks."""
+        with self._state_lock:
+            requested_gpus = list(self._gpu_indices)
         model = None
         tokenizer = None
         generating_args: dict[str, Any] = {}
         error = ""
+        previous_visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES")
         try:
+            # Set visibility before LLaMA-Factory/PyTorch initializes CUDA.
+            if requested_gpus:
+                os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(str(item) for item in requested_gpus)
             if self._teardown_old:
                 self._release_model()
+                with self._state_lock:
+                    self._gpu_indices = requested_gpus
             from llamafactory.chat import ChatModel
             from llamafactory.hparams import get_infer_args
 
@@ -138,6 +154,12 @@ class InferenceSession:
                 generating_args = dict(generating_args)
         except Exception as exc:  # noqa: BLE001 - surface load failure via status
             error = str(exc)
+        finally:
+            if requested_gpus:
+                if previous_visible_devices is None:
+                    os.environ.pop("CUDA_VISIBLE_DEVICES", None)
+                else:
+                    os.environ["CUDA_VISIBLE_DEVICES"] = previous_visible_devices
         with self._state_lock:
             if error:
                 self._model = None
@@ -152,6 +174,7 @@ class InferenceSession:
                 self._model = None
                 self._tokenizer = None
                 self._status = "idle"
+                self._gpu_indices = []
                 return
             self._model = model
             self._tokenizer = tokenizer
@@ -189,6 +212,7 @@ class InferenceSession:
                 self._adapter_path = ""
                 self._loaded_at = 0.0
                 self._error = ""
+                self._gpu_indices = []
 
     def unload(self) -> dict[str, Any]:
         with self._state_lock:
@@ -208,6 +232,7 @@ class InferenceSession:
                 self._adapter_path = ""
                 self._loaded_at = 0.0
                 self._error = ""
+                self._gpu_indices = []
         return {"unloaded": True, "status": "idle"}
 
     def chat(self, messages, temperature=0.95, top_p=0.7, max_new_tokens=1024, do_sample=True, **kwargs) -> dict[str, Any]: