add mpirun mode

flagscale/auto_tuner/tuner.py

@@ -64,7 +64,9 @@ def __init__(self, config: DictConfig):
         # Set platform envs
         if "platform" not in self.config.auto_tuner:
             self.config.auto_tuner.platform = {}
-        if os.environ.get("AIRS_SWITCH", None) != "False":
+
+        # As long as AIRS_SWITCH has value it means running on the platform
+        if os.environ.get("AIRS_SWITCH", None):
             self.config.auto_tuner.platform.airs_switch = True
 
             if os.environ.get("AIRS_SIZE", None):

flagscale/auto_tuner/utils.py

@@ -1,3 +1,7 @@
+import re
+import socket
+
+
 def divisible(x, y):
     if x % y == 0:
         return True
@@ -25,6 +29,7 @@ def beside(keys, strategy, history):
 
 
 def sort_by_memory(strategy):
+    """Sort strategy by memory."""
     return (
         -strategy["use_recompute"],
         -strategy["tensor_model_parallel_size"],
@@ -45,6 +50,7 @@ def sort_by_memory(strategy):
 
 
 def sort_by_performance(strategy):
+    """Sort strategy by performance potentially."""
     return (
         strategy["use_recompute"],
         -strategy["tensor_model_parallel_size"],
@@ -62,3 +68,62 @@ def sort_by_performance(strategy):
             else float("inf")
         ),
     )
+
+
+def is_ip_addr(master):
+    """Check if master is ip address."""
+
+    if not isinstance(master, str):
+        return False
+    pattern = r"^((25[0-5]|2[0-4]\d|1\d{2}|[1-9]?\d)\.){3}(25[0-5]|2[0-4]\d|1\d{2}|[1-9]?\d)$"
+    result = re.match(pattern, master)
+    if result:
+        return True
+    else:
+        return False
+
+def get_ip_addr():
+    """Get ip address."""
+    try:
+        hostname = socket.gethostname()
+        ip = socket.gethostbyname(socket.getfqdn(hostname))
+    except:
+        ip = '127.0.0.1'
+    return ip
+
+
+def is_master(config):
+
+    # nnodes > 1 & ip or hostname in hostfile
+    # nnodes > 1:
+    from flagscale.launcher.runner import parse_hostfile
+    multi_nodes = False
+    if config.experiment.runner.get("nnodes", 1) > 1:
+        multi_nodes = True
+
+    hostfile = None
+    if config.experiment.runner.get("hostfile", None):
+        hostfile = config.experiment.runner["hostfile"]
+    if os.environ.get("AIRS_SWITCH", None):
+        if os.environ.get("AIRS_HOSTFILE_PATH", None):
+            hostfile = os.environ["AIRS_HOSTFILE_PATH"]
+    resources = parse_hostfile(hostfile)
+
+    if not resources and multi_nodes:
+        raise ValueError("In the multi-node mode, please set the hostfile")
+
+    if not resources and not multi_nodes:
+        return True
+
+    if resources and multi_nodes:
+        master = resources.keys()[0]
+        if is_ip_addr(master):
+            return get_ip_addr() == master
+        else:
+            output = subprocess.run("hostname",
+                                    check=True,
+                                    shell=True,
+                                    text=True,
+                                    capture_output=True)
+            hostname = output.stdout.strip()
+            return hostname == master

run.py

@@ -9,9 +9,12 @@
 def main(config: DictConfig) -> None:
     if config.action == "auto_tune":
         from flagscale.auto_tuner import AutoTuner
-
-        tuner = AutoTuner(config)
-        tuner.tune()
+        # For MPIRUN scene, just one autotuner process.
+        # NOTE: This is a temporary solution and will be updated with cloud runner.
+        from flagscale.auto_tuner.utils import is_master
+        if is_master(config):
+            tuner = AutoTuner(config)
+            tuner.tune()
     else:
         if config.experiment.runner.get("type", "ssh") == "ssh":
             runner = SSHRunner(config)