diff --git a/.github/scripts/update_disabled_issues.py b/.github/scripts/update_disabled_issues.py index 631727fde7..00125d0db3 100755 --- a/.github/scripts/update_disabled_issues.py +++ b/.github/scripts/update_disabled_issues.py @@ -28,7 +28,10 @@ def main() -> None: with urlopen( Request( f"{HUD_URL}/api/flaky-tests/getDisabledTestsAndJobs", - headers={"Authorization": os.environ["FLAKY_TEST_BOT_KEY"]}, + headers={ + "Authorization": os.environ["FLAKY_TEST_BOT_KEY"], + "x-hud-internal-bot": os.environ["HUD_API_TOKEN"], + }, ) ) as result: if result.status != 200: diff --git a/.github/workflows/check-alerts.yml b/.github/workflows/check-alerts.yml index 4b6595e5e4..488d005e53 100644 --- a/.github/workflows/check-alerts.yml +++ b/.github/workflows/check-alerts.yml @@ -47,6 +47,7 @@ jobs: env: # NOTE: Should be a blank string for pull requests GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + HUD_API_TOKEN: ${{ secrets.HUD_API_TOKEN }} update-queue-alert: env: @@ -66,3 +67,4 @@ jobs: env: # NOTE: Should be a blank string for pull requests GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + HUD_API_TOKEN: ${{ secrets.HUD_API_TOKEN }} diff --git a/.github/workflows/update-queue-times.yml b/.github/workflows/update-queue-times.yml index 0ab6d8add9..27b0b07d5b 100644 --- a/.github/workflows/update-queue-times.yml +++ b/.github/workflows/update-queue-times.yml @@ -24,3 +24,5 @@ jobs: role-to-assume: arn:aws:iam::308535385114:role/gha_workflow_update_queue_times aws-region: us-east-1 - run: yarn node scripts/updateQueueTimes.mjs + env: + HUD_API_TOKEN: ${{ secrets.HUD_API_TOKEN }} diff --git a/.github/workflows/update_disabled_tests.yml b/.github/workflows/update_disabled_tests.yml index fcbf627b65..abc6fca965 100644 --- a/.github/workflows/update_disabled_tests.yml +++ b/.github/workflows/update_disabled_tests.yml @@ -31,6 +31,7 @@ jobs: # environment, we do not have access to this token so fall back to the # GITHUB_TOKEN. FLAKY_TEST_BOT_KEY: ${{ secrets.FLAKY_TEST_BOT_KEY }} + HUD_API_TOKEN: ${{ secrets.HUD_API_TOKEN }} run: | python3 -m pip install GitPython==3.1.44 python3 .github/scripts/update_disabled_issues.py diff --git a/.github/workflows/update_test_file_report.yml b/.github/workflows/update_test_file_report.yml index 237d564cc9..ee38419850 100644 --- a/.github/workflows/update_test_file_report.yml +++ b/.github/workflows/update_test_file_report.yml @@ -53,6 +53,7 @@ jobs: GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} AWS_INFRA_ALERTS_LAMBDA_URL: ${{ secrets.AWS_INFRA_ALERTS_LAMBDA_URL }} TEST_REPORT_AWS_LAMBDA_TOKEN: ${{ secrets.TEST_REPORT_AWS_LAMBDA_TOKEN }} + HUD_API_TOKEN: ${{ secrets.HUD_API_TOKEN }} run: | cd test-infra python3 tools/torchci/test_insights/daily_regression.py diff --git a/tools/torchci/check_alerts.py b/tools/torchci/check_alerts.py index a9e4311d2a..6b1161f045 100755 --- a/tools/torchci/check_alerts.py +++ b/tools/torchci/check_alerts.py @@ -12,6 +12,7 @@ import requests from setuptools import distutils # type: ignore[import] +from torchci.utils import get_hud_headers FAILURE_CHAIN_THRESHOLD = 2 @@ -351,7 +352,10 @@ def create_issue(issue: Dict, dry_run: bool) -> Dict: def fetch_hud_data(repo: str, branch: str) -> Tuple[List[str], list[list[JobData]]]: - response = requests.get(f"https://hud.pytorch.org/api/hud/{repo}/{branch}/0") + response = requests.get( + f"https://hud.pytorch.org/api/hud/{repo}/{branch}/0", + headers=get_hud_headers(), + ) response.raise_for_status() hud_data = json.loads(response.text) diff --git a/tools/torchci/queue_alert.py b/tools/torchci/queue_alert.py index 7fdcf8b7f9..0ed99b7f40 100644 --- a/tools/torchci/queue_alert.py +++ b/tools/torchci/queue_alert.py @@ -13,6 +13,7 @@ fetch_alerts, update_issue, ) +from torchci.utils import get_hud_headers REPO_ROOT = Path(__file__).resolve().parent.parent.parent @@ -109,7 +110,7 @@ def queuing_alert(dry_run: bool) -> None: url = ( "https://hud.pytorch.org/api/clickhouse/queued_jobs_by_label?parameters=%7B%7D" ) - response = requests.get(url).json() + response = requests.get(url, headers=get_hud_headers()).json() large_queue = filter_long_queues(response) diff --git a/tools/torchci/test_insights/daily_regression.py b/tools/torchci/test_insights/daily_regression.py index e6a881379f..910a7a8b77 100644 --- a/tools/torchci/test_insights/daily_regression.py +++ b/tools/torchci/test_insights/daily_regression.py @@ -5,6 +5,7 @@ import requests from torchci.test_insights.file_report_generator import FileReportGenerator from torchci.test_insights.weekly_notification import send_to_aws_alerting_lambda +from torchci.utils import get_hud_headers FILE_REPORT_URL = "https://hud.pytorch.org/tests/fileReport" @@ -211,7 +212,8 @@ def get_representative_data_for_time( self, start_date, stop_date ) -> list[dict[str, Any]]: response = requests.get( - f"https://hud.pytorch.org/api/flaky-tests/fileReport?startDate={start_date}&endDate={stop_date}" + f"https://hud.pytorch.org/api/flaky-tests/fileReport?startDate={start_date}&endDate={stop_date}", + headers=get_hud_headers(), ) if response.status_code != 200: diff --git a/tools/torchci/utils.py b/tools/torchci/utils.py index 6789d90060..346dddad07 100644 --- a/tools/torchci/utils.py +++ b/tools/torchci/utils.py @@ -4,7 +4,7 @@ import pathlib import subprocess from hashlib import sha256 -from typing import List, Union +from typing import Any, List, Union FILE_CACHE_LIFESPAN_SECONDS = 60 * 60 * 24 # 1 day @@ -66,3 +66,18 @@ def wrapper(*args, **kwargs): return res return wrapper + + +def get_hud_headers() -> dict[str, Any]: + """ + Get headers for requests to the HUD API. This includes the + x-hud-internal-bot header which is required for authentication. If the + HUD_API_TOKEN environment variable is not set, this function returns an + empty dictionary. + """ + if "HUD_API_TOKEN" not in os.environ: + return {} + return { + # Looks like a real browser instead of python-requests + "x-hud-internal-bot": os.environ["HUD_API_TOKEN"], + } diff --git a/torchci/scripts/updateQueueTimes.mjs b/torchci/scripts/updateQueueTimes.mjs index 0a4d019c73..8bce92d10d 100644 --- a/torchci/scripts/updateQueueTimes.mjs +++ b/torchci/scripts/updateQueueTimes.mjs @@ -14,7 +14,12 @@ const s3client = getS3Client(); // %7B%7D = encoded {} const response = await fetch( - "https://hud.pytorch.org/api/clickhouse/queued_jobs_by_label?parameters=%7B%7D" + "https://hud.pytorch.org/api/clickhouse/queued_jobs_by_label?parameters=%7B%7D", + { + headers: { + "x-hud-internal-bot": process.env.HUD_API_TOKEN, + }, + } ).then((r) => r.json()); for (const r of response) { const unixTime = parseInt((new Date(r.time).getTime() / 1000).toFixed(0));