From 6a274dfd76b9e5808275e9bb7f147cda40ea6b1c Mon Sep 17 00:00:00 2001
From: Chen Shen <scv119@gmail.com>
Date: Thu, 9 Dec 2021 13:41:46 -0800
Subject: [PATCH] CI][Chaos-test] chaos test now can set max-nodes-to-kill
 #20962

---
 python/ray/_private/test_utils.py    | 7 +++++--
 release/nightly_tests/setup_chaos.py | 4 +++-
 2 files changed, 8 insertions(+), 3 deletions(-)

diff --git a/python/ray/_private/test_utils.py b/python/ray/_private/test_utils.py
index 0a6eb1761..46c52fd6a 100644
--- a/python/ray/_private/test_utils.py
+++ b/python/ray/_private/test_utils.py
@@ -957,7 +957,8 @@ def teardown_tls(key_filepath, cert_filepath, temp_dir):
 def get_and_run_node_killer(node_kill_interval_s,
                             namespace=None,
                             lifetime=None,
-                            no_start=False):
+                            no_start=False,
+                            max_nodes_to_kill=2):
     assert ray.is_initialized(), (
         "The API is only available when Ray is initialized.")
 
@@ -1058,7 +1059,9 @@ def get_and_run_node_killer(node_kill_interval_s,
         namespace=namespace,
         name="node_killer",
         lifetime=lifetime).remote(
-            head_node_id, node_kill_interval_s=node_kill_interval_s)
+            head_node_id,
+            node_kill_interval_s=node_kill_interval_s,
+            max_nodes_to_kill=max_nodes_to_kill)
     print("Waiting for node killer actor to be ready...")
     ray.get(node_killer.ready.remote())
     print("Node killer actor is ready now.")
diff --git a/release/nightly_tests/setup_chaos.py b/release/nightly_tests/setup_chaos.py
index 21a623e0a..0c80ef38b 100644
--- a/release/nightly_tests/setup_chaos.py
+++ b/release/nightly_tests/setup_chaos.py
@@ -7,6 +7,7 @@ from ray._private.test_utils import get_and_run_node_killer
 def parse_script_args():
     parser = argparse.ArgumentParser()
     parser.add_argument("--node-kill-interval", type=int, default=60)
+    parser.add_argument("--max-nodes-to-kill", type=int, default=2)
     parser.add_argument(
         "--no-start",
         action="store_true",
@@ -31,7 +32,8 @@ def main():
         args.node_kill_interval,
         namespace="release_test_namespace",
         lifetime="detached",
-        no_start=args.no_start)
+        no_start=args.no_start,
+        max_nodes_to_kill=args.max_nodes_to_kill)
     print("Successfully deployed a node killer.")