From 314578590a19506bae35620c04f3fb7acc26c9d7 Mon Sep 17 00:00:00 2001 From: Teakowa Date: Sat, 26 Sep 2026 19:11:13 +0800 Subject: [PATCH 01/15] feat(training): add Colab GPU recognition runner Fixes #20 --- training/README.md | 62 +++- training/bootstrap.sh | 7 +- training/colab_remote.py | 287 ++++++++++++++++++ training/run_rec_colab.py | 471 ++++++++++++++++++++++++++++++ training/run_rec_smoke.sh | 27 +- training/setup_rec_environment.sh | 87 ++++-- 6 files changed, 916 insertions(+), 25 deletions(-) create mode 100644 training/colab_remote.py create mode 100644 training/run_rec_colab.py diff --git a/training/README.md b/training/README.md index d6521c7..90b7f85 100644 --- a/training/README.md +++ b/training/README.md @@ -398,10 +398,12 @@ Prepare the offline environment once, then run the CPU recognition Smoke: `run_rec_smoke.sh` accepts `--labels-dir`, `--output-dir`, `--epochs` (the target total epoch), and `--resume-checkpoint` (a checkpoint base path without -`.pdparams`, `.pdopt`, or `.states`). It validates both label files, fine-tunes -recognition only, and leaves `latest` plus `best_accuracy` under -`training/.work/`. Per-epoch `iter_epoch_*` dumps and PaddleOCR's duplicate -`best_model/` copy are pruned after training. To reclaim space from older runs: +`.pdparams`, `.pdopt`, or `.states`). `--device cpu|cuda` selects PaddleOCR's +device and defaults to `cpu`; `--evaluation-dir` selects the evaluation output +directory. It validates both label files, fine-tunes recognition only, and +leaves `latest` plus `best_accuracy` under `training/.work/`. Per-epoch +`iter_epoch_*` dumps and PaddleOCR's duplicate `best_model/` copy are pruned +after training. To reclaim space from older runs: ```bash uv run python training/scripts/prune_rec_checkpoints.py --root training/.work @@ -419,6 +421,58 @@ The current release gate is field accuracy at least `364/379` (`0.9604221635883905`). A failed Smoke keeps the checkpoint and evaluation report for inspection but returns a non-zero status. +### Run recognition training on Colab GPU + +Install the [Google Colab CLI](https://github.com/googlecolab/google-colab-cli) +on the local Mac and prepare the reviewed/materialized dataset plus the local +PP-OCRv6 small base checkpoint: + +```bash +uv tool install google-colab-cli +./training/setup_rec_environment.sh +``` + +The first CLI request can prompt for Google OAuth authentication in the +terminal. The CLI keeps those credentials locally; OCRKit does not send +platform, R2, or release credentials to Colab. + +Start training and the existing checkpoint evaluation with one command: + +```bash +uv run python training/run_rec_colab.py +``` + +The default dataset is `datasets/labeled/rec`. To train from a materialized +platform snapshot, select that snapshot's output directory explicitly: + +```bash +uv run python training/run_rec_colab.py \ + --labels-dir datasets/labeled/rec/platform/@ \ + --gpu L4 \ + --epochs 10 +``` + +`--gpu` is a Colab allocation preference (default `L4`), not a model or +training requirement. PaddlePaddle's CUDA runtime and device are checked before +training; an unavailable or unsupported GPU request fails without falling back +to CPU. The local CPU command remains `./training/run_rec_smoke.sh`. + +The runner transfers only the selected train/holdout labels and referenced +crops, available review/snapshot provenance files, the base recognition +checkpoint, OCRKit training/evaluation code, and the fixture images needed by +the existing evaluation. It records source revisions, input checksums, the +effective training configuration, PaddleOCR revision, allocated GPU details, +checkpoint, and evaluation report in the returned `run.json`. + +Checkpoints, `fixture_report.json`, provenance, the remote training log, the +Colab CLI log, and `status.json` are stored below the ignored +`training/.work/colab-runs//` directory. A successful run does not +publish a candidate or change the stable model channel. Provisioning, +staging, training, evaluation, retrieval, and handled failures all stop the +Colab runtime after it has been allocated. Failed runs keep diagnostics and +any partial output under `partial/`; if teardown itself fails, `status.json` +includes the named `colab stop` command to release that session. + To evaluate a checkpoint explicitly, use a new output directory: ```bash diff --git a/training/bootstrap.sh b/training/bootstrap.sh index 02b1de5..9a474f6 100644 --- a/training/bootstrap.sh +++ b/training/bootstrap.sh @@ -15,7 +15,12 @@ printf 'PaddleOCR checkout: %s\n' "${paddleocr_dir}" venv_dir="${work_dir}/venv" if [[ ! -x "${venv_dir}/bin/python" ]]; then - python3.12 -m venv "${venv_dir}" + python_bootstrap="${OCRKIT_TRAINING_PYTHON:-python3.12}" + if ! command -v "${python_bootstrap}" >/dev/null; then + printf 'Python 3.12 or OCRKIT_TRAINING_PYTHON is required to create the training environment.\n' >&2 + exit 1 + fi + "${python_bootstrap}" -m venv "${venv_dir}" fi printf 'Training environment: %s\n' "${venv_dir}" diff --git a/training/colab_remote.py b/training/colab_remote.py new file mode 100644 index 0000000..6174453 --- /dev/null +++ b/training/colab_remote.py @@ -0,0 +1,287 @@ +#!/usr/bin/env python3 +from __future__ import annotations + +import hashlib +import json +import os +import shlex +import shutil +import subprocess +import sys +import tarfile +import traceback +from datetime import UTC, datetime +from pathlib import Path +from typing import Any + +COLAB_ROOT = Path("/content/ocrkit-colab") +INPUT_ARCHIVE = Path("/content/ocrkit-input.tar.gz") +RESULT_ARCHIVE = Path("/content/ocrkit-result.tar.gz") +REPO = COLAB_ROOT / "repo" +DATASET = COLAB_ROOT / "dataset" +RESULTS = COLAB_ROOT / "results" +CHECKPOINTS = RESULTS / "checkpoint" +EVALUATION = RESULTS / "evaluation" +RUN_METADATA = RESULTS / "run.json" +REMOTE_LOG = RESULTS / "remote.log" + + +def sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def safe_extract(archive_path: Path, destination: Path) -> None: + root = destination.resolve() + with tarfile.open(archive_path, "r:gz") as archive: + members = archive.getmembers() + for member in members: + target = (destination / member.name).resolve() + if not target.is_relative_to(root) or not (member.isfile() or member.isdir()): + raise RuntimeError("Colab input archive contains an unsupported path or file type") + archive.extractall(destination) + + +def run_logged(command: list[str], *, cwd: Path, log: Any, env: dict[str, str] | None = None) -> None: + rendered = shlex.join(command) + print(f"$ {rendered}", flush=True) + log.write(f"$ {rendered}\n") + log.flush() + with subprocess.Popen( + command, + cwd=cwd, + env=env, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + bufsize=1, + ) as process: + assert process.stdout is not None + for line in process.stdout: + print(line, end="", flush=True) + log.write(line) + status = process.wait() + log.flush() + if status: + raise RuntimeError(f"remote command exited with status {status}: {rendered}") + + +def checked_gpu() -> dict[str, Any]: + if not shutil.which("nvidia-smi"): + raise RuntimeError("Colab did not provide an NVIDIA GPU runtime") + result = subprocess.run( + [ + "nvidia-smi", + "--query-gpu=name,compute_cap,driver_version,memory.total", + "--format=csv,noheader,nounits", + ], + check=False, + capture_output=True, + text=True, + ) + if result.returncode or not result.stdout.strip(): + raise RuntimeError("Colab could not report an allocated NVIDIA GPU") + devices = [] + for line in result.stdout.splitlines(): + name, capability, driver, memory = (part.strip() for part in line.split(",", 3)) + major, minor = (int(part) for part in capability.split(".", 1)) + devices.append( + { + "name": name, + "compute_capability": capability, + "driver_version": driver, + "memory_mib": int(memory), + "supported_by_paddle": (major, minor) > (7, 5), + } + ) + if not devices or not devices[0]["supported_by_paddle"]: + allocated = devices[0]["compute_capability"] if devices else "unknown" + raise RuntimeError( + f"allocated GPU compute capability {allocated} is unsupported by PaddlePaddle 3.3.1; " + "a capability greater than 7.5 is required" + ) + return {"devices": devices, "nvidia_smi": subprocess.run( + ["nvidia-smi"], check=False, capture_output=True, text=True + ).stdout} + + +def verify_inputs(request: dict[str, Any], root: Path) -> None: + for record in request["input_files"]: + path = root / record["path"] + if not path.is_file() or path.stat().st_size != record["size_bytes"] or sha256(path) != record["sha256"]: + raise RuntimeError(f"staged input failed checksum verification: {record['path']}") + + +def output_files() -> list[dict[str, Any]]: + records = [] + for directory in (CHECKPOINTS, EVALUATION): + if not directory.is_dir(): + continue + for path in sorted(directory.rglob("*")): + if path.is_file(): + records.append( + { + "path": path.relative_to(RESULTS).as_posix(), + "size_bytes": path.stat().st_size, + "sha256": sha256(path), + } + ) + return records + + +def write_result_archive() -> None: + RESULT_ARCHIVE.unlink(missing_ok=True) + with tarfile.open(RESULT_ARCHIVE, "w:gz") as archive: + for path in (RUN_METADATA, REMOTE_LOG): + if path.is_file(): + archive.add(path, arcname=path.relative_to(COLAB_ROOT)) + for directory in (CHECKPOINTS, EVALUATION): + if directory.is_dir(): + archive.add(directory, arcname=directory.relative_to(COLAB_ROOT)) + + +def main() -> int: + RESULTS.mkdir(parents=True, exist_ok=True) + result: dict[str, Any] = { + "schema_version": 1, + "status": "failed", + "completed_at": datetime.now(UTC).isoformat(), + } + try: + with REMOTE_LOG.open("w", encoding="utf-8") as log: + try: + if not INPUT_ARCHIVE.is_file(): + raise RuntimeError("OCRKit input archive was not uploaded to the Colab runtime") + COLAB_ROOT.mkdir(parents=True, exist_ok=True) + safe_extract(INPUT_ARCHIVE, COLAB_ROOT) + request_path = COLAB_ROOT / "request.json" + request = json.loads(request_path.read_text(encoding="utf-8")) + result["request"] = request["run"] + verify_inputs(request, COLAB_ROOT) + gpu = checked_gpu() + REPO.mkdir(parents=True, exist_ok=True) + DATASET.mkdir(parents=True, exist_ok=True) + CHECKPOINTS.mkdir(parents=True, exist_ok=True) + + if not shutil.which("uv"): + run_logged( + [sys.executable, "-m", "pip", "install", "uv"], + cwd=COLAB_ROOT, + log=log, + ) + run_logged( + ["uv", "sync", "--locked", "--no-dev", "--python", sys.executable], + cwd=REPO, + log=log, + ) + environment = os.environ.copy() + environment["OCRKIT_TRAINING_PYTHON"] = sys.executable + run_logged( + ["bash", "training/setup_rec_environment.sh", "--device", "cuda"], + cwd=REPO, + env=environment, + log=log, + ) + + paddle = subprocess.run( + [ + str(REPO / "training/.work/venv/bin/python"), + "-c", + "import json, paddle; paddle.device.set_device('gpu:0'); paddle.to_tensor([1.0]).numpy(); print(json.dumps({'version': paddle.__version__, 'cuda': paddle.is_compiled_with_cuda(), 'device': paddle.device.get_device()}))", + ], + check=True, + capture_output=True, + text=True, + ) + paddle_info = json.loads(paddle.stdout.strip().splitlines()[-1]) + if not paddle_info["cuda"] or paddle_info["device"] != "gpu:0": + raise RuntimeError("PaddlePaddle did not select the allocated CUDA device") + + epochs = str(request["run"]["epochs"]) + run_logged( + [ + "bash", + "training/run_rec_smoke.sh", + "--labels-dir", + str(DATASET), + "--output-dir", + str(CHECKPOINTS), + "--evaluation-dir", + str(EVALUATION), + "--epochs", + epochs, + "--device", + "cuda", + ], + cwd=REPO, + log=log, + ) + best_checkpoint = CHECKPOINTS / "best_accuracy.pdparams" + report_path = EVALUATION / "fixture_report.json" + if not best_checkpoint.is_file() or best_checkpoint.stat().st_size == 0: + raise RuntimeError("training did not produce the best-accuracy recognition checkpoint") + if not report_path.is_file(): + raise RuntimeError("checkpoint evaluation did not produce fixture_report.json") + evaluation_report = json.loads(report_path.read_text(encoding="utf-8")) + paddleocr_revision = subprocess.run( + ["git", "-C", str(REPO / "training/.work/PaddleOCR"), "rev-parse", "HEAD"], + check=True, + capture_output=True, + text=True, + ).stdout.strip() + result.update( + { + "status": "success", + "runtime": { + **gpu, + "python": sys.version, + "paddle": paddle_info, + "paddleocr_revision": paddleocr_revision, + }, + "evaluation": { + "report": "evaluation/fixture_report.json", + "field_accuracy": evaluation_report.get("field_accuracy"), + "run_code_accuracy": evaluation_report.get("run_code", {}).get("field_accuracy"), + }, + "outputs": output_files(), + } + ) + except BaseException as exc: + result["error"] = f"{type(exc).__name__}: {exc}" + traceback.print_exc(file=log) + log.flush() + finally: + result["completed_at"] = datetime.now(UTC).isoformat() + result["outputs"] = output_files() + RUN_METADATA.write_text( + json.dumps(result, ensure_ascii=False, indent=2) + "\n", + encoding="utf-8", + ) + write_result_archive() + except BaseException as exc: + print(f"OCRKit Colab run failed: {type(exc).__name__}: {exc}", file=sys.stderr, flush=True) + try: + RUN_METADATA.write_text( + json.dumps( + {"schema_version": 1, "status": "failed", "error": f"{type(exc).__name__}: {exc}"}, + ensure_ascii=False, + indent=2, + ) + + "\n", + encoding="utf-8", + ) + write_result_archive() + except OSError: + pass + return 1 + if result.get("status") != "success": + print(f"OCRKit Colab run failed: {result.get('error', 'remote step failed')}", file=sys.stderr, flush=True) + return 1 + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/training/run_rec_colab.py b/training/run_rec_colab.py new file mode 100644 index 0000000..1b8bbcd --- /dev/null +++ b/training/run_rec_colab.py @@ -0,0 +1,471 @@ +#!/usr/bin/env python3 +from __future__ import annotations + +import argparse +import hashlib +import json +import os +import shlex +import shutil +import subprocess +import sys +import tarfile +from datetime import UTC, datetime +from pathlib import Path, PurePosixPath, PureWindowsPath +from typing import Any + +ROOT = Path(__file__).resolve().parents[1] +RUNS = ROOT / "training/.work/colab-runs" +REMOTE_RUNNER = ROOT / "training/colab_remote.py" +PRETRAINED_CHECKPOINT = ROOT / "training/.work/pretrained/PP-OCRv6_small_rec_pretrained.pdparams" +SOURCE_FILES = ( + "pyproject.toml", + "uv.lock", + "scripts/batch_eval.py", + "training/bootstrap.sh", + "training/setup_rec_environment.sh", + "training/run_rec_smoke.sh", + "training/evaluate_rec_checkpoint.sh", + "training/configs/rec_pp_ocrv6_small.yaml", + "training/configs/pp_ocrv6_small_det.lock.json", + "training/scripts/prepare_detector.py", + "training/scripts/prepare_rapidocr_config.py", + "training/scripts/prune_rec_checkpoints.py", + "training/scripts/validate_annotations.py", + "training/colab_remote.py", +) + + +def sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def add_file(files: dict[str, tuple[Path, str]], archive_path: str, source: Path, kind: str) -> None: + if source.is_symlink() or not source.is_file(): + raise ValueError(f"required Colab input is missing or is a symbolic link: {source}") + if archive_path in files: + if files[archive_path][0] != source: + raise ValueError(f"two Colab inputs map to the same path: {archive_path}") + return + files[archive_path] = (source, kind) + + +def safe_relative(value: str) -> Path: + path = Path(value) + if ( + not value + or "\\" in value + or path.is_absolute() + or PureWindowsPath(value).is_absolute() + or any(part in ("", ".", "..") for part in PurePosixPath(value).parts) + ): + raise ValueError(f"training input contains an unsafe image path: {value!r}") + return path + + +def validate_labels(label_path: Path) -> tuple[int, list[str]]: + result = subprocess.run( + [sys.executable, str(ROOT / "training/scripts/validate_annotations.py"), "rec", str(label_path)], + cwd=ROOT, + check=True, + capture_output=True, + text=True, + ) + sample_count = int(json.loads(result.stdout)["valid_samples"]) + if sample_count < 1: + raise ValueError(f"recognition split is empty: {label_path.name}") + images = [] + for line in label_path.read_text(encoding="utf-8").splitlines(): + if not line.strip(): + continue + image_name = line.split("\t", 1)[0] + relative = safe_relative(image_name) + candidates = (label_path.parent / "images" / relative, label_path.parent.parent / relative) + image_path = next((candidate for candidate in candidates if candidate.is_file()), None) + if image_path is None: + raise ValueError(f"validated recognition image disappeared: {image_name}") + image_path = image_path.resolve() + dataset_root = label_path.parent.parent.resolve() + if not image_path.is_relative_to(dataset_root): + raise ValueError(f"recognition image resolves outside the selected dataset: {image_name}") + images.append(image_name) + return sample_count, images + + +def add_fixture_set(files: dict[str, tuple[Path, str]], relative_cases: str) -> None: + cases_path = ROOT / relative_cases + add_file(files, f"repo/{relative_cases}", cases_path, "evaluation-fixture") + cases = json.loads(cases_path.read_text(encoding="utf-8")) + for case in cases: + image_path = safe_relative(case["image"]) + source = cases_path.parent / image_path + add_file( + files, + f"repo/{cases_path.parent.relative_to(ROOT).as_posix()}/{image_path.as_posix()}", + source, + "evaluation-fixture", + ) + + +def git_metadata(path: Path) -> tuple[str | None, bool]: + revision = subprocess.run( + ["git", "-C", str(path), "rev-parse", "HEAD"], + capture_output=True, + text=True, + check=False, + ) + if revision.returncode: + return None, False + status = subprocess.run( + ["git", "-C", str(path), "status", "--porcelain", "--untracked-files=normal"], + capture_output=True, + text=True, + check=False, + ) + return revision.stdout.strip(), bool(status.stdout.strip()) + + +def read_dataset_provenance(dataset_root: Path) -> dict[str, Any]: + for name in ("provenance.json", "snapshot.json"): + path = dataset_root / name + if not path.is_file(): + continue + data = json.loads(path.read_text(encoding="utf-8")) + snapshot = data.get("snapshot", {}) + snapshot_id = snapshot.get("snapshot_id") or data.get("snapshot_id") + if snapshot_id: + return { + "snapshot_id": snapshot_id, + "snapshot_version": snapshot.get("version") or data.get("snapshot_version") or data.get("version"), + "code_revision": data.get("code_revision"), + } + return {} + + +def source_files(root: Path, files: dict[str, tuple[Path, str]]) -> None: + for directory in ("app", "configs"): + for path in sorted((root / directory).rglob("*")): + if path.is_file() and not path.is_symlink() and "__pycache__" not in path.parts: + add_file(files, f"repo/{path.relative_to(root).as_posix()}", path, "ocrkit-source") + for relative in SOURCE_FILES: + add_file(files, f"repo/{relative}", root / relative, "ocrkit-source") + add_fixture_set(files, "datasets/fixtures/challenge/cases.json") + add_fixture_set(files, "tests/fixtures/run_code/cases.json") + + +def stage_inputs( + archive_path: Path, + run_request: dict[str, Any], + dataset_root: Path, + checkpoint_path: Path, + train_images: list[str], + holdout_images: list[str], +) -> None: + files: dict[str, tuple[Path, str]] = {} + source_files(ROOT, files) + + for relative in ("labels/train.txt", "labels/holdout.txt"): + add_file(files, f"dataset/{relative}", dataset_root / relative, "reviewed-labels") + image_names = set(train_images + holdout_images) + for image_name in sorted(image_names): + relative = safe_relative(image_name) + label_parent = dataset_root / "labels" + candidates = (label_parent / "images" / relative, dataset_root / relative) + image_path = next((candidate for candidate in candidates if candidate.is_file()), None) + if image_path is None: + raise ValueError(f"training crop is missing: {image_name}") + add_file(files, f"dataset/{relative.as_posix()}", image_path, "reviewed-crop") + + for relative in ( + "provenance.json", + "snapshot.json", + "crop_manifest.json", + "review/train.jsonl", + "review/holdout.jsonl", + ): + path = dataset_root / relative + if path.is_file(): + add_file(files, f"dataset/{relative}", path, "dataset-provenance") + add_file( + files, + "repo/training/.work/pretrained/PP-OCRv6_small_rec_pretrained.pdparams", + checkpoint_path, + "base-recognition-checkpoint", + ) + + records = [] + with tarfile.open(archive_path, "w:gz") as archive: + for path, (source, kind) in sorted(files.items()): + archive.add(source, arcname=path, recursive=False) + records.append( + { + "path": path, + "kind": kind, + "size_bytes": source.stat().st_size, + "sha256": sha256(source), + } + ) + request = {"run": run_request, "input_files": records} + request_path = archive_path.parent / "request.json" + request_path.write_text(json.dumps(request, ensure_ascii=False, indent=2) + "\n", encoding="utf-8") + archive.add(request_path, arcname="request.json", recursive=False) + + +def stream_colab(command: list[str], log_path: Path) -> int: + with log_path.open("a", encoding="utf-8") as log: + rendered = shlex.join(command) + print(f"$ {rendered}", flush=True) + log.write(f"$ {rendered}\n") + with subprocess.Popen( + command, + cwd=ROOT, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + bufsize=1, + ) as process: + assert process.stdout is not None + try: + for line in process.stdout: + print(line, end="", flush=True) + log.write(line) + except KeyboardInterrupt: + process.terminate() + process.wait() + raise + return process.wait() + + +def safe_extract_result(archive_path: Path, destination: Path) -> None: + destination.mkdir(parents=True, exist_ok=True) + root = destination.resolve() + with tarfile.open(archive_path, "r:gz") as archive: + members = archive.getmembers() + for member in members: + target = (destination / member.name).resolve() + if not target.is_relative_to(root) or not (member.isfile() or member.isdir()): + raise ValueError("Colab result archive contains an unsupported path or file type") + archive.extractall(destination) + + +def copy_remote_result(remote_root: Path, run_dir: Path, *, success: bool) -> dict[str, Any]: + results = remote_root / "results" + metadata_path = results / "run.json" + remote_metadata = json.loads(metadata_path.read_text(encoding="utf-8")) if metadata_path.is_file() else None + remote_log = results / "remote.log" + if success: + if remote_metadata is None or remote_metadata.get("status") != "success": + raise ValueError("Colab did not return successful run metadata") + for name in ("checkpoint", "evaluation"): + if not (results / name).is_dir(): + raise ValueError(f"Colab did not return the {name} artifacts") + output_records = remote_metadata.get("outputs") + required_outputs = { + "checkpoint/best_accuracy.pdparams", + "evaluation/fixture_report.json", + } + if not isinstance(output_records, list) or not required_outputs.issubset( + {record.get("path") for record in output_records} + ): + raise ValueError("Colab did not checksum the required checkpoint and evaluation report") + for record in output_records: + output = results / safe_relative(record["path"]) + if ( + not output.is_file() + or output.stat().st_size != record["size_bytes"] + or sha256(output) != record["sha256"] + ): + raise ValueError(f"Colab result failed checksum verification: {record['path']}") + if not (results / "checkpoint/best_accuracy.pdparams").is_file(): + raise ValueError("Colab result is missing the best-accuracy checkpoint") + report_path = results / "evaluation/fixture_report.json" + if not report_path.is_file(): + raise ValueError("Colab result is missing the fixture evaluation report") + report = json.loads(report_path.read_text(encoding="utf-8")) + if ( + report.get("field_accuracy", 0.0) < 0.9604221635883905 + or report.get("run_code", {}).get("field_accuracy", 0.0) < 1.0 + ): + raise ValueError("Colab result does not meet the recognition evaluation contract") + for name in ("checkpoint", "evaluation"): + shutil.copytree(results / name, run_dir / name) + shutil.copy2(metadata_path, run_dir / "run.json") + if remote_log.is_file(): + shutil.copy2(remote_log, run_dir / "remote.log") + return remote_metadata + + partial = run_dir / "partial" + partial.mkdir(exist_ok=True) + for name in ("checkpoint", "evaluation"): + source = results / name + if source.is_dir(): + shutil.copytree(source, partial / name) + if metadata_path.is_file(): + shutil.copy2(metadata_path, partial / "run.json") + if remote_log.is_file(): + shutil.copy2(remote_log, partial / "remote.log") + return remote_metadata or {} + + +def main() -> int: + parser = argparse.ArgumentParser(description="Run the OCRKit recognition training and evaluation workflow on Colab GPU.") + parser.add_argument("--labels-dir", type=Path, default=ROOT / "datasets/labeled/rec") + parser.add_argument("--pretrained-checkpoint", type=Path, default=PRETRAINED_CHECKPOINT) + parser.add_argument("--epochs", type=int, default=10) + parser.add_argument("--gpu", default="L4", help="Colab GPU preference; no accelerator fallback is attempted") + args = parser.parse_args() + if args.epochs < 1: + parser.error("--epochs must be a positive integer") + colab = shutil.which("colab") + if not colab: + parser.error("Google Colab CLI is required; install it with uv tool install google-colab-cli") + + dataset_root = args.labels_dir.resolve() + train_label = dataset_root / "labels/train.txt" + holdout_label = dataset_root / "labels/holdout.txt" + for path in (train_label, holdout_label, args.pretrained_checkpoint): + if not path.is_file(): + parser.error(f"required training input is missing: {path}") + train_count, train_images = validate_labels(train_label) + holdout_count, holdout_images = validate_labels(holdout_label) + + run_id = datetime.now(UTC).strftime("%Y%m%dT%H%M%SZ") + f"-{os.getpid():x}" + run_dir = RUNS / run_id + run_dir.mkdir(parents=True, exist_ok=False) + input_archive = run_dir / "ocrkit-input.tar.gz" + result_archive = run_dir / "ocrkit-result.tar.gz" + colab_log = run_dir / "colab.log" + retrieved = run_dir / "retrieved" + + source_revision, source_dirty = git_metadata(ROOT) + dataset_revision, dataset_dirty = git_metadata(dataset_root) + try: + dataset_source = dataset_root.relative_to(ROOT).as_posix() + except ValueError: + dataset_source = "external" + checkpoint_path = args.pretrained_checkpoint.resolve() + run_request = { + "run_id": run_id, + "ocrkit_revision": source_revision, + "ocrkit_worktree_dirty": source_dirty, + "dataset": { + **read_dataset_provenance(dataset_root), + "source": dataset_source, + "revision": dataset_revision, + "worktree_dirty": dataset_dirty, + "train_samples": train_count, + "holdout_samples": holdout_count, + }, + "base_checkpoint": { + "model": "PP-OCRv6_small_rec_pretrained", + "source_name": checkpoint_path.name, + "sha256": sha256(checkpoint_path), + }, + "training": { + "epochs": args.epochs, + "device": "cuda", + "gpu_preference": args.gpu, + "recipe": "training/configs/rec_pp_ocrv6_small.yaml", + "paddleocr_recipe": "configs/rec/PP-OCRv6/PP-OCRv6_small_rec.yml", + "batch_size_per_card": 8, + "first_batch_size": 8, + "evaluation_batch_step": [0, 1], + "save_epoch_step": args.epochs + 1, + "minimum_field_accuracy": 0.9604221635883905, + "minimum_run_code_accuracy": 1.0, + }, + } + try: + stage_inputs(input_archive, run_request, dataset_root, checkpoint_path, train_images, holdout_images) + except Exception as exc: + input_archive.unlink(missing_ok=True) + (run_dir / "status.json").write_text( + json.dumps({"status": "failed", "error": f"{type(exc).__name__}: {exc}"}, indent=2) + "\n", + encoding="utf-8", + ) + raise + + session = f"ocrkit-rec-{run_id.lower()}" + session_attempted = False + remote_metadata: dict[str, Any] = {} + error: str | None = None + stop_status: int | None = None + exec_status: int | None = None + + try: + session_attempted = True + provision_status = stream_colab([colab, "new", "-s", session, "--gpu", args.gpu], colab_log) + if provision_status: + raise RuntimeError( + f"Colab could not provision the requested GPU {args.gpu}; no fallback accelerator was selected." + ) + upload_status = stream_colab( + [colab, "upload", "-s", session, str(input_archive), "/content/ocrkit-input.tar.gz"], + colab_log, + ) + input_archive.unlink(missing_ok=True) + if upload_status: + raise RuntimeError("Colab failed to stage OCRKit training inputs") + exec_status = stream_colab( + [colab, "exec", "-s", session, "-f", str(REMOTE_RUNNER)], + colab_log, + ) + download_status = stream_colab( + [colab, "download", "-s", session, "/content/ocrkit-result.tar.gz", str(result_archive)], + colab_log, + ) + if download_status: + raise RuntimeError("Colab failed to retrieve the remote run logs and artifacts") + safe_extract_result(result_archive, retrieved) + try: + remote_metadata = copy_remote_result(retrieved, run_dir, success=exec_status == 0) + except Exception: + try: + copy_remote_result(retrieved, run_dir, success=False) + except Exception: + pass + raise + if exec_status: + raise RuntimeError(f"Colab training or evaluation failed with exit status {exec_status}.") + if remote_metadata.get("status") != "success": + raise RuntimeError("Colab training or evaluation did not return a successful status.") + except KeyboardInterrupt: + error = "Colab run interrupted by the operator." + except Exception as exc: + error = f"{type(exc).__name__}: {exc}" + finally: + input_archive.unlink(missing_ok=True) + if session_attempted: + try: + stop_status = stream_colab([colab, "stop", "-s", session], colab_log) + except Exception as exc: + stop_status = -1 + stop_error = f"Colab runtime teardown command failed: {type(exc).__name__}: {exc}" + error = f"{error}; {stop_error}" if error else stop_error + if stop_status and error is None: + error = f"training completed, but Colab runtime teardown failed; run colab stop -s {session}" + + status = { + "status": "success" if error is None and stop_status in (None, 0) else "failed", + "runtime_stopped": stop_status == 0, + "colab_session": session, + "error": error, + } + (run_dir / "status.json").write_text(json.dumps(status, ensure_ascii=False, indent=2) + "\n", encoding="utf-8") + result_archive.unlink(missing_ok=True) + if retrieved.exists(): + shutil.rmtree(retrieved) + if status["status"] != "success": + print(f"Colab OCRKit run failed: {error or 'runtime did not stop successfully'}", file=sys.stderr) + print(f"Local run logs and any partial outputs: {run_dir}", file=sys.stderr) + return 1 + print(f"Colab OCRKit run completed. Checkpoint, evaluation, provenance, and logs: {run_dir}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/training/run_rec_smoke.sh b/training/run_rec_smoke.sh index 887f35e..673289b 100755 --- a/training/run_rec_smoke.sh +++ b/training/run_rec_smoke.sh @@ -10,6 +10,8 @@ config_path="${paddleocr_dir}/configs/rec/PP-OCRv6/PP-OCRv6_small_rec.yml" output_dir="${work_dir}/checkpoints/rec_pp_ocrv6_small" epoch_num=10 resume_checkpoint="" +device=cpu +evaluation_dir="" while [[ $# -gt 0 ]]; do case "$1" in @@ -29,13 +31,30 @@ while [[ $# -gt 0 ]]; do resume_checkpoint="$2" shift 2 ;; + --device) + device="$2" + shift 2 + ;; + --evaluation-dir) + evaluation_dir="$2" + shift 2 + ;; *) - printf 'usage: %s [--labels-dir ] [--output-dir ] [--epochs ] [--resume-checkpoint ]\n' "$0" >&2 + printf 'usage: %s [--labels-dir ] [--output-dir ] [--epochs ] [--resume-checkpoint ] [--device cpu|cuda] [--evaluation-dir ]\n' "$0" >&2 exit 2 ;; esac done +case "$device" in + cpu) use_gpu=False ;; + cuda) use_gpu=True ;; + *) + printf 'device must be cpu or cuda\n' >&2 + exit 2 + ;; +esac + if ! [[ "${epoch_num}" =~ ^[1-9][0-9]*$ ]]; then printf 'epochs must be a positive integer\n' >&2 exit 2 @@ -64,7 +83,7 @@ fi cd "${paddleocr_dir}" "${python_bin}" tools/train.py -c "${config_path}" -o \ - Global.use_gpu=False \ + Global.use_gpu="$use_gpu" \ "Global.epoch_num=${epoch_num}" \ Global.save_model_dir="${output_dir}" \ Global.save_epoch_step="$((epoch_num + 1))" \ @@ -85,7 +104,9 @@ cd "${paddleocr_dir}" cd "${root_dir}" "${python_bin}" "${root_dir}/training/scripts/prune_rec_checkpoints.py" "${output_dir}" -evaluation_dir="${work_dir}/evaluations/rec_pp_ocrv6_small/$(date -u +%Y.%m.%d-%H%M%S)-$$" +if [[ -z "$evaluation_dir" ]]; then + evaluation_dir="$work_dir/evaluations/rec_pp_ocrv6_small/$(date -u +%Y.%m.%d-%H%M%S)-$$" +fi "${root_dir}/training/evaluate_rec_checkpoint.sh" \ "${output_dir}/best_accuracy" \ "${evaluation_dir}" diff --git a/training/setup_rec_environment.sh b/training/setup_rec_environment.sh index 5bbb089..69153c6 100755 --- a/training/setup_rec_environment.sh +++ b/training/setup_rec_environment.sh @@ -1,16 +1,29 @@ #!/usr/bin/env bash set -euo pipefail -root_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" -work_dir="${root_dir}/training/.work" -paddleocr_dir="${work_dir}/PaddleOCR" -python_bin="${work_dir}/venv/bin/python" -pretrained_dir="${work_dir}/pretrained" -pretrained_model="${pretrained_dir}/PP-OCRv6_small_rec_pretrained.pdparams" +root_dir="$(cd "$(dirname "$0")/.." && pwd)" +work_dir="$root_dir/training/.work" +paddleocr_dir="$work_dir/PaddleOCR" +python_bin="$work_dir/venv/bin/python" +pretrained_dir="$work_dir/pretrained" +pretrained_model="$pretrained_dir/PP-OCRv6_small_rec_pretrained.pdparams" +device=cpu -bash "${root_dir}/training/bootstrap.sh" +if [[ $# -gt 0 ]]; then + if [[ $# -ne 2 || "$1" != "--device" ]]; then + printf 'usage: %s [--device cpu|cuda]\n' "$0" >&2 + exit 2 + fi + device="$2" +fi +if [[ "$device" != cpu && "$device" != cuda ]]; then + printf 'device must be cpu or cuda\n' >&2 + exit 2 +fi -if [[ "$(uname -s)" == "Darwin" ]]; then +bash "$root_dir/training/bootstrap.sh" + +if [[ "$device" == cpu && "$(uname -s)" == Darwin ]]; then if ! command -v brew >/dev/null; then printf 'Homebrew is required to install ccache on macOS. Install Homebrew, then rerun this script.\n' >&2 exit 1 @@ -24,16 +37,56 @@ if [[ "$(uname -s)" == "Darwin" ]]; then fi fi -"${python_bin}" -m pip install --upgrade pip -"${python_bin}" -m pip install "paddlepaddle==3.3.1" -i https://www.paddlepaddle.org.cn/packages/stable/cpu/ -"${python_bin}" -m pip install -r "${paddleocr_dir}/requirements.txt" -"${python_bin}" -m pip install "paddle2onnx==2.1.0" +paddle_package="paddlepaddle==3.3.1" +paddle_index="https://www.paddlepaddle.org.cn/packages/stable/cpu/" +if [[ "$device" == cuda ]]; then + if [[ "$(uname -s)" != Linux ]] || ! command -v nvidia-smi >/dev/null; then + printf 'CUDA training requires a Linux runtime with an NVIDIA GPU and nvidia-smi.\n' >&2 + exit 1 + fi + + cuda_version="$(nvidia-smi 2>/dev/null | awk -F 'CUDA Version: ' 'NF > 1 {split($2, version, " "); print version[1]; exit}')" + if [[ -z "$cuda_version" ]]; then + printf 'could not determine the NVIDIA driver CUDA version from nvidia-smi.\n' >&2 + exit 1 + fi + cuda_major="$(cut -d. -f1 <<< "$cuda_version")" + cuda_minor="$(cut -d. -f2 <<< "$cuda_version")" + if (( cuda_major > 12 || (cuda_major == 12 && cuda_minor >= 9) )); then + paddle_index="https://www.paddlepaddle.org.cn/packages/stable/cu129/" + elif (( cuda_major > 12 || (cuda_major == 12 && cuda_minor >= 6) )); then + paddle_index="https://www.paddlepaddle.org.cn/packages/stable/cu126/" + elif (( cuda_major > 11 || (cuda_major == 11 && cuda_minor >= 8) )); then + paddle_index="https://www.paddlepaddle.org.cn/packages/stable/cu118/" + else + printf 'NVIDIA CUDA %s is unsupported; PaddlePaddle 3.3.1 needs CUDA 11.8 or newer.\n' "$cuda_version" >&2 + exit 1 + fi + + gpu_capabilities="$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader 2>/dev/null || true)" + if [[ -z "$gpu_capabilities" ]] || ! awk -F. 'NF == 2 && (($1 + 0) > 7 || (($1 + 0) == 7 && ($2 + 0) > 5)) {supported=1} END {exit !supported}' <<< "$gpu_capabilities"; then + printf 'the allocated GPU is not supported by PaddlePaddle 3.3.1 (compute capability must be greater than 7.5).\n' >&2 + exit 1 + fi + paddle_package="paddlepaddle-gpu==3.3.1" +fi + +"$python_bin" -m pip install --upgrade pip +"$python_bin" -m pip install "$paddle_package" -i "$paddle_index" +"$python_bin" -m pip install -r "$paddleocr_dir/requirements.txt" +"$python_bin" -m pip install "paddle2onnx==2.1.0" + +if [[ "$device" == cpu && "$(uname -s)" == Darwin ]]; then + "$python_bin" -c 'import platform; import paddle; assert platform.machine() == "arm64"; assert paddle.device.get_device() == "cpu"; assert not paddle.is_compiled_with_cuda(); print(f"PaddlePaddle {paddle.__version__}: {platform.machine()} {paddle.device.get_device()}")' +fi -if [[ "$(uname -s)" == "Darwin" ]]; then - "${python_bin}" -c 'import platform; import paddle; assert platform.machine() == "arm64"; assert paddle.device.get_device() == "cpu"; assert not paddle.is_compiled_with_cuda(); print(f"PaddlePaddle {paddle.__version__}: {platform.machine()} {paddle.device.get_device()}")' +if [[ "$device" == cuda ]]; then + "$python_bin" -m pip uninstall -y paddlepaddle paddlepaddle-gpu >/dev/null 2>&1 || true + "$python_bin" -m pip install --no-deps "$paddle_package" -i "$paddle_index" + "$python_bin" -c 'import paddle; assert paddle.is_compiled_with_cuda(), "installed PaddlePaddle is not CUDA-enabled"; assert paddle.device.cuda.device_count() > 0, "no CUDA device is visible to PaddlePaddle"; paddle.device.set_device("gpu:0"); value = paddle.to_tensor([1.0]).numpy()[0]; assert value == 1.0; print(f"PaddlePaddle {paddle.__version__}: CUDA {paddle.device.get_device()}")' fi -mkdir -p "${pretrained_dir}" -if [[ ! -f "${pretrained_model}" ]]; then - "${python_bin}" -c "from urllib.request import urlretrieve; urlretrieve('https://paddle-model-ecology.bj.bcebos.com/paddlex/official_pretrained_model/PP-OCRv6_small_rec_pretrained.pdparams', '${pretrained_model}')" +mkdir -p "$pretrained_dir" +if [[ ! -f "$pretrained_model" ]]; then + "$python_bin" -c "from urllib.request import urlretrieve; urlretrieve('https://paddle-model-ecology.bj.bcebos.com/paddlex/official_pretrained_model/PP-OCRv6_small_rec_pretrained.pdparams', '$pretrained_model')" fi From c5bea27a0908322600a4aafd028b3e3065c4018d Mon Sep 17 00:00:00 2001 From: e54-bot Date: Sat, 26 Sep 2026 20:25:16 +0800 Subject: [PATCH 02/15] fix(training): address Colab runner review findings Stage returned artifacts and promote them to the run directory only after runtime teardown succeeds; otherwise keep them under partial/. Drop the duplicated accuracy gates and hardcoded training values so the shared evaluation script stays the single source of truth. Restore the unrelated setup script churn, install Paddle once for CUDA, and surface the exec failure when result retrieval also fails. Add mocked runner tests. Refs #20 --- tests/test_run_rec_colab.py | 143 ++++++++++++++++++++++++++++++ training/run_rec_colab.py | 41 +++++---- training/run_rec_smoke.sh | 6 +- training/setup_rec_environment.sh | 60 ++++++------- 4 files changed, 199 insertions(+), 51 deletions(-) create mode 100644 tests/test_run_rec_colab.py diff --git a/tests/test_run_rec_colab.py b/tests/test_run_rec_colab.py new file mode 100644 index 0000000..8b5f0bc --- /dev/null +++ b/tests/test_run_rec_colab.py @@ -0,0 +1,143 @@ +from __future__ import annotations + +import hashlib +import json +import sys +import tarfile +from pathlib import Path + +import pytest + +from training import run_rec_colab + +SESSION_STOP = "stop" + + +def build_result_archive(path: Path, *, status: str = "success") -> None: + files = { + "results/checkpoint/best_accuracy.pdparams": b"weights", + "results/evaluation/fixture_report.json": b'{"field_accuracy": 0.99}', + "results/remote.log": b"log", + } + outputs = [ + { + "path": name.removeprefix("results/"), + "size_bytes": len(data), + "sha256": hashlib.sha256(data).hexdigest(), + } + for name, data in files.items() + if not name.endswith("remote.log") + ] + files["results/run.json"] = json.dumps({"status": status, "outputs": outputs}).encode() + with tarfile.open(path, "w:gz") as archive: + for name, data in files.items(): + source = path.parent / "member" + source.write_bytes(data) + archive.add(source, arcname=name) + + +@pytest.fixture +def colab_run(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + labels = tmp_path / "labels-dir" + (labels / "labels").mkdir(parents=True) + for name in ("train.txt", "holdout.txt"): + (labels / "labels" / name).write_text("a.png\ta\n", encoding="utf-8") + checkpoint = tmp_path / "base.pdparams" + checkpoint.write_bytes(b"base") + runs = tmp_path / "runs" + calls: list[list[str]] = [] + behavior = {"exec_status": 0, "result_status": "success", "stop_status": 0} + + def fake_stream(command: list[str], _log: Path) -> int: + calls.append(command) + verb = command[1] + if verb == "exec": + return behavior["exec_status"] + if verb == "download": + build_result_archive(Path(command[-1]), status=behavior["result_status"]) + return 0 + if verb == SESSION_STOP: + return behavior["stop_status"] + return 0 + + monkeypatch.setattr(run_rec_colab, "RUNS", runs) + monkeypatch.setattr(run_rec_colab.shutil, "which", lambda _name: "colab") + monkeypatch.setattr(run_rec_colab, "validate_labels", lambda _path: (1, ["a.png"])) + monkeypatch.setattr(run_rec_colab, "stage_inputs", lambda archive, *_a: archive.write_bytes(b"input")) + monkeypatch.setattr(run_rec_colab, "stream_colab", fake_stream) + monkeypatch.setattr( + sys, "argv", ["run_rec_colab.py", "--labels-dir", str(labels), "--pretrained-checkpoint", str(checkpoint)] + ) + return runs, calls, behavior + + +def only_run(runs: Path) -> Path: + (run_dir,) = runs.iterdir() + return run_dir + + +def test_success_keeps_verified_artifacts_and_stops_runtime(colab_run) -> None: + runs, calls, _ = colab_run + + assert run_rec_colab.main() == 0 + + run_dir = only_run(runs) + assert (run_dir / "checkpoint/best_accuracy.pdparams").is_file() + assert (run_dir / "evaluation/fixture_report.json").is_file() + assert json.loads((run_dir / "status.json").read_text())["runtime_stopped"] is True + assert calls[-1][1] == SESSION_STOP + assert not (run_dir / "accepted").exists() + + +def test_training_failure_keeps_partial_output_and_stops_runtime(colab_run) -> None: + runs, calls, behavior = colab_run + behavior.update(exec_status=1, result_status="failed") + + assert run_rec_colab.main() == 1 + + run_dir = only_run(runs) + assert not (run_dir / "checkpoint").exists() + assert not (run_dir / "run.json").exists() + assert (run_dir / "partial/checkpoint/best_accuracy.pdparams").is_file() + assert calls[-1][1] == SESSION_STOP + assert json.loads((run_dir / "status.json").read_text())["status"] == "failed" + + +def test_teardown_failure_demotes_accepted_artifacts_to_partial(colab_run) -> None: + runs, _, behavior = colab_run + behavior["stop_status"] = 1 + + assert run_rec_colab.main() == 1 + + run_dir = only_run(runs) + assert not (run_dir / "checkpoint").exists() + assert (run_dir / "partial/checkpoint/best_accuracy.pdparams").is_file() + assert "colab stop" in json.loads((run_dir / "status.json").read_text())["error"] + + +def test_provisioning_failure_stops_runtime_and_uploads_nothing(colab_run) -> None: + runs, calls, _ = colab_run + original = run_rec_colab.stream_colab + run_rec_colab.stream_colab = lambda command, log: 1 if command[1] == "new" else original(command, log) + try: + assert run_rec_colab.main() == 1 + finally: + run_rec_colab.stream_colab = original + + assert [command[1] for command in calls] == [SESSION_STOP] + assert not (only_run(runs) / "checkpoint").exists() + + +def test_tampered_result_is_rejected(tmp_path: Path) -> None: + remote = tmp_path / "retrieved" + archive = tmp_path / "result.tar.gz" + build_result_archive(archive) + run_rec_colab.safe_extract_result(archive, remote) + (remote / "results/checkpoint/best_accuracy.pdparams").write_bytes(b"tampered") + run_dir = tmp_path / "run" + run_dir.mkdir() + + with pytest.raises(ValueError, match="checksum"): + run_rec_colab.copy_remote_result(remote, run_dir, success=True) + + assert not (run_dir / "accepted").exists() diff --git a/training/run_rec_colab.py b/training/run_rec_colab.py index 1b8bbcd..1b9abe2 100644 --- a/training/run_rec_colab.py +++ b/training/run_rec_colab.py @@ -16,6 +16,7 @@ ROOT = Path(__file__).resolve().parents[1] RUNS = ROOT / "training/.work/colab-runs" +ACCEPTED_STAGING = "accepted" REMOTE_RUNNER = ROOT / "training/colab_remote.py" PRETRAINED_CHECKPOINT = ROOT / "training/.work/pretrained/PP-OCRv6_small_rec_pretrained.pdparams" SOURCE_FILES = ( @@ -285,17 +286,12 @@ def copy_remote_result(remote_root: Path, run_dir: Path, *, success: bool) -> di report_path = results / "evaluation/fixture_report.json" if not report_path.is_file(): raise ValueError("Colab result is missing the fixture evaluation report") - report = json.loads(report_path.read_text(encoding="utf-8")) - if ( - report.get("field_accuracy", 0.0) < 0.9604221635883905 - or report.get("run_code", {}).get("field_accuracy", 0.0) < 1.0 - ): - raise ValueError("Colab result does not meet the recognition evaluation contract") + accepted = run_dir / ACCEPTED_STAGING for name in ("checkpoint", "evaluation"): - shutil.copytree(results / name, run_dir / name) - shutil.copy2(metadata_path, run_dir / "run.json") + shutil.copytree(results / name, accepted / name) + shutil.copy2(metadata_path, accepted / "run.json") if remote_log.is_file(): - shutil.copy2(remote_log, run_dir / "remote.log") + shutil.copy2(remote_log, accepted / "remote.log") return remote_metadata partial = run_dir / "partial" @@ -371,12 +367,6 @@ def main() -> int: "gpu_preference": args.gpu, "recipe": "training/configs/rec_pp_ocrv6_small.yaml", "paddleocr_recipe": "configs/rec/PP-OCRv6/PP-OCRv6_small_rec.yml", - "batch_size_per_card": 8, - "first_batch_size": 8, - "evaluation_batch_step": [0, 1], - "save_epoch_step": args.epochs + 1, - "minimum_field_accuracy": 0.9604221635883905, - "minimum_run_code_accuracy": 1.0, }, } try: @@ -419,11 +409,15 @@ def main() -> int: colab_log, ) if download_status: - raise RuntimeError("Colab failed to retrieve the remote run logs and artifacts") + raise RuntimeError( + "Colab failed to retrieve the remote run logs and artifacts" + + (f" after the remote run exited with status {exec_status}" if exec_status else "") + ) safe_extract_result(result_archive, retrieved) try: remote_metadata = copy_remote_result(retrieved, run_dir, success=exec_status == 0) except Exception: + shutil.rmtree(run_dir / ACCEPTED_STAGING, ignore_errors=True) try: copy_remote_result(retrieved, run_dir, success=False) except Exception: @@ -449,8 +443,21 @@ def main() -> int: if stop_status and error is None: error = f"training completed, but Colab runtime teardown failed; run colab stop -s {session}" + succeeded = error is None and stop_status in (None, 0) + accepted = run_dir / ACCEPTED_STAGING + if accepted.is_dir(): + if succeeded: + for child in accepted.iterdir(): + child.rename(run_dir / child.name) + accepted.rmdir() + else: + partial = run_dir / "partial" + partial.mkdir(exist_ok=True) + for child in accepted.iterdir(): + child.rename(partial / child.name) + accepted.rmdir() status = { - "status": "success" if error is None and stop_status in (None, 0) else "failed", + "status": "success" if succeeded else "failed", "runtime_stopped": stop_status == 0, "colab_session": session, "error": error, diff --git a/training/run_rec_smoke.sh b/training/run_rec_smoke.sh index 673289b..115e444 100755 --- a/training/run_rec_smoke.sh +++ b/training/run_rec_smoke.sh @@ -83,7 +83,7 @@ fi cd "${paddleocr_dir}" "${python_bin}" tools/train.py -c "${config_path}" -o \ - Global.use_gpu="$use_gpu" \ + "Global.use_gpu=${use_gpu}" \ "Global.epoch_num=${epoch_num}" \ Global.save_model_dir="${output_dir}" \ Global.save_epoch_step="$((epoch_num + 1))" \ @@ -104,8 +104,8 @@ cd "${paddleocr_dir}" cd "${root_dir}" "${python_bin}" "${root_dir}/training/scripts/prune_rec_checkpoints.py" "${output_dir}" -if [[ -z "$evaluation_dir" ]]; then - evaluation_dir="$work_dir/evaluations/rec_pp_ocrv6_small/$(date -u +%Y.%m.%d-%H%M%S)-$$" +if [[ -z "${evaluation_dir}" ]]; then + evaluation_dir="${work_dir}/evaluations/rec_pp_ocrv6_small/$(date -u +%Y.%m.%d-%H%M%S)-$$" fi "${root_dir}/training/evaluate_rec_checkpoint.sh" \ "${output_dir}/best_accuracy" \ diff --git a/training/setup_rec_environment.sh b/training/setup_rec_environment.sh index 69153c6..e1776f9 100755 --- a/training/setup_rec_environment.sh +++ b/training/setup_rec_environment.sh @@ -1,14 +1,14 @@ #!/usr/bin/env bash set -euo pipefail -root_dir="$(cd "$(dirname "$0")/.." && pwd)" -work_dir="$root_dir/training/.work" -paddleocr_dir="$work_dir/PaddleOCR" -python_bin="$work_dir/venv/bin/python" -pretrained_dir="$work_dir/pretrained" -pretrained_model="$pretrained_dir/PP-OCRv6_small_rec_pretrained.pdparams" -device=cpu +root_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +work_dir="${root_dir}/training/.work" +paddleocr_dir="${work_dir}/PaddleOCR" +python_bin="${work_dir}/venv/bin/python" +pretrained_dir="${work_dir}/pretrained" +pretrained_model="${pretrained_dir}/PP-OCRv6_small_rec_pretrained.pdparams" +device=cpu if [[ $# -gt 0 ]]; then if [[ $# -ne 2 || "$1" != "--device" ]]; then printf 'usage: %s [--device cpu|cuda]\n' "$0" >&2 @@ -16,14 +16,14 @@ if [[ $# -gt 0 ]]; then fi device="$2" fi -if [[ "$device" != cpu && "$device" != cuda ]]; then +if [[ "${device}" != cpu && "${device}" != cuda ]]; then printf 'device must be cpu or cuda\n' >&2 exit 2 fi -bash "$root_dir/training/bootstrap.sh" +bash "${root_dir}/training/bootstrap.sh" -if [[ "$device" == cpu && "$(uname -s)" == Darwin ]]; then +if [[ "${device}" == cpu && "$(uname -s)" == "Darwin" ]]; then if ! command -v brew >/dev/null; then printf 'Homebrew is required to install ccache on macOS. Install Homebrew, then rerun this script.\n' >&2 exit 1 @@ -39,54 +39,52 @@ fi paddle_package="paddlepaddle==3.3.1" paddle_index="https://www.paddlepaddle.org.cn/packages/stable/cpu/" -if [[ "$device" == cuda ]]; then +if [[ "${device}" == cuda ]]; then if [[ "$(uname -s)" != Linux ]] || ! command -v nvidia-smi >/dev/null; then printf 'CUDA training requires a Linux runtime with an NVIDIA GPU and nvidia-smi.\n' >&2 exit 1 fi cuda_version="$(nvidia-smi 2>/dev/null | awk -F 'CUDA Version: ' 'NF > 1 {split($2, version, " "); print version[1]; exit}')" - if [[ -z "$cuda_version" ]]; then + if [[ -z "${cuda_version}" ]]; then printf 'could not determine the NVIDIA driver CUDA version from nvidia-smi.\n' >&2 exit 1 fi - cuda_major="$(cut -d. -f1 <<< "$cuda_version")" - cuda_minor="$(cut -d. -f2 <<< "$cuda_version")" + cuda_major="$(cut -d. -f1 <<< "${cuda_version}")" + cuda_minor="$(cut -d. -f2 <<< "${cuda_version}")" if (( cuda_major > 12 || (cuda_major == 12 && cuda_minor >= 9) )); then paddle_index="https://www.paddlepaddle.org.cn/packages/stable/cu129/" - elif (( cuda_major > 12 || (cuda_major == 12 && cuda_minor >= 6) )); then + elif (( cuda_major == 12 && cuda_minor >= 6 )); then paddle_index="https://www.paddlepaddle.org.cn/packages/stable/cu126/" - elif (( cuda_major > 11 || (cuda_major == 11 && cuda_minor >= 8) )); then + elif (( cuda_major == 11 && cuda_minor >= 8 )) || (( cuda_major == 12 )); then paddle_index="https://www.paddlepaddle.org.cn/packages/stable/cu118/" else - printf 'NVIDIA CUDA %s is unsupported; PaddlePaddle 3.3.1 needs CUDA 11.8 or newer.\n' "$cuda_version" >&2 + printf 'NVIDIA CUDA %s is unsupported; PaddlePaddle 3.3.1 needs CUDA 11.8 or newer.\n' "${cuda_version}" >&2 exit 1 fi gpu_capabilities="$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader 2>/dev/null || true)" - if [[ -z "$gpu_capabilities" ]] || ! awk -F. 'NF == 2 && (($1 + 0) > 7 || (($1 + 0) == 7 && ($2 + 0) > 5)) {supported=1} END {exit !supported}' <<< "$gpu_capabilities"; then + if [[ -z "${gpu_capabilities}" ]] || ! awk -F. 'NF == 2 && (($1 + 0) > 7 || (($1 + 0) == 7 && ($2 + 0) > 5)) {supported=1} END {exit !supported}' <<< "${gpu_capabilities}"; then printf 'the allocated GPU is not supported by PaddlePaddle 3.3.1 (compute capability must be greater than 7.5).\n' >&2 exit 1 fi paddle_package="paddlepaddle-gpu==3.3.1" fi -"$python_bin" -m pip install --upgrade pip -"$python_bin" -m pip install "$paddle_package" -i "$paddle_index" -"$python_bin" -m pip install -r "$paddleocr_dir/requirements.txt" -"$python_bin" -m pip install "paddle2onnx==2.1.0" +"${python_bin}" -m pip install --upgrade pip +"${python_bin}" -m pip install "${paddle_package}" -i "${paddle_index}" +"${python_bin}" -m pip install -r "${paddleocr_dir}/requirements.txt" +"${python_bin}" -m pip install "paddle2onnx==2.1.0" -if [[ "$device" == cpu && "$(uname -s)" == Darwin ]]; then - "$python_bin" -c 'import platform; import paddle; assert platform.machine() == "arm64"; assert paddle.device.get_device() == "cpu"; assert not paddle.is_compiled_with_cuda(); print(f"PaddlePaddle {paddle.__version__}: {platform.machine()} {paddle.device.get_device()}")' +if [[ "${device}" == cpu && "$(uname -s)" == "Darwin" ]]; then + "${python_bin}" -c 'import platform; import paddle; assert platform.machine() == "arm64"; assert paddle.device.get_device() == "cpu"; assert not paddle.is_compiled_with_cuda(); print(f"PaddlePaddle {paddle.__version__}: {platform.machine()} {paddle.device.get_device()}")' fi -if [[ "$device" == cuda ]]; then - "$python_bin" -m pip uninstall -y paddlepaddle paddlepaddle-gpu >/dev/null 2>&1 || true - "$python_bin" -m pip install --no-deps "$paddle_package" -i "$paddle_index" - "$python_bin" -c 'import paddle; assert paddle.is_compiled_with_cuda(), "installed PaddlePaddle is not CUDA-enabled"; assert paddle.device.cuda.device_count() > 0, "no CUDA device is visible to PaddlePaddle"; paddle.device.set_device("gpu:0"); value = paddle.to_tensor([1.0]).numpy()[0]; assert value == 1.0; print(f"PaddlePaddle {paddle.__version__}: CUDA {paddle.device.get_device()}")' +if [[ "${device}" == cuda ]]; then + "${python_bin}" -c 'import paddle; assert paddle.is_compiled_with_cuda(), "installed PaddlePaddle is not CUDA-enabled"; assert paddle.device.cuda.device_count() > 0, "no CUDA device is visible to PaddlePaddle"; paddle.device.set_device("gpu:0"); assert paddle.to_tensor([1.0]).numpy()[0] == 1.0; print(f"PaddlePaddle {paddle.__version__}: CUDA {paddle.device.get_device()}")' fi -mkdir -p "$pretrained_dir" -if [[ ! -f "$pretrained_model" ]]; then - "$python_bin" -c "from urllib.request import urlretrieve; urlretrieve('https://paddle-model-ecology.bj.bcebos.com/paddlex/official_pretrained_model/PP-OCRv6_small_rec_pretrained.pdparams', '$pretrained_model')" +mkdir -p "${pretrained_dir}" +if [[ ! -f "${pretrained_model}" ]]; then + "${python_bin}" -c "from urllib.request import urlretrieve; urlretrieve('https://paddle-model-ecology.bj.bcebos.com/paddlex/official_pretrained_model/PP-OCRv6_small_rec_pretrained.pdparams', '${pretrained_model}')" fi From 80a6463ae735c5a21944c618997b5101eceb774b Mon Sep 17 00:00:00 2001 From: e54-bot Date: Sat, 26 Sep 2026 20:28:37 +0800 Subject: [PATCH 03/15] fix(training): bound the Colab exec timeout colab exec defaults to a 30 second timeout, which would abort training. Pass an explicit --timeout-seconds bound and report the remote error when the run metadata is not successful, since colab exec does not map a remote SystemExit to its exit status. Refs #20 --- tests/test_run_rec_colab.py | 2 ++ training/README.md | 3 ++- training/run_rec_colab.py | 6 ++++-- 3 files changed, 8 insertions(+), 3 deletions(-) diff --git a/tests/test_run_rec_colab.py b/tests/test_run_rec_colab.py index 8b5f0bc..b9f73db 100644 --- a/tests/test_run_rec_colab.py +++ b/tests/test_run_rec_colab.py @@ -81,6 +81,8 @@ def test_success_keeps_verified_artifacts_and_stops_runtime(colab_run) -> None: assert run_rec_colab.main() == 0 + exec_command = next(command for command in calls if command[1] == "exec") + assert float(exec_command[exec_command.index("--timeout") + 1]) == 6 * 3600 run_dir = only_run(runs) assert (run_dir / "checkpoint/best_accuracy.pdparams").is_file() assert (run_dir / "evaluation/fixture_report.json").is_file() diff --git a/training/README.md b/training/README.md index 90b7f85..ca4702a 100644 --- a/training/README.md +++ b/training/README.md @@ -455,7 +455,8 @@ uv run python training/run_rec_colab.py \ `--gpu` is a Colab allocation preference (default `L4`), not a model or training requirement. PaddlePaddle's CUDA runtime and device are checked before training; an unavailable or unsupported GPU request fails without falling back -to CPU. The local CPU command remains `./training/run_rec_smoke.sh`. +to CPU. `--timeout-seconds` (default 6 hours) bounds the remote run; the Colab +CLI's own `exec` default of 30 seconds is always overridden. The local CPU command remains `./training/run_rec_smoke.sh`. The runner transfers only the selected train/holdout labels and referenced crops, available review/snapshot provenance files, the base recognition diff --git a/training/run_rec_colab.py b/training/run_rec_colab.py index 1b9abe2..74f7c5d 100644 --- a/training/run_rec_colab.py +++ b/training/run_rec_colab.py @@ -260,7 +260,8 @@ def copy_remote_result(remote_root: Path, run_dir: Path, *, success: bool) -> di remote_log = results / "remote.log" if success: if remote_metadata is None or remote_metadata.get("status") != "success": - raise ValueError("Colab did not return successful run metadata") + remote_error = (remote_metadata or {}).get("error", "no run metadata") + raise ValueError(f"Colab run did not succeed: {remote_error}") for name in ("checkpoint", "evaluation"): if not (results / name).is_dir(): raise ValueError(f"Colab did not return the {name} artifacts") @@ -313,6 +314,7 @@ def main() -> int: parser.add_argument("--pretrained-checkpoint", type=Path, default=PRETRAINED_CHECKPOINT) parser.add_argument("--epochs", type=int, default=10) parser.add_argument("--gpu", default="L4", help="Colab GPU preference; no accelerator fallback is attempted") + parser.add_argument("--timeout-seconds", type=float, default=6 * 3600, help="Upper bound for the remote training and evaluation run") args = parser.parse_args() if args.epochs < 1: parser.error("--epochs must be a positive integer") @@ -401,7 +403,7 @@ def main() -> int: if upload_status: raise RuntimeError("Colab failed to stage OCRKit training inputs") exec_status = stream_colab( - [colab, "exec", "-s", session, "-f", str(REMOTE_RUNNER)], + [colab, "exec", "-s", session, "--timeout", str(args.timeout_seconds), "-f", str(REMOTE_RUNNER)], colab_log, ) download_status = stream_colab( From 1439e7d2c7302608ba0a22808b2733036b5d3659 Mon Sep 17 00:00:00 2001 From: e54-bot Date: Sat, 26 Sep 2026 21:34:31 +0800 Subject: [PATCH 04/15] fix(training): allow Turing GPUs and create the Colab venv with uv PaddlePaddle 3.3.1 cu126 and cu129 run convolution forward/backward and matmul on a Colab T4 (compute capability 7.5), so drop the capability gate and rely on the runtime device check. Colab ships without python3-venv, so bootstrap creates the venv with uv when OCRKIT_TRAINING_PYTHON is set. Default the GPU preference to T4. Refs #20 --- training/README.md | 4 ++-- training/bootstrap.sh | 7 ++++++- training/colab_remote.py | 8 -------- training/run_rec_colab.py | 2 +- training/setup_rec_environment.sh | 5 ----- 5 files changed, 9 insertions(+), 17 deletions(-) diff --git a/training/README.md b/training/README.md index ca4702a..eca0d9e 100644 --- a/training/README.md +++ b/training/README.md @@ -448,11 +448,11 @@ platform snapshot, select that snapshot's output directory explicitly: ```bash uv run python training/run_rec_colab.py \ --labels-dir datasets/labeled/rec/platform/@ \ - --gpu L4 \ + --gpu T4 \ --epochs 10 ``` -`--gpu` is a Colab allocation preference (default `L4`), not a model or +`--gpu` is a Colab allocation preference (default `T4`), not a model or training requirement. PaddlePaddle's CUDA runtime and device are checked before training; an unavailable or unsupported GPU request fails without falling back to CPU. `--timeout-seconds` (default 6 hours) bounds the remote run; the Colab diff --git a/training/bootstrap.sh b/training/bootstrap.sh index 9a474f6..75e63eb 100644 --- a/training/bootstrap.sh +++ b/training/bootstrap.sh @@ -20,7 +20,12 @@ if [[ ! -x "${venv_dir}/bin/python" ]]; then printf 'Python 3.12 or OCRKIT_TRAINING_PYTHON is required to create the training environment.\n' >&2 exit 1 fi - "${python_bootstrap}" -m venv "${venv_dir}" + if [[ -n "${OCRKIT_TRAINING_PYTHON:-}" ]]; then + # Managed runtimes such as Colab ship without python3-venv/ensurepip. + uv venv --seed --python "${python_bootstrap}" "${venv_dir}" + else + "${python_bootstrap}" -m venv "${venv_dir}" + fi fi printf 'Training environment: %s\n' "${venv_dir}" diff --git a/training/colab_remote.py b/training/colab_remote.py index 6174453..3deacf4 100644 --- a/training/colab_remote.py +++ b/training/colab_remote.py @@ -87,22 +87,14 @@ def checked_gpu() -> dict[str, Any]: devices = [] for line in result.stdout.splitlines(): name, capability, driver, memory = (part.strip() for part in line.split(",", 3)) - major, minor = (int(part) for part in capability.split(".", 1)) devices.append( { "name": name, "compute_capability": capability, "driver_version": driver, "memory_mib": int(memory), - "supported_by_paddle": (major, minor) > (7, 5), } ) - if not devices or not devices[0]["supported_by_paddle"]: - allocated = devices[0]["compute_capability"] if devices else "unknown" - raise RuntimeError( - f"allocated GPU compute capability {allocated} is unsupported by PaddlePaddle 3.3.1; " - "a capability greater than 7.5 is required" - ) return {"devices": devices, "nvidia_smi": subprocess.run( ["nvidia-smi"], check=False, capture_output=True, text=True ).stdout} diff --git a/training/run_rec_colab.py b/training/run_rec_colab.py index 74f7c5d..b502990 100644 --- a/training/run_rec_colab.py +++ b/training/run_rec_colab.py @@ -313,7 +313,7 @@ def main() -> int: parser.add_argument("--labels-dir", type=Path, default=ROOT / "datasets/labeled/rec") parser.add_argument("--pretrained-checkpoint", type=Path, default=PRETRAINED_CHECKPOINT) parser.add_argument("--epochs", type=int, default=10) - parser.add_argument("--gpu", default="L4", help="Colab GPU preference; no accelerator fallback is attempted") + parser.add_argument("--gpu", default="T4", help="Colab GPU preference; no accelerator fallback is attempted") parser.add_argument("--timeout-seconds", type=float, default=6 * 3600, help="Upper bound for the remote training and evaluation run") args = parser.parse_args() if args.epochs < 1: diff --git a/training/setup_rec_environment.sh b/training/setup_rec_environment.sh index e1776f9..d946bfb 100755 --- a/training/setup_rec_environment.sh +++ b/training/setup_rec_environment.sh @@ -63,11 +63,6 @@ if [[ "${device}" == cuda ]]; then exit 1 fi - gpu_capabilities="$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader 2>/dev/null || true)" - if [[ -z "${gpu_capabilities}" ]] || ! awk -F. 'NF == 2 && (($1 + 0) > 7 || (($1 + 0) == 7 && ($2 + 0) > 5)) {supported=1} END {exit !supported}' <<< "${gpu_capabilities}"; then - printf 'the allocated GPU is not supported by PaddlePaddle 3.3.1 (compute capability must be greater than 7.5).\n' >&2 - exit 1 - fi paddle_package="paddlepaddle-gpu==3.3.1" fi From 34bfc4aee163c7e37ba52445a524e2d491d0a32b Mon Sep 17 00:00:00 2001 From: e54-bot Date: Sat, 26 Sep 2026 21:41:07 +0800 Subject: [PATCH 05/15] fix(training): transfer Colab archives in bounded parts colab upload and download send each file as a single base64 JSON request, and a 145 MB input archive was rejected with HTTP 400. Send the input in 32 MiB parts that the runtime joins, and return the result as checksummed parts listed in an index. Refs #20 --- tests/test_run_rec_colab.py | 49 ++++++++++++++++++++++++++++-- training/colab_remote.py | 33 ++++++++++++++++++--- training/run_rec_colab.py | 59 ++++++++++++++++++++++++++++++++----- 3 files changed, 127 insertions(+), 14 deletions(-) diff --git a/tests/test_run_rec_colab.py b/tests/test_run_rec_colab.py index b9f73db..8c48a0f 100644 --- a/tests/test_run_rec_colab.py +++ b/tests/test_run_rec_colab.py @@ -54,7 +54,27 @@ def fake_stream(command: list[str], _log: Path) -> int: if verb == "exec": return behavior["exec_status"] if verb == "download": - build_result_archive(Path(command[-1]), status=behavior["result_status"]) + remote, local = command[-2], Path(command[-1]) + if "remote_result" not in behavior: + archive = local.parent / "remote-result.tar.gz" + build_result_archive(archive, status=behavior["result_status"]) + behavior["remote_result"] = archive.read_bytes() + archive.unlink() + data = behavior["remote_result"] + parts = [data[i : i + 2000] for i in range(0, len(data), 2000)] + if remote.endswith("index.json"): + local.write_text( + json.dumps( + { + "parts": [ + {"name": f"ocrkit-result.part{i:04d}", "sha256": hashlib.sha256(part).hexdigest()} + for i, part in enumerate(parts) + ] + } + ) + ) + else: + local.write_bytes(parts[int(remote.rsplit("part", 1)[1])]) return 0 if verb == SESSION_STOP: return behavior["stop_status"] @@ -63,7 +83,8 @@ def fake_stream(command: list[str], _log: Path) -> int: monkeypatch.setattr(run_rec_colab, "RUNS", runs) monkeypatch.setattr(run_rec_colab.shutil, "which", lambda _name: "colab") monkeypatch.setattr(run_rec_colab, "validate_labels", lambda _path: (1, ["a.png"])) - monkeypatch.setattr(run_rec_colab, "stage_inputs", lambda archive, *_a: archive.write_bytes(b"input")) + monkeypatch.setattr(run_rec_colab, "PART_BYTES", 4) + monkeypatch.setattr(run_rec_colab, "stage_inputs", lambda archive, *_a: archive.write_bytes(b"input-archive")) monkeypatch.setattr(run_rec_colab, "stream_colab", fake_stream) monkeypatch.setattr( sys, "argv", ["run_rec_colab.py", "--labels-dir", str(labels), "--pretrained-checkpoint", str(checkpoint)] @@ -81,6 +102,8 @@ def test_success_keeps_verified_artifacts_and_stops_runtime(colab_run) -> None: assert run_rec_colab.main() == 0 + uploads = [command[-1] for command in calls if command[1] == "upload"] + assert uploads == [f"/content/ocrkit-input.part{i:04d}" for i in range(4)] exec_command = next(command for command in calls if command[1] == "exec") assert float(exec_command[exec_command.index("--timeout") + 1]) == 6 * 3600 run_dir = only_run(runs) @@ -143,3 +166,25 @@ def test_tampered_result_is_rejected(tmp_path: Path) -> None: run_rec_colab.copy_remote_result(remote, run_dir, success=True) assert not (run_dir / "accepted").exists() + + +def test_remote_parts_round_trip(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + from training import colab_remote + + monkeypatch.setattr(colab_remote, "CONTENT", tmp_path) + monkeypatch.setattr(colab_remote, "INPUT_ARCHIVE", tmp_path / "in.tar.gz") + monkeypatch.setattr(colab_remote, "RESULT_ARCHIVE", tmp_path / "out.tar.gz") + monkeypatch.setattr(colab_remote, "RESULT_INDEX", tmp_path / "out.index.json") + monkeypatch.setattr(colab_remote, "PART_BYTES", 3) + payload = b"0123456789" + for index in range(2): + (tmp_path / f"ocrkit-input.part{index:04d}").write_bytes(payload[index * 5 : index * 5 + 5]) + colab_remote.join_input_parts() + assert (tmp_path / "in.tar.gz").read_bytes() == payload + assert not list(tmp_path.glob("ocrkit-input.part*")) + + (tmp_path / "out.tar.gz").write_bytes(payload) + colab_remote.split_result_archive() + index = json.loads((tmp_path / "out.index.json").read_text()) + assert b"".join((tmp_path / part["name"]).read_bytes() for part in index["parts"]) == payload + assert len(index["parts"]) == 4 diff --git a/training/colab_remote.py b/training/colab_remote.py index 3deacf4..0798f45 100644 --- a/training/colab_remote.py +++ b/training/colab_remote.py @@ -15,8 +15,11 @@ from typing import Any COLAB_ROOT = Path("/content/ocrkit-colab") -INPUT_ARCHIVE = Path("/content/ocrkit-input.tar.gz") -RESULT_ARCHIVE = Path("/content/ocrkit-result.tar.gz") +CONTENT = Path("/content") +INPUT_ARCHIVE = CONTENT / "ocrkit-input.tar.gz" +RESULT_ARCHIVE = CONTENT / "ocrkit-result.tar.gz" +RESULT_INDEX = CONTENT / "ocrkit-result.index.json" +PART_BYTES = 32 * 1024 * 1024 REPO = COLAB_ROOT / "repo" DATASET = COLAB_ROOT / "dataset" RESULTS = COLAB_ROOT / "results" @@ -124,6 +127,28 @@ def output_files() -> list[dict[str, Any]]: return records +def join_input_parts() -> None: + parts = sorted(CONTENT.glob("ocrkit-input.part*")) + if not parts: + raise RuntimeError("OCRKit input parts were not uploaded to the Colab runtime") + with INPUT_ARCHIVE.open("wb") as archive: + for part in parts: + archive.write(part.read_bytes()) + part.unlink() + + +def split_result_archive() -> None: + for stale in (*CONTENT.glob("ocrkit-result.part*"), RESULT_INDEX): + stale.unlink(missing_ok=True) + parts = [] + with RESULT_ARCHIVE.open("rb") as archive: + for index, chunk in enumerate(iter(lambda: archive.read(PART_BYTES), b"")): + name = f"ocrkit-result.part{index:04d}" + (CONTENT / name).write_bytes(chunk) + parts.append({"name": name, "sha256": hashlib.sha256(chunk).hexdigest()}) + RESULT_INDEX.write_text(json.dumps({"parts": parts}), encoding="utf-8") + + def write_result_archive() -> None: RESULT_ARCHIVE.unlink(missing_ok=True) with tarfile.open(RESULT_ARCHIVE, "w:gz") as archive: @@ -133,6 +158,7 @@ def write_result_archive() -> None: for directory in (CHECKPOINTS, EVALUATION): if directory.is_dir(): archive.add(directory, arcname=directory.relative_to(COLAB_ROOT)) + split_result_archive() def main() -> int: @@ -145,8 +171,7 @@ def main() -> int: try: with REMOTE_LOG.open("w", encoding="utf-8") as log: try: - if not INPUT_ARCHIVE.is_file(): - raise RuntimeError("OCRKit input archive was not uploaded to the Colab runtime") + join_input_parts() COLAB_ROOT.mkdir(parents=True, exist_ok=True) safe_extract(INPUT_ARCHIVE, COLAB_ROOT) request_path = COLAB_ROOT / "request.json" diff --git a/training/run_rec_colab.py b/training/run_rec_colab.py index b502990..5847f2f 100644 --- a/training/run_rec_colab.py +++ b/training/run_rec_colab.py @@ -17,6 +17,7 @@ ROOT = Path(__file__).resolve().parents[1] RUNS = ROOT / "training/.work/colab-runs" ACCEPTED_STAGING = "accepted" +PART_BYTES = 32 * 1024 * 1024 REMOTE_RUNNER = ROOT / "training/colab_remote.py" PRETRAINED_CHECKPOINT = ROOT / "training/.work/pretrained/PP-OCRv6_small_rec_pretrained.pdparams" SOURCE_FILES = ( @@ -241,6 +242,54 @@ def stream_colab(command: list[str], log_path: Path) -> int: return process.wait() +def upload_archive(colab: str, session: str, archive: Path, run_dir: Path, log_path: Path) -> int: + """Colab uploads are single JSON requests, so send the archive in bounded parts.""" + parts_dir = run_dir / "input-parts" + parts_dir.mkdir() + try: + with archive.open("rb") as stream: + for index, chunk in enumerate(iter(lambda: stream.read(PART_BYTES), b"")): + part = parts_dir / f"part{index:04d}" + part.write_bytes(chunk) + status = stream_colab( + [colab, "upload", "-s", session, str(part), f"/content/ocrkit-input.part{index:04d}"], + log_path, + ) + part.unlink() + if status: + return status + return 0 + finally: + shutil.rmtree(parts_dir, ignore_errors=True) + + +def download_archive(colab: str, session: str, destination: Path, run_dir: Path, log_path: Path) -> int: + parts_dir = run_dir / "result-parts" + parts_dir.mkdir() + try: + index_path = parts_dir / "index.json" + status = stream_colab( + [colab, "download", "-s", session, "/content/ocrkit-result.index.json", str(index_path)], log_path + ) + if status: + return status + with destination.open("wb") as result: + for record in json.loads(index_path.read_text(encoding="utf-8"))["parts"]: + part = parts_dir / safe_relative(record["name"]) + status = stream_colab( + [colab, "download", "-s", session, f"/content/{record['name']}", str(part)], log_path + ) + if status: + return status + if sha256(part) != record["sha256"]: + raise ValueError(f"Colab result part failed checksum verification: {record['name']}") + result.write(part.read_bytes()) + part.unlink() + return 0 + finally: + shutil.rmtree(parts_dir, ignore_errors=True) + + def safe_extract_result(archive_path: Path, destination: Path) -> None: destination.mkdir(parents=True, exist_ok=True) root = destination.resolve() @@ -395,10 +444,7 @@ def main() -> int: raise RuntimeError( f"Colab could not provision the requested GPU {args.gpu}; no fallback accelerator was selected." ) - upload_status = stream_colab( - [colab, "upload", "-s", session, str(input_archive), "/content/ocrkit-input.tar.gz"], - colab_log, - ) + upload_status = upload_archive(colab, session, input_archive, run_dir, colab_log) input_archive.unlink(missing_ok=True) if upload_status: raise RuntimeError("Colab failed to stage OCRKit training inputs") @@ -406,10 +452,7 @@ def main() -> int: [colab, "exec", "-s", session, "--timeout", str(args.timeout_seconds), "-f", str(REMOTE_RUNNER)], colab_log, ) - download_status = stream_colab( - [colab, "download", "-s", session, "/content/ocrkit-result.tar.gz", str(result_archive)], - colab_log, - ) + download_status = download_archive(colab, session, result_archive, run_dir, colab_log) if download_status: raise RuntimeError( "Colab failed to retrieve the remote run logs and artifacts" From d15103e09afb1238b742d2525863e05145bcd12a Mon Sep 17 00:00:00 2001 From: e54-bot Date: Sun, 27 Sep 2026 02:36:53 +0800 Subject: [PATCH 06/15] fix(training): build the Colab training venv on Python 3.12 paddle2onnx 2.1.0 has no Python 3.13 wheel, and Colab now ships 3.13. Create the training venv with uv on Python 3.12, matching the local environment, and drop the OCRKIT_TRAINING_PYTHON override. Refs #20 --- training/bootstrap.sh | 16 +++++++--------- training/colab_remote.py | 4 ---- 2 files changed, 7 insertions(+), 13 deletions(-) diff --git a/training/bootstrap.sh b/training/bootstrap.sh index 75e63eb..a25b25b 100644 --- a/training/bootstrap.sh +++ b/training/bootstrap.sh @@ -15,16 +15,14 @@ printf 'PaddleOCR checkout: %s\n' "${paddleocr_dir}" venv_dir="${work_dir}/venv" if [[ ! -x "${venv_dir}/bin/python" ]]; then - python_bootstrap="${OCRKIT_TRAINING_PYTHON:-python3.12}" - if ! command -v "${python_bootstrap}" >/dev/null; then - printf 'Python 3.12 or OCRKIT_TRAINING_PYTHON is required to create the training environment.\n' >&2 - exit 1 - fi - if [[ -n "${OCRKIT_TRAINING_PYTHON:-}" ]]; then - # Managed runtimes such as Colab ship without python3-venv/ensurepip. - uv venv --seed --python "${python_bootstrap}" "${venv_dir}" + if command -v python3.12 >/dev/null; then + python3.12 -m venv "${venv_dir}" + elif command -v uv >/dev/null; then + # Managed runtimes such as Colab have neither python3.12 nor python3-venv. + uv venv --seed --python 3.12 "${venv_dir}" else - "${python_bootstrap}" -m venv "${venv_dir}" + printf 'Python 3.12 or uv is required to create the training environment.\n' >&2 + exit 1 fi fi diff --git a/training/colab_remote.py b/training/colab_remote.py index 0798f45..e9b42ab 100644 --- a/training/colab_remote.py +++ b/training/colab_remote.py @@ -3,7 +3,6 @@ import hashlib import json -import os import shlex import shutil import subprocess @@ -194,12 +193,9 @@ def main() -> int: cwd=REPO, log=log, ) - environment = os.environ.copy() - environment["OCRKIT_TRAINING_PYTHON"] = sys.executable run_logged( ["bash", "training/setup_rec_environment.sh", "--device", "cuda"], cwd=REPO, - env=environment, log=log, ) From 25ed112474feb8eec4e7997ea8d283a5fbac6c3c Mon Sep 17 00:00:00 2001 From: e54-bot Date: Sun, 27 Sep 2026 02:49:39 +0800 Subject: [PATCH 07/15] fix(training): fall back to uv when python3.12 lacks venv support Colab provides python3.12 without ensurepip, so the plain venv step fails there. Keep python3.12 -m venv as the first choice and fall back to a uv-managed Python 3.12 venv. Refs #20 --- training/bootstrap.sh | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/training/bootstrap.sh b/training/bootstrap.sh index a25b25b..09a029a 100644 --- a/training/bootstrap.sh +++ b/training/bootstrap.sh @@ -15,13 +15,13 @@ printf 'PaddleOCR checkout: %s\n' "${paddleocr_dir}" venv_dir="${work_dir}/venv" if [[ ! -x "${venv_dir}/bin/python" ]]; then - if command -v python3.12 >/dev/null; then - python3.12 -m venv "${venv_dir}" + if command -v python3.12 >/dev/null && python3.12 -m venv "${venv_dir}"; then + : elif command -v uv >/dev/null; then - # Managed runtimes such as Colab have neither python3.12 nor python3-venv. - uv venv --seed --python 3.12 "${venv_dir}" + # Managed runtimes such as Colab ship python3.12 without python3-venv/ensurepip. + uv venv --clear --seed --python 3.12 "${venv_dir}" else - printf 'Python 3.12 or uv is required to create the training environment.\n' >&2 + printf 'Python 3.12 with venv support, or uv, is required to create the training environment.\n' >&2 exit 1 fi fi From 8f6325ce43e20356bf1ca867064c2cf2c2dddc88 Mon Sep 17 00:00:00 2001 From: e54-bot Date: Sun, 27 Sep 2026 04:11:51 +0800 Subject: [PATCH 08/15] feat(training): install the Colab Paddle wheel from a checksummed mirror The official cu129 PaddlePaddle wheel is 2.9 GB and downloads at about 0.4 MB/s from a Colab runtime. Install a mirror of the identical file from the OWBastion CDN, verify its sha256, and keep the official index as the fallback for other CUDA versions. Refs #20 --- training/README.md | 5 +++++ training/colab_remote.py | 7 +++++++ training/run_rec_colab.py | 5 +++++ training/setup_rec_environment.sh | 15 ++++++++++++++- 4 files changed, 31 insertions(+), 1 deletion(-) diff --git a/training/README.md b/training/README.md index eca0d9e..6ba2bff 100644 --- a/training/README.md +++ b/training/README.md @@ -458,6 +458,11 @@ training; an unavailable or unsupported GPU request fails without falling back to CPU. `--timeout-seconds` (default 6 hours) bounds the remote run; the Colab CLI's own `exec` default of 30 seconds is always overridden. The local CPU command remains `./training/run_rec_smoke.sh`. +The 2.9 GB `paddlepaddle-gpu` wheel is slow to fetch from the official +CDN outside China, so the runner installs a checksummed mirror of the official +cu129 build (`PADDLE_WHEEL` in `run_rec_colab.py`) when the runtime selects the +cu129 index, and otherwise falls back to the official index. + The runner transfers only the selected train/holdout labels and referenced crops, available review/snapshot provenance files, the base recognition checkpoint, OCRKit training/evaluation code, and the fixture images needed by diff --git a/training/colab_remote.py b/training/colab_remote.py index e9b42ab..58309a9 100644 --- a/training/colab_remote.py +++ b/training/colab_remote.py @@ -3,6 +3,7 @@ import hashlib import json +import os import shlex import shutil import subprocess @@ -193,9 +194,15 @@ def main() -> int: cwd=REPO, log=log, ) + mirror = request["run"]["paddle_wheel_mirror"] run_logged( ["bash", "training/setup_rec_environment.sh", "--device", "cuda"], cwd=REPO, + env={ + **os.environ, + "OCRKIT_PADDLE_WHEEL_URL": mirror["url"], + "OCRKIT_PADDLE_WHEEL_SHA256": mirror["sha256"], + }, log=log, ) diff --git a/training/run_rec_colab.py b/training/run_rec_colab.py index 5847f2f..58ad560 100644 --- a/training/run_rec_colab.py +++ b/training/run_rec_colab.py @@ -17,6 +17,10 @@ ROOT = Path(__file__).resolve().parents[1] RUNS = ROOT / "training/.work/colab-runs" ACCEPTED_STAGING = "accepted" +PADDLE_WHEEL = { + "url": "https://cdn.owbastion.codes/ocrkit/wheels/cu129/paddlepaddle_gpu-3.3.1-cp312-cp312-linux_x86_64.whl", + "sha256": "03fc5211183ba20ef71e63a35e589fd386a8805c1011c3a409a0e5b118be4668", +} PART_BYTES = 32 * 1024 * 1024 REMOTE_RUNNER = ROOT / "training/colab_remote.py" PRETRAINED_CHECKPOINT = ROOT / "training/.work/pretrained/PP-OCRv6_small_rec_pretrained.pdparams" @@ -412,6 +416,7 @@ def main() -> int: "source_name": checkpoint_path.name, "sha256": sha256(checkpoint_path), }, + "paddle_wheel_mirror": PADDLE_WHEEL, "training": { "epochs": args.epochs, "device": "cuda", diff --git a/training/setup_rec_environment.sh b/training/setup_rec_environment.sh index d946bfb..bd89bd5 100755 --- a/training/setup_rec_environment.sh +++ b/training/setup_rec_environment.sh @@ -67,7 +67,20 @@ if [[ "${device}" == cuda ]]; then fi "${python_bin}" -m pip install --upgrade pip -"${python_bin}" -m pip install "${paddle_package}" -i "${paddle_index}" +if [[ "${device}" == cuda && "${paddle_index}" == */cu129/ && -n "${OCRKIT_PADDLE_WHEEL_URL:-}" && -n "${OCRKIT_PADDLE_WHEEL_SHA256:-}" ]]; then + # Optional mirror of the official cu129 wheel for runtimes far from the official CDN. + mkdir -p "${work_dir}/wheels" + wheel_path="${work_dir}/wheels/${OCRKIT_PADDLE_WHEEL_URL##*/}" + curl -L --fail --retry 3 -o "${wheel_path}" "${OCRKIT_PADDLE_WHEEL_URL}" + wheel_sha256="$("${python_bin}" -c 'import hashlib, sys; print(hashlib.sha256(open(sys.argv[1], "rb").read()).hexdigest())' "${wheel_path}")" + if [[ "${wheel_sha256}" != "${OCRKIT_PADDLE_WHEEL_SHA256}" ]]; then + printf 'mirrored PaddlePaddle wheel failed checksum verification: %s\n' "${wheel_path}" >&2 + exit 1 + fi + "${python_bin}" -m pip install "${wheel_path}" +else + "${python_bin}" -m pip install "${paddle_package}" -i "${paddle_index}" +fi "${python_bin}" -m pip install -r "${paddleocr_dir}/requirements.txt" "${python_bin}" -m pip install "paddle2onnx==2.1.0" From ca5a2d2125be80cba2cb104afc8599b9f0f8fc16 Mon Sep 17 00:00:00 2001 From: e54-bot Date: Sun, 27 Sep 2026 04:25:22 +0800 Subject: [PATCH 09/15] fix(training): read the epoch count from the training request section The remote runner read request.run.epochs, but the request stores it under run.training. Assert the request keys the remote runner reads. Refs #20 --- tests/test_run_rec_colab.py | 11 +++++++++-- training/colab_remote.py | 2 +- 2 files changed, 10 insertions(+), 3 deletions(-) diff --git a/tests/test_run_rec_colab.py b/tests/test_run_rec_colab.py index 8c48a0f..bbcd421 100644 --- a/tests/test_run_rec_colab.py +++ b/tests/test_run_rec_colab.py @@ -84,7 +84,11 @@ def fake_stream(command: list[str], _log: Path) -> int: monkeypatch.setattr(run_rec_colab.shutil, "which", lambda _name: "colab") monkeypatch.setattr(run_rec_colab, "validate_labels", lambda _path: (1, ["a.png"])) monkeypatch.setattr(run_rec_colab, "PART_BYTES", 4) - monkeypatch.setattr(run_rec_colab, "stage_inputs", lambda archive, *_a: archive.write_bytes(b"input-archive")) + def fake_stage(archive: Path, run_request: dict, *_rest) -> None: + behavior["run_request"] = run_request + archive.write_bytes(b"input-archive") + + monkeypatch.setattr(run_rec_colab, "stage_inputs", fake_stage) monkeypatch.setattr(run_rec_colab, "stream_colab", fake_stream) monkeypatch.setattr( sys, "argv", ["run_rec_colab.py", "--labels-dir", str(labels), "--pretrained-checkpoint", str(checkpoint)] @@ -98,10 +102,13 @@ def only_run(runs: Path) -> Path: def test_success_keeps_verified_artifacts_and_stops_runtime(colab_run) -> None: - runs, calls, _ = colab_run + runs, calls, behavior = colab_run assert run_rec_colab.main() == 0 + request = behavior["run_request"] + assert request["training"]["epochs"] == 10 + assert set(request["paddle_wheel_mirror"]) == {"url", "sha256"} uploads = [command[-1] for command in calls if command[1] == "upload"] assert uploads == [f"/content/ocrkit-input.part{i:04d}" for i in range(4)] exec_command = next(command for command in calls if command[1] == "exec") diff --git a/training/colab_remote.py b/training/colab_remote.py index 58309a9..0a1e185 100644 --- a/training/colab_remote.py +++ b/training/colab_remote.py @@ -220,7 +220,7 @@ def main() -> int: if not paddle_info["cuda"] or paddle_info["device"] != "gpu:0": raise RuntimeError("PaddlePaddle did not select the allocated CUDA device") - epochs = str(request["run"]["epochs"]) + epochs = str(request["run"]["training"]["epochs"]) run_logged( [ "bash", From c8abf0d1a1240b2c511b81eb6ca3bae54d227bf4 Mon Sep 17 00:00:00 2001 From: e54-bot Date: Sun, 27 Sep 2026 05:04:21 +0800 Subject: [PATCH 10/15] fix(training): export the recognition checkpoint on CPU A checkpoint trained on a CUDA device was exported with a GPU-specific inference graph that paddle2onnx aborts on. Exporting the same checkpoint on CPU converts and passes the fixture gate, and matches the existing CPU-only local flow, so make the export device independent. Refs #20 --- training/evaluate_rec_checkpoint.sh | 1 + 1 file changed, 1 insertion(+) diff --git a/training/evaluate_rec_checkpoint.sh b/training/evaluate_rec_checkpoint.sh index cde4fd1..84fe30a 100755 --- a/training/evaluate_rec_checkpoint.sh +++ b/training/evaluate_rec_checkpoint.sh @@ -36,6 +36,7 @@ uv run python training/scripts/prepare_detector.py \ cd "${paddleocr_dir}" "${training_python}" tools/export_model.py -c "${config_path}" -o \ + Global.use_gpu=False \ Global.pretrained_model="${checkpoint}" \ Global.save_inference_dir="${artifact_dir}/paddle_rec" \ Global.character_dict_path="${paddleocr_dir}/ppocr/utils/dict/ppocrv6_dict.txt" From 4650325c6fd757bb584b46ff932fec95ec1dc0cd Mon Sep 17 00:00:00 2001 From: e54-bot Date: Sun, 27 Sep 2026 05:28:38 +0800 Subject: [PATCH 11/15] fix(training): evaluate the Colab-trained checkpoint locally CUDA builds of PaddlePaddle export nn.Linear as linear_v2, which paddle2onnx 2.1.0 aborts on, so the evaluation export cannot run on the Colab GPU runtime. Train on Colab with --train-only, stop the runtime after retrieval, and run the unchanged evaluate_rec_checkpoint.sh locally. The run only succeeds when that evaluation passes; otherwise the checkpoint is kept under partial/. Stop uploading the evaluation-only inputs. Refs #20 --- tests/test_run_rec_colab.py | 39 ++++++++-- training/README.md | 25 +++--- training/colab_remote.py | 27 +------ training/evaluate_rec_checkpoint.sh | 1 - training/run_rec_colab.py | 117 +++++++++++++--------------- training/run_rec_smoke.sh | 16 ++-- 6 files changed, 113 insertions(+), 112 deletions(-) diff --git a/tests/test_run_rec_colab.py b/tests/test_run_rec_colab.py index bbcd421..cb6ac25 100644 --- a/tests/test_run_rec_colab.py +++ b/tests/test_run_rec_colab.py @@ -16,7 +16,6 @@ def build_result_archive(path: Path, *, status: str = "success") -> None: files = { "results/checkpoint/best_accuracy.pdparams": b"weights", - "results/evaluation/fixture_report.json": b'{"field_accuracy": 0.99}', "results/remote.log": b"log", } outputs = [ @@ -26,7 +25,7 @@ def build_result_archive(path: Path, *, status: str = "success") -> None: "sha256": hashlib.sha256(data).hexdigest(), } for name, data in files.items() - if not name.endswith("remote.log") + if name.startswith("results/checkpoint/") ] files["results/run.json"] = json.dumps({"status": status, "outputs": outputs}).encode() with tarfile.open(path, "w:gz") as archive: @@ -46,10 +45,17 @@ def colab_run(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): checkpoint.write_bytes(b"base") runs = tmp_path / "runs" calls: list[list[str]] = [] - behavior = {"exec_status": 0, "result_status": "success", "stop_status": 0} + behavior = {"exec_status": 0, "result_status": "success", "stop_status": 0, "eval_status": 0} def fake_stream(command: list[str], _log: Path) -> int: calls.append(command) + if command[0].endswith("evaluate_rec_checkpoint.sh"): + if behavior["eval_status"]: + return behavior["eval_status"] + output = Path(command[-1]) + output.mkdir(parents=True) + (output / "fixture_report.json").write_text('{"field_accuracy": 0.99, "run_code": {"field_accuracy": 1.0}}') + return 0 verb = command[1] if verb == "exec": return behavior["exec_status"] @@ -89,7 +95,7 @@ def fake_stage(archive: Path, run_request: dict, *_rest) -> None: archive.write_bytes(b"input-archive") monkeypatch.setattr(run_rec_colab, "stage_inputs", fake_stage) - monkeypatch.setattr(run_rec_colab, "stream_colab", fake_stream) + monkeypatch.setattr(run_rec_colab, "stream_command", fake_stream) monkeypatch.setattr( sys, "argv", ["run_rec_colab.py", "--labels-dir", str(labels), "--pretrained-checkpoint", str(checkpoint)] ) @@ -116,8 +122,11 @@ def test_success_keeps_verified_artifacts_and_stops_runtime(colab_run) -> None: run_dir = only_run(runs) assert (run_dir / "checkpoint/best_accuracy.pdparams").is_file() assert (run_dir / "evaluation/fixture_report.json").is_file() + assert json.loads((run_dir / "run.json").read_text())["evaluation"]["field_accuracy"] == 0.99 assert json.loads((run_dir / "status.json").read_text())["runtime_stopped"] is True - assert calls[-1][1] == SESSION_STOP + stop_index = next(i for i, command in enumerate(calls) if command[1] == SESSION_STOP) + evaluation_index = next(i for i, command in enumerate(calls) if command[0].endswith("evaluate_rec_checkpoint.sh")) + assert stop_index < evaluation_index assert not (run_dir / "accepted").exists() @@ -135,6 +144,20 @@ def test_training_failure_keeps_partial_output_and_stops_runtime(colab_run) -> N assert json.loads((run_dir / "status.json").read_text())["status"] == "failed" +def test_local_evaluation_failure_demotes_checkpoint_to_partial(colab_run) -> None: + runs, calls, behavior = colab_run + behavior["eval_status"] = 1 + + assert run_rec_colab.main() == 1 + + run_dir = only_run(runs) + assert not (run_dir / "checkpoint").exists() + assert not (run_dir / "run.json").exists() + assert (run_dir / "partial/checkpoint/best_accuracy.pdparams").is_file() + assert "local checkpoint evaluation failed" in json.loads((run_dir / "status.json").read_text())["error"] + assert json.loads((run_dir / "status.json").read_text())["runtime_stopped"] is True + + def test_teardown_failure_demotes_accepted_artifacts_to_partial(colab_run) -> None: runs, _, behavior = colab_run behavior["stop_status"] = 1 @@ -149,12 +172,12 @@ def test_teardown_failure_demotes_accepted_artifacts_to_partial(colab_run) -> No def test_provisioning_failure_stops_runtime_and_uploads_nothing(colab_run) -> None: runs, calls, _ = colab_run - original = run_rec_colab.stream_colab - run_rec_colab.stream_colab = lambda command, log: 1 if command[1] == "new" else original(command, log) + original = run_rec_colab.stream_command + run_rec_colab.stream_command = lambda command, log: 1 if command[1] == "new" else original(command, log) try: assert run_rec_colab.main() == 1 finally: - run_rec_colab.stream_colab = original + run_rec_colab.stream_command = original assert [command[1] for command in calls] == [SESSION_STOP] assert not (only_run(runs) / "checkpoint").exists() diff --git a/training/README.md b/training/README.md index 6ba2bff..e4b49db 100644 --- a/training/README.md +++ b/training/README.md @@ -399,8 +399,8 @@ Prepare the offline environment once, then run the CPU recognition Smoke: `run_rec_smoke.sh` accepts `--labels-dir`, `--output-dir`, `--epochs` (the target total epoch), and `--resume-checkpoint` (a checkpoint base path without `.pdparams`, `.pdopt`, or `.states`). `--device cpu|cuda` selects PaddleOCR's -device and defaults to `cpu`; `--evaluation-dir` selects the evaluation output -directory. It validates both label files, fine-tunes recognition only, and +device and defaults to `cpu`; `--train-only` stops after training and pruning +without running the evaluation. It validates both label files, fine-tunes recognition only, and leaves `latest` plus `best_accuracy` under `training/.work/`. Per-epoch `iter_epoch_*` dumps and PaddleOCR's duplicate `best_model/` copy are pruned after training. To reclaim space from older runs: @@ -436,7 +436,7 @@ The first CLI request can prompt for Google OAuth authentication in the terminal. The CLI keeps those credentials locally; OCRKit does not send platform, R2, or release credentials to Colab. -Start training and the existing checkpoint evaluation with one command: +Train on Colab and evaluate the retrieved checkpoint locally with one command: ```bash uv run python training/run_rec_colab.py @@ -463,19 +463,26 @@ CDN outside China, so the runner installs a checksummed mirror of the official cu129 build (`PADDLE_WHEEL` in `run_rec_colab.py`) when the runtime selects the cu129 index, and otherwise falls back to the official index. +Only training runs on Colab. CUDA builds of PaddlePaddle export `nn.Linear` as +`linear_v2`, which `paddle2onnx` cannot convert, so the runner retrieves the +checkpoint, stops the runtime, and then runs the unchanged +`training/evaluate_rec_checkpoint.sh` on your machine (the local training +environment from `setup_rec_environment.sh` is required). The run only +succeeds if that evaluation passes the same gate as a local run. + The runner transfers only the selected train/holdout labels and referenced crops, available review/snapshot provenance files, the base recognition -checkpoint, OCRKit training/evaluation code, and the fixture images needed by -the existing evaluation. It records source revisions, input checksums, the -effective training configuration, PaddleOCR revision, allocated GPU details, -checkpoint, and evaluation report in the returned `run.json`. +checkpoint, and the training scripts. It records source revisions, input +checksums, the effective training configuration, PaddleOCR revision, allocated +GPU details, checkpoint checksums, and the local evaluation summary in the +returned `run.json`. Checkpoints, `fixture_report.json`, provenance, the remote training log, the Colab CLI log, and `status.json` are stored below the ignored `training/.work/colab-runs//` directory. A successful run does not publish a candidate or change the stable model channel. Provisioning, -staging, training, evaluation, retrieval, and handled failures all stop the -Colab runtime after it has been allocated. Failed runs keep diagnostics and +staging, training, retrieval, and handled failures all stop the Colab runtime +after it has been allocated; it is stopped before the local evaluation starts. Failed runs keep diagnostics and any partial output under `partial/`; if teardown itself fails, `status.json` includes the named `colab stop` command to release that session. diff --git a/training/colab_remote.py b/training/colab_remote.py index 0a1e185..b754391 100644 --- a/training/colab_remote.py +++ b/training/colab_remote.py @@ -24,7 +24,6 @@ DATASET = COLAB_ROOT / "dataset" RESULTS = COLAB_ROOT / "results" CHECKPOINTS = RESULTS / "checkpoint" -EVALUATION = RESULTS / "evaluation" RUN_METADATA = RESULTS / "run.json" REMOTE_LOG = RESULTS / "remote.log" @@ -112,7 +111,7 @@ def verify_inputs(request: dict[str, Any], root: Path) -> None: def output_files() -> list[dict[str, Any]]: records = [] - for directory in (CHECKPOINTS, EVALUATION): + for directory in (CHECKPOINTS,): if not directory.is_dir(): continue for path in sorted(directory.rglob("*")): @@ -155,9 +154,8 @@ def write_result_archive() -> None: for path in (RUN_METADATA, REMOTE_LOG): if path.is_file(): archive.add(path, arcname=path.relative_to(COLAB_ROOT)) - for directory in (CHECKPOINTS, EVALUATION): - if directory.is_dir(): - archive.add(directory, arcname=directory.relative_to(COLAB_ROOT)) + if CHECKPOINTS.is_dir(): + archive.add(CHECKPOINTS, arcname=CHECKPOINTS.relative_to(COLAB_ROOT)) split_result_archive() @@ -189,11 +187,6 @@ def main() -> int: cwd=COLAB_ROOT, log=log, ) - run_logged( - ["uv", "sync", "--locked", "--no-dev", "--python", sys.executable], - cwd=REPO, - log=log, - ) mirror = request["run"]["paddle_wheel_mirror"] run_logged( ["bash", "training/setup_rec_environment.sh", "--device", "cuda"], @@ -229,23 +222,16 @@ def main() -> int: str(DATASET), "--output-dir", str(CHECKPOINTS), - "--evaluation-dir", - str(EVALUATION), "--epochs", epochs, "--device", "cuda", + "--train-only", ], cwd=REPO, log=log, ) best_checkpoint = CHECKPOINTS / "best_accuracy.pdparams" - report_path = EVALUATION / "fixture_report.json" - if not best_checkpoint.is_file() or best_checkpoint.stat().st_size == 0: - raise RuntimeError("training did not produce the best-accuracy recognition checkpoint") - if not report_path.is_file(): - raise RuntimeError("checkpoint evaluation did not produce fixture_report.json") - evaluation_report = json.loads(report_path.read_text(encoding="utf-8")) paddleocr_revision = subprocess.run( ["git", "-C", str(REPO / "training/.work/PaddleOCR"), "rev-parse", "HEAD"], check=True, @@ -261,11 +247,6 @@ def main() -> int: "paddle": paddle_info, "paddleocr_revision": paddleocr_revision, }, - "evaluation": { - "report": "evaluation/fixture_report.json", - "field_accuracy": evaluation_report.get("field_accuracy"), - "run_code_accuracy": evaluation_report.get("run_code", {}).get("field_accuracy"), - }, "outputs": output_files(), } ) diff --git a/training/evaluate_rec_checkpoint.sh b/training/evaluate_rec_checkpoint.sh index 84fe30a..cde4fd1 100755 --- a/training/evaluate_rec_checkpoint.sh +++ b/training/evaluate_rec_checkpoint.sh @@ -36,7 +36,6 @@ uv run python training/scripts/prepare_detector.py \ cd "${paddleocr_dir}" "${training_python}" tools/export_model.py -c "${config_path}" -o \ - Global.use_gpu=False \ Global.pretrained_model="${checkpoint}" \ Global.save_inference_dir="${artifact_dir}/paddle_rec" \ Global.character_dict_path="${paddleocr_dir}/ppocr/utils/dict/ppocrv6_dict.txt" diff --git a/training/run_rec_colab.py b/training/run_rec_colab.py index 58ad560..3ade5f0 100644 --- a/training/run_rec_colab.py +++ b/training/run_rec_colab.py @@ -25,20 +25,12 @@ REMOTE_RUNNER = ROOT / "training/colab_remote.py" PRETRAINED_CHECKPOINT = ROOT / "training/.work/pretrained/PP-OCRv6_small_rec_pretrained.pdparams" SOURCE_FILES = ( - "pyproject.toml", - "uv.lock", - "scripts/batch_eval.py", "training/bootstrap.sh", "training/setup_rec_environment.sh", "training/run_rec_smoke.sh", - "training/evaluate_rec_checkpoint.sh", "training/configs/rec_pp_ocrv6_small.yaml", - "training/configs/pp_ocrv6_small_det.lock.json", - "training/scripts/prepare_detector.py", - "training/scripts/prepare_rapidocr_config.py", "training/scripts/prune_rec_checkpoints.py", "training/scripts/validate_annotations.py", - "training/colab_remote.py", ) @@ -102,21 +94,6 @@ def validate_labels(label_path: Path) -> tuple[int, list[str]]: return sample_count, images -def add_fixture_set(files: dict[str, tuple[Path, str]], relative_cases: str) -> None: - cases_path = ROOT / relative_cases - add_file(files, f"repo/{relative_cases}", cases_path, "evaluation-fixture") - cases = json.loads(cases_path.read_text(encoding="utf-8")) - for case in cases: - image_path = safe_relative(case["image"]) - source = cases_path.parent / image_path - add_file( - files, - f"repo/{cases_path.parent.relative_to(ROOT).as_posix()}/{image_path.as_posix()}", - source, - "evaluation-fixture", - ) - - def git_metadata(path: Path) -> tuple[str | None, bool]: revision = subprocess.run( ["git", "-C", str(path), "rev-parse", "HEAD"], @@ -153,14 +130,8 @@ def read_dataset_provenance(dataset_root: Path) -> dict[str, Any]: def source_files(root: Path, files: dict[str, tuple[Path, str]]) -> None: - for directory in ("app", "configs"): - for path in sorted((root / directory).rglob("*")): - if path.is_file() and not path.is_symlink() and "__pycache__" not in path.parts: - add_file(files, f"repo/{path.relative_to(root).as_posix()}", path, "ocrkit-source") for relative in SOURCE_FILES: add_file(files, f"repo/{relative}", root / relative, "ocrkit-source") - add_fixture_set(files, "datasets/fixtures/challenge/cases.json") - add_fixture_set(files, "tests/fixtures/run_code/cases.json") def stage_inputs( @@ -221,7 +192,7 @@ def stage_inputs( archive.add(request_path, arcname="request.json", recursive=False) -def stream_colab(command: list[str], log_path: Path) -> int: +def stream_command(command: list[str], log_path: Path) -> int: with log_path.open("a", encoding="utf-8") as log: rendered = shlex.join(command) print(f"$ {rendered}", flush=True) @@ -255,7 +226,7 @@ def upload_archive(colab: str, session: str, archive: Path, run_dir: Path, log_p for index, chunk in enumerate(iter(lambda: stream.read(PART_BYTES), b"")): part = parts_dir / f"part{index:04d}" part.write_bytes(chunk) - status = stream_colab( + status = stream_command( [colab, "upload", "-s", session, str(part), f"/content/ocrkit-input.part{index:04d}"], log_path, ) @@ -272,7 +243,7 @@ def download_archive(colab: str, session: str, destination: Path, run_dir: Path, parts_dir.mkdir() try: index_path = parts_dir / "index.json" - status = stream_colab( + status = stream_command( [colab, "download", "-s", session, "/content/ocrkit-result.index.json", str(index_path)], log_path ) if status: @@ -280,7 +251,7 @@ def download_archive(colab: str, session: str, destination: Path, run_dir: Path, with destination.open("wb") as result: for record in json.loads(index_path.read_text(encoding="utf-8"))["parts"]: part = parts_dir / safe_relative(record["name"]) - status = stream_colab( + status = stream_command( [colab, "download", "-s", session, f"/content/{record['name']}", str(part)], log_path ) if status: @@ -315,18 +286,13 @@ def copy_remote_result(remote_root: Path, run_dir: Path, *, success: bool) -> di if remote_metadata is None or remote_metadata.get("status") != "success": remote_error = (remote_metadata or {}).get("error", "no run metadata") raise ValueError(f"Colab run did not succeed: {remote_error}") - for name in ("checkpoint", "evaluation"): - if not (results / name).is_dir(): - raise ValueError(f"Colab did not return the {name} artifacts") + if not (results / "checkpoint").is_dir(): + raise ValueError("Colab did not return the checkpoint artifacts") output_records = remote_metadata.get("outputs") - required_outputs = { - "checkpoint/best_accuracy.pdparams", - "evaluation/fixture_report.json", - } - if not isinstance(output_records, list) or not required_outputs.issubset( - {record.get("path") for record in output_records} - ): - raise ValueError("Colab did not checksum the required checkpoint and evaluation report") + if not isinstance(output_records, list) or "checkpoint/best_accuracy.pdparams" not in { + record.get("path") for record in output_records + }: + raise ValueError("Colab did not checksum the best-accuracy checkpoint") for record in output_records: output = results / safe_relative(record["path"]) if ( @@ -335,14 +301,8 @@ def copy_remote_result(remote_root: Path, run_dir: Path, *, success: bool) -> di or sha256(output) != record["sha256"] ): raise ValueError(f"Colab result failed checksum verification: {record['path']}") - if not (results / "checkpoint/best_accuracy.pdparams").is_file(): - raise ValueError("Colab result is missing the best-accuracy checkpoint") - report_path = results / "evaluation/fixture_report.json" - if not report_path.is_file(): - raise ValueError("Colab result is missing the fixture evaluation report") accepted = run_dir / ACCEPTED_STAGING - for name in ("checkpoint", "evaluation"): - shutil.copytree(results / name, accepted / name) + shutil.copytree(results / "checkpoint", accepted / "checkpoint") shutil.copy2(metadata_path, accepted / "run.json") if remote_log.is_file(): shutil.copy2(remote_log, accepted / "remote.log") @@ -350,10 +310,8 @@ def copy_remote_result(remote_root: Path, run_dir: Path, *, success: bool) -> di partial = run_dir / "partial" partial.mkdir(exist_ok=True) - for name in ("checkpoint", "evaluation"): - source = results / name - if source.is_dir(): - shutil.copytree(source, partial / name) + if (results / "checkpoint").is_dir(): + shutil.copytree(results / "checkpoint", partial / "checkpoint") if metadata_path.is_file(): shutil.copy2(metadata_path, partial / "run.json") if remote_log.is_file(): @@ -361,13 +319,37 @@ def copy_remote_result(remote_root: Path, run_dir: Path, *, success: bool) -> di return remote_metadata or {} +def evaluate_locally(accepted: Path, log_path: Path) -> None: + """Run the existing local evaluation contract on the retrieved checkpoint.""" + evaluation = accepted / "evaluation" + status = stream_command( + [ + str(ROOT / "training/evaluate_rec_checkpoint.sh"), + str((accepted / "checkpoint/best_accuracy").resolve()), + str(evaluation.resolve()), + ], + log_path, + ) + if status: + raise RuntimeError(f"local checkpoint evaluation failed with exit status {status}") + report = json.loads((evaluation / "fixture_report.json").read_text(encoding="utf-8")) + metadata = json.loads((accepted / "run.json").read_text(encoding="utf-8")) + metadata["evaluation"] = { + "location": "local", + "report": "evaluation/fixture_report.json", + "field_accuracy": report.get("field_accuracy"), + "run_code_accuracy": report.get("run_code", {}).get("field_accuracy"), + } + (accepted / "run.json").write_text(json.dumps(metadata, ensure_ascii=False, indent=2) + "\n", encoding="utf-8") + + def main() -> int: - parser = argparse.ArgumentParser(description="Run the OCRKit recognition training and evaluation workflow on Colab GPU.") + parser = argparse.ArgumentParser(description="Train the OCRKit recognition model on a Colab GPU, then evaluate it locally.") parser.add_argument("--labels-dir", type=Path, default=ROOT / "datasets/labeled/rec") parser.add_argument("--pretrained-checkpoint", type=Path, default=PRETRAINED_CHECKPOINT) parser.add_argument("--epochs", type=int, default=10) parser.add_argument("--gpu", default="T4", help="Colab GPU preference; no accelerator fallback is attempted") - parser.add_argument("--timeout-seconds", type=float, default=6 * 3600, help="Upper bound for the remote training and evaluation run") + parser.add_argument("--timeout-seconds", type=float, default=6 * 3600, help="Upper bound for the remote training run") args = parser.parse_args() if args.epochs < 1: parser.error("--epochs must be a positive integer") @@ -444,7 +426,7 @@ def main() -> int: try: session_attempted = True - provision_status = stream_colab([colab, "new", "-s", session, "--gpu", args.gpu], colab_log) + provision_status = stream_command([colab, "new", "-s", session, "--gpu", args.gpu], colab_log) if provision_status: raise RuntimeError( f"Colab could not provision the requested GPU {args.gpu}; no fallback accelerator was selected." @@ -453,7 +435,7 @@ def main() -> int: input_archive.unlink(missing_ok=True) if upload_status: raise RuntimeError("Colab failed to stage OCRKit training inputs") - exec_status = stream_colab( + exec_status = stream_command( [colab, "exec", "-s", session, "--timeout", str(args.timeout_seconds), "-f", str(REMOTE_RUNNER)], colab_log, ) @@ -474,9 +456,9 @@ def main() -> int: pass raise if exec_status: - raise RuntimeError(f"Colab training or evaluation failed with exit status {exec_status}.") + raise RuntimeError(f"Colab training failed with exit status {exec_status}.") if remote_metadata.get("status") != "success": - raise RuntimeError("Colab training or evaluation did not return a successful status.") + raise RuntimeError("Colab training did not return a successful status.") except KeyboardInterrupt: error = "Colab run interrupted by the operator." except Exception as exc: @@ -485,7 +467,7 @@ def main() -> int: input_archive.unlink(missing_ok=True) if session_attempted: try: - stop_status = stream_colab([colab, "stop", "-s", session], colab_log) + stop_status = stream_command([colab, "stop", "-s", session], colab_log) except Exception as exc: stop_status = -1 stop_error = f"Colab runtime teardown command failed: {type(exc).__name__}: {exc}" @@ -493,8 +475,15 @@ def main() -> int: if stop_status and error is None: error = f"training completed, but Colab runtime teardown failed; run colab stop -s {session}" - succeeded = error is None and stop_status in (None, 0) accepted = run_dir / ACCEPTED_STAGING + if error is None and stop_status in (None, 0) and accepted.is_dir(): + try: + evaluate_locally(accepted, colab_log) + except KeyboardInterrupt: + error = "Local checkpoint evaluation interrupted by the operator." + except Exception as exc: + error = f"{type(exc).__name__}: {exc}" + succeeded = error is None and stop_status in (None, 0) if accepted.is_dir(): if succeeded: for child in accepted.iterdir(): @@ -520,7 +509,7 @@ def main() -> int: print(f"Colab OCRKit run failed: {error or 'runtime did not stop successfully'}", file=sys.stderr) print(f"Local run logs and any partial outputs: {run_dir}", file=sys.stderr) return 1 - print(f"Colab OCRKit run completed. Checkpoint, evaluation, provenance, and logs: {run_dir}") + print(f"Colab OCRKit run completed. Checkpoint, local evaluation, provenance, and logs: {run_dir}") return 0 diff --git a/training/run_rec_smoke.sh b/training/run_rec_smoke.sh index 115e444..015ae6f 100755 --- a/training/run_rec_smoke.sh +++ b/training/run_rec_smoke.sh @@ -11,7 +11,7 @@ output_dir="${work_dir}/checkpoints/rec_pp_ocrv6_small" epoch_num=10 resume_checkpoint="" device=cpu -evaluation_dir="" +train_only=false while [[ $# -gt 0 ]]; do case "$1" in @@ -35,12 +35,12 @@ while [[ $# -gt 0 ]]; do device="$2" shift 2 ;; - --evaluation-dir) - evaluation_dir="$2" - shift 2 + --train-only) + train_only=true + shift ;; *) - printf 'usage: %s [--labels-dir ] [--output-dir ] [--epochs ] [--resume-checkpoint ] [--device cpu|cuda] [--evaluation-dir ]\n' "$0" >&2 + printf 'usage: %s [--labels-dir ] [--output-dir ] [--epochs ] [--resume-checkpoint ] [--device cpu|cuda] [--train-only]\n' "$0" >&2 exit 2 ;; esac @@ -104,9 +104,11 @@ cd "${paddleocr_dir}" cd "${root_dir}" "${python_bin}" "${root_dir}/training/scripts/prune_rec_checkpoints.py" "${output_dir}" -if [[ -z "${evaluation_dir}" ]]; then - evaluation_dir="${work_dir}/evaluations/rec_pp_ocrv6_small/$(date -u +%Y.%m.%d-%H%M%S)-$$" +if [[ "${train_only}" == true ]]; then + exit 0 fi + +evaluation_dir="${work_dir}/evaluations/rec_pp_ocrv6_small/$(date -u +%Y.%m.%d-%H%M%S)-$$" "${root_dir}/training/evaluate_rec_checkpoint.sh" \ "${output_dir}/best_accuracy" \ "${evaluation_dir}" From 857fb972921b20b17f5dbb7e4ae222190277cd7d Mon Sep 17 00:00:00 2001 From: e54-bot Date: Sun, 27 Sep 2026 10:28:50 +0800 Subject: [PATCH 12/15] fix(training): return only the best checkpoint and retry Colab transfers Return best_accuracy.pdparams, config and the training log instead of the full checkpoint directory. The optimizer state and latest checkpoint are not needed until resume is supported and made the result about five times larger. Retry each part transfer up to three times. Refs #20 --- tests/test_run_rec_colab.py | 14 ++++++++++++++ training/colab_remote.py | 28 ++++++++++++++-------------- training/run_rec_colab.py | 15 ++++++++++++--- 3 files changed, 40 insertions(+), 17 deletions(-) diff --git a/tests/test_run_rec_colab.py b/tests/test_run_rec_colab.py index cb6ac25..a1019fb 100644 --- a/tests/test_run_rec_colab.py +++ b/tests/test_run_rec_colab.py @@ -218,3 +218,17 @@ def test_remote_parts_round_trip(tmp_path: Path, monkeypatch: pytest.MonkeyPatch index = json.loads((tmp_path / "out.index.json").read_text()) assert b"".join((tmp_path / part["name"]).read_bytes() for part in index["parts"]) == payload assert len(index["parts"]) == 4 + + +def test_transfer_retries_a_transient_failure(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + statuses = iter([1, 1, 0]) + attempts: list[list[str]] = [] + + def flaky(command: list[str], _log: Path) -> int: + attempts.append(command) + return next(statuses) + + monkeypatch.setattr(run_rec_colab, "stream_command", flaky) + + assert run_rec_colab.transfer_with_retry(["colab", "download"], tmp_path / "log") == 0 + assert len(attempts) == 3 diff --git a/training/colab_remote.py b/training/colab_remote.py index b754391..dd4855a 100644 --- a/training/colab_remote.py +++ b/training/colab_remote.py @@ -24,6 +24,7 @@ DATASET = COLAB_ROOT / "dataset" RESULTS = COLAB_ROOT / "results" CHECKPOINTS = RESULTS / "checkpoint" +RETURNED_CHECKPOINT_FILES = ("best_accuracy.pdparams", "config.yml", "train.log") RUN_METADATA = RESULTS / "run.json" REMOTE_LOG = RESULTS / "remote.log" @@ -111,18 +112,16 @@ def verify_inputs(request: dict[str, Any], root: Path) -> None: def output_files() -> list[dict[str, Any]]: records = [] - for directory in (CHECKPOINTS,): - if not directory.is_dir(): - continue - for path in sorted(directory.rglob("*")): - if path.is_file(): - records.append( - { - "path": path.relative_to(RESULTS).as_posix(), - "size_bytes": path.stat().st_size, - "sha256": sha256(path), - } - ) + for name in RETURNED_CHECKPOINT_FILES: + path = CHECKPOINTS / name + if path.is_file(): + records.append( + { + "path": path.relative_to(RESULTS).as_posix(), + "size_bytes": path.stat().st_size, + "sha256": sha256(path), + } + ) return records @@ -154,8 +153,9 @@ def write_result_archive() -> None: for path in (RUN_METADATA, REMOTE_LOG): if path.is_file(): archive.add(path, arcname=path.relative_to(COLAB_ROOT)) - if CHECKPOINTS.is_dir(): - archive.add(CHECKPOINTS, arcname=CHECKPOINTS.relative_to(COLAB_ROOT)) + for name in RETURNED_CHECKPOINT_FILES: + if (CHECKPOINTS / name).is_file(): + archive.add(CHECKPOINTS / name, arcname=(CHECKPOINTS / name).relative_to(COLAB_ROOT)) split_result_archive() diff --git a/training/run_rec_colab.py b/training/run_rec_colab.py index 3ade5f0..22e9b71 100644 --- a/training/run_rec_colab.py +++ b/training/run_rec_colab.py @@ -217,6 +217,15 @@ def stream_command(command: list[str], log_path: Path) -> int: return process.wait() +def transfer_with_retry(command: list[str], log_path: Path, attempts: int = 3) -> int: + status = 1 + for _ in range(attempts): + status = stream_command(command, log_path) + if not status: + break + return status + + def upload_archive(colab: str, session: str, archive: Path, run_dir: Path, log_path: Path) -> int: """Colab uploads are single JSON requests, so send the archive in bounded parts.""" parts_dir = run_dir / "input-parts" @@ -226,7 +235,7 @@ def upload_archive(colab: str, session: str, archive: Path, run_dir: Path, log_p for index, chunk in enumerate(iter(lambda: stream.read(PART_BYTES), b"")): part = parts_dir / f"part{index:04d}" part.write_bytes(chunk) - status = stream_command( + status = transfer_with_retry( [colab, "upload", "-s", session, str(part), f"/content/ocrkit-input.part{index:04d}"], log_path, ) @@ -243,7 +252,7 @@ def download_archive(colab: str, session: str, destination: Path, run_dir: Path, parts_dir.mkdir() try: index_path = parts_dir / "index.json" - status = stream_command( + status = transfer_with_retry( [colab, "download", "-s", session, "/content/ocrkit-result.index.json", str(index_path)], log_path ) if status: @@ -251,7 +260,7 @@ def download_archive(colab: str, session: str, destination: Path, run_dir: Path, with destination.open("wb") as result: for record in json.loads(index_path.read_text(encoding="utf-8"))["parts"]: part = parts_dir / safe_relative(record["name"]) - status = stream_command( + status = transfer_with_retry( [colab, "download", "-s", session, f"/content/{record['name']}", str(part)], log_path ) if status: From e48d8e5b7994bbf0a5d214b6b71fecffa726f825 Mon Sep 17 00:00:00 2001 From: e54-bot Date: Sun, 27 Sep 2026 15:01:04 +0800 Subject: [PATCH 13/15] feat(training): retrieve Colab checkpoints through R2, not the Colab CLI The Colab CLI transfers files as single base64-encoded JSON requests, which measured well under 1 MB/s for large binaries. Retrieving a 119 MB checkpoint through it took several minutes and repeatedly lost the connection when the Colab runtime was recycled mid-transfer, corrupting or losing the run. Have Colab upload the trained checkpoint directly to a private R2 bucket using a short-lived, single-object presigned PUT URL the runner generates locally with existing OCRKIT_R2_* credentials; Colab never receives R2 credentials, and the URL is redacted before it is logged or persisted. The runner downloads the checkpoint from R2 after stopping the Colab runtime, verifies its checksum, and deletes the object whether or not the run succeeded. Only the small run.json and remote.log now travel through the Colab CLI. The 125 MB official base checkpoint no longer needs to be uploaded from the local machine either: Colab fetches it directly from PaddlePaddle's public model CDN and verifies it against a pinned sha256. A checkpoint passed via --pretrained-checkpoint continues to upload through the CLI, since it is not publicly hosted. Add per-stage timing to the remote run for future capacity planning, and require the R2 environment variables up front with a clear error instead of failing partway through a run. Verified live: presigned PUT from a Colab T4 session to R2 uploaded 120 MB in about 4 seconds. Retrieval and checksum verification back through boto3 are covered by new tests, falsified by disabling the delete-on-finally guarantee, the early-record-before-later-failure guarantee, and the R2 error-code mapping, each of which failed the corresponding test before the fix. Refs #20 --- app/storage/r2_client.py | 31 +++ tests/test_r2_client_checkpoint_transfer.py | 75 ++++++ tests/test_run_rec_colab.py | 221 +++++++++++------- training/README.md | 46 ++-- training/colab_remote.py | 192 +++++++++------- training/run_rec_colab.py | 238 ++++++++++---------- 6 files changed, 514 insertions(+), 289 deletions(-) create mode 100644 tests/test_r2_client_checkpoint_transfer.py diff --git a/app/storage/r2_client.py b/app/storage/r2_client.py index 17afe56..cc03368 100644 --- a/app/storage/r2_client.py +++ b/app/storage/r2_client.py @@ -1,6 +1,7 @@ from __future__ import annotations from dataclasses import dataclass +from pathlib import Path from socket import timeout as SocketTimeout import boto3 @@ -135,3 +136,33 @@ def get_object_bytes( if code in {"AccessDenied", "403"}: raise ObjectAccessDeniedError("Object access denied") from exc raise ObjectDownloadError("Object download failed") from exc + + def download_object(self, bucket: str, object_key: str, destination: Path) -> None: + try: + self._client.download_file(bucket, object_key, str(destination)) + except (ConnectTimeoutError, ReadTimeoutError, SocketTimeout) as exc: + raise ObjectTimeoutError("Object download timed out") from exc + except ClientError as exc: + code = exc.response.get("Error", {}).get("Code", "") + if code in {"NoSuchKey", "NoSuchBucket", "404"}: + raise ObjectNotFoundError("Object not found") from exc + if code in {"AccessDenied", "403"}: + raise ObjectAccessDeniedError("Object access denied") from exc + raise ObjectDownloadError("Object download failed") from exc + + def delete_object(self, bucket: str, object_key: str) -> None: + try: + self._client.delete_object(Bucket=bucket, Key=object_key) + except ClientError as exc: + code = exc.response.get("Error", {}).get("Code", "") + if code in {"AccessDenied", "403"}: + raise ObjectAccessDeniedError("Object delete access denied") from exc + raise ObjectDownloadError("Object delete failed") from exc + + def generate_presigned_put_url(self, bucket: str, object_key: str, expires_in_seconds: int) -> str: + """A short-lived, single-object write URL a remote runtime can use without holding credentials.""" + return self._client.generate_presigned_url( + "put_object", + Params={"Bucket": bucket, "Key": object_key}, + ExpiresIn=expires_in_seconds, + ) diff --git a/tests/test_r2_client_checkpoint_transfer.py b/tests/test_r2_client_checkpoint_transfer.py new file mode 100644 index 0000000..bab720e --- /dev/null +++ b/tests/test_r2_client_checkpoint_transfer.py @@ -0,0 +1,75 @@ +from __future__ import annotations + +from pathlib import Path + +import pytest +from botocore.exceptions import ClientError + +from app.storage.r2_client import ObjectAccessDeniedError, ObjectNotFoundError, R2ObjectStore + + +def _client_error(code: str) -> ClientError: + return ClientError({"Error": {"Code": code, "Message": code}}, "operation") + + +@pytest.fixture +def store() -> R2ObjectStore: + return R2ObjectStore.from_settings( + endpoint_url="https://example.invalid", + access_key_id="key", + secret_access_key="secret", + region_name="auto", + default_bucket="bucket", + allowed_buckets_raw="", + read_timeout_seconds=5, + ) + + +def test_generate_presigned_put_url_targets_the_given_bucket_and_key(store: R2ObjectStore) -> None: + url = store.generate_presigned_put_url("bucket", "colab-runs/run-1/checkpoint.pdparams", expires_in_seconds=900) + + assert "colab-runs/run-1/checkpoint.pdparams" in url + assert "X-Amz-Signature" in url + + +def test_download_object_writes_the_destination_file(store: R2ObjectStore, tmp_path: Path) -> None: + destination = tmp_path / "checkpoint.pdparams" + + def fake_download_file(bucket: str, key: str, filename: str) -> None: + assert (bucket, key) == ("bucket", "checkpoint.pdparams") + Path(filename).write_bytes(b"weights") + + store._client.download_file = fake_download_file # type: ignore[attr-defined] + + store.download_object("bucket", "checkpoint.pdparams", destination) + + assert destination.read_bytes() == b"weights" + + +def test_download_object_missing_key_raises_not_found(store: R2ObjectStore, tmp_path: Path) -> None: + def fake_download_file(*_args: object) -> None: + raise _client_error("NoSuchKey") + + store._client.download_file = fake_download_file # type: ignore[attr-defined] + + with pytest.raises(ObjectNotFoundError): + store.download_object("bucket", "missing.pdparams", tmp_path / "out") + + +def test_delete_object_access_denied_raises_access_denied_error(store: R2ObjectStore) -> None: + def fake_delete_object(**_kwargs: object) -> None: + raise _client_error("AccessDenied") + + store._client.delete_object = fake_delete_object # type: ignore[attr-defined] + + with pytest.raises(ObjectAccessDeniedError): + store.delete_object("bucket", "checkpoint.pdparams") + + +def test_delete_object_succeeds_without_error(store: R2ObjectStore) -> None: + calls: list[dict[str, object]] = [] + store._client.delete_object = lambda **kwargs: calls.append(kwargs) # type: ignore[attr-defined] + + store.delete_object("bucket", "checkpoint.pdparams") + + assert calls == [{"Bucket": "bucket", "Key": "checkpoint.pdparams"}] diff --git a/tests/test_run_rec_colab.py b/tests/test_run_rec_colab.py index a1019fb..005c56a 100644 --- a/tests/test_run_rec_colab.py +++ b/tests/test_run_rec_colab.py @@ -8,31 +8,32 @@ import pytest +from app.storage.r2_client import ObjectNotFoundError from training import run_rec_colab SESSION_STOP = "stop" -def build_result_archive(path: Path, *, status: str = "success") -> None: - files = { - "results/checkpoint/best_accuracy.pdparams": b"weights", - "results/remote.log": b"log", - } - outputs = [ - { - "path": name.removeprefix("results/"), - "size_bytes": len(data), - "sha256": hashlib.sha256(data).hexdigest(), - } - for name, data in files.items() - if name.startswith("results/checkpoint/") - ] - files["results/run.json"] = json.dumps({"status": status, "outputs": outputs}).encode() - with tarfile.open(path, "w:gz") as archive: - for name, data in files.items(): - source = path.parent / "member" - source.write_bytes(data) - archive.add(source, arcname=name) +class FakeR2Store: + """Stands in for R2ObjectStore: an in-memory object map keyed by (bucket, key).""" + + def __init__(self) -> None: + self.default_bucket = "test-bucket" + self.objects: dict[tuple[str, str], bytes] = {} + self.deleted: list[tuple[str, str]] = [] + + def generate_presigned_put_url(self, bucket: str, key: str, expires_in_seconds: int) -> str: + return f"https://example.invalid/put/{bucket}/{key}?expires={expires_in_seconds}" + + def download_object(self, bucket: str, key: str, destination: Path) -> None: + data = self.objects.get((bucket, key)) + if data is None: + raise ObjectNotFoundError(f"no such object: {bucket}/{key}") + destination.write_bytes(data) + + def delete_object(self, bucket: str, key: str) -> None: + self.objects.pop((bucket, key), None) + self.deleted.append((bucket, key)) @pytest.fixture @@ -45,7 +46,20 @@ def colab_run(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): checkpoint.write_bytes(b"base") runs = tmp_path / "runs" calls: list[list[str]] = [] - behavior = {"exec_status": 0, "result_status": "success", "stop_status": 0, "eval_status": 0} + checkpoint_bytes = b"weights" + behavior = { + "exec_status": 0, + "remote_status": "success", + "stop_status": 0, + "eval_status": 0, + "upload_checkpoint": True, + "run_request": None, + "remote_log_text": "remote log contents\n", + "checkpoint_config": "Global: {}\n", + "train_log": "epoch 1 done\n", + "metadata_download_status": 0, + } + fake_r2 = FakeR2Store() def fake_stream(command: list[str], _log: Path) -> int: calls.append(command) @@ -58,30 +72,36 @@ def fake_stream(command: list[str], _log: Path) -> int: return 0 verb = command[1] if verb == "exec": + # Stands in for colab_remote.py actually running: populate the fake R2 object and + # the small metadata payload that fetch_remote_metadata/retrieve_checkpoint consume. + run_request = behavior["run_request"] + metadata: dict[str, object] = {"status": behavior["remote_status"]} + if behavior["remote_status"] != "success": + metadata["error"] = "RuntimeError: training crashed" + if behavior["upload_checkpoint"]: + upload = run_request["checkpoint_upload"] + fake_r2.objects[(upload["bucket"], upload["key"])] = checkpoint_bytes + metadata["checkpoint"] = { + "bucket": upload["bucket"], + "key": upload["key"], + "size_bytes": len(checkpoint_bytes), + "sha256": hashlib.sha256(checkpoint_bytes).hexdigest(), + } + metadata["checkpoint_config"] = behavior["checkpoint_config"] + metadata["train_log"] = behavior["train_log"] + behavior["remote_metadata"] = metadata return behavior["exec_status"] if verb == "download": - remote, local = command[-2], Path(command[-1]) - if "remote_result" not in behavior: - archive = local.parent / "remote-result.tar.gz" - build_result_archive(archive, status=behavior["result_status"]) - behavior["remote_result"] = archive.read_bytes() - archive.unlink() - data = behavior["remote_result"] - parts = [data[i : i + 2000] for i in range(0, len(data), 2000)] - if remote.endswith("index.json"): - local.write_text( - json.dumps( - { - "parts": [ - {"name": f"ocrkit-result.part{i:04d}", "sha256": hashlib.sha256(part).hexdigest()} - for i, part in enumerate(parts) - ] - } - ) - ) - else: - local.write_bytes(parts[int(remote.rsplit("part", 1)[1])]) - return 0 + remote_path, local_path = command[4], Path(command[-1]) + if remote_path.endswith("run.json"): + if behavior["metadata_download_status"]: + return behavior["metadata_download_status"] + local_path.write_text(json.dumps(behavior["remote_metadata"]), encoding="utf-8") + return 0 + if remote_path.endswith("remote.log"): + local_path.write_text(behavior["remote_log_text"], encoding="utf-8") + return 0 + raise AssertionError(f"unexpected download path: {remote_path}") if verb == SESSION_STOP: return behavior["stop_status"] return 0 @@ -90,6 +110,8 @@ def fake_stream(command: list[str], _log: Path) -> int: monkeypatch.setattr(run_rec_colab.shutil, "which", lambda _name: "colab") monkeypatch.setattr(run_rec_colab, "validate_labels", lambda _path: (1, ["a.png"])) monkeypatch.setattr(run_rec_colab, "PART_BYTES", 4) + monkeypatch.setattr(run_rec_colab, "require_r2_store", lambda _parser: fake_r2) + def fake_stage(archive: Path, run_request: dict, *_rest) -> None: behavior["run_request"] = run_request archive.write_bytes(b"input-archive") @@ -99,7 +121,7 @@ def fake_stage(archive: Path, run_request: dict, *_rest) -> None: monkeypatch.setattr( sys, "argv", ["run_rec_colab.py", "--labels-dir", str(labels), "--pretrained-checkpoint", str(checkpoint)] ) - return runs, calls, behavior + return runs, calls, behavior, fake_r2 def only_run(runs: Path) -> Path: @@ -107,20 +129,24 @@ def only_run(runs: Path) -> Path: return run_dir -def test_success_keeps_verified_artifacts_and_stops_runtime(colab_run) -> None: - runs, calls, behavior = colab_run +def test_success_retrieves_checkpoint_from_r2_and_deletes_it(colab_run) -> None: + runs, calls, behavior, fake_r2 = colab_run assert run_rec_colab.main() == 0 request = behavior["run_request"] assert request["training"]["epochs"] == 10 assert set(request["paddle_wheel_mirror"]) == {"url", "sha256"} + assert request["checkpoint_upload"]["bucket"] == fake_r2.default_bucket uploads = [command[-1] for command in calls if command[1] == "upload"] assert uploads == [f"/content/ocrkit-input.part{i:04d}" for i in range(4)] exec_command = next(command for command in calls if command[1] == "exec") assert float(exec_command[exec_command.index("--timeout") + 1]) == 6 * 3600 + run_dir = only_run(runs) - assert (run_dir / "checkpoint/best_accuracy.pdparams").is_file() + assert (run_dir / "checkpoint/best_accuracy.pdparams").read_bytes() == b"weights" + assert (run_dir / "checkpoint/config.yml").read_text() == "Global: {}\n" + assert (run_dir / "remote.log").read_text() == "remote log contents\n" assert (run_dir / "evaluation/fixture_report.json").is_file() assert json.loads((run_dir / "run.json").read_text())["evaluation"]["field_accuracy"] == 0.99 assert json.loads((run_dir / "status.json").read_text())["runtime_stopped"] is True @@ -128,11 +154,13 @@ def test_success_keeps_verified_artifacts_and_stops_runtime(colab_run) -> None: evaluation_index = next(i for i, command in enumerate(calls) if command[0].endswith("evaluate_rec_checkpoint.sh")) assert stop_index < evaluation_index assert not (run_dir / "accepted").exists() + assert fake_r2.deleted == [(request["checkpoint_upload"]["bucket"], request["checkpoint_upload"]["key"])] + assert not fake_r2.objects -def test_training_failure_keeps_partial_output_and_stops_runtime(colab_run) -> None: - runs, calls, behavior = colab_run - behavior.update(exec_status=1, result_status="failed") +def test_training_failure_keeps_partial_checkpoint_and_stops_runtime(colab_run) -> None: + runs, calls, behavior, fake_r2 = colab_run + behavior.update(exec_status=1, remote_status="failed") assert run_rec_colab.main() == 1 @@ -142,10 +170,11 @@ def test_training_failure_keeps_partial_output_and_stops_runtime(colab_run) -> N assert (run_dir / "partial/checkpoint/best_accuracy.pdparams").is_file() assert calls[-1][1] == SESSION_STOP assert json.loads((run_dir / "status.json").read_text())["status"] == "failed" + assert not fake_r2.objects # still deleted even though the run overall failed def test_local_evaluation_failure_demotes_checkpoint_to_partial(colab_run) -> None: - runs, calls, behavior = colab_run + runs, _calls, behavior, _fake_r2 = colab_run behavior["eval_status"] = 1 assert run_rec_colab.main() == 1 @@ -159,7 +188,7 @@ def test_local_evaluation_failure_demotes_checkpoint_to_partial(colab_run) -> No def test_teardown_failure_demotes_accepted_artifacts_to_partial(colab_run) -> None: - runs, _, behavior = colab_run + runs, _calls, behavior, _fake_r2 = colab_run behavior["stop_status"] = 1 assert run_rec_colab.main() == 1 @@ -171,7 +200,7 @@ def test_teardown_failure_demotes_accepted_artifacts_to_partial(colab_run) -> No def test_provisioning_failure_stops_runtime_and_uploads_nothing(colab_run) -> None: - runs, calls, _ = colab_run + runs, calls, _behavior, fake_r2 = colab_run original = run_rec_colab.stream_command run_rec_colab.stream_command = lambda command, log: 1 if command[1] == "new" else original(command, log) try: @@ -181,43 +210,50 @@ def test_provisioning_failure_stops_runtime_and_uploads_nothing(colab_run) -> No assert [command[1] for command in calls] == [SESSION_STOP] assert not (only_run(runs) / "checkpoint").exists() + assert not fake_r2.objects + + +def test_remote_metadata_download_failure_is_reported_clearly(colab_run) -> None: + runs, _calls, behavior, _fake_r2 = colab_run + behavior["metadata_download_status"] = 1 + + assert run_rec_colab.main() == 1 + + error = json.loads((only_run(runs) / "status.json").read_text())["error"] + assert "run metadata" in error -def test_tampered_result_is_rejected(tmp_path: Path) -> None: - remote = tmp_path / "retrieved" - archive = tmp_path / "result.tar.gz" - build_result_archive(archive) - run_rec_colab.safe_extract_result(archive, remote) - (remote / "results/checkpoint/best_accuracy.pdparams").write_bytes(b"tampered") +def test_tampered_checkpoint_is_rejected_and_still_deleted(tmp_path: Path) -> None: + fake_r2 = FakeR2Store() + fake_r2.objects[("bucket", "key")] = b"tampered-bytes" + remote_metadata = { + "status": "success", + "checkpoint": {"bucket": "bucket", "key": "key", "size_bytes": 999, "sha256": "0" * 64}, + } run_dir = tmp_path / "run" run_dir.mkdir() with pytest.raises(ValueError, match="checksum"): - run_rec_colab.copy_remote_result(remote, run_dir, success=True) + run_rec_colab.retrieve_checkpoint(fake_r2, remote_metadata, None, run_dir) - assert not (run_dir / "accepted").exists() + assert not fake_r2.objects # cleaned up even though verification failed + assert json.loads((run_dir / "accepted" / "run.json").read_text())["status"] == "success" -def test_remote_parts_round_trip(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: - from training import colab_remote +def test_stage_inputs_skips_uploading_the_official_checkpoint(tmp_path: Path) -> None: + dataset_root = tmp_path / "dataset" + (dataset_root / "labels" / "images").mkdir(parents=True) + (dataset_root / "labels" / "train.txt").write_text("a.png\ta\n", encoding="utf-8") + (dataset_root / "labels" / "holdout.txt").write_text("", encoding="utf-8") + (dataset_root / "labels" / "images" / "a.png").write_bytes(b"png") + archive_path = tmp_path / "input.tar.gz" - monkeypatch.setattr(colab_remote, "CONTENT", tmp_path) - monkeypatch.setattr(colab_remote, "INPUT_ARCHIVE", tmp_path / "in.tar.gz") - monkeypatch.setattr(colab_remote, "RESULT_ARCHIVE", tmp_path / "out.tar.gz") - monkeypatch.setattr(colab_remote, "RESULT_INDEX", tmp_path / "out.index.json") - monkeypatch.setattr(colab_remote, "PART_BYTES", 3) - payload = b"0123456789" - for index in range(2): - (tmp_path / f"ocrkit-input.part{index:04d}").write_bytes(payload[index * 5 : index * 5 + 5]) - colab_remote.join_input_parts() - assert (tmp_path / "in.tar.gz").read_bytes() == payload - assert not list(tmp_path.glob("ocrkit-input.part*")) + run_rec_colab.stage_inputs(archive_path, {"run_id": "x"}, dataset_root, None, ["a.png"], []) - (tmp_path / "out.tar.gz").write_bytes(payload) - colab_remote.split_result_archive() - index = json.loads((tmp_path / "out.index.json").read_text()) - assert b"".join((tmp_path / part["name"]).read_bytes() for part in index["parts"]) == payload - assert len(index["parts"]) == 4 + with tarfile.open(archive_path) as archive: + names = archive.getnames() + assert not any("pretrained" in name for name in names) + assert "dataset/a.png" in names def test_transfer_retries_a_transient_failure(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: @@ -232,3 +268,30 @@ def flaky(command: list[str], _log: Path) -> int: assert run_rec_colab.transfer_with_retry(["colab", "download"], tmp_path / "log") == 0 assert len(attempts) == 3 + + +def test_remote_join_input_parts_round_trip(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + from training import colab_remote + + monkeypatch.setattr(colab_remote, "CONTENT", tmp_path) + monkeypatch.setattr(colab_remote, "INPUT_ARCHIVE", tmp_path / "in.tar.gz") + payload = b"0123456789" + for index in range(2): + (tmp_path / f"ocrkit-input.part{index:04d}").write_bytes(payload[index * 5 : index * 5 + 5]) + + colab_remote.join_input_parts() + + assert (tmp_path / "in.tar.gz").read_bytes() == payload + assert not list(tmp_path.glob("ocrkit-input.part*")) + + +def test_remote_upload_checkpoint_raises_on_transport_failure(tmp_path: Path) -> None: + from training import colab_remote + + checkpoint = tmp_path / "best_accuracy.pdparams" + checkpoint.write_bytes(b"weights") + + with pytest.raises(RuntimeError, match="uploading the checkpoint"): + colab_remote.upload_checkpoint( + {"bucket": "b", "key": "k", "url": "http://127.0.0.1:1/unreachable"}, checkpoint + ) diff --git a/training/README.md b/training/README.md index e4b49db..c0da217 100644 --- a/training/README.md +++ b/training/README.md @@ -424,17 +424,24 @@ report for inspection but returns a non-zero status. ### Run recognition training on Colab GPU Install the [Google Colab CLI](https://github.com/googlecolab/google-colab-cli) -on the local Mac and prepare the reviewed/materialized dataset plus the local -PP-OCRv6 small base checkpoint: +and prepare the reviewed/materialized dataset. `setup_rec_environment.sh` is +only needed locally if you also want CPU Smoke or a non-default (custom) base +checkpoint; the default official base checkpoint is fetched by Colab itself. ```bash uv tool install google-colab-cli -./training/setup_rec_environment.sh ``` The first CLI request can prompt for Google OAuth authentication in the terminal. The CLI keeps those credentials locally; OCRKit does not send -platform, R2, or release credentials to Colab. +platform or release credentials to Colab. The runner also requires +`OCRKIT_R2_ENDPOINT_URL`, `OCRKIT_R2_ACCESS_KEY_ID`, `OCRKIT_R2_SECRET_ACCESS_KEY`, +and `OCRKIT_R2_DEFAULT_BUCKET` (see `.env.model.example`): the trained +checkpoint is far too large to transfer efficiently through the Colab CLI, so +Colab uploads it straight to that private R2 bucket using a short-lived, +single-object presigned URL that the runner generates locally and never +writes to disk or a log; Colab never receives R2 credentials. The runner +downloads the checkpoint from R2 and deletes the object once it has done so. Train on Colab and evaluate the retrieved checkpoint locally with one command: @@ -470,21 +477,26 @@ checkpoint, stops the runtime, and then runs the unchanged environment from `setup_rec_environment.sh` is required). The run only succeeds if that evaluation passes the same gate as a local run. -The runner transfers only the selected train/holdout labels and referenced -crops, available review/snapshot provenance files, the base recognition -checkpoint, and the training scripts. It records source revisions, input -checksums, the effective training configuration, PaddleOCR revision, allocated -GPU details, checkpoint checksums, and the local evaluation summary in the -returned `run.json`. - -Checkpoints, `fixture_report.json`, provenance, the remote training log, the -Colab CLI log, and `status.json` are stored below the ignored +Through the Colab CLI, the runner transfers only the selected train/holdout +labels and referenced crops, available review/snapshot provenance files, the +training scripts, and (when `--pretrained-checkpoint` names a checkpoint other +than the official default) that custom checkpoint. The official default base +checkpoint is instead fetched by Colab directly from its public URL and +checksum-verified there, and the trained checkpoint returns through R2 rather +than the CLI. It records source revisions, input checksums, the effective +training configuration, PaddleOCR revision, allocated GPU details, checkpoint +checksums, and the local evaluation summary in the returned `run.json`. + +The checkpoint, `fixture_report.json`, provenance, the remote training log, +the Colab CLI log, and `status.json` are stored below the ignored `training/.work/colab-runs//` directory. A successful run does not publish a candidate or change the stable model channel. Provisioning, -staging, training, retrieval, and handled failures all stop the Colab runtime -after it has been allocated; it is stopped before the local evaluation starts. Failed runs keep diagnostics and -any partial output under `partial/`; if teardown itself fails, `status.json` -includes the named `colab stop` command to release that session. +staging, training, metadata retrieval, and handled failures all stop the +Colab runtime after it has been allocated; it is stopped before the local +evaluation starts. Failed runs keep diagnostics and any partial output under +`partial/`; if teardown itself fails, `status.json` includes the named +`colab stop` command to release that session. The R2 checkpoint object is +deleted once retrieved, on both success and failure. To evaluate a checkpoint explicitly, use a new output directory: diff --git a/training/colab_remote.py b/training/colab_remote.py index dd4855a..43cbddf 100644 --- a/training/colab_remote.py +++ b/training/colab_remote.py @@ -9,7 +9,9 @@ import subprocess import sys import tarfile +import time import traceback +import urllib.request from datetime import UTC, datetime from pathlib import Path from typing import Any @@ -17,16 +19,14 @@ COLAB_ROOT = Path("/content/ocrkit-colab") CONTENT = Path("/content") INPUT_ARCHIVE = CONTENT / "ocrkit-input.tar.gz" -RESULT_ARCHIVE = CONTENT / "ocrkit-result.tar.gz" -RESULT_INDEX = CONTENT / "ocrkit-result.index.json" PART_BYTES = 32 * 1024 * 1024 REPO = COLAB_ROOT / "repo" DATASET = COLAB_ROOT / "dataset" RESULTS = COLAB_ROOT / "results" CHECKPOINTS = RESULTS / "checkpoint" -RETURNED_CHECKPOINT_FILES = ("best_accuracy.pdparams", "config.yml", "train.log") RUN_METADATA = RESULTS / "run.json" REMOTE_LOG = RESULTS / "remote.log" +BASE_CHECKPOINT_PATH = REPO / "training/.work/pretrained/PP-OCRv6_small_rec_pretrained.pdparams" def sha256(path: Path) -> str: @@ -53,6 +53,7 @@ def run_logged(command: list[str], *, cwd: Path, log: Any, env: dict[str, str] | print(f"$ {rendered}", flush=True) log.write(f"$ {rendered}\n") log.flush() + started = time.monotonic() with subprocess.Popen( command, cwd=cwd, @@ -67,11 +68,25 @@ def run_logged(command: list[str], *, cwd: Path, log: Any, env: dict[str, str] | print(line, end="", flush=True) log.write(line) status = process.wait() + elapsed = round(time.monotonic() - started, 1) + log.write(f"$ {rendered} -> exit {status} in {elapsed}s\n") log.flush() if status: raise RuntimeError(f"remote command exited with status {status}: {rendered}") +def timed(stages: dict[str, float], name: str): + class _Timer: + def __enter__(self): + self._start = time.monotonic() + return self + + def __exit__(self, *_exc): + stages[name] = round(time.monotonic() - self._start, 1) + + return _Timer() + + def checked_gpu() -> dict[str, Any]: if not shutil.which("nvidia-smi"): raise RuntimeError("Colab did not provide an NVIDIA GPU runtime") @@ -110,21 +125,6 @@ def verify_inputs(request: dict[str, Any], root: Path) -> None: raise RuntimeError(f"staged input failed checksum verification: {record['path']}") -def output_files() -> list[dict[str, Any]]: - records = [] - for name in RETURNED_CHECKPOINT_FILES: - path = CHECKPOINTS / name - if path.is_file(): - records.append( - { - "path": path.relative_to(RESULTS).as_posix(), - "size_bytes": path.stat().st_size, - "sha256": sha256(path), - } - ) - return records - - def join_input_parts() -> None: parts = sorted(CONTENT.glob("ocrkit-input.part*")) if not parts: @@ -135,34 +135,52 @@ def join_input_parts() -> None: part.unlink() -def split_result_archive() -> None: - for stale in (*CONTENT.glob("ocrkit-result.part*"), RESULT_INDEX): - stale.unlink(missing_ok=True) - parts = [] - with RESULT_ARCHIVE.open("rb") as archive: - for index, chunk in enumerate(iter(lambda: archive.read(PART_BYTES), b"")): - name = f"ocrkit-result.part{index:04d}" - (CONTENT / name).write_bytes(chunk) - parts.append({"name": name, "sha256": hashlib.sha256(chunk).hexdigest()}) - RESULT_INDEX.write_text(json.dumps({"parts": parts}), encoding="utf-8") +def fetch_base_checkpoint(base_checkpoint: dict[str, Any]) -> None: + """A checkpoint the runner did not upload; Colab downloads it directly instead.""" + BASE_CHECKPOINT_PATH.parent.mkdir(parents=True, exist_ok=True) + urllib.request.urlretrieve(base_checkpoint["url"], BASE_CHECKPOINT_PATH) + if sha256(BASE_CHECKPOINT_PATH) != base_checkpoint["sha256"]: + raise RuntimeError("downloaded base checkpoint failed checksum verification") -def write_result_archive() -> None: - RESULT_ARCHIVE.unlink(missing_ok=True) - with tarfile.open(RESULT_ARCHIVE, "w:gz") as archive: - for path in (RUN_METADATA, REMOTE_LOG): - if path.is_file(): - archive.add(path, arcname=path.relative_to(COLAB_ROOT)) - for name in RETURNED_CHECKPOINT_FILES: - if (CHECKPOINTS / name).is_file(): - archive.add(CHECKPOINTS / name, arcname=(CHECKPOINTS / name).relative_to(COLAB_ROOT)) - split_result_archive() +def upload_checkpoint(upload: dict[str, Any], checkpoint_path: Path) -> dict[str, Any]: + """PUTs the trained checkpoint straight to R2; the Colab CLI's own transfer is far slower.""" + result = subprocess.run( + [ + "curl", + "-fsS", + "-X", + "PUT", + "--data-binary", + f"@{checkpoint_path}", + "-H", + "Content-Type: application/octet-stream", + "--max-time", + "1800", + "-o", + "/dev/null", + "-w", + "%{http_code}", + upload["url"], + ], + capture_output=True, + text=True, + ) + if result.returncode or result.stdout.strip() not in {"200", "201"}: + raise RuntimeError(f"uploading the checkpoint to R2 failed (curl exit {result.returncode}: {result.stderr.strip()[-300:]})") + return { + "bucket": upload["bucket"], + "key": upload["key"], + "size_bytes": checkpoint_path.stat().st_size, + "sha256": sha256(checkpoint_path), + } def main() -> int: RESULTS.mkdir(parents=True, exist_ok=True) + stages: dict[str, float] = {} result: dict[str, Any] = { - "schema_version": 1, + "schema_version": 2, "status": "failed", "completed_at": datetime.now(UTC).isoformat(), } @@ -174,30 +192,34 @@ def main() -> int: safe_extract(INPUT_ARCHIVE, COLAB_ROOT) request_path = COLAB_ROOT / "request.json" request = json.loads(request_path.read_text(encoding="utf-8")) - result["request"] = request["run"] + run_request = request["run"] + # The presigned checkpoint-upload URL is a write credential; never persist or log it. + result["request"] = { + key: (value if key != "checkpoint_upload" else {k: v for k, v in value.items() if k != "url"}) + for key, value in run_request.items() + } verify_inputs(request, COLAB_ROOT) gpu = checked_gpu() REPO.mkdir(parents=True, exist_ok=True) DATASET.mkdir(parents=True, exist_ok=True) CHECKPOINTS.mkdir(parents=True, exist_ok=True) + if run_request["base_checkpoint"]["source"] == "official-download": + with timed(stages, "fetch_base_checkpoint"): + fetch_base_checkpoint(run_request["base_checkpoint"]) + elif not BASE_CHECKPOINT_PATH.is_file(): + raise RuntimeError("uploaded base checkpoint was not staged at the expected path") + if not shutil.which("uv"): + run_logged([sys.executable, "-m", "pip", "install", "uv"], cwd=COLAB_ROOT, log=log) + mirror = run_request["paddle_wheel_mirror"] + with timed(stages, "setup_environment"): run_logged( - [sys.executable, "-m", "pip", "install", "uv"], - cwd=COLAB_ROOT, + ["bash", "training/setup_rec_environment.sh", "--device", "cuda"], + cwd=REPO, + env={**os.environ, "OCRKIT_PADDLE_WHEEL_URL": mirror["url"], "OCRKIT_PADDLE_WHEEL_SHA256": mirror["sha256"]}, log=log, ) - mirror = request["run"]["paddle_wheel_mirror"] - run_logged( - ["bash", "training/setup_rec_environment.sh", "--device", "cuda"], - cwd=REPO, - env={ - **os.environ, - "OCRKIT_PADDLE_WHEEL_URL": mirror["url"], - "OCRKIT_PADDLE_WHEEL_SHA256": mirror["sha256"], - }, - log=log, - ) paddle = subprocess.run( [ @@ -213,31 +235,43 @@ def main() -> int: if not paddle_info["cuda"] or paddle_info["device"] != "gpu:0": raise RuntimeError("PaddlePaddle did not select the allocated CUDA device") - epochs = str(request["run"]["training"]["epochs"]) - run_logged( - [ - "bash", - "training/run_rec_smoke.sh", - "--labels-dir", - str(DATASET), - "--output-dir", - str(CHECKPOINTS), - "--epochs", - epochs, - "--device", - "cuda", - "--train-only", - ], - cwd=REPO, - log=log, - ) + epochs = str(run_request["training"]["epochs"]) + with timed(stages, "train"): + run_logged( + [ + "bash", + "training/run_rec_smoke.sh", + "--labels-dir", + str(DATASET), + "--output-dir", + str(CHECKPOINTS), + "--epochs", + epochs, + "--device", + "cuda", + "--train-only", + ], + cwd=REPO, + log=log, + ) best_checkpoint = CHECKPOINTS / "best_accuracy.pdparams" + if not best_checkpoint.is_file() or best_checkpoint.stat().st_size == 0: + raise RuntimeError("training did not produce the best-accuracy recognition checkpoint") + + with timed(stages, "upload_checkpoint"): + checkpoint_record = upload_checkpoint(run_request["checkpoint_upload"], best_checkpoint) + # Recorded immediately: if a later step fails, the caller still knows this object + # exists in R2 and can retrieve or delete it instead of leaking it silently. + result["checkpoint"] = checkpoint_record + paddleocr_revision = subprocess.run( ["git", "-C", str(REPO / "training/.work/PaddleOCR"), "rev-parse", "HEAD"], check=True, capture_output=True, text=True, ).stdout.strip() + config_path = CHECKPOINTS / "config.yml" + train_log_path = CHECKPOINTS / "train.log" result.update( { "status": "success", @@ -247,34 +281,32 @@ def main() -> int: "paddle": paddle_info, "paddleocr_revision": paddleocr_revision, }, - "outputs": output_files(), + "stage_seconds": stages, + "checkpoint": checkpoint_record, + "checkpoint_config": config_path.read_text(encoding="utf-8") if config_path.is_file() else None, + "train_log": train_log_path.read_text(encoding="utf-8", errors="replace") if train_log_path.is_file() else None, } ) except BaseException as exc: result["error"] = f"{type(exc).__name__}: {exc}" + result["stage_seconds"] = stages traceback.print_exc(file=log) log.flush() finally: result["completed_at"] = datetime.now(UTC).isoformat() - result["outputs"] = output_files() - RUN_METADATA.write_text( - json.dumps(result, ensure_ascii=False, indent=2) + "\n", - encoding="utf-8", - ) - write_result_archive() + RUN_METADATA.write_text(json.dumps(result, ensure_ascii=False, indent=2) + "\n", encoding="utf-8") except BaseException as exc: print(f"OCRKit Colab run failed: {type(exc).__name__}: {exc}", file=sys.stderr, flush=True) try: RUN_METADATA.write_text( json.dumps( - {"schema_version": 1, "status": "failed", "error": f"{type(exc).__name__}: {exc}"}, + {"schema_version": 2, "status": "failed", "error": f"{type(exc).__name__}: {exc}", "stage_seconds": stages}, ensure_ascii=False, indent=2, ) + "\n", encoding="utf-8", ) - write_result_archive() except OSError: pass return 1 diff --git a/training/run_rec_colab.py b/training/run_rec_colab.py index 22e9b71..9226523 100644 --- a/training/run_rec_colab.py +++ b/training/run_rec_colab.py @@ -14,6 +14,9 @@ from pathlib import Path, PurePosixPath, PureWindowsPath from typing import Any +from app.core.config import settings +from app.storage.r2_client import ObjectNotFoundError, R2ObjectStore + ROOT = Path(__file__).resolve().parents[1] RUNS = ROOT / "training/.work/colab-runs" ACCEPTED_STAGING = "accepted" @@ -21,6 +24,13 @@ "url": "https://cdn.owbastion.codes/ocrkit/wheels/cu129/paddlepaddle_gpu-3.3.1-cp312-cp312-linux_x86_64.whl", "sha256": "03fc5211183ba20ef71e63a35e589fd386a8805c1011c3a409a0e5b118be4668", } +OFFICIAL_BASE_CHECKPOINT = { + "model": "PP-OCRv6_small_rec_pretrained", + "url": "https://paddle-model-ecology.bj.bcebos.com/paddlex/official_pretrained_model/PP-OCRv6_small_rec_pretrained.pdparams", + "sha256": "25c9bd54b0e5900916e8bb6ada938abeffb1eac1baedac0ca54a45b1c9310825", +} +R2_KEY_PREFIX = "colab-runs" +R2_UPLOAD_URL_BUFFER_SECONDS = 900 PART_BYTES = 32 * 1024 * 1024 REMOTE_RUNNER = ROOT / "training/colab_remote.py" PRETRAINED_CHECKPOINT = ROOT / "training/.work/pretrained/PP-OCRv6_small_rec_pretrained.pdparams" @@ -34,6 +44,29 @@ ) +def require_r2_store(parser: argparse.ArgumentParser) -> R2ObjectStore: + """The Colab backend retrieves checkpoints through R2 rather than the slow Colab CLI transfer.""" + if not ( + settings.r2_endpoint_url + and settings.r2_access_key_id + and settings.r2_secret_access_key + and settings.r2_default_bucket + ): + parser.error( + "OCRKIT_R2_ENDPOINT_URL, OCRKIT_R2_ACCESS_KEY_ID, OCRKIT_R2_SECRET_ACCESS_KEY, and " + "OCRKIT_R2_DEFAULT_BUCKET are required to run training on Colab" + ) + return R2ObjectStore.from_settings( + endpoint_url=settings.r2_endpoint_url, + access_key_id=settings.r2_access_key_id, + secret_access_key=settings.r2_secret_access_key, + region_name=settings.r2_region_name, + default_bucket=settings.r2_default_bucket, + allowed_buckets_raw=settings.r2_allowed_buckets, + read_timeout_seconds=settings.r2_read_timeout_seconds, + ) + + def sha256(path: Path) -> str: digest = hashlib.sha256() with path.open("rb") as stream: @@ -138,7 +171,7 @@ def stage_inputs( archive_path: Path, run_request: dict[str, Any], dataset_root: Path, - checkpoint_path: Path, + checkpoint_path: Path | None, train_images: list[str], holdout_images: list[str], ) -> None: @@ -167,12 +200,13 @@ def stage_inputs( path = dataset_root / relative if path.is_file(): add_file(files, f"dataset/{relative}", path, "dataset-provenance") - add_file( - files, - "repo/training/.work/pretrained/PP-OCRv6_small_rec_pretrained.pdparams", - checkpoint_path, - "base-recognition-checkpoint", - ) + if checkpoint_path is not None: + add_file( + files, + "repo/training/.work/pretrained/PP-OCRv6_small_rec_pretrained.pdparams", + checkpoint_path, + "base-recognition-checkpoint", + ) records = [] with tarfile.open(archive_path, "w:gz") as archive: @@ -247,85 +281,64 @@ def upload_archive(colab: str, session: str, archive: Path, run_dir: Path, log_p shutil.rmtree(parts_dir, ignore_errors=True) -def download_archive(colab: str, session: str, destination: Path, run_dir: Path, log_path: Path) -> int: - parts_dir = run_dir / "result-parts" - parts_dir.mkdir() - try: - index_path = parts_dir / "index.json" - status = transfer_with_retry( - [colab, "download", "-s", session, "/content/ocrkit-result.index.json", str(index_path)], log_path - ) - if status: - return status - with destination.open("wb") as result: - for record in json.loads(index_path.read_text(encoding="utf-8"))["parts"]: - part = parts_dir / safe_relative(record["name"]) - status = transfer_with_retry( - [colab, "download", "-s", session, f"/content/{record['name']}", str(part)], log_path - ) - if status: - return status - if sha256(part) != record["sha256"]: - raise ValueError(f"Colab result part failed checksum verification: {record['name']}") - result.write(part.read_bytes()) - part.unlink() - return 0 - finally: - shutil.rmtree(parts_dir, ignore_errors=True) +def fetch_remote_metadata(colab: str, session: str, run_dir: Path, log_path: Path) -> tuple[dict[str, Any], Path | None]: + """Only run.json and remote.log travel through the Colab CLI; both are small, unlike the checkpoint.""" + metadata_path = run_dir / "remote-run.json" + status = transfer_with_retry( + [colab, "download", "-s", session, "/content/ocrkit-colab/results/run.json", str(metadata_path)], log_path + ) + if status: + raise RuntimeError("Colab did not return run metadata (results/run.json)") + remote_log_path = run_dir / "remote-log.txt" + if transfer_with_retry( + [colab, "download", "-s", session, "/content/ocrkit-colab/results/remote.log", str(remote_log_path)], log_path + ): + remote_log_path = None # best-effort: diagnostics, not required for the run's outcome + return json.loads(metadata_path.read_text(encoding="utf-8")), remote_log_path -def safe_extract_result(archive_path: Path, destination: Path) -> None: - destination.mkdir(parents=True, exist_ok=True) - root = destination.resolve() - with tarfile.open(archive_path, "r:gz") as archive: - members = archive.getmembers() - for member in members: - target = (destination / member.name).resolve() - if not target.is_relative_to(root) or not (member.isfile() or member.isdir()): - raise ValueError("Colab result archive contains an unsupported path or file type") - archive.extractall(destination) - - -def copy_remote_result(remote_root: Path, run_dir: Path, *, success: bool) -> dict[str, Any]: - results = remote_root / "results" - metadata_path = results / "run.json" - remote_metadata = json.loads(metadata_path.read_text(encoding="utf-8")) if metadata_path.is_file() else None - remote_log = results / "remote.log" - if success: - if remote_metadata is None or remote_metadata.get("status") != "success": - remote_error = (remote_metadata or {}).get("error", "no run metadata") - raise ValueError(f"Colab run did not succeed: {remote_error}") - if not (results / "checkpoint").is_dir(): - raise ValueError("Colab did not return the checkpoint artifacts") - output_records = remote_metadata.get("outputs") - if not isinstance(output_records, list) or "checkpoint/best_accuracy.pdparams" not in { - record.get("path") for record in output_records - }: - raise ValueError("Colab did not checksum the best-accuracy checkpoint") - for record in output_records: - output = results / safe_relative(record["path"]) - if ( - not output.is_file() - or output.stat().st_size != record["size_bytes"] - or sha256(output) != record["sha256"] - ): - raise ValueError(f"Colab result failed checksum verification: {record['path']}") - accepted = run_dir / ACCEPTED_STAGING - shutil.copytree(results / "checkpoint", accepted / "checkpoint") - shutil.copy2(metadata_path, accepted / "run.json") - if remote_log.is_file(): - shutil.copy2(remote_log, accepted / "remote.log") - return remote_metadata - - partial = run_dir / "partial" - partial.mkdir(exist_ok=True) - if (results / "checkpoint").is_dir(): - shutil.copytree(results / "checkpoint", partial / "checkpoint") - if metadata_path.is_file(): - shutil.copy2(metadata_path, partial / "run.json") - if remote_log.is_file(): - shutil.copy2(remote_log, partial / "remote.log") - return remote_metadata or {} +def retrieve_checkpoint( + r2: R2ObjectStore, remote_metadata: dict[str, Any], remote_log_path: Path | None, run_dir: Path +) -> None: + """The checkpoint travels through R2 (uploaded by colab_remote.py), not the Colab CLI.""" + accepted = run_dir / ACCEPTED_STAGING + accepted.mkdir(parents=True, exist_ok=True) + (accepted / "run.json").write_text( + json.dumps(remote_metadata, ensure_ascii=False, indent=2) + "\n", encoding="utf-8" + ) + if remote_log_path is not None: + shutil.move(str(remote_log_path), accepted / "remote.log") + + checkpoint = remote_metadata.get("checkpoint") + if checkpoint is None: + return + checkpoint_dir = accepted / "checkpoint" + checkpoint_dir.mkdir(parents=True, exist_ok=True) + destination = checkpoint_dir / "best_accuracy.pdparams" + try: + try: + r2.download_object(checkpoint["bucket"], checkpoint["key"], destination) + except ObjectNotFoundError as exc: + raise ValueError( + f"Colab did not upload the checkpoint to R2: {checkpoint['bucket']}/{checkpoint['key']}" + ) from exc + if destination.stat().st_size != checkpoint["size_bytes"] or sha256(destination) != checkpoint["sha256"]: + raise ValueError("Colab checkpoint failed checksum verification after download from R2") + config_text = remote_metadata.get("checkpoint_config") + if config_text: + (checkpoint_dir / "config.yml").write_text(config_text, encoding="utf-8") + train_log_text = remote_metadata.get("train_log") + if train_log_text: + (checkpoint_dir / "train.log").write_text(train_log_text, encoding="utf-8") + finally: + try: + r2.delete_object(checkpoint["bucket"], checkpoint["key"]) + except Exception as exc: + print( + f"warning: failed to delete the Colab checkpoint object from R2 " + f"({checkpoint['bucket']}/{checkpoint['key']}): {exc}", + file=sys.stderr, + ) def evaluate_locally(accepted: Path, log_path: Path) -> None: @@ -366,12 +379,17 @@ def main() -> int: if not colab: parser.error("Google Colab CLI is required; install it with uv tool install google-colab-cli") + r2 = require_r2_store(parser) + dataset_root = args.labels_dir.resolve() train_label = dataset_root / "labels/train.txt" holdout_label = dataset_root / "labels/holdout.txt" - for path in (train_label, holdout_label, args.pretrained_checkpoint): + for path in (train_label, holdout_label): if not path.is_file(): parser.error(f"required training input is missing: {path}") + using_official_checkpoint = args.pretrained_checkpoint == PRETRAINED_CHECKPOINT + if not using_official_checkpoint and not args.pretrained_checkpoint.is_file(): + parser.error(f"required training input is missing: {args.pretrained_checkpoint}") train_count, train_images = validate_labels(train_label) holdout_count, holdout_images = validate_labels(holdout_label) @@ -379,9 +397,7 @@ def main() -> int: run_dir = RUNS / run_id run_dir.mkdir(parents=True, exist_ok=False) input_archive = run_dir / "ocrkit-input.tar.gz" - result_archive = run_dir / "ocrkit-result.tar.gz" colab_log = run_dir / "colab.log" - retrieved = run_dir / "retrieved" source_revision, source_dirty = git_metadata(ROOT) dataset_revision, dataset_dirty = git_metadata(dataset_root) @@ -389,7 +405,21 @@ def main() -> int: dataset_source = dataset_root.relative_to(ROOT).as_posix() except ValueError: dataset_source = "external" - checkpoint_path = args.pretrained_checkpoint.resolve() + if using_official_checkpoint: + checkpoint_path = None + base_checkpoint = {**OFFICIAL_BASE_CHECKPOINT, "source": "official-download"} + else: + checkpoint_path = args.pretrained_checkpoint.resolve() + base_checkpoint = { + "model": "custom", + "source": "uploaded", + "source_name": checkpoint_path.name, + "sha256": sha256(checkpoint_path), + } + checkpoint_upload_key = f"{R2_KEY_PREFIX}/{run_id}/checkpoint.pdparams" + checkpoint_upload_url = r2.generate_presigned_put_url( + r2.default_bucket, checkpoint_upload_key, expires_in_seconds=int(args.timeout_seconds) + R2_UPLOAD_URL_BUFFER_SECONDS + ) run_request = { "run_id": run_id, "ocrkit_revision": source_revision, @@ -402,12 +432,9 @@ def main() -> int: "train_samples": train_count, "holdout_samples": holdout_count, }, - "base_checkpoint": { - "model": "PP-OCRv6_small_rec_pretrained", - "source_name": checkpoint_path.name, - "sha256": sha256(checkpoint_path), - }, + "base_checkpoint": base_checkpoint, "paddle_wheel_mirror": PADDLE_WHEEL, + "checkpoint_upload": {"bucket": r2.default_bucket, "key": checkpoint_upload_key, "url": checkpoint_upload_url}, "training": { "epochs": args.epochs, "device": "cuda", @@ -448,26 +475,14 @@ def main() -> int: [colab, "exec", "-s", session, "--timeout", str(args.timeout_seconds), "-f", str(REMOTE_RUNNER)], colab_log, ) - download_status = download_archive(colab, session, result_archive, run_dir, colab_log) - if download_status: - raise RuntimeError( - "Colab failed to retrieve the remote run logs and artifacts" - + (f" after the remote run exited with status {exec_status}" if exec_status else "") - ) - safe_extract_result(result_archive, retrieved) - try: - remote_metadata = copy_remote_result(retrieved, run_dir, success=exec_status == 0) - except Exception: - shutil.rmtree(run_dir / ACCEPTED_STAGING, ignore_errors=True) - try: - copy_remote_result(retrieved, run_dir, success=False) - except Exception: - pass - raise + remote_metadata, remote_log_path = fetch_remote_metadata(colab, session, run_dir, colab_log) + retrieve_checkpoint(r2, remote_metadata, remote_log_path, run_dir) if exec_status: raise RuntimeError(f"Colab training failed with exit status {exec_status}.") if remote_metadata.get("status") != "success": - raise RuntimeError("Colab training did not return a successful status.") + raise RuntimeError( + f"Colab training did not return a successful status: {remote_metadata.get('error', 'unknown error')}" + ) except KeyboardInterrupt: error = "Colab run interrupted by the operator." except Exception as exc: @@ -511,9 +526,6 @@ def main() -> int: "error": error, } (run_dir / "status.json").write_text(json.dumps(status, ensure_ascii=False, indent=2) + "\n", encoding="utf-8") - result_archive.unlink(missing_ok=True) - if retrieved.exists(): - shutil.rmtree(retrieved) if status["status"] != "success": print(f"Colab OCRKit run failed: {error or 'runtime did not stop successfully'}", file=sys.stderr) print(f"Local run logs and any partial outputs: {run_dir}", file=sys.stderr) From 91950a5ca8c8e117867550641f0d40fcebe01c1a Mon Sep 17 00:00:00 2001 From: e54-bot Date: Sun, 27 Sep 2026 15:13:25 +0800 Subject: [PATCH 14/15] fix(training): make run_rec_colab.py importable of `app` when run directly Running `python training/run_rec_colab.py` puts training/, not the repo root, on sys.path, so the new `from app...` imports failed with ModuleNotFoundError. Insert the repo root before importing, matching the sys.path shim already used by training/scripts/*.py. Refs #20 --- training/run_rec_colab.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/training/run_rec_colab.py b/training/run_rec_colab.py index 9226523..a13553d 100644 --- a/training/run_rec_colab.py +++ b/training/run_rec_colab.py @@ -14,10 +14,14 @@ from pathlib import Path, PurePosixPath, PureWindowsPath from typing import Any -from app.core.config import settings -from app.storage.r2_client import ObjectNotFoundError, R2ObjectStore - ROOT = Path(__file__).resolve().parents[1] +if str(ROOT) not in sys.path: + # Running this file directly (python training/run_rec_colab.py) puts training/, not the + # repo root, on sys.path; the repo's own `app` package needs the root added explicitly. + sys.path.insert(0, str(ROOT)) + +from app.core.config import settings # noqa: E402 +from app.storage.r2_client import ObjectNotFoundError, R2ObjectStore # noqa: E402 RUNS = ROOT / "training/.work/colab-runs" ACCEPTED_STAGING = "accepted" PADDLE_WHEEL = { From d6818fca9ba27ff6bccbda1fd8778b3d9e19d01a Mon Sep 17 00:00:00 2001 From: e54-bot Date: Sun, 27 Sep 2026 15:53:09 +0800 Subject: [PATCH 15/15] fix(training): treat a Colab teardown 404 as already-released colab stop can fail with Not Found simply because Colab already reclaimed an idle or finished runtime on its own; a completed, checksum-verified training run was being reported as failed and its checkpoint demoted to partial/ solely because of this race. Verify with `colab sessions` before treating a failed stop as a real teardown failure. Verified live: a full T4 run (base checkpoint fetched by Colab, trained, checksum-verified checkpoint retrieved through R2) hit exactly this race; colab sessions confirmed no session was actually left running. Refs #20 --- tests/test_run_rec_colab.py | 29 ++++++++++++++++++++++++++++- training/run_rec_colab.py | 17 +++++++++++++++++ 2 files changed, 45 insertions(+), 1 deletion(-) diff --git a/tests/test_run_rec_colab.py b/tests/test_run_rec_colab.py index 005c56a..5d924fd 100644 --- a/tests/test_run_rec_colab.py +++ b/tests/test_run_rec_colab.py @@ -5,6 +5,7 @@ import sys import tarfile from pathlib import Path +from types import SimpleNamespace import pytest @@ -58,9 +59,17 @@ def colab_run(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): "checkpoint_config": "Global: {}\n", "train_log": "epoch 1 done\n", "metadata_download_status": 0, + "session_active_after_stop_failure": False, } fake_r2 = FakeR2Store() + def fake_capture(command: list[str]) -> SimpleNamespace: + assert command[1] == "sessions" + if behavior["session_active_after_stop_failure"]: + session_name = next(c[3] for c in calls if len(c) > 1 and c[1] == "new") + return SimpleNamespace(returncode=0, stdout=f"[{session_name}] fake-endpoint | Hardware: T4\n", stderr="") + return SimpleNamespace(returncode=0, stdout="[colab] No active sessions found on server.\n", stderr="") + def fake_stream(command: list[str], _log: Path) -> int: calls.append(command) if command[0].endswith("evaluate_rec_checkpoint.sh"): @@ -118,6 +127,7 @@ def fake_stage(archive: Path, run_request: dict, *_rest) -> None: monkeypatch.setattr(run_rec_colab, "stage_inputs", fake_stage) monkeypatch.setattr(run_rec_colab, "stream_command", fake_stream) + monkeypatch.setattr(run_rec_colab, "capture_command", fake_capture) monkeypatch.setattr( sys, "argv", ["run_rec_colab.py", "--labels-dir", str(labels), "--pretrained-checkpoint", str(checkpoint)] ) @@ -187,9 +197,10 @@ def test_local_evaluation_failure_demotes_checkpoint_to_partial(colab_run) -> No assert json.loads((run_dir / "status.json").read_text())["runtime_stopped"] is True -def test_teardown_failure_demotes_accepted_artifacts_to_partial(colab_run) -> None: +def test_teardown_failure_with_a_still_active_session_demotes_checkpoint_to_partial(colab_run) -> None: runs, _calls, behavior, _fake_r2 = colab_run behavior["stop_status"] = 1 + behavior["session_active_after_stop_failure"] = True assert run_rec_colab.main() == 1 @@ -199,6 +210,22 @@ def test_teardown_failure_demotes_accepted_artifacts_to_partial(colab_run) -> No assert "colab stop" in json.loads((run_dir / "status.json").read_text())["error"] +def test_teardown_failure_when_colab_already_released_the_session_still_succeeds(colab_run) -> None: + """`colab stop` can 404 simply because Colab already reclaimed a finished runtime.""" + runs, _calls, behavior, _fake_r2 = colab_run + behavior["stop_status"] = 1 + behavior["session_active_after_stop_failure"] = False + + assert run_rec_colab.main() == 0 + + run_dir = only_run(runs) + assert (run_dir / "checkpoint/best_accuracy.pdparams").is_file() + status = json.loads((run_dir / "status.json").read_text()) + assert status["status"] == "success" + assert status["runtime_stopped"] is True + assert status["error"] is None + + def test_provisioning_failure_stops_runtime_and_uploads_nothing(colab_run) -> None: runs, calls, _behavior, fake_r2 = colab_run original = run_rec_colab.stream_command diff --git a/training/run_rec_colab.py b/training/run_rec_colab.py index a13553d..18e7f48 100644 --- a/training/run_rec_colab.py +++ b/training/run_rec_colab.py @@ -255,6 +255,20 @@ def stream_command(command: list[str], log_path: Path) -> int: return process.wait() +def capture_command(command: list[str]) -> subprocess.CompletedProcess[str]: + return subprocess.run(command, cwd=ROOT, capture_output=True, text=True) + + +def session_still_active(colab: str, session: str, log_path: Path) -> bool: + """`colab stop` can fail simply because Colab already reclaimed an idle/finished runtime.""" + result = capture_command([colab, "sessions"]) + with log_path.open("a", encoding="utf-8") as log: + log.write(f"$ {colab} sessions\n{result.stdout}{result.stderr}\n") + if result.returncode: + return True # could not verify; assume active so a real leak is not silently dropped + return session in result.stdout + + def transfer_with_retry(command: list[str], log_path: Path, attempts: int = 3) -> int: status = 1 for _ in range(attempts): @@ -500,6 +514,9 @@ def main() -> int: stop_status = -1 stop_error = f"Colab runtime teardown command failed: {type(exc).__name__}: {exc}" error = f"{error}; {stop_error}" if error else stop_error + if stop_status and not session_still_active(colab, session, colab_log): + # Colab had already released the runtime on its own; nothing was left running. + stop_status = 0 if stop_status and error is None: error = f"training completed, but Colab runtime teardown failed; run colab stop -s {session}"