From f711c492e59c90e883b59a2ab69c343be6631418 Mon Sep 17 00:00:00 2001
From: "pierre.delaunay" <delaunap@rtx5.server.mila.quebec>
Date: Tue, 6 Aug 2024 15:04:06 -0400
Subject: [PATCH] update main attr lookup

---
 milabench/remote.py             | 8 ++++----
 scripts/article/run_cuda_dev.sh | 2 ++
 2 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/milabench/remote.py b/milabench/remote.py
index 725c57b45..ce22ed7dd 100644
--- a/milabench/remote.py
+++ b/milabench/remote.py
@@ -70,9 +70,9 @@ def milabench_remote_sync(pack, worker):
 
 def should_run_for(worker, setup_for):
     if setup_for == "worker":
-        return not worker["main"]
+        return not worker.get("main", False)
 
-    return worker["main"]
+    return worker.get("main", False)
 
 
 def worker_commands(pack, worker_plan, setup_for="worker"):
@@ -178,7 +178,7 @@ def is_multinode(pack):
     count = 0
     nodes = pack.config["system"]["nodes"]
     for node in nodes:
-        if not node["main"]:
+        if not node.get("main", False):
             count += 1
     return count > 0
 
@@ -191,7 +191,7 @@ def is_remote(pack):
 def is_main_local(pack):
     """Only the local main can send remote commands to remote"""
     self = pack.config["system"]["self"]
-    return self is not None and self["local"] and self["main"]
+    return self is not None and self["local"] and self.get("main", False)
 
 
 def is_worker(pack):
diff --git a/scripts/article/run_cuda_dev.sh b/scripts/article/run_cuda_dev.sh
index 37729f4b7..0f52a1724 100644
--- a/scripts/article/run_cuda_dev.sh
+++ b/scripts/article/run_cuda_dev.sh
@@ -99,6 +99,8 @@ if [ "$MILABENCH_PREPARE" -eq 0 ]; then
     cd $MILABENCH_WORDIR
     
 
+    milabench prepare --system $MILABENCH_WORDIR/system.yaml "$@"
+    
     # milabench prepare "$@"
     #
     #   Run the benchmakrs