fc-data 0.2.0__py3-none-any.whl

datasmith/update/__init__.py

@@ -0,0 +1,5 @@
+"""ds.update — Pipeline orchestrator and CLI."""
+
+from datasmith.update.pipeline import Pipeline
+
+__all__ = ["Pipeline"]

datasmith/update/cli.py

@@ -0,0 +1,169 @@
+from __future__ import annotations
+
+import argparse
+import asyncio
+import os
+import re
+import signal
+import sys
+
+from datasmith.utils import get_logger
+
+logger = get_logger("update.cli")
+
+_DATE_RE = re.compile(r"^\d{4}-\d{2}-\d{2}$")
+
+
+_STAGE_DESCRIPTIONS = {
+    1: "scrape_repos      — Fetch repository metadata from GitHub for all tracked repos",
+    2: "scrape_commits    — Scrape merged PR commits and patches from each repository",
+    3: "classify_prs      — Use LLM agents to classify PRs as performance-related",
+    4: "resolve_packages  — Resolve Python dependencies for performance commits via uv",
+    5: "render_problems   — Scrape linked issues and render deconstructed problem contexts",
+    6: "synthesize_images — Generate Docker build contexts for confirmed performance commits",
+    7: "publish           — Build, verify, and publish Docker images to DockerHub",
+}
+
+
+def _stages_epilog() -> str:
+    lines = ["pipeline stages (run in order by default):"]
+    for num, desc in _STAGE_DESCRIPTIONS.items():
+        lines.append(f"  {num}. {desc}")
+    return "\n".join(lines)
+
+
+def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        prog="fc-data",
+        description="Run the FormulaCode update pipeline — discovers performance-improving "
+        "commits from GitHub, classifies them with LLM agents, and builds Docker images "
+        "for benchmark evaluation.",
+        epilog=_stages_epilog(),
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+    parser.add_argument("--start-date", required=True, help="Start of the date range to scan for commits (YYYY-MM-DD)")
+    parser.add_argument("--end-date", required=True, help="End of the date range to scan for commits (YYYY-MM-DD)")
+    parser.add_argument(
+        "--resume",
+        action="store_true",
+        help="Skip stages already marked complete and resume from the next pending stage",
+    )
+    parser.add_argument(
+        "--stage",
+        type=int,
+        default=None,
+        action="append",
+        metavar="N",
+        help="Run only stage N (1-7); repeat to run multiple stages (e.g. --stage 1 --stage 2)",
+    )
+    parser.add_argument("--dry-run", action="store_true", help="Log what each stage would do without executing")
+    parser.add_argument(
+        "--n-concurrent", type=int, default=None, metavar="N", help="Max concurrent items per runner stage"
+    )
+    parser.add_argument(
+        "--tasks-per-repo",
+        type=int,
+        default=None,
+        metavar="N",
+        help="Max tasks per repo for stages 5 (render_problems) and 6 (synthesize_images). Ignored by other stages.",
+    )
+    parser.add_argument(
+        "--agent",
+        type=str,
+        default=None,
+        choices=["claude", "codex", "gemini", "none"],
+        help="CLI agent to use for stage 6 synthesis (default: auto-detect first available). "
+        "'none' skips LLM generation and relies only on similar-context matching.",
+    )
+    parser.add_argument(
+        "--force",
+        action="store_true",
+        help="Re-run even for tasks already processed; applies to stages 5 (render_problems) and 6 (synthesize_images)",
+    )
+    parser.add_argument(
+        "--offline-source",
+        type=str,
+        default=None,
+        metavar="PATH",
+        help="Path to a parquet file with offline PR data to import into stages 1 and 2",
+    )
+    parser.add_argument(
+        "--min-stars",
+        type=int,
+        default=500,
+        metavar="N",
+        help="Minimum stars for GitHub code search repo discovery in stage 1 (default: 500)",
+    )
+    return parser.parse_args(argv)
+
+
+def validate_dates(args: argparse.Namespace) -> None:
+    for name in ("start_date", "end_date"):
+        val = getattr(args, name)
+        if not _DATE_RE.match(val):
+            print(f"Error: {name.replace('_', '-')} must be YYYY-MM-DD, got '{val}'", file=sys.stderr)
+            sys.exit(1)
+
+
+_sigint_count = 0
+
+
+def _sigint_handler(signum: int, frame: object) -> None:
+    """Handle CTRL+C by killing agent subprocesses so worker threads unblock.
+
+    Agent subprocesses run in their own sessions (``start_new_session=True``)
+    and don't receive SIGINT from the terminal.  Without this handler the
+    threads blocked on ``proc.communicate()`` never return and the process
+    hangs.
+
+    First CTRL+C: SIGTERM all tracked agent processes, raise KeyboardInterrupt.
+    Second CTRL+C: SIGKILL all tracked agent processes, force-exit immediately.
+    """
+    global _sigint_count
+    _sigint_count += 1
+
+    from datasmith.agents.installed.base import terminate_all_agents
+
+    if _sigint_count >= 2:
+        logger.info("Force-killing all agent subprocesses")
+        terminate_all_agents(force=True)
+        os._exit(1)
+
+    logger.info("Interrupted — terminating agent subprocesses (press Ctrl+C again to force-quit)")
+    terminate_all_agents(force=False)
+    raise KeyboardInterrupt
+
+
+def main(argv: list[str] | None = None) -> None:
+    args = parse_args(argv)
+    validate_dates(args)
+
+    from datasmith.update.pipeline import Pipeline
+
+    signal.signal(signal.SIGINT, _sigint_handler)
+
+    pipeline = Pipeline(
+        dry_run=args.dry_run,
+        n_concurrent=args.n_concurrent,
+        tasks_per_repo=args.tasks_per_repo,
+        agent=args.agent,
+        force=args.force,
+        offline_source=args.offline_source,
+        min_stars=args.min_stars,
+    )
+    try:
+        asyncio.run(
+            pipeline.run(
+                start_date=args.start_date,
+                end_date=args.end_date,
+                resume=args.resume,
+                stage=args.stage,  # None or list[int] from append action
+            )
+        )
+    except KeyboardInterrupt:
+        logger.info("Pipeline interrupted by user")
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

datasmith/update/offline.py

@@ -0,0 +1,173 @@
+"""Convert rows from an offline parquet source into ``pull_requests`` records."""
+
+from __future__ import annotations
+
+import re
+from typing import Any
+
+import pandas as pd
+
+from datasmith.filters import symbolic_compliance
+from datasmith.utils import get_logger
+
+logger = get_logger("update.offline")
+
+# Matches a data row in the markdown file-change table:
+#   | filename | additions | deletions | ... |
+_TABLE_ROW_RE = re.compile(r"^\|\s*(?P<filename>[^|]+?)\s*\|\s*(?P<additions>\d+)\s*\|\s*(?P<deletions>\d+)\s*\|")
+
+
+def _sanitize_text(value: object) -> str:
+    """Return a clean string, handling NaN/None and Postgres-illegal null bytes."""
+    if value is None or (isinstance(value, float) and pd.isna(value)):
+        return ""
+    return str(value).replace("\u0000", "")
+
+
+def parse_file_change_summary(summary: object) -> list[dict[str, Any]] | None:
+    """Parse a markdown file-change table into the ``file_changes`` list format.
+
+    Expected input::
+
+        | File | Lines Added | Lines Removed | Total Changes |
+        |------|-------------|----------------|----------------|
+        | foo.py | 10 | 3 | 13 |
+
+    Returns ``None`` when *summary* is empty/NaN or contains no parseable rows.
+    """
+    text = _sanitize_text(summary)
+    if not text:
+        return None
+    changes: list[dict[str, Any]] = []
+    for line in text.splitlines():
+        # Skip header and separator rows
+        if line.startswith("|--") or "Lines Added" in line or "File" in line:
+            continue
+        m = _TABLE_ROW_RE.match(line)
+        if m:
+            changes.append({
+                "filename": m.group("filename").strip(),
+                "additions": int(m.group("additions")),
+                "deletions": int(m.group("deletions")),
+            })
+    return changes or None
+
+
+def _extract_labels(raw_labels: Any) -> list[str]:
+    """Extract label name strings from the parquet labels column.
+
+    The column stores a numpy array of label dicts (each with a ``name`` key),
+    or an empty array.
+    """
+    if raw_labels is None:
+        return []
+    try:
+        return [label["name"] for label in raw_labels if isinstance(label, dict) and "name" in label]
+    except (TypeError, KeyError):
+        return []
+
+
+def _safe_str(value: object) -> str:
+    if value is None or (isinstance(value, float) and pd.isna(value)):
+        return ""
+    return str(value)
+
+
+def _dict_sha(value: object) -> str:
+    """Extract ``sha`` from a dict-like value (pr_head / pr_base columns)."""
+    if isinstance(value, dict):
+        return str(value.get("sha", ""))
+    return ""
+
+
+def load_offline_repo_names(path: str) -> list[tuple[str, str]]:
+    """Return unique ``(owner, repo)`` pairs from an offline parquet file."""
+    df = pd.read_parquet(path, columns=["repo_name"])
+    pairs: list[tuple[str, str]] = []
+    seen: set[tuple[str, str]] = set()
+    for name in df["repo_name"].dropna().unique():
+        name = str(name).strip()
+        if "/" not in name:
+            continue
+        owner, repo = name.split("/", 1)
+        pair = (owner, repo)
+        if pair not in seen:
+            seen.add(pair)
+            pairs.append(pair)
+    return pairs
+
+
+def _row_to_record(row: pd.Series) -> dict[str, Any]:
+    """Convert a single parquet row into a ``pull_requests`` upsert record."""
+    repo_name = str(row["repo_name"])
+    owner, repo = repo_name.split("/", 1)
+
+    title = _sanitize_text(row.get("pr_title"))
+    body = _sanitize_text(row.get("pr_body"))
+    patch = _sanitize_text(row.get("original_patch"))
+    file_changes = parse_file_change_summary(row.get("file_change_summary"))
+
+    record: dict[str, Any] = {
+        "owner": owner,
+        "repo": repo,
+        "issue_number": int(row["pr_number"]),
+        "title": title,
+        "body": body,
+        "state": _safe_str(row.get("pr_state")),
+        "created_at": _safe_str(row.get("pr_created_at")) or None,
+        "merged_at": _safe_str(row.get("pr_merged_at")) or None,
+        "closed_at": _safe_str(row.get("pr_closed_at")) or None,
+        "merge_commit_sha": _safe_str(row.get("pr_merge_commit_sha")),
+        "base_sha": _dict_sha(row.get("pr_base")),
+        "head_sha": _dict_sha(row.get("pr_head")),
+        "labels": _extract_labels(row.get("pr_labels")),
+        "is_performance_commit_symbolic": symbolic_compliance(
+            title=title,
+            patch=patch or None,
+            file_changes=file_changes,
+        ),
+    }
+    if patch:
+        record["patch"] = patch
+    if file_changes:
+        record["file_changes"] = file_changes
+    return record
+
+
+def load_offline_pull_requests(
+    path: str,
+    since: str | None = None,
+    until: str | None = None,
+) -> list[dict[str, Any]]:
+    """Load and convert parquet rows into ``pull_requests`` upsert records.
+
+    Filters by ``pr_merged_at`` using the same ``[since, until)`` semantics
+    as :class:`~datasmith.runners.scrape_commits.ScrapeCommitsRunner`.
+    """
+    df = pd.read_parquet(path)
+    logger.info("Loaded %d rows from offline source %s", len(df), path)
+
+    # Date filtering on pr_merged_at
+    if since or until:
+        merged: pd.Series[Any] = pd.to_datetime(df["pr_merged_at"], utc=True, errors="coerce")
+        if since:
+            df = df[merged >= pd.Timestamp(since, tz="UTC")]
+            merged = merged.loc[df.index]
+        if until:
+            df = df[merged < pd.Timestamp(until, tz="UTC")]
+        logger.info("After date filtering: %d rows", len(df))
+
+    records: list[dict[str, Any]] = []
+    for _, row in df.iterrows():
+        try:
+            records.append(_row_to_record(row))
+        except Exception:
+            logger.warning(
+                "Skipping row %s/%s#%s: conversion error",
+                row.get("repo_name", "?"),
+                row.get("pr_number", "?"),
+                row.get("pr_merge_commit_sha", "?"),
+                exc_info=True,
+            )
+    logger.info("Converted %d records for upsert", len(records))
+    return records