test: increase timeouts for running experiments on k8s after env split (

#9530)
determined-ai · Jun 17, 2024 · edbeee9 · edbeee9
1 parent 0f6eb24
commit edbeee9
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 5 deletions.
diff --git a/e2e_tests/tests/cluster/managed_cluster_k8s.py b/e2e_tests/tests/cluster/managed_cluster_k8s.py
@@ -17,8 +17,18 @@ def __init__(self) -> None:
 
         # Verify we have pulled our image.
         # TODO this won't work if we have multiple nodes.
-        utils.wait_for_command_state(sess, utils.run_command(sess, 0, slots=0), "TERMINATED", 300)
-        utils.wait_for_command_state(sess, utils.run_command(sess, 0, slots=1), "TERMINATED", 300)
+        utils.wait_for_command_state(
+            sess,
+            utils.run_command(sess, 0, slots=0),
+            "TERMINATED",
+            utils.KUBERNETES_EXPERIMENT_TIMEOUT,
+        )
+        utils.wait_for_command_state(
+            sess,
+            utils.run_command(sess, 0, slots=1),
+            "TERMINATED",
+            utils.KUBERNETES_EXPERIMENT_TIMEOUT,
+        )
 
     def kill_master(self) -> None:
         self._scale_master(up=False)

diff --git a/e2e_tests/tests/cluster/test_agent_disable.py b/e2e_tests/tests/cluster/test_agent_disable.py
@@ -92,7 +92,10 @@ def test_disable_agent_experiment_resume() -> None:
         ["--config", "max_restarts=0"],
     )
     exp.wait_for_experiment_state(
-        sess, exp_id, bindings.experimentv1State.RUNNING, max_wait_secs=300
+        sess,
+        exp_id,
+        bindings.experimentv1State.RUNNING,
+        max_wait_secs=utils.KUBERNETES_EXPERIMENT_TIMEOUT,
     )
 
     with _disable_agent(admin, agent_id):
@@ -123,7 +126,7 @@ def test_disable_agent_zero_slots() -> None:
 
     command_id = utils.run_zero_slot_command(sess, sleep=180)
     # Wait for it to run.
-    utils.wait_for_command_state(sess, command_id, "RUNNING", 300)
+    utils.wait_for_command_state(sess, command_id, "RUNNING", utils.KUBERNETES_EXPERIMENT_TIMEOUT)
 
     try:
         with _disable_agent(admin, agent_id):
@@ -155,7 +158,10 @@ def test_drain_agent() -> None:
         ["--config", "hyperparameters.training_batch_seconds=0.15"],  # Take 15 seconds.
     )
     exp.wait_for_experiment_state(
-        sess, experiment_id, bindings.experimentv1State.RUNNING, max_wait_secs=300
+        sess,
+        experiment_id,
+        bindings.experimentv1State.RUNNING,
+        max_wait_secs=utils.KUBERNETES_EXPERIMENT_TIMEOUT,
     )
     exp.wait_for_experiment_active_workload(sess, experiment_id)
     exp.wait_for_experiment_workload_progress(sess, experiment_id)

diff --git a/e2e_tests/tests/cluster/utils.py b/e2e_tests/tests/cluster/utils.py
@@ -15,6 +15,8 @@
 from tests import config as conf
 from tests import detproc
 
+KUBERNETES_EXPERIMENT_TIMEOUT = 600
+
 
 class _HTTPServerWithRequest(http.server.HTTPServer):
     def __init__(