Refactor timeout handling in context manager; update test to reduce avg_step_timeout and rename test function

Author: recursix

src/agentlab/experiments/exp_utils.py

@@ -49,8 +49,9 @@ def _episode_timeout(exp_arg: ExpArgs, avg_step_timeout=60):
 def timeout_manager(seconds: int = None):
     """Context manager to handle timeouts."""
 
-    seconds = max(1, int(seconds))  # make sure seconds is at least 1
-
+    if isinstance(seconds, float):
+        seconds = max(1, int(seconds))  # make sure seconds is at least 1
+        
     if seconds is None or sys.platform == "win32":
         try:
             logger.warning("Timeouts are not supported on Windows.")

tests/experiments/test_launch_exp.py

@@ -31,7 +31,7 @@ def _test_launch_system(backend="ray", cause_timeout=False):
         avg_step_timeout = 0.5
     else:
         wait_time = 0
-        avg_step_timeout = 30
+        avg_step_timeout = 10
 
     exp_args_list = []
     for seed in range(3):
@@ -91,7 +91,7 @@ def test_launch_system_ray():
     _test_launch_system(backend="ray")
 
 
-def test_timeout_ray():
+def _test_timeout_ray():
     _test_launch_system(backend="ray", cause_timeout=True)
 
 
@@ -120,7 +120,7 @@ def test_4o_mini_on_miniwob_tiny_test():
 
 
 if __name__ == "__main__":
-    test_timeout_ray()
+    _test_timeout_ray()
     # test_4o_mini_on_miniwob_tiny_test()
     # test_launch_system_ray()
     # test_launch_system_sequntial()