cleanup and format

Author: pmcurtin

tests/test_CI.py

@@ -1,14 +1,11 @@
 import os
-import sys
 import tempfile
 
 import pytest
 import torch
 import torch.distributed as dist
 
-sys.path.append("../src")
-
-import torchrunx  # noqa: I001
+import torchrunx as trx
 
 
 def test_simple_localhost():
@@ -31,7 +28,7 @@ def dist_func():
 
         return o.detach()
 
-    r = torchrunx.launch(
+    r = trx.launch(
         func=dist_func, func_kwargs={}, workers_per_host=2, backend="gloo", log_dir="./test_logs"
     )
 
@@ -44,9 +41,7 @@ def dist_func():
         print(f"worker rank: {rank}")
 
     tmp = tempfile.mkdtemp()
-    torchrunx.launch(
-        func=dist_func, func_kwargs={}, workers_per_host=2, backend="gloo", log_dir=tmp
-    )
+    trx.launch(func=dist_func, func_kwargs={}, workers_per_host=2, backend="gloo", log_dir=tmp)
 
     log_files = next(os.walk(tmp), (None, None, []))[2]
 
@@ -69,7 +64,7 @@ def error_func():
         raise ValueError("abcdefg")
 
     with pytest.raises(RuntimeError) as excinfo:
-        torchrunx.launch(
+        trx.launch(
             func=error_func,
             func_kwargs={},
             workers_per_host=1,

tests/test_func.py

@@ -3,14 +3,14 @@
 import torch
 import torch.distributed as dist
 
-import torchrunx
+import torchrunx as trx
 
 
 def test_launch():
-    result = torchrunx.launch(
+    result = trx.launch(
         func=simple_matmul,
-        hostnames=torchrunx.slurm_hosts(),
-        workers_per_host=torchrunx.slurm_workers(),
+        hostnames=trx.slurm_hosts(),
+        workers_per_host=trx.slurm_workers(),
     )
 
     t = True
@@ -19,8 +19,6 @@ def test_launch():
 
     assert t, "Not all tensors equal"
 
-    dist.destroy_process_group()
-
 
 def simple_matmul():
     rank = int(os.environ["RANK"])

tests/test_train.py

@@ -1,11 +1,6 @@
 import os
-import sys
 
-sys.path.append("../src")
-
-import torch.distributed as dist
-
-import torchrunx
+import torchrunx as trx
 
 
 def worker():
@@ -38,10 +33,10 @@ def forward(self, x):
 
 
 def test_distributed_train():
-    torchrunx.launch(
+    trx.launch(
         worker,
-        hostnames=torchrunx.slurm_hosts(),
-        workers_per_host=torchrunx.slurm_workers(),
+        hostnames=trx.slurm_hosts(),
+        workers_per_host=trx.slurm_workers(),
         backend="nccl",
     )