Summarize training steps across all buckets

Author: apstenku123

scripts/report_training_steps.py

@@ -217,6 +217,58 @@ def _format_int(value: int) -> str:
     return f"{value:,}"
 
 
+def _format_short_tokens(value: int) -> str:
+    if value >= 1_000_000_000:
+        return f"{value / 1_000_000_000:.3f}B"
+    if value >= 1_000_000:
+        return f"{value / 1_000_000:.1f}M"
+    return _format_int(value)
+
+
+def _rows_by_length_and_kind(rows: Iterable[ReportRow]) -> dict[int, dict[str, ReportRow]]:
+    result: dict[int, dict[str, ReportRow]] = {}
+    for row in rows:
+        result.setdefault(row.length, {})[row.kind] = row
+    return result
+
+
+def _print_summary(rows: Iterable[ReportRow]) -> None:
+    """Print one unambiguous curriculum row per context length."""
+    print(
+        "| length | bs | tokens/step | code tokens | code steps | commit+PR-doc tokens | commit steps | main tokens | main steps | standalone PR tokens | +PR total steps | skipped files |"
+    )
+    print("|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|")
+    for length, by_kind in sorted(_rows_by_length_and_kind(rows).items()):
+        main = by_kind.get("main_code_plus_commits")
+        if main is None:
+            continue
+        code = by_kind.get("code_only")
+        commits = by_kind.get("commits_with_pr_docstring")
+        pr = by_kind.get("standalone_pr_side_stream")
+        main_pr = by_kind.get("main_plus_standalone_pr")
+        skipped = sum(row.skipped_files for row in by_kind.values())
+        print(
+            "| "
+            + " | ".join(
+                [
+                    str(length),
+                    str(main.batch_size),
+                    _format_int(main.tokens_per_step),
+                    _format_short_tokens(code.trained_tokens) if code else "-",
+                    _format_int(code.steps_by_trained_tokens) if code else "-",
+                    _format_short_tokens(commits.trained_tokens) if commits else "-",
+                    _format_int(commits.steps_by_trained_tokens) if commits else "-",
+                    _format_short_tokens(main.trained_tokens),
+                    _format_int(main.steps_by_trained_tokens),
+                    _format_short_tokens(pr.trained_tokens) if pr else "-",
+                    _format_int(main_pr.steps_by_trained_tokens) if main_pr else "-",
+                    _format_int(skipped),
+                ]
+            )
+            + " |"
+        )
+
+
 def _print_markdown(rows: Iterable[ReportRow]) -> None:
     print(
         "| kind | length | bs | tokens/step | files | skipped | rows | docs | trained tokens | steps(trained) | valid tokens | steps(valid) |"
@@ -272,7 +324,12 @@ def main() -> int:
         action="store_true",
         help="Skip files that disappear or are unreadable during live conveyor writes, and report the skip count.",
     )
-    parser.add_argument("--format", choices=("text", "markdown", "json"), default="text")
+    parser.add_argument(
+        "--format",
+        choices=("summary", "text", "markdown", "json"),
+        default="summary",
+        help="summary prints one row per bucket; text/markdown print every source kind.",
+    )
     args = parser.parse_args()
 
     rows = build_report(
@@ -285,6 +342,8 @@ def main() -> int:
     )
     if args.format == "json":
         print(json.dumps([asdict(row) for row in rows], indent=2, sort_keys=True))
+    elif args.format == "summary":
+        _print_summary(rows)
     elif args.format == "markdown":
         _print_markdown(rows)
     else:

tests/test_report_training_steps.py

@@ -5,7 +5,7 @@
 import pyarrow as pa
 import pyarrow.parquet as pq
 
-from scripts.report_training_steps import build_report, parse_batch_schedule
+from scripts.report_training_steps import build_report, parse_batch_schedule, _print_summary
 
 
 def _write_counter_parquet(path: Path, rows: list[dict[str, int]]) -> None:
@@ -65,6 +65,43 @@ def test_report_keeps_code_commits_main_and_pr_stream_distinct(tmp_path: Path) -
     assert by_kind["main_plus_standalone_pr"].steps_by_trained_tokens == 1
 
 
+def test_summary_reports_every_length_once(tmp_path: Path, capsys) -> None:
+    code = tmp_path / "code"
+    commits = tmp_path / "commits"
+    pr = tmp_path / "pr"
+    _write_counter_parquet(
+        code / "1024" / "code.parquet",
+        [{"valid": 4096, "trained": 4096, "docs": 1}],
+    )
+    _write_counter_parquet(
+        commits / "2048" / "commits.parquet",
+        [{"valid": 8192, "trained": 8192, "docs": 1}],
+    )
+    _write_counter_parquet(
+        pr / "2048" / "pr.parquet",
+        [{"valid": 4096, "trained": 4096, "docs": 1}],
+    )
+
+    rows = build_report(
+        code_root=code,
+        commit_root=commits,
+        pr_root=pr,
+        batch_by_length={1024: 2, 2048: 2},
+        max_workers=1,
+        allow_concurrent_skips=False,
+    )
+
+    _print_summary(rows)
+    out = capsys.readouterr().out
+
+    assert "| 1024 | 2 | 2,048 |" in out
+    assert "| 2048 | 2 | 4,096 |" in out
+    assert out.count("| 1024 |") == 1
+    assert out.count("| 2048 |") == 1
+    assert "main tokens" in out
+    assert "standalone PR tokens" in out
+
+
 def test_parse_batch_schedule_rejects_bad_items() -> None:
     assert parse_batch_schedule("1024=192,2048=96") == {1024: 192, 2048: 96}