switch to LogMap in CI test

apoorvkh · Aug 17, 2024 · ee8e6e8 · ee8e6e8
1 parent ce1c067
commit ee8e6e8
Showing 1 changed file with 3 additions and 5 deletions.
diff --git a/tests/test_CI.py b/tests/test_CI.py
@@ -5,8 +5,8 @@
 import torch
 import torch.distributed as dist
 
-import torchrunx  # noqa: I001
-from torchrunx.logging_utils import DefaultLogSpec
+import torchrunx
+from torchrunx import LogMap
 
 
 def test_simple_localhost():
@@ -56,9 +56,7 @@ def dist_func():
         func_kwargs={},
         workers_per_host=2,
         backend="gloo",
-        log_spec=DefaultLogSpec.basic(
-            hostnames=["localhost"], workers_per_host=[2], log_dir="./test_logs"
-        ),
+        log_map=LogMap.basic(hostnames=["localhost"], workers_per_host=[2], log_dir="./test_logs"),
     )
 
     log_files = next(os.walk("./test_logs"), (None, None, []))[2]