diff --git a/.github/scripts/calculate_run_times_for_label_combinations.py b/.github/scripts/calculate_run_times_for_label_combinations.py index dda8927e43c..9d6e50b020b 100644 --- a/.github/scripts/calculate_run_times_for_label_combinations.py +++ b/.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) ) diff --git a/.github/scripts/helpers.py b/.github/scripts/helpers.py index 746052e19dd..81a525d126d 100644 --- a/.github/scripts/helpers.py +++ b/.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 diff --git a/.github/scripts/helpers_test.py b/.github/scripts/helpers_test.py index 5845ef919af..b6d5fee7e0c 100644 --- a/.github/scripts/helpers_test.py +++ b/.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] diff --git a/.github/scripts/nebius_populate_vms.py b/.github/scripts/nebius_populate_vms.py index 0eb1c4c4e06..39f1884f18c 100644 --- a/.github/scripts/nebius_populate_vms.py +++ b/.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 diff --git a/.github/scripts/nebius_runners_wait_times.py b/.github/scripts/nebius_runners_wait_times.py index b6459608af1..76fe9db03d7 100644 --- a/.github/scripts/nebius_runners_wait_times.py +++ b/.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) diff --git a/.github/scripts/nebius_watch_runners.py b/.github/scripts/nebius_watch_runners.py index 5f3175bc08d..1eb845d7843 100644 --- a/.github/scripts/nebius_watch_runners.py +++ b/.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, diff --git a/.github/scripts/nightly_pr_check.py b/.github/scripts/nightly_pr_check.py index b5cb6d2e3af..c72ca6f53dd 100644 --- a/.github/scripts/nightly_pr_check.py +++ b/.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", ""), ) diff --git a/.github/scripts/tests/finalize_workload_comments.py b/.github/scripts/tests/finalize_workload_comments.py index 2dd9a3ac937..5a8f7f96152 100644 --- a/.github/scripts/tests/finalize_workload_comments.py +++ b/.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, ) diff --git a/.github/scripts/tests/workload_comment.py b/.github/scripts/tests/workload_comment.py index 2ccd5de0780..0e07c33b826 100644 --- a/.github/scripts/tests/workload_comment.py +++ b/.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, diff --git a/.github/scripts/tests/workload_comment_test.py b/.github/scripts/tests/workload_comment_test.py index 84ff0f8f545..4ae4ea9e100 100644 --- a/.github/scripts/tests/workload_comment_test.py +++ b/.github/scripts/tests/workload_comment_test.py @@ -29,15 +29,15 @@ def test_find_current_job_url_falls_back_to_run_url(monkeypatch) -> None: monkeypatch.setenv("GITHUB_REPOSITORY", "org/repo") monkeypatch.setenv("GITHUB_RUN_ID", "123") - def fake_get_jobs_raw(token: str, repo: str, run_id: int) -> list[object]: - assert token == "token" + def fake_get_jobs(github: object, repo: str, run_id: int) -> list[object]: + assert github is not None assert repo == "org/repo" assert run_id == 123 return [] - monkeypatch.setattr(h, "get_jobs_raw", fake_get_jobs_raw) + monkeypatch.setattr(h, "get_jobs", fake_get_jobs) - assert h.find_current_job_url("job", "runner") == ( + assert h.find_current_job_url(object(), "job", "runner") == ( "https://github.com/org/repo/actions/runs/123" ) @@ -47,8 +47,8 @@ def test_find_current_job_url_matches_reusable_workflow_job_name(monkeypatch) -> monkeypatch.setenv("GITHUB_REPOSITORY", "org/repo") monkeypatch.setenv("GITHUB_RUN_ID", "123") - def fake_get_jobs_raw(token: str, repo: str, run_id: int) -> list[object]: - assert token == "token" + def fake_get_jobs(github: object, repo: str, run_id: int) -> list[object]: + assert github is not None assert repo == "org/repo" assert run_id == 123 return [ @@ -64,10 +64,11 @@ def fake_get_jobs_raw(token: str, repo: str, run_id: int) -> list[object]: ) ] - monkeypatch.setattr(h, "get_jobs_raw", fake_get_jobs_raw) + monkeypatch.setattr(h, "get_jobs", fake_get_jobs) assert ( h.find_current_job_url( + object(), "Build and test [build_preset=relwithdebinfo component=blockstore] [id=1 ip=10.0.0.1]", "runner-1", ) @@ -80,8 +81,8 @@ def test_find_current_job_url_prefers_runner_specific_match(monkeypatch) -> None monkeypatch.setenv("GITHUB_REPOSITORY", "org/repo") monkeypatch.setenv("GITHUB_RUN_ID", "123") - def fake_get_jobs_raw(token: str, repo: str, run_id: int) -> list[object]: - assert token == "token" + def fake_get_jobs(github: object, repo: str, run_id: int) -> list[object]: + assert github is not None assert repo == "org/repo" assert run_id == 123 return [ @@ -99,10 +100,11 @@ def fake_get_jobs_raw(token: str, repo: str, run_id: int) -> list[object]: ), ] - monkeypatch.setattr(h, "get_jobs_raw", fake_get_jobs_raw) + monkeypatch.setattr(h, "get_jobs", fake_get_jobs) assert ( h.find_current_job_url( + object(), "Build and test [build_preset=relwithdebinfo component=blockstore]", "runner-b", )