Auto balance CI tests (#4238)

2025-03-09 21:05:55 -07:00
parent 730d084f2a
commit fbd560028a
5 changed files with 151 additions and 84 deletions
--- a/test/srt/test_custom_allreduce.py
+++ b/test/srt/test_custom_allreduce.py
@@ -42,7 +42,8 @@ def multi_process_parallel(
    # as compared to multiprocessing.
    # NOTE: We need to set working_dir for distributed tests,
    # otherwise we may get import errors on ray workers
-    ray.init(log_to_driver=False)
+
+    ray.init(log_to_driver=True)

    distributed_init_port = get_open_port()
    refs = []