Harden corpus scraper + first pair-hunting run

scripts/corpus/fetch_corpus.py

@@ -56,6 +56,7 @@
 import json
 import os
 import re
+import shutil
 import sys
 import time
 from dataclasses import dataclass, field
@@ -144,27 +145,30 @@ class Scheduler:
         extensions={".sh", ".slurm", ".job", ".sbatch", ".batch", ".sl", ".bash", ""},
         directive=True, marker="#SBATCH",
         accept=accept_directive("#SBATCH"),
+        # Directive phrases are quoted: an unquoted "--array" / "-l" leads with a
+        # dash, which GitHub code search parses as a NOT operator (422 fatal).
         queries=[
-            "#SBATCH --array extension:sh",
-            "#SBATCH --gres=gpu extension:sh",
-            "#SBATCH --ntasks-per-node extension:sh",
-            "#SBATCH --dependency extension:sh",
-            "#SBATCH hetjob extension:sh",
-            "#SBATCH --partition extension:slurm",
+            '"#SBATCH --array" extension:sh',
+            '"#SBATCH --gres=gpu" extension:sh',
+            '"#SBATCH --ntasks-per-node" extension:sh',
+            '"#SBATCH --dependency" extension:sh',
+            '"#SBATCH" hetjob extension:sh',
+            '"#SBATCH --partition" extension:slurm',
         ],
     ),
     "pbs": Scheduler(
         name="pbs",
         extensions={".sh", ".pbs", ".job", ".bash", ""},
         directive=True, marker="#PBS",
         accept=accept_directive("#PBS"),
+        # Quoted for the same reason as slurm: "-l"/"-q"/"-W" lead with a dash.
         queries=[
-            "#PBS -l select extension:sh",
-            "#PBS -l nodes extension:sh",
-            "#PBS -q extension:pbs",
-            "#PBS -l walltime extension:sh",
-            "#PBS -J extension:sh",              # array jobs
-            "#PBS -W depend extension:sh",       # dependencies
+            '"#PBS -l select" extension:sh',
+            '"#PBS -l nodes" extension:sh',
+            '"#PBS -q" extension:pbs',
+            '"#PBS -l walltime" extension:sh',
+            '"#PBS -J" extension:sh',            # array jobs
+            '"#PBS -W depend" extension:sh',     # dependencies
         ],
     ),
     "htcondor": Scheduler(
@@ -242,14 +246,29 @@ def make_session(token: str) -> requests.Session:
     return s
 
 
-def get_with_retry(session, url, params=None, retries=3):
-    for _ in range(retries):
-        resp = session.get(url, params=params)
+def get_with_retry(session, url, params=None, retries=5):
+    for attempt in range(retries):
+        try:
+            resp = session.get(url, params=params, timeout=30)
+        except requests.exceptions.RequestException as e:
+            # GitHub drops the connection under abuse-detection throttling; treat
+            # it like a rate-limit and back off rather than crashing the run.
+            wait = 15 * (attempt + 1)
+            print(f"    connection error ({type(e).__name__}); sleeping {wait}s", flush=True)
+            time.sleep(wait)
+            continue
         if resp.status_code == 200:
             return resp
         if resp.status_code in (403, 429):
-            reset = int(resp.headers.get("X-RateLimit-Reset", time.time() + 60))
-            wait = max(5, reset - int(time.time()) + 2)
+            # Secondary (abuse) rate limits set Retry-After and can fire even
+            # while the documented per-minute bucket still shows quota, so honor
+            # Retry-After first, then the reset header, then a safe default.
+            retry_after = resp.headers.get("Retry-After")
+            if retry_after and retry_after.isdigit():
+                wait = int(retry_after) + 2
+            else:
+                reset = int(resp.headers.get("X-RateLimit-Reset", time.time() + 60))
+                wait = max(30, reset - int(time.time()) + 2)
             print(f"    rate-limited; sleeping {wait}s", flush=True)
             time.sleep(wait)
             continue
@@ -258,10 +277,22 @@ def get_with_retry(session, url, params=None, retries=3):
             return None
         if resp.status_code == 404:
             return None
+        if resp.status_code in (408, 500, 502, 503, 504):
+            # GitHub's search backend times out (408) or blips (5xx) on broad
+            # queries; back off and retry rather than crashing the run.
+            wait = 15 * (attempt + 1)
+            print(f"    transient {resp.status_code}; sleeping {wait}s", flush=True)
+            time.sleep(wait)
+            continue
         resp.raise_for_status()
     return None
 
 
+# GitHub code search allows ~30 req/min but trips a stricter secondary limit on
+# bursts; pace pages well under that.
+SEARCH_PAGE_GAP = 3.0
+
+
 def search_code(session, query, per_page=100):
     for page in range(1, 11):  # cap at 1000 results/query
         resp = get_with_retry(session, GITHUB_SEARCH_URL,
@@ -272,7 +303,7 @@ def search_code(session, query, per_page=100):
         yield from items
         if len(items) < per_page:
             break
-        time.sleep(1.2)
+        time.sleep(SEARCH_PAGE_GAP)
 
 
 def fetch_content(session, item):
@@ -370,6 +401,10 @@ def run_pairs(args):
     almost never co-occur across independently-scraped single-scheduler corpora.
     """
     out_dir = Path(args.out or "testdata/corpus/candidate-pairs")
+    # Start clean: pairs mode isn't resumable, and leaving repo dirs from a prior
+    # run would misrepresent this run's results (the manifest is rewritten fresh).
+    if out_dir.exists():
+        shutil.rmtree(out_dir)
     out_dir.mkdir(parents=True, exist_ok=True)
     session = make_session(args.token)
 
@@ -381,18 +416,24 @@ def run_pairs(args):
         for query in sched.queries:
             if found >= args.search_limit:
                 break
-            for item in search_code(session, query):
-                if found >= args.search_limit:
-                    break
-                repo = item["repository"]["full_name"]
-                if is_excluded(repo):
-                    continue
-                if Path(item["path"]).suffix.lower() not in sched.extensions:
-                    continue
-                bucket = repo_map.setdefault(repo, {}).setdefault(name, [])
-                if item["path"] not in {it["path"] for it in bucket}:
-                    bucket.append(item)
-                    found += 1
+            try:
+                for item in search_code(session, query):
+                    if found >= args.search_limit:
+                        break
+                    repo = item["repository"]["full_name"]
+                    if is_excluded(repo):
+                        continue
+                    if Path(item["path"]).suffix.lower() not in sched.extensions:
+                        continue
+                    bucket = repo_map.setdefault(repo, {}).setdefault(name, [])
+                    if item["path"] not in {it["path"] for it in bucket}:
+                        bucket.append(item)
+                        found += 1
+            except requests.exceptions.RequestException as e:
+                # A single query dying must not abort the whole multi-scheduler
+                # sweep; log it and move on to the next query.
+                print(f"    query failed ({type(e).__name__}); skipping", flush=True)
+                continue
 
     # Phase 2 — keep repos spanning >= min_schedulers formats.
     candidates = {r: sm for r, sm in repo_map.items() if len(sm) >= args.min_schedulers}
@@ -409,7 +450,11 @@ def run_pairs(args):
             sched = SCHEDULERS[name]
             for item in items[:args.per_repo_limit]:
                 stars = item["repository"].get("stargazers_count", stars)
-                content = fetch_content(session, item)
+                try:
+                    content = fetch_content(session, item)
+                except requests.exceptions.RequestException as e:
+                    print(f"  SKIP {name}: {item['path']} ({type(e).__name__})")
+                    continue
                 if content is None or len(content.encode()) > args.max_bytes:
                     continue
                 ok, reason = sched.accept(content, item["path"], args)

testdata/corpus/candidate-pairs/manifest.json

@@ -0,0 +1,104 @@
+{
+  "NVIDIA/knavigator": {
+    "repo": "NVIDIA/knavigator",
+    "stars": 0,
+    "schedulers": [
+      "kueue",
+      "volcano"
+    ],
+    "files": {
+      "volcano": [
+        {
+          "path": "resources/benchmarks/templates/volcano/job.yaml",
+          "html_url": "https://github.com/NVIDIA/knavigator/blob/e5f1892df98235ba2c20a49574b6620b7a6f0e94/resources/benchmarks/templates/volcano/job.yaml",
+          "sha": "4c656f0800a2b6f3887cb4a60cf0318b24855bf0"
+        }
+      ],
+      "kueue": [
+        {
+          "path": "resources/benchmarks/templates/kueue/job.yaml",
+          "html_url": "https://github.com/NVIDIA/knavigator/blob/e5f1892df98235ba2c20a49574b6620b7a6f0e94/resources/benchmarks/templates/kueue/job.yaml",
+          "sha": "e1d6230b8a5f8263c27688811334330a67cde63f"
+        }
+      ]
+    }
+  },
+  "Project-HAMi/website": {
+    "repo": "Project-HAMi/website",
+    "stars": 0,
+    "schedulers": [
+      "kueue",
+      "volcano"
+    ],
+    "files": {
+      "volcano": [
+        {
+          "path": "tutorials/labs/examples/08-volcano-vgpu/02-gang-job.yaml",
+          "html_url": "https://github.com/Project-HAMi/website/blob/aa282781239fc86c4f6a59a95dca36c475bfca7d/tutorials/labs/examples/08-volcano-vgpu/02-gang-job.yaml",
+          "sha": "513894bed730e3578d8ebf5bd8be9272fadcbab8"
+        },
+        {
+          "path": "tutorials/labs/examples/08-volcano-vgpu/05-queue-fit-job.yaml",
+          "html_url": "https://github.com/Project-HAMi/website/blob/aa282781239fc86c4f6a59a95dca36c475bfca7d/tutorials/labs/examples/08-volcano-vgpu/05-queue-fit-job.yaml",
+          "sha": "f26c5ccf7815889c25347ef8a58f58a715c41e87"
+        }
+      ],
+      "kueue": [
+        {
+          "path": "tutorials/labs/examples/09-kueue-hami-vgpu/03-jobs.yaml",
+          "html_url": "https://github.com/Project-HAMi/website/blob/aa282781239fc86c4f6a59a95dca36c475bfca7d/tutorials/labs/examples/09-kueue-hami-vgpu/03-jobs.yaml",
+          "sha": "02c8870b549e6ca7a36ace952996ece33114fe67"
+        }
+      ]
+    }
+  },
+  "lasyard/docs": {
+    "repo": "lasyard/docs",
+    "stars": 0,
+    "schedulers": [
+      "kueue",
+      "volcano"
+    ],
+    "files": {
+      "volcano": [
+        {
+          "path": "_files/macos/workspace/k8s/volcano/sleep_vj_high.yaml",
+          "html_url": "https://github.com/lasyard/docs/blob/5dd43cb2ec7a2e1e74e3ed98f01a8c301b769b96/_files/macos/workspace/k8s/volcano/sleep_vj_high.yaml",
+          "sha": "f9d8e3173e24190fc62d5c730671af608c9a02c0"
+        },
+        {
+          "path": "_files/macos/workspace/k8s/volcano/hierarchical_vj_a1.yaml",
+          "html_url": "https://github.com/lasyard/docs/blob/5dd43cb2ec7a2e1e74e3ed98f01a8c301b769b96/_files/macos/workspace/k8s/volcano/hierarchical_vj_a1.yaml",
+          "sha": "433a0324aff9b66f8de0009009558cb54fe45a31"
+        }
+      ],
+      "kueue": [
+        {
+          "path": "_files/macos/workspace/k8s/kueue/tas3_job.yaml",
+          "html_url": "https://github.com/lasyard/docs/blob/5dd43cb2ec7a2e1e74e3ed98f01a8c301b769b96/_files/macos/workspace/k8s/kueue/tas3_job.yaml",
+          "sha": "e6b61c93aeb23bac55863163330e393a210b0869"
+        },
+        {
+          "path": "_files/macos/workspace/k8s/kueue/tas1_job.yaml",
+          "html_url": "https://github.com/lasyard/docs/blob/5dd43cb2ec7a2e1e74e3ed98f01a8c301b769b96/_files/macos/workspace/k8s/kueue/tas1_job.yaml",
+          "sha": "74191c0910ab75630a08829eb48dc3f596564255"
+        },
+        {
+          "path": "_files/macos/workspace/k8s/kueue/tas4_job.yaml",
+          "html_url": "https://github.com/lasyard/docs/blob/5dd43cb2ec7a2e1e74e3ed98f01a8c301b769b96/_files/macos/workspace/k8s/kueue/tas4_job.yaml",
+          "sha": "417e0a32ebd5dd1cb864cbfd05afb66a3b000847"
+        },
+        {
+          "path": "_files/macos/workspace/k8s/kueue/sleep_job_kueue_high.yaml",
+          "html_url": "https://github.com/lasyard/docs/blob/5dd43cb2ec7a2e1e74e3ed98f01a8c301b769b96/_files/macos/workspace/k8s/kueue/sleep_job_kueue_high.yaml",
+          "sha": "a8f6e48f63683735f3facecf931defa32d9712b4"
+        },
+        {
+          "path": "_files/macos/workspace/k8s/sleep_job_kueue.yaml",
+          "html_url": "https://github.com/lasyard/docs/blob/5dd43cb2ec7a2e1e74e3ed98f01a8c301b769b96/_files/macos/workspace/k8s/sleep_job_kueue.yaml",
+          "sha": "cab958a71a498b984f2dd15eb05a04010fda7ba4"
+        }
+      ]
+    }
+  }
+}