Add Hud x-token in lambda (#7511)

Add token in pytorch-lab/pytorch now, modify code to access it

Author: yangw-dev

aws/lambda/benchmark_regression_summary_report/common/benchmark_time_series_api_model.py

@@ -45,7 +45,7 @@ class BenchmarkTimeSeriesApiResponse:
 
     @classmethod
     def from_request(
-        cls, url: str, query: dict, timeout: int = 180
+        cls, url: str, query: dict, access_token: str, timeout: int = 180
     ) -> "BenchmarkTimeSeriesApiResponse":
         """
         Send a POST request and parse into BenchmarkTimeSeriesApiResponse.
@@ -62,14 +62,7 @@ def from_request(
         """
 
         headers = {
-            # Looks like a real browser instead of python-requests
-            "User-Agent": (
-                "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) "
-                "AppleWebKit/537.36 (KHTML, like Gecko) "
-                "Chrome/120.0.0.0 Safari/537.36"
-            ),
-            "Accept": "application/json,text/html;q=0.9,*/*;q=0.8",
-            "Accept-Language": "en-US,en;q=0.9",
+            "x-hud-internal-bot": access_token,
         }
         resp = requests.post(url, json=query, timeout=timeout, headers=headers)
         resp.raise_for_status()

aws/lambda/benchmark_regression_summary_report/lambda_function.py

@@ -32,6 +32,7 @@
     "CLICKHOUSE_ENDPOINT": os.getenv("CLICKHOUSE_ENDPOINT", ""),
     "CLICKHOUSE_PASSWORD": os.getenv("CLICKHOUSE_PASSWORD", ""),
     "CLICKHOUSE_USERNAME": os.getenv("CLICKHOUSE_USERNAME", ""),
+    "HUD_INTERNAL_BOT_TOKEN": os.getenv("HUD_INTERNAL_BOT_TOKEN", ""),
 }
 
 
@@ -71,13 +72,15 @@ def __init__(
         self,
         config_id: str,
         end_time: int,
+        hud_access_token: str = "",
         is_dry_run: bool = False,
         is_pass_check: bool = False,
     ) -> None:
         self.is_dry_run = is_dry_run
         self.is_pass_check = is_pass_check
         self.config_id = config_id
         self.end_time = end_time
+        self.hud_access_token = hud_access_token
 
     def log_info(self, msg: str):
         logger.info("[%s][%s] %s", self.end_time, self.config_id, msg)
@@ -136,13 +139,15 @@ def process(
                 f"with frequency {report_freq.get_text()}..."
             )
 
-        target, ls, le = self.get_target(config, self.end_time)
+        target, ls, le = self.get_target(config, self.end_time, self.hud_access_token)
         if not target.time_series:
             self.log_info(
                 f"no target data found for time range [{ls},{le}] with frequency {report_freq.get_text()}..."
             )
             return
-        baseline, bs, be = self.get_baseline(config, self.end_time)
+        baseline, bs, be = self.get_baseline(
+            config, self.end_time, self.hud_access_token
+        )
 
         if not baseline.time_series:
             self.log_info(
@@ -165,7 +170,7 @@ def process(
         reportManager.run(cc, ENVS["GITHUB_TOKEN"])
         return
 
-    def get_target(self, config: BenchmarkConfig, end_time: int):
+    def get_target(self, config: BenchmarkConfig, end_time: int, hud_access_token: str):
         data_range = config.policy.range
         target_s = end_time - data_range.comparison_timedelta_s()
         target_e = end_time
@@ -178,10 +183,16 @@ def get_target(self, config: BenchmarkConfig, end_time: int):
             start_time=target_s,
             end_time=target_e,
             source=config.source,
+            access_token=hud_access_token,
         )
         self.log_info(
             f"done. found {len(target_data.time_series)} # of data groups, with time range {target_data.time_range}",
         )
+
+        if len(target_data.time_series) > 0:
+            self.log_info(
+                f"peeking the first data: {target_data.time_series[0]}",
+            )
         if not target_data.time_range or not target_data.time_range.end:
             return None, target_s, target_e
 
@@ -190,7 +201,9 @@ def get_target(self, config: BenchmarkConfig, end_time: int):
             return None, target_s, target_e
         return target_data, target_s, target_e
 
-    def get_baseline(self, config: BenchmarkConfig, end_time: int):
+    def get_baseline(
+        self, config: BenchmarkConfig, end_time: int, hud_access_token: str
+    ):
         data_range = config.policy.range
         baseline_s = end_time - data_range.total_timedelta_s()
         baseline_e = end_time - data_range.comparison_timedelta_s()
@@ -204,12 +217,18 @@ def get_baseline(self, config: BenchmarkConfig, end_time: int):
             start_time=baseline_s,
             end_time=baseline_e,
             source=config.source,
+            access_token=hud_access_token,
         )
 
         self.log_info(
             f"Done. found {len(raw_data.time_series)} # of data, with time range {raw_data.time_range}",
         )
 
+        if len(raw_data.time_series) > 0:
+            self.log_info(
+                f"peeking the first data: {raw_data.time_series[0]}",
+            )
+
         baseline_latest_ts = int(isoparse(raw_data.time_range.end).timestamp())
 
         if not self.should_use_data(baseline_latest_ts, baseline_e):
@@ -245,6 +264,7 @@ def _fetch_from_benchmark_ts_api(
         config_id: str,
         end_time: int,
         start_time: int,
+        access_token: str,
         source: BenchmarkApiSource,
     ):
         str_end_time = format_ts_with_t(end_time)
@@ -261,7 +281,7 @@ def _fetch_from_benchmark_ts_api(
         t0 = time.perf_counter()
         try:
             resp: BenchmarkTimeSeriesApiResponse = (
-                BenchmarkTimeSeriesApiResponse.from_request(url, query)
+                BenchmarkTimeSeriesApiResponse.from_request(url, query, access_token)
             )
 
             elapsed_ms = (time.perf_counter() - t0) * 1000.0
@@ -353,6 +373,7 @@ def _get_latest_record_ts(
 def main(
     config_id: str,
     github_access_token: str = "",
+    hud_access_token: str = "",
     args: Optional[argparse.Namespace] = None,
     *,
     is_dry_run: bool = False,
@@ -392,6 +413,7 @@ def main(
             end_time=end_time_ts,
             is_dry_run=is_dry_run,
             is_pass_check=is_forced,
+            hud_access_token=hud_access_token,
         )
         processor.process(args=args)
     except Exception as e:
@@ -411,6 +433,7 @@ def lambda_handler(event: Any, context: Any) -> None:
     main(
         config_id=config_id,
         github_access_token=ENVS["GITHUB_TOKEN"],
+        hud_access_token=ENVS["HUD_INTERNAL_BOT_TOKEN"],
     )
     return
 
@@ -473,6 +496,12 @@ def parse_args() -> argparse.Namespace:
         type=str,
         help="the end time to run, in format of YYYY-MM-DD HH:MM:SS",
     )
+    parser.add_argument(
+        "--hud-internal-bot-token",
+        type=str,
+        default=ENVS["HUD_INTERNAL_BOT_TOKEN"],
+        help="the hud internal bot token to access hud api",
+    )
     parser.set_defaults(dry_run=True)  # default is True
     args, _ = parser.parse_known_args()
     return args
@@ -487,6 +516,7 @@ def local_run() -> None:
     # update environment variables for input parameters
     main(
         config_id=args.config_id,
+        hud_access_token=args.hud_internal_bot_token,
         github_access_token=args.github_access_token,
         args=args,
         is_dry_run=args.dry_run,