Support ranks=-1

machinable-org · Jun 30, 2024 · 6e7baf5 · 6e7baf5
1 parent eafc91c
commit 6e7baf5
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/docs/examples/mpi-execution/mpi.py b/docs/examples/mpi-execution/mpi.py
@@ -1,5 +1,6 @@
 from typing import Literal, Optional, Union
 
+import os
 import sys
 
 from machinable import Execution
@@ -23,7 +24,10 @@ def on_compute_default_resources(self, executable):
 
         ranks = executable.config.get("ranks", False)
         if ranks not in [None, False]:
-            resources["-n"] = int(ranks)
+            if ranks == -1:
+                ranks = os.environ.get("MPI_RANKS", 0)
+            if int(ranks) > 0:
+                resources["-n"] = int(ranks)
 
         return resources