[Runner] Multi-node runner

deepray/__init__.py

@@ -41,13 +41,12 @@
 # _check_tf_version()
 
 logger = logging_util.get_logger()
-
 common_flags.define_common_flags()
-flags.FLAGS(sys.argv, known_only=True)
 
 
 def init():
   logger.debug(f"sys.argv = {sys.argv}")  # sys.argv from Horovod
+  flags.FLAGS(sys.argv, known_only=True)
 
   gpus = tf.config.experimental.list_physical_devices("GPU")
   for gpu in gpus:
@@ -79,24 +78,99 @@ def start_tensorflow_server(cluster_resolver):
   server.join()
 
 
+def get_num_proc(hosts, hostfile):
+  if hosts is not None and hostfile is not None:
+    raise ValueError("Argument hosts and hostfile only allow one provided.")
+  if hosts:
+    # Parse hosts parameter format: "host1:2,host2:4,host3:1"
+    total_slots = 0
+    host_list = hosts.split(",")
+    for host_entry in host_list:
+      if ":" not in host_entry:
+        raise ValueError(f"Invalid host format: '{host_entry}'. Expected 'hostname:slots'")
+
+      hostname, slots_str = host_entry.split(":", 1)
+      if not slots_str.isdigit():
+        raise ValueError(f"Invalid slots value: '{slots_str}'. Must be integer")
+
+      slots = int(slots_str)
+      if slots <= 0:
+        raise ValueError(f"Slots must be positive: {slots}")
+      total_slots += slots
+    num_proc = total_slots
+  elif hostfile:
+    # Parse hostfile with format: "hostname slots=X" per line
+    total_slots = 0
+    with open(hostfile, "r") as f:
+      for line_num, line in enumerate(f, 1):
+        line = line.strip()
+        # Skip empty lines and comments
+        if not line or line.startswith("#"):
+          continue
+
+        if "slots=" not in line:
+          raise ValueError(f"Invalid hostfile format at line {line_num}: '{line}'. Expected 'slots=X'")
+
+        # Parse "slots=X" format
+        parts = line.split()
+        slots_found = False
+        for part in parts:
+          if part.startswith("slots="):
+            slots_str = part.split("=")[1]
+            if not slots_str.isdigit():
+              raise ValueError(f"Invalid slots value at line {line_num}: '{slots_str}'. Must be integer")
+
+            slots = int(slots_str)
+            if slots <= 0:
+              raise ValueError(f"Slots must be positive at line {line_num}: {slots}")
+
+            total_slots += slots
+            slots_found = True
+            break
+
+        if not slots_found:
+          raise ValueError(f"No valid slots entry found at line {line_num}: '{line}'")
+    num_proc = total_slots
+  else:
+    # Use local GPU devices
+    physical_devices = tf.config.list_physical_devices("GPU")
+    num_proc = len(physical_devices)
+
+  logger.debug(f"world_size = {num_proc}")
+  return num_proc
+
+
 def runner(function, verbose=None):
   parser = argparse.ArgumentParser(description="Deepray Runner")
   parser.add_argument("-v", "--version", action="version", version=__version__, help="Shows Deepray version.")
   parser.add_argument(
     "--distribution_strategy", type=str, default="Horovod", help="Whether run distributed training with Horovod."
   )
-
-  physical_devices = tf.config.list_physical_devices("GPU")
-  world_size = len(physical_devices)
-  logger.debug(f"world_size = {world_size}")
+  parser.add_argument(
+    "--hosts",
+    type=str,
+    default=None,
+    help="Path to a host file containing the list of host names and the number of available slots. \
+      Each line of the file must be of the form: <hostname> slots=<slots>",
+  )
+  parser.add_argument(
+    "--hostfile",
+    type=str,
+    default=None,
+    help="Path to a host file containing the list of host names and the number of available slots. \
+      Each line of the file must be of the form: <hostname> slots=<slots>",
+  )
+  parser.add_argument("--ssh_port", type=int, default=None, help="SSH port on all the hosts.")
 
   user_argv = sys.argv  # get user specified args
   args, unknown = parser.parse_known_args()
 
-  if world_size > 1 and args.distribution_strategy == "Horovod":
+  num_proc = get_num_proc(args.hosts, args.hostfile)
+
+  if num_proc > 1 and args.distribution_strategy == "Horovod":
     user_argv.extend([
       "--distribution_strategy=horovod",
-      f"--num_gpus={world_size}",
+      f"--num_gpus={num_proc}",
       "--use_horovod",
     ])
     try:
@@ -114,7 +188,17 @@ def helper(argv, main):
       init()
       main()
 
-    horovod.run(helper, args=(sys.argv,), kwargs={"main": function}, np=world_size, verbose=verbose, use_mpi=True)
+    horovod.run(
+      helper,
+      args=(sys.argv,),
+      kwargs={"main": function},
+      hosts=args.hosts,
+      hostfile=args.hostfile,
+      ssh_port=args.ssh_port,
+      num_proc=num_proc,
+      verbose=verbose,
+      use_mpi=True,
+    )
   elif args.distribution_strategy == "ParameterServer":
     cluster_resolver = tf.distribute.cluster_resolver.TFConfigClusterResolver()
     if cluster_resolver.task_type in ("worker", "ps"):
@@ -125,6 +209,6 @@ def helper(argv, main):
       function()
   else:
     logger.info("Deepray finds only one GPU available, so we turn off distribution_strategy.")
-    user_argv.extend(["--distribution_strategy=off", f"--num_gpus={world_size}"])
+    user_argv.extend(["--distribution_strategy=off", f"--num_gpus={num_proc}"])
     init()
     function()

deepray/version.py

@@ -21,7 +21,7 @@
 # We follow Semantic Versioning (https://semver.org/)
 _MAJOR_VERSION = "0"
 _MINOR_VERSION = "21"
-_PATCH_VERSION = "95"
+_PATCH_VERSION = "96"
 
 # When building releases, we can update this value on the release branch to
 # reflect the current release candidate ('rc0', 'rc1') or, finally, the official