1212
1313
1414@click .command ()
15- @click .argument ("targets" , required = True , type = str )
16- @click .argument ("team" , required = True , type = str )
15+ @click .argument ("targets" , required = True , type = str , nargs = - 1 )
16+ @click .argument ("team" , required = True , type = str , nargs = 1 )
1717@click .option (
1818 "--concurrency" ,
1919 default = 3 ,
2626 type = int ,
2727 help = ("Index of the concurrent shard to run." ),
2828)
29- def main (targets : str , team : str , concurrency : int , shard : int ) -> None :
29+ @click .option (
30+ "--size" ,
31+ default = "small,medium,large" ,
32+ type = str ,
33+ help = ("Size of tests to run." ),
34+ )
35+ @click .option (
36+ "--run-flaky-tests" ,
37+ is_flag = True ,
38+ show_default = True ,
39+ default = False ,
40+ help = ("Run flaky tests." ),
41+ )
42+ def main (
43+ targets : List [str ],
44+ team : str ,
45+ concurrency : int ,
46+ shard : int ,
47+ size : str ,
48+ run_flaky_tests : bool ,
49+ ) -> None :
3050 if not bazel_workspace_dir :
3151 raise Exception ("Please use `bazelisk run //ci/ray_ci`" )
32-
3352 os .chdir (bazel_workspace_dir )
3453
35- test_targets = _get_test_targets (targets , team , concurrency , shard )
54+ if run_flaky_tests :
55+ test_targets = _get_flaky_test_targets (team )
56+ else :
57+ test_targets = _get_test_targets (targets , team , concurrency , shard , size )
3658 if not test_targets :
3759 logging .info ("No tests to run" )
3860 return
@@ -50,13 +72,7 @@ def _run_tests(test_targets: List[str]) -> None:
5072 .split ()
5173 )
5274 subprocess .check_call (
53- [
54- "bazel" ,
55- "test" ,
56- "--config=ci" ,
57- ]
58- + bazel_options
59- + test_targets
75+ ["bazel" , "test" , "--config=ci" ] + bazel_options + test_targets
6076 )
6177
6278
@@ -65,16 +81,14 @@ def _get_test_targets(
6581 team : str ,
6682 concurrency : int ,
6783 shard : int ,
84+ size : str ,
6885 yaml_dir : Optional [str ] = None ,
6986) -> List [str ]:
7087 """
7188 Get test targets to run for a particular shard
7289 """
73- if not yaml_dir :
74- yaml_dir = os .path .join (bazel_workspace_dir , "ci/ray_ci" )
75-
7690 return _chunk_into_n (
77- _get_all_test_targets (targets , team , yaml_dir = yaml_dir ),
91+ _get_all_test_targets (targets , team , size , yaml_dir = yaml_dir ),
7892 concurrency ,
7993 )[shard ]
8094
@@ -84,28 +98,52 @@ def _chunk_into_n(list: List[str], n: int):
8498 return [list [x * size : x * size + size ] for x in range (n )]
8599
86100
87- def _get_all_test_targets (targets : str , team : str , yaml_dir : str ) -> List [str ]:
101+ def _get_all_test_query (targets : List [str ], team : str , size : str ) -> str :
102+ """
103+ Bazel query to get all test targets given a team and test size
104+ """
105+ test_query = " union " .join ([f"tests({ target } for target in targets ])
106+ team_query = f"attr(tags, team:{ team } { test_query }
107+ size_query = " union " .join (
108+ [f"attr(size, { s } { test_query } for s in size .split ("," )]
109+ )
110+ except_query = " union " .join (
111+ [
112+ f"attr(tags, { t } { test_query }
113+ for t in ["debug_tests" , "asan_tests" , "ray_ha" ]
114+ ]
115+ )
116+
117+ return f"({ team_query } { size_query } { except_query }
118+
119+
120+ def _get_all_test_targets (
121+ targets : str , team : str , size : str , yaml_dir : str
122+ ) -> List [str ]:
88123 """
89124 Get all test targets that are not flaky
90125 """
91126
92127 test_targets = (
93128 subprocess .check_output (
94- [
95- "bazel" ,
96- "query" ,
97- f"attr(tags, team:{ team } { targets }
98- # TODO(can): Remove this once we have a better way
99- # to filter out test size
100- f"attr(size, small, tests({ targets }
101- f"attr(size, medium, tests({ targets }
102- ")" ,
103- ]
129+ ["bazel" , "query" , _get_all_test_query (targets , team , size )],
104130 )
105131 .decode ("utf-8" )
106132 .split ("\n " )
107133 )
134+ flaky_tests = _get_flaky_test_targets (team , yaml_dir )
135+
136+ return [test for test in test_targets if test and test not in flaky_tests ]
137+
138+
139+ def _get_flaky_test_targets (team : str , yaml_dir : Optional [str ] = None ) -> List [str ]:
140+ """
141+ Get all test targets that are flaky
142+ """
143+ if not yaml_dir :
144+ yaml_dir = os .path .join (bazel_workspace_dir , "ci/ray_ci" )
145+
108146 with open (f"{ yaml_dir } { team } , "rb" ) as f :
109147 flaky_tests = yaml .safe_load (f )["flaky_tests" ]
110148
111- return [ test for test in test_targets if test and test not in flaky_tests ]
149+ return flaky_tests
0 commit comments