[ci] Restore pytest_checker script, but at correct location (ray-proj…

…ect#34523)

ray-project#34463 removed the scripts under the `scripts/` directory because all of them should have been symlinks. However, `pytest_checker.py` was an actual script that was not symlinked from the `ci/` directory. This PR restores this script at the correct location in `ci/lint` and adjusts all references to it in the codebase.

Signed-off-by: Kai Fricke <[email protected]>
Signed-off-by: Jack He <[email protected]>

build-docker.sh

@@ -2,7 +2,7 @@
 # shellcheck disable=SC2086
 # This script is for users to build docker images locally. It is most useful for users wishing to edit the
 # base-deps, ray-deps, or ray images. This script is *not* tested, so please look at the 
-# scripts/build-docker-images.py if there are problems with using this script.
+# ci/build/build-docker-images.py if there are problems with using this script.
 
 set -x
 

ci/ci.sh

@@ -555,7 +555,7 @@ lint_bazel_pytest() {
     # - find all py_test rules in bazel that have the specified team tag EXCEPT ones with "no_main" tag and outputs them as xml
     # - converts the xml to json
     # - feeds the json into pytest_checker.py
-    bazel query "kind(py_test.*, tests(python/...) intersect attr(tags, \"\b$team\b\", python/...) except attr(tags, \"\bno_main\b\", python/...))" --output xml | xq | python scripts/pytest_checker.py
+    bazel query "kind(py_test.*, tests(python/...) intersect attr(tags, \"\b$team\b\", python/...) except attr(tags, \"\bno_main\b\", python/...))" --output xml | xq | python ci/lint/pytest_checker.py
   done
 }
 

ci/lint/pytest_checker.py

@@ -0,0 +1,126 @@
+import json
+import re
+import sys
+from pathlib import Path
+
+
+def check_file(file_contents: str) -> bool:
+    """Check file for the snippet"""
+    return bool(re.search(r"^if __name__ == \"__main__\":", file_contents, re.M))
+
+
+def parse_json(data: str) -> dict:
+    return json.loads(data)
+
+
+def treat_path(path: str) -> Path:
+    """Treat bazel paths to filesystem paths"""
+    path = path[2:].replace(":", "/")
+    return Path(path)
+
+
+def get_paths_from_parsed_data(parsed_data: dict) -> list:
+    # Example JSON input:
+    # "rule": [
+    #     {
+    #       "@class": "py_test",
+    #       "@location": "/home/ubuntu/ray/python/ray/tests/BUILD:345:8",
+    #       "@name": "//python/ray/tests:test_tracing",
+    #       "string": [
+    #         {
+    #           "@name": "name",
+    #           "@value": "test_tracing"
+    #         },
+    #       ],
+    #       "list": [
+    #         {
+    #           "@name": "srcs",
+    #           "label": [
+    #             {
+    #               "@value": "//python/ray/tests:aws/conftest.py"
+    #             },
+    #             {
+    #               "@value": "//python/ray/tests:conftest.py"
+    #             },
+    #             {
+    #               "@value": "//python/ray/tests:test_tracing.py"
+    #             }
+    #           ]
+    #         }
+    #       ],
+    #       ... other fields ...
+    #       "label": {
+    #         "@name": "main",
+    #         "@value": "//python/ray/tests:test_runtime_env_working_dir_remote_uri.py"
+    #       },
+    #       ... other fields ...
+    #     }
+    # ]
+    #
+    # We want to get the location of the actual test file.
+    # This can be, in order of priority:
+    #   1. Specified as the "main" label
+    #   2. Specified as the ONLY "srcs" label
+    #   3. Specified as the "srcs" label matching the "name" of the test
+    # https://docs.bazel.build/versions/main/be/python.html#py_test
+
+    paths = []
+    for rule in parsed_data["query"]["rule"]:
+        name = rule["@name"]
+        if "label" in rule and rule["label"]["@name"] == "main":
+            paths.append((name, treat_path(rule["label"]["@value"])))
+        else:
+            list_args = {e["@name"]: e for e in rule["list"]}
+            label = list_args["srcs"]["label"]
+            if isinstance(label, dict):
+                paths.append((name, treat_path(label["@value"])))
+            else:
+                # list
+                string_name = next(
+                    x["@value"] for x in rule["string"] if x["@name"] == "name"
+                )
+                main_path = next(
+                    x["@value"] for x in label if string_name in x["@value"]
+                )
+                paths.append((name, treat_path(main_path)))
+    return paths
+
+
+def main(data: str):
+    print("Checking files for the pytest snippet...")
+    parsed_data = parse_json(data)
+    paths = get_paths_from_parsed_data(parsed_data)
+
+    bad_paths = []
+    for name, path in paths:
+        # Special case for myst doc checker
+        if "test_myst_doc" in str(path):
+            continue
+
+        print(f"Checking test '{name}' | file '{path}'...")
+        try:
+            with open(path, "r") as f:
+                if not check_file(f.read()):
+                    print(f"File '{path}' is missing the pytest snippet.")
+                    bad_paths.append(path)
+        except FileNotFoundError:
+            print(f"File '{path}' is missing.")
+            bad_paths.append((path, "path is missing!"))
+    if bad_paths:
+        formatted_bad_paths = "\n".join([str(x) for x in bad_paths])
+        raise RuntimeError(
+            'Found py_test files without `if __name__ == "__main__":` snippet:'
+            f"\n{formatted_bad_paths}\n"
+            "If this is intentional, please add a `no_main` tag to bazel BUILD "
+            "entry for those files."
+        )
+
+
+if __name__ == "__main__":
+    # Expects a json
+    # Invocation from workspace root:
+    # bazel query 'kind(py_test.*, tests(python/...) intersect
+    # attr(tags, "\bteam:ml\b", python/...) except attr(tags, "\bno_main\b",
+    # python/...))' --output xml | xq | python ci/lint/pytest_checker.py
+    data = sys.stdin.read()
+    main(data)