Add headers for hud api auth for some scripts

.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:

.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 }}

.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 }}

.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

.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

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)
 

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)
 

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:

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"],
+    }

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));