248 lines
8.2 KiB
Python
248 lines
8.2 KiB
Python
from __future__ import annotations
|
|
|
|
import re
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from . import db
|
|
|
|
|
|
ANSI_RE = re.compile(r"\x1b\[[0-?]*[ -/]*[@-~]")
|
|
ULTRALYTICS_EPOCH_RE = re.compile(r"^\s*(?P<current>\d+)\s*/\s*(?P<total>\d+)(?=\s)")
|
|
EPOCH_WORD_RE = re.compile(
|
|
r"\bepoch\b[^0-9]{0,12}(?P<current>\d+)\s*(?:/|of)\s*(?P<total>\d+)",
|
|
re.IGNORECASE,
|
|
)
|
|
MMENGINE_EPOCH_RE = re.compile(
|
|
r"Epoch\((?P<stage>train|val|test)\)\s*\[(?P<epoch>\d+)\]\s*\[(?P<current>\d+)\s*/\s*(?P<total>\d+)\]",
|
|
re.IGNORECASE,
|
|
)
|
|
MMENGINE_ITER_RE = re.compile(
|
|
r"Iter\((?P<stage>train|val|test)\)\s*\[(?P<current>\d+)\s*/\s*(?P<total>\d+)\]",
|
|
re.IGNORECASE,
|
|
)
|
|
TQDM_PERCENT_RE = re.compile(r"(?P<percent>\d{1,3}(?:\.\d+)?)%\|")
|
|
COMMAND_EPOCHS_RE = re.compile(r"(?:--epochs|--max-epochs|max_epochs)\D{0,20}(?P<total>\d+)", re.IGNORECASE)
|
|
|
|
|
|
def strip_ansi(value: str) -> str:
|
|
return ANSI_RE.sub("", value.replace("\r", "\n"))
|
|
|
|
|
|
def progress_from_log_path(log_path: str | Path, status: str | None = None, max_bytes: int = 65536) -> dict[str, Any]:
|
|
return parse_progress(db.log_tail(log_path, max_bytes=max_bytes), status=status)
|
|
|
|
|
|
def parse_progress(log_text: str, status: str | None = None) -> dict[str, Any]:
|
|
text = strip_ansi(log_text or "")
|
|
parsed = _parse_training_progress(text)
|
|
if parsed is None:
|
|
parsed = _status_progress(status)
|
|
|
|
if status == "success":
|
|
return {
|
|
**parsed,
|
|
"percent": 100.0,
|
|
"stage": "completed",
|
|
"label": _terminal_label("已完成", parsed),
|
|
"source": "status",
|
|
}
|
|
if status == "failed":
|
|
return {
|
|
**parsed,
|
|
"stage": "failed",
|
|
"label": _terminal_label("失败", parsed),
|
|
"source": parsed.get("source", "status"),
|
|
}
|
|
if status == "cancelled":
|
|
return {
|
|
**parsed,
|
|
"stage": "cancelled",
|
|
"label": _terminal_label("已取消", parsed),
|
|
"source": parsed.get("source", "status"),
|
|
}
|
|
return parsed
|
|
|
|
|
|
def _parse_training_progress(text: str) -> dict[str, Any] | None:
|
|
lines = [line.strip() for line in text.splitlines() if line.strip()]
|
|
if not lines:
|
|
return None
|
|
total_epochs = _find_total_epochs(text)
|
|
|
|
for line in reversed(lines[-400:]):
|
|
mmengine_epoch = MMENGINE_EPOCH_RE.search(line)
|
|
if mmengine_epoch:
|
|
epoch = int(mmengine_epoch.group("epoch"))
|
|
current = int(mmengine_epoch.group("current"))
|
|
total = int(mmengine_epoch.group("total"))
|
|
stage = _stage_name(mmengine_epoch.group("stage"))
|
|
if total_epochs:
|
|
percent = ((max(epoch - 1, 0) + current / max(total, 1)) / total_epochs) * 100
|
|
return _progress(
|
|
percent=percent,
|
|
label=f"Epoch {epoch}/{total_epochs} · iter {current}/{total}",
|
|
stage=stage,
|
|
current=epoch,
|
|
total=total_epochs,
|
|
unit="epoch",
|
|
source="mmengine_epoch",
|
|
)
|
|
return _progress(
|
|
percent=(current / max(total, 1)) * 100,
|
|
label=f"Epoch {epoch} · iter {current}/{total}",
|
|
stage=stage,
|
|
current=current,
|
|
total=total,
|
|
unit="iter",
|
|
source="mmengine_iter_in_epoch",
|
|
)
|
|
|
|
epoch_word = EPOCH_WORD_RE.search(line)
|
|
if epoch_word:
|
|
current = int(epoch_word.group("current"))
|
|
total = int(epoch_word.group("total"))
|
|
return _progress(
|
|
percent=(current / max(total, 1)) * 100,
|
|
label=f"Epoch {current}/{total}",
|
|
stage=_stage_from_line(line),
|
|
current=current,
|
|
total=total,
|
|
unit="epoch",
|
|
source="epoch_word",
|
|
)
|
|
|
|
ultralytics_epoch = ULTRALYTICS_EPOCH_RE.search(line)
|
|
if ultralytics_epoch:
|
|
current = int(ultralytics_epoch.group("current"))
|
|
total = int(ultralytics_epoch.group("total"))
|
|
if total > 1 or "loss" in line.lower() or "gpu" in line.lower():
|
|
return _progress(
|
|
percent=(current / max(total, 1)) * 100,
|
|
label=f"Epoch {current}/{total}",
|
|
stage=_stage_from_line(line),
|
|
current=current,
|
|
total=total,
|
|
unit="epoch",
|
|
source="ultralytics_epoch",
|
|
)
|
|
|
|
mmengine_iter = MMENGINE_ITER_RE.search(line)
|
|
if mmengine_iter:
|
|
current = int(mmengine_iter.group("current"))
|
|
total = int(mmengine_iter.group("total"))
|
|
return _progress(
|
|
percent=(current / max(total, 1)) * 100,
|
|
label=f"Iter {current}/{total}",
|
|
stage=_stage_name(mmengine_iter.group("stage")),
|
|
current=current,
|
|
total=total,
|
|
unit="iter",
|
|
source="mmengine_iter",
|
|
)
|
|
|
|
tqdm = TQDM_PERCENT_RE.search(line)
|
|
if tqdm:
|
|
percent = float(tqdm.group("percent"))
|
|
if 0 <= percent <= 100:
|
|
return _progress(
|
|
percent=percent,
|
|
label=f"{percent:g}%",
|
|
stage=_stage_from_line(line),
|
|
current=None,
|
|
total=None,
|
|
unit="percent",
|
|
source="tqdm_percent",
|
|
)
|
|
|
|
tail = "\n".join(lines[-20:]).lower()
|
|
if "results saved" in tail or "save_dir=" in tail:
|
|
return _progress(
|
|
percent=99.0,
|
|
label="保存结果",
|
|
stage="saving",
|
|
current=None,
|
|
total=None,
|
|
unit=None,
|
|
source="completion_hint",
|
|
)
|
|
if "starting training" in tail or "train:" in tail:
|
|
return _progress(
|
|
percent=None,
|
|
label="训练中",
|
|
stage="training",
|
|
current=None,
|
|
total=total_epochs,
|
|
unit="epoch" if total_epochs else None,
|
|
source="training_hint",
|
|
)
|
|
return None
|
|
|
|
|
|
def _find_total_epochs(text: str) -> int | None:
|
|
values = [int(match.group("total")) for match in COMMAND_EPOCHS_RE.finditer(text)]
|
|
return max(values) if values else None
|
|
|
|
|
|
def _status_progress(status: str | None) -> dict[str, Any]:
|
|
if status == "queued":
|
|
return _progress(0.0, "排队中", "queued", None, None, None, "status")
|
|
if status == "running":
|
|
return _progress(None, "运行中", "running", None, None, None, "status")
|
|
if status == "success":
|
|
return _progress(100.0, "已完成", "completed", None, None, None, "status")
|
|
if status == "failed":
|
|
return _progress(None, "失败", "failed", None, None, None, "status")
|
|
if status == "cancelled":
|
|
return _progress(None, "已取消", "cancelled", None, None, None, "status")
|
|
return _progress(None, "未开始", "unknown", None, None, None, "status")
|
|
|
|
|
|
def _progress(
|
|
percent: float | None,
|
|
label: str,
|
|
stage: str,
|
|
current: int | None,
|
|
total: int | None,
|
|
unit: str | None,
|
|
source: str,
|
|
) -> dict[str, Any]:
|
|
if percent is not None:
|
|
percent = round(max(0.0, min(100.0, float(percent))), 2)
|
|
return {
|
|
"percent": percent,
|
|
"label": label,
|
|
"stage": stage,
|
|
"current": current,
|
|
"total": total,
|
|
"unit": unit,
|
|
"source": source,
|
|
}
|
|
|
|
|
|
def _terminal_label(prefix: str, parsed: dict[str, Any]) -> str:
|
|
label = parsed.get("label")
|
|
if label and label not in {prefix, "运行中", "排队中", "未开始"}:
|
|
return f"{prefix} · {label}"
|
|
return prefix
|
|
|
|
|
|
def _stage_from_line(line: str) -> str:
|
|
lower = line.lower()
|
|
if "val" in lower or "validat" in lower:
|
|
return "validation"
|
|
if "predict" in lower or "infer" in lower:
|
|
return "prediction"
|
|
if "save" in lower:
|
|
return "saving"
|
|
return "training"
|
|
|
|
|
|
def _stage_name(value: str) -> str:
|
|
lower = value.lower()
|
|
if lower == "val":
|
|
return "validation"
|
|
if lower == "test":
|
|
return "evaluation"
|
|
return "training"
|