Consolidate GitHub workflow job retrieval

.github/scripts/calculate_run_times_for_label_combinations.py

@@ -24,7 +24,7 @@
 from collections import defaultdict
 from dataclasses import dataclass
 from datetime import datetime, timedelta, timezone
-from typing import DefaultDict, Dict, Iterable, List, Optional, Tuple, Union
+from typing import DefaultDict, Dict, Iterable, List, Optional, Tuple
 
 import numpy as np
 from tabulate import tabulate
@@ -69,29 +69,6 @@ def label_combo_key(labels: Iterable[str]) -> Tuple[str, ...]:
     return tuple(sorted(labels))
 
 
-#
-# PyGithub requester helper: handle variable return shapes from requestJsonAndCheck
-#
-def _req_json_and_headers(requester, path: str):
-    try:
-        res = requester.requestJsonAndCheck("GET", path, headers={})
-    except Exception:
-        raise
-    if isinstance(res, (tuple, list)):
-        if len(res) == 3:
-            data, _, headers = res
-            return data, headers or {}
-        elif len(res) == 2:
-            data, headers = res
-            return data, headers or {}
-        else:
-            data = res[0]
-            headers = res[-1] if len(res) > 1 else {}
-            return data, headers or {}
-    else:
-        return res, {}
-
-
 def get_workflow_id_by_path(repo, workflow_ref: str) -> int:
     LOG.info("Resolving workflow reference: %r", workflow_ref)
 
@@ -190,47 +167,6 @@ def iter_successful_runs(
     )
 
 
-def iter_jobs_for_run(run) -> Iterable[Union[dict, object]]:
-    """
-    Yield job dicts or PyGithub Job objects for the run.
-    Prefer run.jobs() (PyGithub), else fallback to raw jobs endpoint with pagination.
-    """
-    try:
-        yield from run.jobs()
-        return
-    except Exception as e:
-        LOG.debug(
-            "run.jobs() failed for run %s: %s. Falling back to raw jobs API.",
-            getattr(run, "id", None),
-            e,
-        )
-
-    try:
-        requester = run._requester
-        owner = run.repository.owner.login
-        repo = run.repository.name
-        path = f"/repos/{owner}/{repo}/actions/runs/{run.id}/jobs?per_page=100"
-        while path:
-            data, headers = _req_json_and_headers(requester, path)
-            yield from data.get("jobs", [])
-
-            link = headers.get("link") or headers.get("Link")
-            next_url = None
-            if link:
-                parts = [p.strip() for p in link.split(",")]
-                for p in parts:
-                    if 'rel="next"' in p:
-                        next_url = p.split(";")[0].strip().strip("<>")
-                        break
-
-            if next_url and next_url.startswith("https://api.github.com"):
-                path = next_url.replace("https://api.github.com", "")
-            else:
-                path = None
-    except Exception as e:
-        LOG.error("Raw jobs API failed for run %s: %s", getattr(run, "id", None), e)
-
-
 def _extract_steps_from_job(job_obj) -> List[Dict]:
     if isinstance(job_obj, dict):
         return job_obj.get("steps") or []
@@ -341,7 +277,7 @@ def matches(n: Optional[str]) -> bool:
 
 
 def get_job_duration_seconds(run, job_name: str, match_mode: str) -> Optional[float]:
-    jobs = list(iter_jobs_for_run(run))
+    jobs = list(run.jobs(_filter="latest"))
     LOG.debug(
         "Run %s has %d jobs (via chosen method)", getattr(run, "id", None), len(jobs)
     )

.github/scripts/helpers.py

@@ -137,10 +137,12 @@ def job_name_matches(expected_name: str, actual_name: str) -> bool:
     return False
 
 
-def find_current_job_url(current_job_name: str, runner_name: str) -> str:
+def find_current_job_url(
+    github: Github, current_job_name: str, runner_name: str
+) -> str:
     try:
-        jobs = get_jobs_raw(
-            os.environ["GITHUB_TOKEN"],
+        jobs = get_jobs(
+            github,
             os.environ["GITHUB_REPOSITORY"],
             int(os.environ["GITHUB_RUN_ID"]),
         )
@@ -466,10 +468,14 @@ def extract_github_runner_release(payload: dict) -> GithubRunnerRelease:
     return GithubRunnerRelease(version=version, sha256_by_arch=sha256_by_arch)
 
 
-def github_client(github_token: str | None = None) -> Github:
+def github_client(
+    github_token: str | None = None,
+    *,
+    base_url: str = "https://api.github.com",
+) -> Github:
     if github_token:
-        return Github(auth=GithubAuth.Token(github_token))
-    return Github()
+        return Github(auth=GithubAuth.Token(github_token), base_url=base_url)
+    return Github(base_url=base_url)
 
 
 def github_client_from_env() -> Github:
@@ -826,6 +832,16 @@ def parse_actions_job_url(job_url: str) -> tuple[int, int] | None:
     interval_sec=GITHUB_API_RETRY_INTERVAL_SEC,
     retry_exceptions=PYGITHUB_RETRY_EXCEPTIONS,
 )
-def get_jobs_raw(token, repo_full_name, run_id) -> list[WorkflowJob]:
-    repo = github_client(token).get_repo(repo_full_name)
-    return list(repo.get_workflow_run(run_id).jobs())
+def get_jobs(
+    github: Github,
+    repo_full_name: str,
+    run_id: int,
+    *,
+    run_attempt: int | None = None,
+) -> list[WorkflowJob]:
+    repo = github.get_repo(repo_full_name)
+    run = repo.get_workflow_run(run_id)
+    jobs = list(run.jobs(_filter="all" if run_attempt is not None else "latest"))
+    if run_attempt is not None:
+        jobs = [job for job in jobs if job.run_attempt == run_attempt]
+    return jobs

.github/scripts/helpers_test.py

@@ -448,11 +448,16 @@ def get_repo(self, repo):
     assert release.sha256_by_arch["arm64"] == "b" * 64
 
 
-def test_get_jobs_raw_fetches_workflow_jobs_with_pygithub(monkeypatch):
-    jobs = [SimpleNamespace(name="job-1")]
+def test_get_jobs_fetches_latest_or_one_attempt():
+    jobs = [
+        SimpleNamespace(name="old", run_attempt=1),
+        SimpleNamespace(name="current", run_attempt=2),
+    ]
+    filters = []
 
     class FakeRun:
-        def jobs(self):
+        def jobs(self, *, _filter):
+            filters.append(_filter)
             return jobs
 
     class FakeRepo:
@@ -465,22 +470,22 @@ def get_repo(self, repo):
             assert repo == "owner/repo"
             return FakeRepo()
 
-    def fake_github_client(token):
-        assert token == "token"
-        return FakeGithub()
-
-    monkeypatch.setattr(h, "github_client", fake_github_client)
+    github = FakeGithub()
+    assert h.get_jobs(github, "owner/repo", 123) == jobs
+    assert [
+        job.name for job in h.get_jobs(github, "owner/repo", 123, run_attempt=2)
+    ] == ["current"]
+    assert filters == ["latest", "all"]
 
-    assert h.get_jobs_raw("token", "owner/repo", 123) == jobs
 
-
-def test_get_jobs_raw_retries_pygithub_failures(monkeypatch):
+def test_get_jobs_retries_pygithub_failures(monkeypatch):
     jobs = [SimpleNamespace(name="job-1")]
     attempts = []
     sleeps = []
 
     class FakeRun:
-        def jobs(self):
+        def jobs(self, *, _filter):
+            assert _filter == "latest"
             return jobs
 
     class FakeRepo:
@@ -496,14 +501,9 @@ def get_repo(self, repo):
             assert repo == "owner/repo"
             return FakeRepo()
 
-    def fake_github_client(token):
-        assert token == "token"
-        return FakeGithub()
-
-    monkeypatch.setattr(h, "github_client", fake_github_client)
     monkeypatch.setattr(h.time, "sleep", sleeps.append)
 
-    assert h.get_jobs_raw("token", "owner/repo", 123) == jobs
+    assert h.get_jobs(FakeGithub(), "owner/repo", 123) == jobs
     assert attempts == [123, 123]
     assert sleeps == [h.GITHUB_API_RETRY_INTERVAL_SEC]
 

.github/scripts/nebius_populate_vms.py

@@ -7,7 +7,7 @@
 from grpc import StatusCode
 from github import Github
 from typing import List
-from .helpers import github_client
+from .helpers import get_jobs, github_client
 from nebius.sdk import SDK
 from nebius.aio.cli_config import Config
 from nebius.api.nebius.compute.v1 import (
@@ -320,7 +320,11 @@ async def run(github: Github, sdk: SDK, args: argparse.Namespace) -> bool:
     queued_workflows_count = 0
     for workflow in queued_workflows:
         # search through workflow jobs to get labels for jobs with status queued
-        jobs = workflow.jobs()
+        jobs = get_jobs(
+            github,
+            f"{args.github_repo_owner}/{args.github_repo}",
+            workflow.id,
+        )
         for job in jobs:
             if job.status != "queued":
                 continue

.github/scripts/nebius_runners_wait_times.py

@@ -8,7 +8,7 @@
 from collections import defaultdict
 from dateutil import parser as dateparser
 from dateutil.relativedelta import relativedelta
-from .helpers import setup_logger, github_client, get_jobs_raw, classify_runner
+from .helpers import setup_logger, github_client, get_jobs, classify_runner
 
 logger = setup_logger()
 
@@ -97,7 +97,7 @@ def output_results(all_jobs: list[dict], summary, threshold: int):
     )
 
 
-def main(start, end, threshold, github_token, repo):
+def main(start, end, threshold, github, repo):
     logger.info(f"Fetching workflow runs from {start} to {end}")
     all_jobs = []
     summary = defaultdict(lambda: {"total_wait": 0.0, "count": 0, "waits": []})
@@ -112,7 +112,7 @@ def main(start, end, threshold, github_token, repo):
             continue
 
         try:
-            jobs = get_jobs_raw(github_token, repo.full_name, run.id)
+            jobs = get_jobs(github, repo.full_name, run.id)
         except Exception as e:
             logger.warning(f"Failed to get jobs for run {run.id}: {e}")
             continue
@@ -211,4 +211,4 @@ def main(start, end, threshold, github_token, repo):
     g = github_client(github_token)
     repo = g.get_repo(f"{args.owner}/{args.repo}")
 
-    main(start, end, args.threshold, github_token, repo)
+    main(start, end, args.threshold, g, repo)

.github/scripts/nebius_watch_runners.py

@@ -6,7 +6,7 @@
 from .helpers import (
     setup_logger,
     github_client,
-    get_jobs_raw,
+    get_jobs,
     compact_workflow_name,
     compact_job_name,
     date_to_hms,
@@ -86,7 +86,7 @@ async def main():
         queued_workflows_runs.append(run.id)
 
     for run in workflow_runs:
-        for job in get_jobs_raw(token, repo.full_name, run.id):
+        for job in get_jobs(g, f"{args.owner}/{args.repo}", run.id):
             if job.status in ("in_progress", "queued") and job.runner_name:
                 active_jobs[job.runner_name] = {
                     "job_name": job.name,

.github/scripts/nightly_pr_check.py

@@ -849,6 +849,7 @@ def load_context(gh) -> tuple[
     runs = selected_workflows(labels)
     marker = f"pr-{pr_number}-run-{os.environ['GITHUB_RUN_ID']}-attempt-{os.environ.get('GITHUB_RUN_ATTEMPT', '1')}"
     collector_url = find_current_job_url(
+        gh,
         os.environ.get("GITHUB_JOB", "nightly-builds"),
         os.environ.get("RUNNER_NAME", ""),
     )

.github/scripts/tests/finalize_workload_comments.py

@@ -11,7 +11,7 @@
 from ..helpers import (
     BUILD_AND_TEST_JOB_NAME_PREFIX,
     PYGITHUB_RETRY_EXCEPTIONS,
-    get_jobs_raw,
+    get_jobs,
     github_client_from_env,
     load_github_event,
     parse_actions_job_url,
@@ -162,10 +162,8 @@ def main() -> None:
     ):
         return
 
-    pr = pull_request_from_event(
-        github_client_from_env(),
-        load_github_event(),
-    )
+    github = github_client_from_env()
+    pr = pull_request_from_event(github, load_github_event())
     run_number = int(os.environ.get("GITHUB_RUN_NUMBER", "0"))
     current_run_id = int(os.environ.get("GITHUB_RUN_ID", "0"))
     jobs_cache: JobsCache = {}
@@ -175,8 +173,8 @@ def jobs_for_run(run_id: WorkflowRunId) -> WorkflowJobs:
         if run_id in jobs_cache:
             return jobs_cache[run_id]
         try:
-            jobs = get_jobs_raw(
-                os.environ["GITHUB_TOKEN"],
+            jobs = get_jobs(
+                github,
                 os.environ["GITHUB_REPOSITORY"],
                 run_id,
             )

.github/scripts/tests/workload_comment.py

@@ -82,10 +82,8 @@ def main() -> None:
     ):
         return
 
-    pr = pull_request_from_event(
-        github_client_from_env(),
-        load_github_event(),
-    )
+    github = github_client_from_env()
+    pr = pull_request_from_event(github, load_github_event())
     run_number = int(os.environ.get("GITHUB_RUN_NUMBER", "0"))
 
     if args.command == "init":
@@ -101,7 +99,7 @@ def main() -> None:
 
     job_url = ""
     if args.current_job_name:
-        job_url = find_current_job_url(args.current_job_name, args.runner_name)
+        job_url = find_current_job_url(github, args.current_job_name, args.runner_name)
         write_output(args.job_url_out, job_url)
     gs.update_pr_comment_workload_check(
         run_number=run_number,