feat: 训练任务支持从最新 checkpoint 继续训练
- 后端: 新增 resume/checkpoints 端点,compute 客户端与适配器支持续训 - 计算节点: llama_factory 适配器断点续训支持及测试 - 前端: fine-tune API 封装与列表页续训/断点展示
This commit is contained in:
@@ -521,6 +521,117 @@ async def _prepare_resource_on_node(store: Any, resource_type: str, resource_id:
|
||||
return f"/data/yg-ft/{root_name}/{resource_id}"
|
||||
|
||||
|
||||
async def _archive_training_checkpoints(
|
||||
store: Any,
|
||||
task: dict[str, Any],
|
||||
node: dict[str, Any],
|
||||
job: dict[str, Any],
|
||||
) -> dict[str, Any]:
|
||||
"""Persist checkpoints from a stopped job so resume can use another node."""
|
||||
task_id = str(task.get("id") or "")
|
||||
output_dir = str(job.get("output_dir") or task.get("output_dir") or "")
|
||||
checkpoints = job.get("checkpoints") or store.task_checkpoints(task_id)
|
||||
if not task_id or not output_dir or not checkpoints:
|
||||
return store.update_task(
|
||||
task_id,
|
||||
{
|
||||
"checkpoint_archive_status": "not_available",
|
||||
"checkpoint_archive_error": "停止时没有发现可用 checkpoint",
|
||||
},
|
||||
)
|
||||
if not get_settings().minio_enabled:
|
||||
return store.update_task(
|
||||
task_id,
|
||||
{
|
||||
"checkpoint_archive_status": "local_only",
|
||||
"checkpoint_archive_error": "MinIO 未启用,仅支持原算力节点本地继续",
|
||||
},
|
||||
)
|
||||
try:
|
||||
version_id = str(job.get("id") or task.get("compute_job_id") or task_id)
|
||||
client = ComputeNodeClient(node["api_base_url"], timeout=900)
|
||||
archived = await _archive_node_directory(
|
||||
store,
|
||||
client,
|
||||
node,
|
||||
output_dir,
|
||||
"fine-tune",
|
||||
task_id,
|
||||
version_id,
|
||||
f"training/{task_id}/checkpoints",
|
||||
)
|
||||
return store.update_task(
|
||||
task_id,
|
||||
{
|
||||
"checkpoint_archive_status": "completed" if archived else "empty",
|
||||
"checkpoint_archive_version_id": version_id,
|
||||
"checkpoint_archive_object_ids": [str(item["id"]) for item in archived],
|
||||
"checkpoint_archive_error": "" if archived else "checkpoint 目录为空",
|
||||
},
|
||||
)
|
||||
except Exception as exc: # noqa: BLE001 - stopping must remain successful if archive is delayed
|
||||
return store.update_task(
|
||||
task_id,
|
||||
{
|
||||
"checkpoint_archive_status": "pending",
|
||||
"checkpoint_archive_error": str(exc)[:2000],
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
async def _prepare_resume_checkpoint(
|
||||
store: Any,
|
||||
task: dict[str, Any],
|
||||
checkpoint: dict[str, Any],
|
||||
node: dict[str, Any],
|
||||
) -> str | None:
|
||||
"""Materialize the selected checkpoint on the node used for resuming."""
|
||||
task_id = str(task.get("id") or "")
|
||||
checkpoint_name = str(checkpoint.get("name") or Path(str(checkpoint.get("path") or "")).name)
|
||||
if not task_id or not checkpoint_name:
|
||||
return None
|
||||
|
||||
# Prefer the archived copy. This makes resume independent of the node that
|
||||
# originally ran the task and preserves the MinIO-as-source-of-truth rule.
|
||||
if get_settings().minio_enabled:
|
||||
prefix = f"training/{task_id}/checkpoints/"
|
||||
objects = [
|
||||
item
|
||||
for item in store.storage_objects_for_resource("fine-tune", task_id)
|
||||
if str(item.get("object_key") or "").startswith(prefix)
|
||||
and str(item.get("file_name") or "").startswith(f"{checkpoint_name}/")
|
||||
]
|
||||
if objects:
|
||||
client = ComputeNodeClient(node["api_base_url"], timeout=900)
|
||||
root = f"/data/yg-ft/fine-tunes/{task_id}/resume/{checkpoint_name}"
|
||||
for item in objects:
|
||||
file_name = str(item.get("file_name") or "")
|
||||
relative = file_name[len(checkpoint_name) + 1 :]
|
||||
await client.prepare_cache(
|
||||
{
|
||||
"resource_id": task_id,
|
||||
"version_id": item.get("version_id"),
|
||||
"download_url": get_object_storage().presigned_get(str(item["object_key"])),
|
||||
"checksum_sha256": item.get("checksum_sha256") or "",
|
||||
"byte_size": item.get("byte_size") or 0,
|
||||
"relative_path": f"fine-tunes/{task_id}/resume/{checkpoint_name}/{relative}",
|
||||
}
|
||||
)
|
||||
return root
|
||||
|
||||
# Compatibility fallback for installations without MinIO archives: only
|
||||
# reuse the original path when the same node still owns the files.
|
||||
original_node_id = str(task.get("compute_node_id") or "")
|
||||
original_path = str(checkpoint.get("path") or "")
|
||||
if original_node_id == str(node.get("id") or "") and original_path:
|
||||
check = await ComputeNodeClient(node["api_base_url"]).check_paths(
|
||||
[{"path": original_path, "type": "dir", "required": True}]
|
||||
)
|
||||
if check.get("valid"):
|
||||
return original_path
|
||||
return None
|
||||
|
||||
|
||||
def _build_messages_payload(payload: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Convert frontend inference payload to compute API messages format.
|
||||
|
||||
@@ -2458,8 +2569,28 @@ async def stop_fine_tune(task_id: str, current_user: dict = Depends(get_current_
|
||||
return pending
|
||||
node = _node_for_task(task)
|
||||
if task.get("compute_job_id") and node and get_settings().compute_mode != "simulator":
|
||||
job = await ComputeNodeClient(node["api_base_url"]).stop_job(task["compute_job_id"])
|
||||
return ok(store.apply_compute_job(task_id, job))
|
||||
try:
|
||||
job = await ComputeNodeClient(
|
||||
node["api_base_url"],
|
||||
timeout=max(float(get_settings().compute_request_timeout_seconds), 30.0),
|
||||
).stop_job(task["compute_job_id"])
|
||||
except httpx.TimeoutException as exc:
|
||||
raise fail(504, "停止训练超时,算力节点未及时响应,请稍后查看任务状态") from exc
|
||||
except httpx.HTTPStatusError as exc:
|
||||
if exc.response.status_code == 404:
|
||||
# The process has already disappeared from the node. Make
|
||||
# the local task terminal and let the poller release state.
|
||||
return ok(store.stop_task(task_id))
|
||||
raise fail(502, f"算力节点拒绝停止训练:HTTP {exc.response.status_code}") from exc
|
||||
except httpx.RequestError as exc:
|
||||
raise fail(502, f"算力节点不可达,暂时无法停止训练:{exc}") from exc
|
||||
updated = store.apply_compute_job(task_id, job)
|
||||
# Stop is terminal for scheduling, but checkpoint archiving is
|
||||
# best-effort so a slow MinIO service never turns a successful stop
|
||||
# into a 500 response.
|
||||
return ok(await _archive_training_checkpoints(store, updated, node, job))
|
||||
if task.get("compute_job_id") and not node:
|
||||
raise fail(409, "训练任务关联的算力节点不存在,无法安全停止远程进程")
|
||||
return ok(store.stop_task(task_id))
|
||||
except KeyError:
|
||||
raise fail(404, "fine tune task not found")
|
||||
@@ -2470,6 +2601,76 @@ async def stop_fine_tune_alt(task_id: str, current_user: dict = Depends(get_curr
|
||||
return await stop_fine_tune(task_id, current_user)
|
||||
|
||||
|
||||
@router.post("/fine-tune/{task_id}/resume")
|
||||
@op_log(module=OpModule.FINE_TUNE, action=OpAction.RETRY, target_type="fine_tune", target_name_param="task_id")
|
||||
async def resume_fine_tune(
|
||||
task_id: str,
|
||||
payload: dict[str, Any] | None = Body(default=None),
|
||||
current_user: dict[str, Any] = Depends(get_current_user),
|
||||
) -> dict[str, Any]:
|
||||
"""Resume a stopped/failed training task from its newest checkpoint."""
|
||||
store = get_platform_store()
|
||||
payload = payload or {}
|
||||
try:
|
||||
task = store.task(task_id)
|
||||
except KeyError:
|
||||
raise fail(404, "fine tune task not found")
|
||||
if not has_resource_access("fine-tune", task_id, current_user, "execute"):
|
||||
raise fail(403, "no permission to resume this task")
|
||||
if task.get("status") not in {"stopped", "failed"}:
|
||||
raise fail(409, "只有已停止或失败的训练任务可以继续")
|
||||
|
||||
checkpoints = store.task_checkpoints(task_id)
|
||||
if not checkpoints:
|
||||
raise fail(409, "没有可用的训练断点,请确认任务停止前已经生成 checkpoint")
|
||||
checkpoint = max(checkpoints, key=lambda item: (int(item.get("step") or 0), str(item.get("create_time") or "")))
|
||||
resume_payload = {
|
||||
**task,
|
||||
**payload,
|
||||
"task_id": task_id,
|
||||
"id": task_id,
|
||||
"compute_node_id": payload.get("compute_node_id") or task.get("compute_node_id"),
|
||||
"gpus": payload.get("gpus") or task.get("gpus") or [],
|
||||
"strict_node_selection": True,
|
||||
"resume_checkpoint_id": checkpoint.get("id"),
|
||||
"resume_checkpoint_name": checkpoint.get("name"),
|
||||
"resume_from_checkpoint": checkpoint.get("path"),
|
||||
}
|
||||
try:
|
||||
# Select the requested/original node first, materialize the checkpoint,
|
||||
# then run the normal model/dataset/GPU preflight against that node.
|
||||
node, _ = store.prepare_compute_job_payload(task_id, resume_payload)
|
||||
if get_settings().compute_mode == "simulator":
|
||||
prepared_checkpoint = str(checkpoint.get("path") or "")
|
||||
else:
|
||||
prepared_checkpoint = await _prepare_resume_checkpoint(store, task, checkpoint, node)
|
||||
if not prepared_checkpoint:
|
||||
raise RuntimeError(
|
||||
f"断点 {checkpoint.get('name') or checkpoint.get('path')} 不可用:原算力节点文件已不存在,且 MinIO 中没有可恢复副本"
|
||||
)
|
||||
resume_payload["resume_from_checkpoint"] = prepared_checkpoint
|
||||
resume_payload["compute_node_id"] = node["id"]
|
||||
preflight = await _fine_tune_preflight(
|
||||
store,
|
||||
task_id,
|
||||
resume_payload,
|
||||
validate=True,
|
||||
sync_resources=True,
|
||||
)
|
||||
if not preflight.get("valid"):
|
||||
errors = "; ".join(preflight.get("errors") or ["继续训练预检未通过"])
|
||||
raise RuntimeError(errors)
|
||||
store.reset_task_for_retry(task_id, resume_payload)
|
||||
return ok(await _submit_fine_tune_task(store, resume_payload))
|
||||
except RuntimeError as exc:
|
||||
raise fail(409, f"继续训练预检未通过:{exc}") from exc
|
||||
except httpx.TimeoutException as exc:
|
||||
raise fail(504, "继续训练预检超时,请确认算力节点和 MinIO 服务可达") from exc
|
||||
except Exception as exc: # noqa: BLE001 - keep the resource diagnosis visible to the user
|
||||
store.mark_task_failed(task_id, str(exc))
|
||||
raise fail(502, f"继续训练失败:{exc}") from exc
|
||||
|
||||
|
||||
@router.post("/fine-tune/{task_id}/retry")
|
||||
@op_log(module=OpModule.FINE_TUNE, action=OpAction.RETRY, target_type="fine_tune", target_name_param="task_id")
|
||||
async def retry_fine_tune(task_id: str, payload: dict[str, Any] | None = Body(default=None), current_user: dict = Depends(get_current_user)) -> dict[str, Any]:
|
||||
|
||||
@@ -158,7 +158,12 @@ class ComputeNodeClient:
|
||||
return _unwrap_dict(response.json())
|
||||
|
||||
async def stop_job(self, job_id: str) -> dict[str, Any]:
|
||||
async with httpx.AsyncClient(timeout=self.timeout, headers=self.headers()) as client:
|
||||
# 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())
|
||||
|
||||
Reference in New Issue
Block a user