From 6e7baf50ecc460bd596d946921b79aba94b09de6 Mon Sep 17 00:00:00 2001 From: Frithjof Gressmann Date: Sat, 29 Jun 2024 20:47:10 -0500 Subject: [PATCH] Support ranks=-1 --- docs/examples/mpi-execution/mpi.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/docs/examples/mpi-execution/mpi.py b/docs/examples/mpi-execution/mpi.py index 5729533b..6436eb77 100644 --- 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