From f079e4516c0b706944c14dafe14d7cc8300309c5 Mon Sep 17 00:00:00 2001 From: fuhailin Date: Tue, 14 Oct 2025 14:46:38 +0800 Subject: [PATCH] [Runner] Multi-node runner --- deepray/__init__.py | 104 +++++++++++++++++++++++++++++++++++++++----- deepray/version.py | 2 +- 2 files changed, 95 insertions(+), 11 deletions(-) diff --git a/deepray/__init__.py b/deepray/__init__.py index 9d19928..7f9ba80 100644 --- a/deepray/__init__.py +++ b/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: 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: 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() diff --git a/deepray/version.py b/deepray/version.py index 446cfc0..82586d5 100644 --- a/deepray/version.py +++ b/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