From 0d56d17948d00bc7471c8f67b08d2ed22f020ed8 Mon Sep 17 00:00:00 2001
From: Lonnie Liu <lonnie@anyscale.com>
Date: Sat, 8 Feb 2025 19:21:59 +0000
Subject: [PATCH] [doc] compile api annotation check into single job

this avoids building the ray wheel for many times

Signed-off-by: Lonnie Liu <lonnie@anyscale.com>
---
 .buildkite/lint.rayci.yml                  |  7 +---
 ci/lint/lint.sh                            |  1 +
 ci/ray_ci/doc/cmd_check_api_discrepancy.py | 40 ++++++++++++++++------
 3 files changed, 31 insertions(+), 17 deletions(-)

diff --git a/.buildkite/lint.rayci.yml b/.buildkite/lint.rayci.yml
index c1a812d98ed8b..8cf4844d8e9e0 100644
--- a/.buildkite/lint.rayci.yml
+++ b/.buildkite/lint.rayci.yml
@@ -33,12 +33,7 @@ steps:
       - ./ci/lint/lint.sh {{matrix}}
     matrix:
       - api_annotations
-      - api_policy_check core
-      - api_policy_check serve
-      - api_policy_check data
-      - api_policy_check train
-      - api_policy_check tune
-      - api_policy_check rllib
+      - api_policy_check
 
   - label: ":lint-roller: lint: linkcheck"
     instance_type: medium
diff --git a/ci/lint/lint.sh b/ci/lint/lint.sh
index b68723f90d1cd..2031132184eb3 100755
--- a/ci/lint/lint.sh
+++ b/ci/lint/lint.sh
@@ -94,6 +94,7 @@ api_policy_check() {
   # install ray and compile doc to generate API files
   make -C doc/ html
   RAY_DISABLE_EXTRA_CPP=1 pip install -e "python[all]"
+
   # validate the API files
   bazel run //ci/ray_ci/doc:cmd_check_api_discrepancy -- /ray "$@"
 }
diff --git a/ci/ray_ci/doc/cmd_check_api_discrepancy.py b/ci/ray_ci/doc/cmd_check_api_discrepancy.py
index 49e3b84d2593a..71f1c053f9fe8 100644
--- a/ci/ray_ci/doc/cmd_check_api_discrepancy.py
+++ b/ci/ray_ci/doc/cmd_check_api_discrepancy.py
@@ -65,15 +65,7 @@
 }
 
 
-@click.command()
-@click.argument("ray_checkout_dir", required=True, type=str)
-@click.argument("team", required=True, type=click.Choice(list(TEAM_API_CONFIGS.keys())))
-def main(ray_checkout_dir: str, team: str) -> None:
-    """
-    This script checks for annotated classes and functions in a module, and finds
-    discrepancies between the annotations and the documentation.
-    """
-
+def _check_team(ray_checkout_dir: str, team: str) -> bool:
     # Load all APIs from the codebase
     api_in_codes = {}
     for module in TEAM_API_CONFIGS[team]["head_modules"]:
@@ -103,9 +95,35 @@ def main(ray_checkout_dir: str, team: str) -> None:
     for api in bad_apis:
         logger.info(f"\t{api}")
 
-    assert not bad_apis, "Some public APIs are not documented. Please document them."
+    if not bad_apis:
+        logger.info(
+            f"Some public {team} APIs are not documented. Please document them."
+        )
+        return False
+    return True
+
+
+@click.command()
+@click.argument("ray_checkout_dir", required=True, type=str)
+@click.argument(
+    "team", default="ALL", type=click.Choice(list(TEAM_API_CONFIGS.keys()) + ["ALL"])
+)
+def main(ray_checkout_dir: str, team: str) -> None:
+    """
+    This script checks for annotated classes and functions in a module, and finds
+    discrepancies between the annotations and the documentation.
+    """
+    if team != "ALL":
+        if not _check_team(ray_checkout_dir, team):
+            exit(1)
+        return
 
-    return
+    all_pass = True
+    for team in TEAM_API_CONFIGS:
+        if not _check_team(ray_checkout_dir, team):
+            all_pass = False
+    if not all_pass:
+        exit(1)
 
 
 if __name__ == "__main__":