[ci] add script to filter out flaky tests targets (#42717)

the script takes test target names from stdin, and filters out all flaky ones based on the test
state in the database. this script is required for macos test invocation on the pipeline.

---------

Signed-off-by: khluu <[email protected]>
Signed-off-by: Lonnie Liu <[email protected]>
Co-authored-by: Lonnie Liu <[email protected]>

ci/ray_ci/automation/BUILD.bazel

@@ -17,3 +17,10 @@ py_binary(
  srcs = ["weekly_green_metric.py"],
  deps = ["//ci/ray_ci:ray_ci_lib"],
 )
+
+py_binary(
+ name = "filter_out_flaky_tests",
+ srcs = ["filter_out_flaky_tests.py"],
+ deps = ["//ci/ray_ci:ray_ci_lib"],
+)
+

ci/ray_ci/automation/filter_out_flaky_tests.py

@@ -0,0 +1,30 @@
+import sys
+
+from ci.ray_ci.utils import filter_out_flaky_tests
+from ray_release.configs.global_config import init_global_config
+from ray_release.bazel import bazel_runfile
+
+
+def main():
+ """
+ Filter out flaky tests.
+ Input (stdin) as a list of test names.
+ Output (stdout) as a list of test names without flaky tests.
+
+ Args:
+ prefix: Prefix to query tests with.
+ """
+ # Process arguments
+ if len(sys.argv) != 2:
+ raise ValueError("Invalid number of arguments.")
+
+ prefix = sys.argv[1]
+
+ # Initialize global config
+ init_global_config(bazel_runfile("release/ray_release/configs/oss_config.yaml"))
+
+ filter_out_flaky_tests(sys.stdin, sys.stdout, prefix)
+
+
+if __name__ == "__main__":
+ main()

ci/ray_ci/test_utils.py

@@ -1,10 +1,17 @@
 import base64
+import io
 import sys
 import pytest
 from unittest import mock
 from typing import List
 
-from ci.ray_ci.utils import chunk_into_n, docker_login
+from ray_release.test import Test
+from ci.ray_ci.utils import (
+ chunk_into_n,
+ docker_login,
+ get_flaky_test_names,
+ filter_out_flaky_tests,
+)
 
 
 def test_chunk_into_n() -> None:
@@ -34,5 +41,55 @@ def _mock_subprocess_run(
  docker_login("docker_ecr")
 
 
+def _make_test(name: str, state: str, team: str) -> Test:
+ return Test(
+ {
+ "name": name,
+ "state": state,
+ "team": team,
+ }
+ )
+
+
+@mock.patch("ray_release.test.Test.gen_from_s3")
+def test_get_flaky_test_names(mock_gen_from_s3):
+ mock_gen_from_s3.side_effect = (
+ [
+ _make_test("darwin:https://test_1", "flaky", "core"),
+ _make_test("darwin:https://test_2", "flaky", "ci"),
+ _make_test("darwin:https://test_3", "passing", "core"),
+ ],
+ [
+ _make_test("linux:https://test_1", "flaky", "core"),
+ _make_test("linux:https://test_2", "passing", "ci"),
+ ],
+ )
+ flaky_test_names = get_flaky_test_names(
+ prefix="darwin:",
+ )
+ assert flaky_test_names == ["//test_1", "//test_2"]
+ flaky_test_names = get_flaky_test_names(
+ prefix="linux:",
+ )
+ assert flaky_test_names == ["//test_1"]
+
+
+@mock.patch("ray_release.test.Test.gen_from_s3")
+def test_filter_out_flaky_tests(mock_gen_from_s3):
+ mock_gen_from_s3.side_effect = (
+ [
+ _make_test("darwin:https://test_1", "flaky", "core"),
+ _make_test("darwin:https://test_2", "flaky", "ci"),
+ _make_test("darwin:https://test_3", "passing", "core"),
+ _make_test("darwin:https://test_4", "passing", "ci"),
+ ],
+ )
+
+ test_targets = ["//test_1", "//test_2", "//test_3", "//test_4"]
+ output = io.StringIO()
+ filter_out_flaky_tests(io.StringIO("\n".join(test_targets)), output, "darwin:")
+ assert output.getvalue() == "//test_3\n//test_4\n"
+
+
 if __name__ == "__main__":
  sys.exit(pytest.main(["-v", __file__]))

ci/ray_ci/utils.py

@@ -1,4 +1,5 @@
 import base64
+import io
 import logging
 import subprocess
 import sys
@@ -9,6 +10,7 @@
 from math import ceil
 
 import ci.ray_ci.bazel_sharding as bazel_sharding
+from ray_release.test import Test, TestState
 
 
 POSTMERGE_PIPELINE = "0189e759-8c96-4302-b6b5-b4274406bf89"
@@ -77,6 +79,52 @@ def docker_pull(image: str) -> None:
  )
 
 
+def get_flaky_test_names(prefix: str) -> List[str]:
+ """
+ Query all flaky tests with specified prefix.
+
+ Args:
+ prefix: A prefix to filter by.
+
+ Returns:
+ List[str]: List of test names.
+ """
+ tests = Test.gen_from_s3(prefix)
+ # Filter tests by test state
+ state = TestState.FLAKY
+ test_names = [t.get_name() for t in tests if t.get_state() == state]
+
+ # Remove prefixes.
+ for i in range(len(test_names)):
+ test = test_names[i]
+ if test.startswith(prefix):
+ test_names[i] = test[len(prefix) :]
+
+ return test_names
+
+
+def filter_out_flaky_tests(input: io.TextIOBase, output: io.TextIOBase, prefix: str):
+ """
+ Filter out flaky tests from list of test targets.
+
+ Args:
+ input: Input stream, each test name in one line.
+ output: Output stream, each test name in one line.
+ prefix: Prefix to query tests with.
+ """
+ # Obtain all existing tests with specified test state
+ flaky_tests = set(get_flaky_test_names(prefix))
+
+ # Filter out these test from list of test targets
+ for t in input:
+ t = t.strip()
+ if not t:
+ continue
+ if t in flaky_tests:
+ continue
+ output.write(f"{t}\n")
+
+
 logger = logging.getLogger()
 logger.setLevel(logging.INFO)