ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

[pre-commit.ci] auto fixes from pre-commit.com hooks

Browse source 
for more information, see https://pre-commit.ci
This commit is contained in:
pre-commit-ci[bot] 2026-04-24 13:50:14 +00:00
commit 6c05ba524e
8 changed files with 416 additions and 165 deletions
Download patch file Download diff file Expand all files Collapse all files

28
studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/config.py
View file

@ -14,28 +14,28 @@ class GitHubRepoSeedSource(SeedSource):
seed_type: Literal["github_repo"] = "github_repo"
repos: list[str] = Field(
default_factory=list,
description="List of GitHub repositories to scrape, each in `owner/name` form.",
default_factory = list,
description = "List of GitHub repositories to scrape, each in `owner/name` form.",
)
token: str = Field(
default="",
description="Personal access token. Leave blank to read GH_TOKEN / GITHUB_TOKEN from env at run time.",
default = "",
description = "Personal access token. Leave blank to read GH_TOKEN / GITHUB_TOKEN from env at run time.",
)
item_types: list[Literal["issues", "pulls", "commits"]] = Field(
default=["issues", "pulls"],
description="Which GitHub item types to fetch per repo.",
default = ["issues", "pulls"],
description = "Which GitHub item types to fetch per repo.",
)
limit: int = Field(
default=100,
ge=1,
le=5000,
description="Maximum items per repo per item type (e.g. limit=100 + ['issues','pulls'] => up to 200 items per repo).",
default = 100,
ge = 1,
le = 5000,
description = "Maximum items per repo per item type (e.g. limit=100 + ['issues','pulls'] => up to 200 items per repo).",
)
include_comments: bool = Field(
default=True,
description="Fetch the first N comments of each issue/PR and include them in the `comments` column.",
default = True,
description = "Fetch the first N comments of each issue/PR and include them in the `comments` column.",
)
max_comments_per_item: int = Field(default=30, ge=0, le=200)
max_comments_per_item: int = Field(default = 30, ge = 0, le = 200)
@field_validator("repos")
@classmethod
@ -57,7 +57,7 @@ class GitHubRepoSeedSource(SeedSource):
raise ValueError("item_types must not be empty")
return list(dict.fromkeys(v))
@model_validator(mode="after")
@model_validator(mode = "after")
def _ensure_repos(self) -> "GitHubRepoSeedSource":
if not self.repos:
raise ValueError("At least one repo is required")

12
studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/impl.py
View file

@ -20,12 +20,12 @@ class GitHubRepoSeedReader(SeedReader[GitHubRepoSeedSource]):
def get_dataset_uri(self) -> str:
out_dir = Path(tempfile.gettempdir()) / "studio-github-repo-seed"
cfg = ScrapeConfig(
repos=list(self.source.repos),
token=self.source.token,
item_types=list(self.source.item_types),
limit=self.source.limit,
include_comments=self.source.include_comments,
max_comments_per_item=self.source.max_comments_per_item,
repos = list(self.source.repos),
token = self.source.token,
item_types = list(self.source.item_types),
limit = self.source.limit,
include_comments = self.source.include_comments,
max_comments_per_item = self.source.max_comments_per_item,
)
path = materialize_to_jsonl(cfg, out_dir)
return str(path)

6
studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/plugin.py
View file

@ -4,7 +4,7 @@
from data_designer.plugins.plugin import Plugin, PluginType
github_repo_seed_plugin = Plugin(
impl_qualified_name="data_designer_github_repo_seed.impl.GitHubRepoSeedReader",
config_qualified_name="data_designer_github_repo_seed.config.GitHubRepoSeedSource",
plugin_type=PluginType.SEED_READER,
impl_qualified_name = "data_designer_github_repo_seed.impl.GitHubRepoSeedReader",
config_qualified_name = "data_designer_github_repo_seed.config.GitHubRepoSeedSource",
plugin_type = PluginType.SEED_READER,
)

47
studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper.py
View file

@ -9,6 +9,7 @@ After scraping, we read the per-resource JSONL shards and flatten them into
a single unified JSONL with stable columns (`item_type`, `repo`, `number`,
`title`, `body`, ...).
"""
from __future__ import annotations
import json
@ -59,7 +60,7 @@ def _resolve_token(token: str) -> str:
def _read_jsonl(path: Path, max_rows: int | None = None):
if not path.exists():
return
with path.open(encoding="utf-8") as f:
with path.open(encoding = "utf-8") as f:
for i, line in enumerate(f):
if not line.strip():
continue
@ -72,7 +73,11 @@ def _read_jsonl(path: Path, max_rows: int | None = None):
def _flatten_issue_row(r: dict, repo: str, include_comments: bool, max_c: int) -> dict:
labels = [l.get("name") for l in (r.get("labels", {}) or {}).get("nodes", []) if l.get("name")]
labels = [
l.get("name")
for l in (r.get("labels", {}) or {}).get("nodes", [])
if l.get("name")
]
comments_nodes = (r.get("comments") or {}).get("nodes") or []
comments_text = ""
if include_comments and comments_nodes:
@ -98,7 +103,11 @@ def _flatten_issue_row(r: dict, repo: str, include_comments: bool, max_c: int) -
def _flatten_pr_row(r: dict, repo: str, include_comments: bool, max_c: int) -> dict:
labels = [l.get("name") for l in (r.get("labels", {}) or {}).get("nodes", []) if l.get("name")]
labels = [
l.get("name")
for l in (r.get("labels", {}) or {}).get("nodes", [])
if l.get("name")
]
comments_nodes = (r.get("comments") or {}).get("nodes") or []
comments_text = ""
if include_comments and comments_nodes:
@ -126,7 +135,7 @@ def _flatten_pr_row(r: dict, repo: str, include_comments: bool, max_c: int) -> d
def _flatten_commit_row(r: dict, repo: str) -> dict:
msg = r.get("messageHeadline") or r.get("message") or ""
body = r.get("messageBody") or r.get("message") or msg
author = (r.get("author") or {})
author = r.get("author") or {}
return {
"item_type": "commit",
"repo": repo,
@ -148,7 +157,7 @@ def scrape(cfg: ScrapeConfig, base_dir: Path):
os.environ["GH_TOKEN"] = token
GitHubClient, RepoScraper = _load_impl()
client = GitHubClient()
base_dir.mkdir(parents=True, exist_ok=True)
base_dir.mkdir(parents = True, exist_ok = True)
# Per-resource trial limits. limit <= 0 means "all": use a very large cap.
effective_limit = cfg.limit if cfg.limit and cfg.limit > 0 else 1_000_000
@ -164,11 +173,11 @@ def scrape(cfg: ScrapeConfig, base_dir: Path):
for repo in cfg.repos:
owner, name = repo.split("/", 1)
scraper = RepoScraper(
owner=owner,
name=name,
base_dir=base_dir,
client=client,
trial_limits=trial_limits,
owner = owner,
name = name,
base_dir = base_dir,
client = client,
trial_limits = trial_limits,
)
try:
scraper.scrape_repo_meta()
@ -185,10 +194,18 @@ def scrape(cfg: ScrapeConfig, base_dir: Path):
repo_dir = base_dir / f"{owner}__{name}"
if "issues" in cfg.item_types:
for row in _read_jsonl(repo_dir / "issues.jsonl", read_cap):
all_rows.append(_flatten_issue_row(row, repo, cfg.include_comments, cfg.max_comments_per_item))
all_rows.append(
_flatten_issue_row(
row, repo, cfg.include_comments, cfg.max_comments_per_item
)
)
if "pulls" in cfg.item_types:
for row in _read_jsonl(repo_dir / "pull_requests.jsonl", read_cap):
all_rows.append(_flatten_pr_row(row, repo, cfg.include_comments, cfg.max_comments_per_item))
all_rows.append(
_flatten_pr_row(
row, repo, cfg.include_comments, cfg.max_comments_per_item
)
)
if "commits" in cfg.item_types:
for row in _read_jsonl(repo_dir / "commits.jsonl", read_cap):
all_rows.append(_flatten_commit_row(row, repo))
@ -197,13 +214,13 @@ def scrape(cfg: ScrapeConfig, base_dir: Path):
def materialize_to_jsonl(cfg: ScrapeConfig, out_dir: Path) -> Path:
out_dir.mkdir(parents=True, exist_ok=True)
out_dir.mkdir(parents = True, exist_ok = True)
tag = "-".join(r.replace("/", "__") for r in cfg.repos)[:120]
kinds = "-".join(cfg.item_types)
fname = f"github_{tag}__{kinds}__{cfg.limit}_{int(time.time())}.jsonl"
out = out_dir / fname
rows = scrape(cfg, out_dir / "raw")
with out.open("w", encoding="utf-8") as f:
with out.open("w", encoding = "utf-8") as f:
for r in rows:
f.write(json.dumps(r, ensure_ascii=False) + "\n")
f.write(json.dumps(r, ensure_ascii = False) + "\n")
return out

42
studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/gh_client.py
View file

@ -2,6 +2,7 @@
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""GitHub API client with rate-limit awareness, retry, and dual REST/GraphQL support."""
from __future__ import annotations
import json
@ -71,8 +72,12 @@ class GitHubClient:
else:
self.rest_remaining = None
def graphql(self, query: str, variables: Optional[Dict[str, Any]] = None,
max_retries: int = 20) -> Dict[str, Any]:
def graphql(
self,
query: str,
variables: Optional[Dict[str, Any]] = None,
max_retries: int = 20,
) -> Dict[str, Any]:
self._check_rate_and_wait("graphql")
backoff = 2
last_err = None
@ -80,8 +85,8 @@ class GitHubClient:
try:
r = self.session.post(
GRAPHQL_URL,
json={"query": query, "variables": variables or {}},
timeout=120,
json = {"query": query, "variables": variables or {}},
timeout = 120,
)
self.calls_graphql += 1
# Update rate info from response headers
@ -123,7 +128,9 @@ class GitHubClient:
# Retry on RATE_LIMITED
for e in errs:
if e.get("type") == "RATE_LIMITED":
self._sleep_until((self.graphql_reset or int(time.time()) + 60))
self._sleep_until(
(self.graphql_reset or int(time.time()) + 60)
)
break
else:
# No rate-limit error, log and return partial
@ -138,8 +145,14 @@ class GitHubClient:
backoff = min(backoff * 2, 60)
raise RuntimeError(f"GraphQL failed after {max_retries} retries: {last_err}")
def rest(self, method: str, path: str, params: Optional[Dict[str, Any]] = None,
json_body: Optional[Dict[str, Any]] = None, max_retries: int = 6) -> requests.Response:
def rest(
self,
method: str,
path: str,
params: Optional[Dict[str, Any]] = None,
json_body: Optional[Dict[str, Any]] = None,
max_retries: int = 6,
) -> requests.Response:
self._check_rate_and_wait("rest")
if path.startswith("http"):
url = path
@ -149,7 +162,9 @@ class GitHubClient:
last_err = None
for attempt in range(max_retries):
try:
r = self.session.request(method, url, params=params, json=json_body, timeout=120)
r = self.session.request(
method, url, params = params, json = json_body, timeout = 120
)
self.calls_rest += 1
rem = r.headers.get("X-RateLimit-Remaining")
rst = r.headers.get("X-RateLimit-Reset")
@ -190,15 +205,18 @@ class GitHubClient:
backoff = min(backoff * 2, 60)
raise RuntimeError(f"REST failed after {max_retries} retries: {last_err}")
def rest_paginate(self, path: str, params: Optional[Dict[str, Any]] = None,
per_page: int = 100) -> Iterator[dict]:
def rest_paginate(
self, path: str, params: Optional[Dict[str, Any]] = None, per_page: int = 100
) -> Iterator[dict]:
params = dict(params or {})
params.setdefault("per_page", per_page)
url = path
while True:
r = self.rest("GET", url, params=params if url == path else None)
r = self.rest("GET", url, params = params if url == path else None)
if r.status_code != 200:
log.error("REST paginate got %s at %s: %s", r.status_code, url, r.text[:200])
log.error(
"REST paginate got %s at %s: %s", r.status_code, url, r.text[:200]
)
return
items = r.json()
if isinstance(items, dict):

77
studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/queries.py
View file

@ -131,7 +131,9 @@ def _q(parts: list[str], body: str) -> str:
return "\n".join(parts + [body])
ISSUES_PAGE_QUERY = _q([F_ACTOR, F_LABEL, F_TIMELINE], """
ISSUES_PAGE_QUERY = _q(
[F_ACTOR, F_LABEL, F_TIMELINE],
"""
query IssuesPage($owner: String!, $name: String!, $first: Int!, $after: String) {
repository(owner: $owner, name: $name) {
issues(first: $first, after: $after, orderBy: {field: CREATED_AT, direction: ASC}) {
@ -169,9 +171,12 @@ query IssuesPage($owner: String!, $name: String!, $first: Int!, $after: String)
}
rateLimit { cost remaining resetAt }
}
""")
""",
)
PRS_PAGE_QUERY = _q([F_ACTOR, F_LABEL, F_PR_TIMELINE], """
PRS_PAGE_QUERY = _q(
[F_ACTOR, F_LABEL, F_PR_TIMELINE],
"""
query PRsPage($owner: String!, $name: String!, $first: Int!, $after: String) {
repository(owner: $owner, name: $name) {
pullRequests(first: $first, after: $after, orderBy: {field: CREATED_AT, direction: ASC}) {
@ -265,9 +270,12 @@ query PRsPage($owner: String!, $name: String!, $first: Int!, $after: String) {
}
rateLimit { cost remaining resetAt }
}
""")
""",
)
ISSUE_COMMENTS_QUERY = _q([F_ACTOR], """
ISSUE_COMMENTS_QUERY = _q(
[F_ACTOR],
"""
query IssueComments($owner: String!, $name: String!, $number: Int!, $after: String) {
repository(owner: $owner, name: $name) {
issueOrPullRequest(number: $number) {
@ -298,9 +306,12 @@ query IssueComments($owner: String!, $name: String!, $number: Int!, $after: Stri
}
rateLimit { cost remaining resetAt }
}
""")
""",
)
ISSUE_TIMELINE_QUERY = _q([F_ACTOR, F_TIMELINE], """
ISSUE_TIMELINE_QUERY = _q(
[F_ACTOR, F_TIMELINE],
"""
query IssueTimeline($owner: String!, $name: String!, $number: Int!, $after: String) {
repository(owner: $owner, name: $name) {
issue(number: $number) {
@ -312,9 +323,12 @@ query IssueTimeline($owner: String!, $name: String!, $number: Int!, $after: Stri
}
rateLimit { cost remaining resetAt }
}
""")
""",
)
PR_TIMELINE_QUERY = _q([F_ACTOR, F_PR_TIMELINE], """
PR_TIMELINE_QUERY = _q(
[F_ACTOR, F_PR_TIMELINE],
"""
query PRTimeline($owner: String!, $name: String!, $number: Int!, $after: String) {
repository(owner: $owner, name: $name) {
pullRequest(number: $number) {
@ -326,7 +340,8 @@ query PRTimeline($owner: String!, $name: String!, $number: Int!, $after: String)
}
rateLimit { cost remaining resetAt }
}
""")
""",
)
PR_COMMITS_QUERY = """
query PRCommits($owner: String!, $name: String!, $number: Int!, $after: String) {
@ -364,7 +379,9 @@ query PRFiles($owner: String!, $name: String!, $number: Int!, $after: String) {
}
"""
PR_REVIEW_THREADS_QUERY = _q([F_ACTOR], """
PR_REVIEW_THREADS_QUERY = _q(
[F_ACTOR],
"""
query PRReviewThreads($owner: String!, $name: String!, $number: Int!, $after: String) {
repository(owner: $owner, name: $name) {
pullRequest(number: $number) {
@ -389,9 +406,12 @@ query PRReviewThreads($owner: String!, $name: String!, $number: Int!, $after: St
}
rateLimit { cost remaining resetAt }
}
""")
""",
)
DISCUSSIONS_PAGE_QUERY = _q([F_ACTOR, F_LABEL], """
DISCUSSIONS_PAGE_QUERY = _q(
[F_ACTOR, F_LABEL],
"""
query DiscussionsPage($owner: String!, $name: String!, $first: Int!, $after: String) {
repository(owner: $owner, name: $name) {
discussions(first: $first, after: $after, orderBy: {field: CREATED_AT, direction: ASC}) {
@ -437,9 +457,12 @@ query DiscussionsPage($owner: String!, $name: String!, $first: Int!, $after: Str
}
rateLimit { cost remaining resetAt }
}
""")
""",
)
DISCUSSION_COMMENTS_QUERY = _q([F_ACTOR], """
DISCUSSION_COMMENTS_QUERY = _q(
[F_ACTOR],
"""
query DiscussionComments($owner: String!, $name: String!, $number: Int!, $after: String) {
repository(owner: $owner, name: $name) {
discussion(number: $number) {
@ -467,9 +490,12 @@ query DiscussionComments($owner: String!, $name: String!, $number: Int!, $after:
}
rateLimit { cost remaining resetAt }
}
""")
""",
)
DISCUSSION_REPLIES_QUERY = _q([F_ACTOR], """
DISCUSSION_REPLIES_QUERY = _q(
[F_ACTOR],
"""
query DiscussionReplies($commentId: ID!, $after: String) {
node(id: $commentId) {
... on DiscussionComment {
@ -486,7 +512,8 @@ query DiscussionReplies($commentId: ID!, $after: String) {
}
rateLimit { cost remaining resetAt }
}
""")
""",
)
COMMITS_PAGE_QUERY = """
query CommitsPage($owner: String!, $name: String!, $first: Int!, $after: String, $branch: String!) {
@ -519,7 +546,9 @@ query CommitsPage($owner: String!, $name: String!, $first: Int!, $after: String,
}
"""
RELEASES_QUERY = _q([F_ACTOR], """
RELEASES_QUERY = _q(
[F_ACTOR],
"""
query Releases($owner: String!, $name: String!, $first: Int!, $after: String) {
repository(owner: $owner, name: $name) {
releases(first: $first, after: $after, orderBy: {field: CREATED_AT, direction: ASC}) {
@ -540,9 +569,12 @@ query Releases($owner: String!, $name: String!, $first: Int!, $after: String) {
}
rateLimit { cost remaining resetAt }
}
""")
""",
)
LABELS_QUERY = _q([F_LABEL], """
LABELS_QUERY = _q(
[F_LABEL],
"""
query LabelsList($owner: String!, $name: String!, $first: Int!, $after: String) {
repository(owner: $owner, name: $name) {
labels(first: $first, after: $after) {
@ -552,7 +584,8 @@ query LabelsList($owner: String!, $name: String!, $first: Int!, $after: String)
}
rateLimit { cost remaining resetAt }
}
""")
""",
)
MILESTONES_QUERY = """
query Milestones($owner: String!, $name: String!, $first: Int!, $after: String) {

357
studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/scraper.py
View file

@ -5,6 +5,7 @@
Resumable via state file. Writes JSONL shards under data/{repo}/{resource}.jsonl.
"""
from __future__ import annotations
import argparse
@ -34,25 +35,42 @@ def ts() -> str:
class RepoScraper:
def __init__(self, owner: str, name: str, base_dir: Path, client: GitHubClient,
trial_limits: Optional[Dict[str, int]] = None):
def __init__(
self,
owner: str,
name: str,
base_dir: Path,
client: GitHubClient,
trial_limits: Optional[Dict[str, int]] = None,
):
self.owner = owner
self.name = name
self.base_dir = base_dir
self.client = client
self.trial_limits = trial_limits or {}
self.repo_dir = base_dir / f"{owner}__{name}"
self.repo_dir.mkdir(parents=True, exist_ok=True)
self.repo_dir.mkdir(parents = True, exist_ok = True)
self.state = StateStore(base_dir / "state" / f"{owner}__{name}.json")
# Writers
self.writers: Dict[str, JsonlWriter] = {}
for key in ("issues", "pull_requests", "discussions",
"commits", "releases", "labels", "milestones",
"pr_extra_comments", "pr_extra_timeline", "pr_extra_reviews",
"issue_extra_comments", "issue_extra_timeline",
"discussion_extra_comments", "discussion_extra_replies",
"repo_meta"):
for key in (
"issues",
"pull_requests",
"discussions",
"commits",
"releases",
"labels",
"milestones",
"pr_extra_comments",
"pr_extra_timeline",
"pr_extra_reviews",
"issue_extra_comments",
"issue_extra_timeline",
"discussion_extra_comments",
"discussion_extra_replies",
"repo_meta",
):
self.writers[key] = JsonlWriter(self.repo_dir / f"{key}.jsonl")
# ----- helpers -----
@ -63,14 +81,25 @@ class RepoScraper:
return counter >= lim
def _log_rate(self, where: str, data: Dict[str, Any]) -> None:
rl = data.get("data", {}).get("rateLimit") if isinstance(data.get("data"), dict) else None
rl = (
data.get("data", {}).get("rateLimit")
if isinstance(data.get("data"), dict)
else None
)
if rl:
log.debug("[%s] rate cost=%s remaining=%s resetAt=%s",
where, rl.get("cost"), rl.get("remaining"), rl.get("resetAt"))
log.debug(
"[%s] rate cost=%s remaining=%s resetAt=%s",
where,
rl.get("cost"),
rl.get("remaining"),
rl.get("resetAt"),
)
# ----- repo meta -----
def scrape_repo_meta(self) -> None:
data = self.client.graphql(Q.REPO_META_QUERY, {"owner": self.owner, "name": self.name})
data = self.client.graphql(
Q.REPO_META_QUERY, {"owner": self.owner, "name": self.name}
)
self._log_rate("repo_meta", data)
repo = data.get("data", {}).get("repository") or {}
repo["_fetchedAt"] = ts()
@ -89,11 +118,16 @@ class RepoScraper:
per_page = 15 # conservative for heavy nested query
while True:
page += 1
vars_ = {"owner": self.owner, "name": self.name, "first": per_page, "after": cursor}
vars_ = {
"owner": self.owner,
"name": self.name,
"first": per_page,
"after": cursor,
}
data = self.client.graphql(Q.ISSUES_PAGE_QUERY, vars_)
self._log_rate("issues", data)
repo = (data.get("data") or {}).get("repository") or {}
issues = (repo.get("issues") or {})
issues = repo.get("issues") or {}
nodes = issues.get("nodes") or []
for it in nodes:
it["_owner"] = self.owner
@ -101,17 +135,27 @@ class RepoScraper:
it["_fetchedAt"] = ts()
# Paginate nested comments/timeline if more exist
if it.get("comments", {}).get("pageInfo", {}).get("hasNextPage"):
self._paginate_issue_comments(it["number"], it["comments"]["pageInfo"]["endCursor"])
self._paginate_issue_comments(
it["number"], it["comments"]["pageInfo"]["endCursor"]
)
if it.get("timelineItems", {}).get("pageInfo", {}).get("hasNextPage"):
self._paginate_issue_timeline(it["number"], it["timelineItems"]["pageInfo"]["endCursor"])
self._paginate_issue_timeline(
it["number"], it["timelineItems"]["pageInfo"]["endCursor"]
)
if self.writers[key].write(it):
total_new += 1
info = issues.get("pageInfo") or {}
cursor = info.get("endCursor")
self.state.set(f"{key}_cursor", cursor)
log.info("[%s/%s] issues page %d (+%d) cursor=%s remaining=%s",
self.owner, self.name, page, len(nodes), str(cursor)[:20],
self.client.graphql_remaining)
log.info(
"[%s/%s] issues page %d (+%d) cursor=%s remaining=%s",
self.owner,
self.name,
page,
len(nodes),
str(cursor)[:20],
self.client.graphql_remaining,
)
if self._trial_stop(key, total_new):
log.info("Trial limit reached for issues (%d)", total_new)
return total_new
@ -123,10 +167,17 @@ class RepoScraper:
def _paginate_issue_comments(self, number: int, after: str) -> None:
cur = after
while cur:
vars_ = {"owner": self.owner, "name": self.name, "number": number, "after": cur}
vars_ = {
"owner": self.owner,
"name": self.name,
"number": number,
"after": cur,
}
data = self.client.graphql(Q.ISSUE_COMMENTS_QUERY, vars_)
item = (((data.get("data") or {}).get("repository") or {}).get("issueOrPullRequest") or {})
comments = (item.get("comments") or {})
item = ((data.get("data") or {}).get("repository") or {}).get(
"issueOrPullRequest"
) or {}
comments = item.get("comments") or {}
for c in comments.get("nodes") or []:
c["_owner"] = self.owner
c["_repo"] = self.name
@ -138,10 +189,15 @@ class RepoScraper:
def _paginate_issue_timeline(self, number: int, after: str) -> None:
cur = after
while cur:
vars_ = {"owner": self.owner, "name": self.name, "number": number, "after": cur}
vars_ = {
"owner": self.owner,
"name": self.name,
"number": number,
"after": cur,
}
data = self.client.graphql(Q.ISSUE_TIMELINE_QUERY, vars_)
item = (((data.get("data") or {}).get("repository") or {}).get("issue") or {})
tl = (item.get("timelineItems") or {})
item = ((data.get("data") or {}).get("repository") or {}).get("issue") or {}
tl = item.get("timelineItems") or {}
for ev in tl.get("nodes") or []:
ev["_owner"] = self.owner
ev["_repo"] = self.name
@ -163,11 +219,16 @@ class RepoScraper:
per_page = 3 # PR query is heavy; keep small so huge PRs don't OOM GraphQL
while True:
page += 1
vars_ = {"owner": self.owner, "name": self.name, "first": per_page, "after": cursor}
vars_ = {
"owner": self.owner,
"name": self.name,
"first": per_page,
"after": cursor,
}
data = self.client.graphql(Q.PRS_PAGE_QUERY, vars_)
self._log_rate("prs", data)
repo = (data.get("data") or {}).get("repository") or {}
prs = (repo.get("pullRequests") or {})
prs = repo.get("pullRequests") or {}
nodes = prs.get("nodes") or []
for pr in nodes:
pr["_owner"] = self.owner
@ -175,23 +236,37 @@ class RepoScraper:
pr["_fetchedAt"] = ts()
num = pr["number"]
if pr.get("comments", {}).get("pageInfo", {}).get("hasNextPage"):
self._paginate_pr_comments(num, pr["comments"]["pageInfo"]["endCursor"])
self._paginate_pr_comments(
num, pr["comments"]["pageInfo"]["endCursor"]
)
if pr.get("timelineItems", {}).get("pageInfo", {}).get("hasNextPage"):
self._paginate_pr_timeline(num, pr["timelineItems"]["pageInfo"]["endCursor"])
self._paginate_pr_timeline(
num, pr["timelineItems"]["pageInfo"]["endCursor"]
)
if pr.get("commits", {}).get("pageInfo", {}).get("hasNextPage"):
self._paginate_pr_commits(num, pr["commits"]["pageInfo"]["endCursor"])
self._paginate_pr_commits(
num, pr["commits"]["pageInfo"]["endCursor"]
)
if pr.get("files", {}).get("pageInfo", {}).get("hasNextPage"):
self._paginate_pr_files(num, pr["files"]["pageInfo"]["endCursor"])
if pr.get("reviewThreads", {}).get("pageInfo", {}).get("hasNextPage"):
self._paginate_pr_review_threads(num, pr["reviewThreads"]["pageInfo"]["endCursor"])
self._paginate_pr_review_threads(
num, pr["reviewThreads"]["pageInfo"]["endCursor"]
)
if self.writers[key].write(pr):
total_new += 1
info = prs.get("pageInfo") or {}
cursor = info.get("endCursor")
self.state.set(f"{key}_cursor", cursor)
log.info("[%s/%s] PRs page %d (+%d) cursor=%s remaining=%s",
self.owner, self.name, page, len(nodes), str(cursor)[:20],
self.client.graphql_remaining)
log.info(
"[%s/%s] PRs page %d (+%d) cursor=%s remaining=%s",
self.owner,
self.name,
page,
len(nodes),
str(cursor)[:20],
self.client.graphql_remaining,
)
if self._trial_stop(key, total_new):
log.info("Trial limit reached for PRs (%d)", total_new)
return total_new
@ -203,10 +278,17 @@ class RepoScraper:
def _paginate_pr_comments(self, number: int, after: str) -> None:
cur = after
while cur:
vars_ = {"owner": self.owner, "name": self.name, "number": number, "after": cur}
vars_ = {
"owner": self.owner,
"name": self.name,
"number": number,
"after": cur,
}
data = self.client.graphql(Q.ISSUE_COMMENTS_QUERY, vars_)
item = (((data.get("data") or {}).get("repository") or {}).get("issueOrPullRequest") or {})
comments = (item.get("comments") or {})
item = ((data.get("data") or {}).get("repository") or {}).get(
"issueOrPullRequest"
) or {}
comments = item.get("comments") or {}
for c in comments.get("nodes") or []:
c["_owner"] = self.owner
c["_repo"] = self.name
@ -218,10 +300,17 @@ class RepoScraper:
def _paginate_pr_timeline(self, number: int, after: str) -> None:
cur = after
while cur:
vars_ = {"owner": self.owner, "name": self.name, "number": number, "after": cur}
vars_ = {
"owner": self.owner,
"name": self.name,
"number": number,
"after": cur,
}
data = self.client.graphql(Q.PR_TIMELINE_QUERY, vars_)
item = (((data.get("data") or {}).get("repository") or {}).get("pullRequest") or {})
tl = (item.get("timelineItems") or {})
item = ((data.get("data") or {}).get("repository") or {}).get(
"pullRequest"
) or {}
tl = item.get("timelineItems") or {}
for ev in tl.get("nodes") or []:
ev["_owner"] = self.owner
ev["_repo"] = self.name
@ -236,10 +325,17 @@ class RepoScraper:
if out_key not in self.writers:
self.writers[out_key] = JsonlWriter(self.repo_dir / f"{out_key}.jsonl")
while cur:
vars_ = {"owner": self.owner, "name": self.name, "number": number, "after": cur}
vars_ = {
"owner": self.owner,
"name": self.name,
"number": number,
"after": cur,
}
data = self.client.graphql(Q.PR_COMMITS_QUERY, vars_)
item = (((data.get("data") or {}).get("repository") or {}).get("pullRequest") or {})
cc = (item.get("commits") or {})
item = ((data.get("data") or {}).get("repository") or {}).get(
"pullRequest"
) or {}
cc = item.get("commits") or {}
for c in cc.get("nodes") or []:
c["_owner"] = self.owner
c["_repo"] = self.name
@ -254,10 +350,17 @@ class RepoScraper:
if out_key not in self.writers:
self.writers[out_key] = JsonlWriter(self.repo_dir / f"{out_key}.jsonl")
while cur:
vars_ = {"owner": self.owner, "name": self.name, "number": number, "after": cur}
vars_ = {
"owner": self.owner,
"name": self.name,
"number": number,
"after": cur,
}
data = self.client.graphql(Q.PR_FILES_QUERY, vars_)
item = (((data.get("data") or {}).get("repository") or {}).get("pullRequest") or {})
ff = (item.get("files") or {})
item = ((data.get("data") or {}).get("repository") or {}).get(
"pullRequest"
) or {}
ff = item.get("files") or {}
for f in ff.get("nodes") or []:
f["_owner"] = self.owner
f["_repo"] = self.name
@ -274,10 +377,17 @@ class RepoScraper:
if out_key not in self.writers:
self.writers[out_key] = JsonlWriter(self.repo_dir / f"{out_key}.jsonl")
while cur:
vars_ = {"owner": self.owner, "name": self.name, "number": number, "after": cur}
vars_ = {
"owner": self.owner,
"name": self.name,
"number": number,
"after": cur,
}
data = self.client.graphql(Q.PR_REVIEW_THREADS_QUERY, vars_)
item = (((data.get("data") or {}).get("repository") or {}).get("pullRequest") or {})
rt = (item.get("reviewThreads") or {})
item = ((data.get("data") or {}).get("repository") or {}).get(
"pullRequest"
) or {}
rt = item.get("reviewThreads") or {}
for th in rt.get("nodes") or []:
th["_owner"] = self.owner
th["_repo"] = self.name
@ -299,11 +409,16 @@ class RepoScraper:
per_page = 15
while True:
page += 1
vars_ = {"owner": self.owner, "name": self.name, "first": per_page, "after": cursor}
vars_ = {
"owner": self.owner,
"name": self.name,
"first": per_page,
"after": cursor,
}
data = self.client.graphql(Q.DISCUSSIONS_PAGE_QUERY, vars_)
self._log_rate("discussions", data)
repo = (data.get("data") or {}).get("repository") or {}
dd = (repo.get("discussions") or {})
dd = repo.get("discussions") or {}
nodes = dd.get("nodes") or []
for d in nodes:
d["_owner"] = self.owner
@ -311,19 +426,29 @@ class RepoScraper:
d["_fetchedAt"] = ts()
num = d["number"]
if d.get("comments", {}).get("pageInfo", {}).get("hasNextPage"):
self._paginate_discussion_comments(num, d["comments"]["pageInfo"]["endCursor"])
self._paginate_discussion_comments(
num, d["comments"]["pageInfo"]["endCursor"]
)
# paginate replies per comment if needed
for c in d.get("comments", {}).get("nodes", []) or []:
if c.get("replies", {}).get("pageInfo", {}).get("hasNextPage"):
self._paginate_discussion_replies(c["id"], c["replies"]["pageInfo"]["endCursor"], num)
self._paginate_discussion_replies(
c["id"], c["replies"]["pageInfo"]["endCursor"], num
)
if self.writers[key].write(d):
total_new += 1
info = dd.get("pageInfo") or {}
cursor = info.get("endCursor")
self.state.set(f"{key}_cursor", cursor)
log.info("[%s/%s] discussions page %d (+%d) cursor=%s remaining=%s",
self.owner, self.name, page, len(nodes), str(cursor)[:20],
self.client.graphql_remaining)
log.info(
"[%s/%s] discussions page %d (+%d) cursor=%s remaining=%s",
self.owner,
self.name,
page,
len(nodes),
str(cursor)[:20],
self.client.graphql_remaining,
)
if self._trial_stop(key, total_new):
return total_new
if not info.get("hasNextPage"):
@ -334,10 +459,17 @@ class RepoScraper:
def _paginate_discussion_comments(self, number: int, after: str) -> None:
cur = after
while cur:
vars_ = {"owner": self.owner, "name": self.name, "number": number, "after": cur}
vars_ = {
"owner": self.owner,
"name": self.name,
"number": number,
"after": cur,
}
data = self.client.graphql(Q.DISCUSSION_COMMENTS_QUERY, vars_)
disc = (((data.get("data") or {}).get("repository") or {}).get("discussion") or {})
cc = (disc.get("comments") or {})
disc = ((data.get("data") or {}).get("repository") or {}).get(
"discussion"
) or {}
cc = disc.get("comments") or {}
for c in cc.get("nodes") or []:
c["_owner"] = self.owner
c["_repo"] = self.name
@ -346,10 +478,17 @@ class RepoScraper:
info = cc.get("pageInfo") or {}
cur = info.get("endCursor") if info.get("hasNextPage") else None
def _paginate_discussion_replies(self, comment_id: str, after: str, disc_number: int) -> None:
def _paginate_discussion_replies(
self, comment_id: str, after: str, disc_number: int
) -> None:
cur = after
while cur:
vars_ = {"owner": self.owner, "name": self.name, "commentId": comment_id, "after": cur}
vars_ = {
"owner": self.owner,
"name": self.name,
"commentId": comment_id,
"after": cur,
}
data = self.client.graphql(Q.DISCUSSION_REPLIES_QUERY, vars_)
node = (data.get("data") or {}).get("node") or {}
replies = node.get("replies") or {}
@ -374,11 +513,16 @@ class RepoScraper:
per_page = 100
while True:
page += 1
vars_ = {"owner": self.owner, "name": self.name,
"first": per_page, "after": cursor, "branch": branch}
vars_ = {
"owner": self.owner,
"name": self.name,
"first": per_page,
"after": cursor,
"branch": branch,
}
data = self.client.graphql(Q.COMMITS_PAGE_QUERY, vars_)
self._log_rate("commits", data)
ref = (((data.get("data") or {}).get("repository") or {}).get("ref") or {})
ref = ((data.get("data") or {}).get("repository") or {}).get("ref") or {}
tgt = ref.get("target") or {}
hist = tgt.get("history") or {}
nodes = hist.get("nodes") or []
@ -391,8 +535,14 @@ class RepoScraper:
info = hist.get("pageInfo") or {}
cursor = info.get("endCursor")
self.state.set(f"{key}_cursor", cursor)
log.info("[%s/%s] commits page %d (+%d) remaining=%s",
self.owner, self.name, page, len(nodes), self.client.graphql_remaining)
log.info(
"[%s/%s] commits page %d (+%d) remaining=%s",
self.owner,
self.name,
page,
len(nodes),
self.client.graphql_remaining,
)
if self._trial_stop(key, total_new):
return total_new
if not info.get("hasNextPage"):
@ -417,10 +567,15 @@ class RepoScraper:
return 0
total_new = 0
while True:
vars_ = {"owner": self.owner, "name": self.name, "first": 50, "after": cursor}
vars_ = {
"owner": self.owner,
"name": self.name,
"first": 50,
"after": cursor,
}
data = self.client.graphql(query, vars_)
repo = (data.get("data") or {}).get("repository") or {}
col = (repo.get(field) or {})
col = repo.get(field) or {}
for it in col.get("nodes") or []:
it["_owner"] = self.owner
it["_repo"] = self.name
@ -447,47 +602,70 @@ class RepoScraper:
def setup_logging(log_file: Path) -> None:
log_file.parent.mkdir(parents=True, exist_ok=True)
log_file.parent.mkdir(parents = True, exist_ok = True)
fmt = "%(asctime)s %(levelname)s [%(name)s] %(message)s"
handlers = [
logging.StreamHandler(sys.stdout),
logging.FileHandler(log_file, mode="a", encoding="utf-8"),
logging.FileHandler(log_file, mode = "a", encoding = "utf-8"),
]
logging.basicConfig(level=logging.INFO, format=fmt, handlers=handlers, force=True)
logging.basicConfig(level = logging.INFO, format = fmt, handlers = handlers, force = True)
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--base-dir", default="/mnt/disks/unslothai/ubuntu/workspace_34/github_scraper")
ap.add_argument("--repos", nargs="+", default=["unslothai/unsloth", "unslothai/unsloth-zoo"])
ap.add_argument("--trial", action="store_true", help="Small trial run")
ap.add_argument("--only", nargs="+", default=None,
help="Only run these resource keys: issues,pulls,discussions,commits,releases,labels,milestones,meta")
ap.add_argument("--hf-upload-interval", type=int, default=900,
help="Seconds between HF uploads (0 to disable)")
ap.add_argument(
"--base-dir", default = "/mnt/disks/unslothai/ubuntu/workspace_34/github_scraper"
)
ap.add_argument(
"--repos", nargs = "+", default = ["unslothai/unsloth", "unslothai/unsloth-zoo"]
)
ap.add_argument("--trial", action = "store_true", help = "Small trial run")
ap.add_argument(
"--only",
nargs = "+",
default = None,
help = "Only run these resource keys: issues,pulls,discussions,commits,releases,labels,milestones,meta",
)
ap.add_argument(
"--hf-upload-interval",
type = int,
default = 900,
help = "Seconds between HF uploads (0 to disable)",
)
args = ap.parse_args()
base = Path(args.base_dir)
data_dir = base / "data"
data_dir.mkdir(parents=True, exist_ok=True)
data_dir.mkdir(parents = True, exist_ok = True)
setup_logging(base / "logs" / f"scraper_{time.strftime('%Y%m%d_%H%M%S')}.log")
log.info("Scraper starting: repos=%s trial=%s", args.repos, args.trial)
client = GitHubClient(min_remaining_graphql=80, min_remaining_rest=80)
client = GitHubClient(min_remaining_graphql = 80, min_remaining_rest = 80)
rl = client.rate_snapshot()
log.info("Rate limit snapshot: %s", json.dumps(rl.get("resources", {}), default=str)[:400])
log.info(
"Rate limit snapshot: %s",
json.dumps(rl.get("resources", {}), default = str)[:400],
)
# Start HF uploader in background if requested
uploader = None
if args.hf_upload_interval > 0:
from hf_uploader import HFUploader
uploader = HFUploader(data_dir, interval_s=args.hf_upload_interval)
uploader = HFUploader(data_dir, interval_s = args.hf_upload_interval)
uploader.start()
trial_limits = None
if args.trial:
trial_limits = {"issues": 5, "pull_requests": 5, "discussions": 3,
"commits": 20, "releases": 3, "labels": 20, "milestones": 20}
trial_limits = {
"issues": 5,
"pull_requests": 5,
"discussions": 3,
"commits": 20,
"releases": 3,
"labels": 20,
"milestones": 20,
}
only = set(args.only or [])
@ -517,9 +695,12 @@ def main():
finally:
if uploader:
log.info("Stopping uploader and final sync...")
uploader.stop(final_upload=True)
log.info("Scraper complete. GraphQL calls=%d REST calls=%d",
client.calls_graphql, client.calls_rest)
uploader.stop(final_upload = True)
log.info(
"Scraper complete. GraphQL calls=%d REST calls=%d",
client.calls_graphql,
client.calls_rest,
)
if __name__ == "__main__":

12
studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/state_store.py
View file

@ -2,6 +2,7 @@
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Checkpoint state management for resumable scraping."""
from __future__ import annotations
import json
@ -14,7 +15,7 @@ from typing import Any, Dict
class StateStore:
def __init__(self, path: str | Path):
self.path = Path(path)
self.path.parent.mkdir(parents=True, exist_ok=True)
self.path.parent.mkdir(parents = True, exist_ok = True)
self._lock = threading.Lock()
self._data: Dict[str, Any] = {}
if self.path.exists():
@ -47,17 +48,18 @@ class StateStore:
def _flush(self) -> None:
tmp = self.path.with_suffix(self.path.suffix + ".tmp")
with tmp.open("w") as f:
json.dump(self._data, f, indent=2, default=str)
json.dump(self._data, f, indent = 2, default = str)
os.replace(tmp, self.path)
class JsonlWriter:
"""Append-only JSONL writer, thread-safe, with line buffering."""
def __init__(self, path: str | Path):
self.path = Path(path)
self.path.parent.mkdir(parents=True, exist_ok=True)
self.path.parent.mkdir(parents = True, exist_ok = True)
self._lock = threading.Lock()
self._fh = self.path.open("a", buffering=1)
self._fh = self.path.open("a", buffering = 1)
self._count_seen_keys: set[str] = set()
# Preload seen keys if file exists (for dedup across resumes)
if self.path.exists() and self.path.stat().st_size > 0:
@ -91,7 +93,7 @@ class JsonlWriter:
return False
if k is not None:
self._count_seen_keys.add(k)
self._fh.write(json.dumps(obj, default=str, ensure_ascii=False))
self._fh.write(json.dumps(obj, default = str, ensure_ascii = False))
self._fh.write("\n")
self._fh.flush()
return True