Add docker to ForeachNode

mila-iqia · Jan 17, 2025 · 40c35bd · 40c35bd
1 parent d73af7d
commit 40c35bd
Showing 1 changed file with 7 additions and 3 deletions.
diff --git a/milabench/commands/__init__.py b/milabench/commands/__init__.py
@@ -635,12 +635,17 @@ def executors(self):
                 )
 
             print(rank, node, node_address(node))
+
+            bench_cmd = self.make_new_node_executor(rank, node, self.executor)
+
+            docker_cmd = DockerRunCommand(bench_cmd, DockerConfig(**config["system"].get("docker")))
+
             worker = SSHCommand(
                 host=node_address(node),
                 user=node["user"],
                 key=key,
                 port=node.get("sshport", 22),
-                executor=self.make_new_node_executor(rank, node, self.executor),
+                executor=docker_cmd,
                 **options
             )
             executors.append(worker)
@@ -706,8 +711,7 @@ def make_new_node_executor(self, rank, node, base):
         ]
         executor.wrapper_argv = new_args
 
-        config = executor.pack.config
-        return DockerRunCommand(executor, DockerConfig(**config["system"].get("docker")))
+        return executor
 
     def __init__(self, executor: Command, *args, **kwargs) -> None:
         base_exec = TorchrunAllNodes.make_base_executor(