0b59c1ebee
- 后端: 新增 resume/checkpoints 端点，compute 客户端与适配器支持续训 - 计算节点: llama_factory 适配器断点续训支持及测试 - 前端: fine-tune API 封装与列表页续训/断点展示
283 lines
12 KiB
Python
283 lines
12 KiB
Python
from __future__ import annotations
|
|
|
|
import time
|
|
from typing import Any
|
|
from urllib.parse import urljoin
|
|
|
|
import httpx
|
|
|
|
from app.core.config import get_settings
|
|
|
|
|
|
def _join_url(base_url: str, path: str) -> str:
|
|
return urljoin(base_url.rstrip("/") + "/", path.lstrip("/"))
|
|
|
|
|
|
def _unwrap_items(payload: Any) -> list[dict[str, Any]]:
|
|
if isinstance(payload, list):
|
|
return [item for item in payload if isinstance(item, dict)]
|
|
if isinstance(payload, dict):
|
|
data = payload.get("data")
|
|
if isinstance(data, dict) and isinstance(data.get("items"), list):
|
|
return [item for item in data["items"] if isinstance(item, dict)]
|
|
if isinstance(payload.get("items"), list):
|
|
return [item for item in payload["items"] if isinstance(item, dict)]
|
|
if isinstance(data, list):
|
|
return [item for item in data if isinstance(item, dict)]
|
|
return []
|
|
|
|
|
|
def _unwrap_dict(payload: Any) -> dict[str, Any]:
|
|
if isinstance(payload, dict) and isinstance(payload.get("data"), dict):
|
|
return payload["data"]
|
|
return payload if isinstance(payload, dict) else {}
|
|
|
|
|
|
# Inference calls are intentionally short-timeout:
|
|
# - load dispatch only confirms the compute node accepted the request
|
|
# (the actual model load now runs asynchronously on the node).
|
|
# - status/unload must never block the platform for long when a node is
|
|
# unreachable but still marked online.
|
|
INFERENCE_LOAD_TIMEOUT = httpx.Timeout(30, connect=10)
|
|
INFERENCE_STATUS_TIMEOUT = httpx.Timeout(30, connect=5)
|
|
INFERENCE_UNLOAD_TIMEOUT = httpx.Timeout(30, connect=5)
|
|
|
|
|
|
class ComputeNodeClient:
|
|
"""Application-side client for one compute node.
|
|
|
|
The client accepts both current YG Compute API responses and common
|
|
wrapper shapes such as `{code,message,data}` to make future engine/node
|
|
adapters less brittle.
|
|
"""
|
|
|
|
def __init__(self, api_base_url: str, token: str | None = None, timeout: float | None = None) -> None:
|
|
settings = get_settings()
|
|
self.api_base_url = api_base_url.rstrip("/")
|
|
self.token = token or settings.compute_service_token
|
|
self.timeout = timeout or settings.compute_request_timeout_seconds
|
|
self.route_prefix = settings.route_prefix.rstrip("/") or "/modelTF"
|
|
|
|
def headers(self) -> dict[str, str]:
|
|
if not self.token:
|
|
return {}
|
|
return {"X-Compute-Token": self.token}
|
|
|
|
async def test_connection(self) -> dict[str, Any]:
|
|
started = time.perf_counter()
|
|
health = await self.health()
|
|
gpus = await self.gpus()
|
|
return {
|
|
"success": True,
|
|
"latency_ms": int((time.perf_counter() - started) * 1000),
|
|
"health": health,
|
|
"gpus": gpus,
|
|
}
|
|
|
|
async def health(self) -> dict[str, Any]:
|
|
paths = [f"{self.route_prefix}/v1/compute/health", f"{self.route_prefix}/health", "/health"]
|
|
last_error = ""
|
|
async with httpx.AsyncClient(timeout=self.timeout, headers=self.headers()) as client:
|
|
for path in paths:
|
|
try:
|
|
response = await client.get(_join_url(self.api_base_url, path))
|
|
response.raise_for_status()
|
|
return _unwrap_dict(response.json())
|
|
except Exception as exc: # noqa: BLE001 - keep endpoint compatibility fallback broad
|
|
last_error = str(exc)
|
|
raise RuntimeError(last_error or "compute health check failed")
|
|
|
|
async def gpus(self) -> list[dict[str, Any]]:
|
|
paths = [
|
|
f"{self.route_prefix}/compute/resources/gpus",
|
|
f"{self.route_prefix}/v1/compute/resources/gpus",
|
|
"/compute/resources/gpus",
|
|
]
|
|
last_error = ""
|
|
async with httpx.AsyncClient(timeout=self.timeout, headers=self.headers()) as client:
|
|
for path in paths:
|
|
try:
|
|
response = await client.get(_join_url(self.api_base_url, path))
|
|
response.raise_for_status()
|
|
return _unwrap_items(response.json())
|
|
except Exception as exc: # noqa: BLE001
|
|
last_error = str(exc)
|
|
raise RuntimeError(last_error or "compute gpu discovery failed")
|
|
|
|
async def create_job(self, payload: dict[str, Any]) -> dict[str, Any]:
|
|
async with httpx.AsyncClient(timeout=self.timeout, headers=self.headers()) as client:
|
|
response = await client.post(_join_url(self.api_base_url, f"{self.route_prefix}/compute/jobs"), json=payload)
|
|
response.raise_for_status()
|
|
return _unwrap_dict(response.json())
|
|
|
|
async def preview_job(self, payload: dict[str, Any]) -> dict[str, Any]:
|
|
async with httpx.AsyncClient(timeout=self.timeout, headers=self.headers()) as client:
|
|
response = await client.post(
|
|
_join_url(self.api_base_url, f"{self.route_prefix}/compute/jobs/preview"),
|
|
json=payload,
|
|
)
|
|
response.raise_for_status()
|
|
return _unwrap_dict(response.json())
|
|
|
|
async def validate_job(self, payload: dict[str, Any]) -> dict[str, Any]:
|
|
async with httpx.AsyncClient(timeout=self.timeout, headers=self.headers()) as client:
|
|
response = await client.post(
|
|
_join_url(self.api_base_url, f"{self.route_prefix}/compute/jobs/validate"),
|
|
json=payload,
|
|
)
|
|
response.raise_for_status()
|
|
return _unwrap_dict(response.json())
|
|
|
|
async def check_paths(self, paths: list[dict[str, Any]]) -> dict[str, Any]:
|
|
async with httpx.AsyncClient(timeout=self.timeout, headers=self.headers()) as client:
|
|
response = await client.post(
|
|
_join_url(self.api_base_url, f"{self.route_prefix}/compute/files/check-paths"),
|
|
json={"paths": paths},
|
|
)
|
|
response.raise_for_status()
|
|
return _unwrap_dict(response.json())
|
|
|
|
async def list_files(
|
|
self,
|
|
root: str = "data",
|
|
relative_path: str = "",
|
|
directories_only: bool = False,
|
|
) -> dict[str, Any]:
|
|
async with httpx.AsyncClient(timeout=self.timeout, headers=self.headers()) as client:
|
|
response = await client.get(
|
|
_join_url(self.api_base_url, f"{self.route_prefix}/compute/files/list"),
|
|
params={"root": root, "relative_path": relative_path, "directories_only": directories_only},
|
|
)
|
|
response.raise_for_status()
|
|
return _unwrap_dict(response.json())
|
|
|
|
async def get_job(self, job_id: str) -> dict[str, Any]:
|
|
async with httpx.AsyncClient(timeout=self.timeout, headers=self.headers()) as client:
|
|
response = await client.get(_join_url(self.api_base_url, f"{self.route_prefix}/compute/jobs/{job_id}"))
|
|
response.raise_for_status()
|
|
return _unwrap_dict(response.json())
|
|
|
|
async def stop_job(self, job_id: str) -> dict[str, Any]:
|
|
# Compute API waits for the child process to exit (up to 10 seconds)
|
|
# before returning. The normal polling timeout is intentionally short,
|
|
# but is too aggressive for a stop request and used to surface as a
|
|
# platform 500 even when the node eventually stopped the job.
|
|
stop_timeout = max(float(self.timeout), 30.0)
|
|
async with httpx.AsyncClient(timeout=stop_timeout, headers=self.headers()) as client:
|
|
response = await client.post(_join_url(self.api_base_url, f"{self.route_prefix}/compute/jobs/{job_id}/stop"))
|
|
response.raise_for_status()
|
|
return _unwrap_dict(response.json())
|
|
|
|
async def job_logs(
|
|
self,
|
|
job_id: str,
|
|
tail_lines: int | None = None,
|
|
offset: int | None = None,
|
|
limit: int | None = None,
|
|
) -> dict[str, Any]:
|
|
params = {
|
|
key: value
|
|
for key, value in {"tail_lines": tail_lines, "offset": offset, "limit": limit}.items()
|
|
if value is not None
|
|
}
|
|
async with httpx.AsyncClient(timeout=self.timeout, headers=self.headers()) as client:
|
|
response = await client.get(
|
|
_join_url(self.api_base_url, f"{self.route_prefix}/compute/jobs/{job_id}/logs"),
|
|
params=params,
|
|
)
|
|
response.raise_for_status()
|
|
return _unwrap_dict(response.json())
|
|
|
|
async def import_local_file(self, payload: dict[str, Any]) -> dict[str, Any]:
|
|
async with httpx.AsyncClient(timeout=self.timeout, headers=self.headers()) as client:
|
|
response = await client.post(
|
|
_join_url(self.api_base_url, f"{self.route_prefix}/compute/files/import-local"),
|
|
json=payload,
|
|
)
|
|
response.raise_for_status()
|
|
return _unwrap_dict(response.json())
|
|
|
|
async def prepare_cache(self, payload: dict[str, Any]) -> dict[str, Any]:
|
|
async with httpx.AsyncClient(timeout=httpx.Timeout(900, connect=30), headers=self.headers()) as client:
|
|
response = await client.post(
|
|
_join_url(self.api_base_url, f"{self.route_prefix}/compute/cache/prepare"),
|
|
json=payload,
|
|
)
|
|
response.raise_for_status()
|
|
return _unwrap_dict(response.json())
|
|
|
|
async def cache_status(self, resource_id: str, version_id: str | None = None) -> dict[str, Any]:
|
|
params = {"resource_id": resource_id}
|
|
if version_id:
|
|
params["version_id"] = version_id
|
|
async with httpx.AsyncClient(timeout=self.timeout, headers=self.headers()) as client:
|
|
response = await client.get(
|
|
_join_url(self.api_base_url, f"{self.route_prefix}/compute/cache/status"),
|
|
params=params,
|
|
)
|
|
response.raise_for_status()
|
|
return _unwrap_dict(response.json())
|
|
|
|
async def _request(
|
|
self,
|
|
method: str,
|
|
path: str,
|
|
json_data: dict[str, Any] | None = None,
|
|
timeout: float | None = None,
|
|
) -> dict[str, Any]:
|
|
"""Generic request method for compute API endpoints."""
|
|
url = _join_url(self.api_base_url, f"{self.route_prefix}{path}")
|
|
async with httpx.AsyncClient(timeout=timeout or 300, headers=self.headers()) as client:
|
|
if method.upper() == "GET":
|
|
response = await client.get(url)
|
|
else:
|
|
response = await client.post(url, json=json_data)
|
|
response.raise_for_status()
|
|
return _unwrap_dict(response.json())
|
|
|
|
# ── Inference helpers (short timeouts — see module constants) ──────────
|
|
|
|
async def inference_load(self, payload: dict[str, Any]) -> dict[str, Any]:
|
|
"""Dispatch a model load. Returns as soon as the node accepts the
|
|
request; the node now loads asynchronously (status goes 'loading')."""
|
|
return await self._request("POST", "/inference/load", json_data=payload, timeout=INFERENCE_LOAD_TIMEOUT)
|
|
|
|
async def inference_status(self) -> dict[str, Any]:
|
|
return await self._request("GET", "/inference/status", timeout=INFERENCE_STATUS_TIMEOUT)
|
|
|
|
async def gpu_resources(self) -> list[dict[str, Any]]:
|
|
"""Read live per-GPU metrics from this compute node."""
|
|
return await self.gpus()
|
|
|
|
async def inference_unload(self) -> dict[str, Any]:
|
|
return await self._request("POST", "/inference/unload", json_data={}, timeout=INFERENCE_UNLOAD_TIMEOUT)
|
|
|
|
async def upload_file(
|
|
self,
|
|
filename: str,
|
|
content: bytes,
|
|
target_relative_path: str,
|
|
resource_type: str | None = None,
|
|
resource_id: str | None = None,
|
|
) -> dict[str, Any]:
|
|
data = {
|
|
"target_relative_path": target_relative_path,
|
|
"resource_type": resource_type or "",
|
|
"resource_id": resource_id or "",
|
|
}
|
|
files = {"file": (filename, content)}
|
|
timeout = httpx.Timeout(max(self.timeout, 60), connect=self.timeout)
|
|
async with httpx.AsyncClient(timeout=timeout, headers=self.headers()) as client:
|
|
response = await client.post(
|
|
_join_url(self.api_base_url, f"{self.route_prefix}/compute/files/upload"),
|
|
data=data,
|
|
files=files,
|
|
)
|
|
response.raise_for_status()
|
|
return _unwrap_dict(response.json())
|
|
|
|
async def upload_file_to_url(self, path: str, upload_url: str, object_key: str = "", content_type: str = "application/octet-stream") -> dict[str, Any]:
|
|
return await self._request("POST", "/compute/files/upload-to-url", json_data={
|
|
"path": path, "upload_url": upload_url, "object_key": object_key, "content_type": content_type,
|
|
}, timeout=900)
|