handle older PyTorch versions

.github/workflows/main.yml

@@ -55,13 +55,17 @@ jobs:
 
   test:
     runs-on: ubuntu-latest
+    strategy: 
+      matrix:
+        python: [py32, py39, py310, py311] 
+        pytorch: [torch20, torch21, torch22, torch23, torch24]
     steps:
       - uses: actions/checkout@v4
       - uses: prefix-dev/[email protected]
         with:
           pixi-version: v0.25.0
-          cache: true
+          cache: false
           frozen: true
-          environments: default
-          activate-environment: default
+          environments: ${{ matrix.python }}-${{ matrix.pytorch }}
+          activate-environment: ${{ matrix.python }}-${{ matrix.pytorch }}
       - run: pytest tests

examples/slurm_poc.py

@@ -21,7 +21,7 @@ def test_launch():
     print("PASS")
 
 
-def simple_matmul(test):
+def simple_matmul():
     rank = int(os.environ["RANK"])
     local_rank = int(os.environ["LOCAL_RANK"])
     device = torch.device(local_rank) if torch.cuda.is_available() else torch.device("cpu")
@@ -35,7 +35,6 @@ def simple_matmul(test):
 
     i = torch.rand((500, 100), device=device)  # batch, dim
     o = torch.matmul(i, w)
-    print(test)
     dist.all_reduce(o, op=dist.ReduceOp.SUM)
     print(i)
     return o.detach().cpu()

src/torchrunx/agent.py

@@ -4,14 +4,14 @@
 import os
 import socket
 import sys
+import tempfile
 from dataclasses import dataclass
 from typing import Callable, Literal
 
 import cloudpickle
 import torch
 import torch.distributed as dist
-from torch.distributed.elastic.multiprocessing import DefaultLogsSpecs
-from torch.distributed.elastic.multiprocessing.api import MultiprocessContext, Std
+from torch.distributed.elastic.multiprocessing import start_processes
 from typing_extensions import Self
 
 from .utils import (
@@ -108,7 +108,7 @@ def main(launcher_agent_group: LauncherAgentGroup):
         port=get_open_port(),
         process_id=os.getpid(),
     )
-
+    # DefaultLogsSpecs(log_dir=None, tee=Std.ALL, local_ranks_filter={0}),
     all_payloads = launcher_agent_group.sync_payloads(payload=payload)
     launcher_payload: LauncherPayload = all_payloads[0]  # pyright: ignore[reportAssignmentType]
     main_agent_payload: AgentPayload = all_payloads[1]  # pyright: ignore[reportAssignmentType]
@@ -119,36 +119,40 @@ def main(launcher_agent_group: LauncherAgentGroup):
     worker_log_files = launcher_payload.worker_log_files[agent_rank]
     num_workers = len(worker_global_ranks)
 
-    # spawn workers
+    if torch.__version__ >= '2.3':
+        # DefaultLogsSpecs only exists in torch >= 2.3
+        from torch.distributed.elastic.multiprocessing import DefaultLogsSpecs
+        log_arg = DefaultLogsSpecs(log_dir=tempfile.mkdtemp())
+    else:
+        log_arg = tempfile.mkdtemp()
 
-    ctx = MultiprocessContext(
-        name=f"{hostname}_",
-        entrypoint=entrypoint,
-        args={
-            i: (
-                WorkerArgs(
-                    function=launcher_payload.fn,
-                    master_hostname=main_agent_payload.hostname,
-                    master_port=main_agent_payload.port,
-                    backend=launcher_payload.backend,
-                    rank=worker_global_ranks[i],
-                    local_rank=i,
-                    local_world_size=num_workers,
-                    world_size=worker_world_size,
-                    log_file=worker_log_files[i],
-                    timeout=launcher_payload.timeout,
-                ).to_bytes(),
+    # spawn workers
+
+    ctx = start_processes(
+            f"{hostname}_",
+            entrypoint,
+            {
+                i: (
+                    WorkerArgs(
+                        function=launcher_payload.fn,
+                        master_hostname=main_agent_payload.hostname,
+                        master_port=main_agent_payload.port,
+                        backend=launcher_payload.backend,
+                        rank=worker_global_ranks[i],
+                        local_rank=i,
+                        local_world_size=num_workers,
+                        world_size=worker_world_size,
+                        log_file=worker_log_files[i],
+                        timeout=launcher_payload.timeout,
+                    ).to_bytes(),
+                )
+                for i in range(num_workers)
+            },
+            {i: {} for i in range(num_workers)},
+            log_arg # type: ignore
             )
-            for i in range(num_workers)
-        },
-        envs={i: {} for i in range(num_workers)},
-        logs_specs=DefaultLogsSpecs(log_dir=None, tee=Std.ALL, local_ranks_filter={0}),
-        start_method="spawn",
-    )
-
+
     try:
-        ctx.start()
-
         status = AgentStatus()
         while True:
             if status.is_running():
@@ -163,7 +167,6 @@ def main(launcher_agent_group: LauncherAgentGroup):
 
             if any(s.is_failed() for s in agent_statuses):
                 raise RuntimeError()
-
     except:
         raise
     finally: