chore: import upstream snapshot with attribution

python/ray/scripts/symmetric_run.py

@@ -0,0 +1,274 @@
+"""Symmetric Run for Ray."""
+
+import socket
+import subprocess
+import sys
+import time
+from typing import List
+
+import click
+
+import ray
+from ray._common.network_utils import is_localhost
+from ray._private.ray_constants import env_integer
+from ray._raylet import GcsClient
+from ray.exceptions import RpcError
+
+import psutil
+
+CLUSTER_WAIT_TIMEOUT = env_integer("RAY_SYMMETRIC_RUN_CLUSTER_WAIT_TIMEOUT", 30)
+
+
+def check_ray_already_started() -> bool:
+    import ray._private.services as services
+
+    # Try auto-detecting the Ray instance.
+    running_gcs_addresses = services.find_gcs_addresses()
+    return len(running_gcs_addresses) > 0
+
+
+def check_cluster_ready(nnodes, timeout=CLUSTER_WAIT_TIMEOUT):
+    """Wait for all nodes to start.
+
+    Raises an exception if the nodes don't start in time.
+    """
+    start_time = time.time()
+    current_nodes = 1
+    ray.init(ignore_reinit_error=True)
+
+    while time.time() - start_time < timeout:
+        time.sleep(5)
+        current_nodes = len(ray.nodes())
+        if current_nodes == nnodes:
+            return True
+        else:
+            click.echo(
+                f"Waiting for nodes to start... {current_nodes}/{nnodes} nodes started"
+            )
+    return False
+
+
+def check_head_node_ready(address: str, timeout=CLUSTER_WAIT_TIMEOUT):
+    start_time = time.time()
+    gcs_client = GcsClient(address=address)
+    while time.time() - start_time < timeout:
+        try:
+            gcs_client.check_alive([], timeout=1)
+            click.echo("Ray cluster is ready!")
+            return True
+        except RpcError:
+            pass
+        time.sleep(5)
+    return False
+
+
+def curate_and_validate_ray_start_args(run_and_start_args: List[str]) -> List[str]:
+    # Reparse the arguments to remove symmetric_run arguments.
+    ctx = symmetric_run.make_context("_", run_and_start_args, resilient_parsing=True)
+    cleaned_args = list(ctx.params["ray_args_and_entrypoint"])
+
+    for arg in cleaned_args:
+        if arg == "--head":
+            raise click.ClickException("Cannot use --head option in symmetric_run.")
+        if arg == "--node-ip-address":
+            raise click.ClickException(
+                "Cannot use --node-ip-address option in symmetric_run."
+            )
+        if arg == "--port":
+            raise click.ClickException("Cannot use --port option in symmetric_run.")
+        if arg == "--block":
+            raise click.ClickException("Cannot use --block option in symmetric_run.")
+
+    return cleaned_args
+
+
+@click.command(
+    name="symmetric_run",
+    context_settings={"ignore_unknown_options": True, "allow_extra_args": True},
+    help="""Command to start Ray across all nodes and execute an entrypoint command.
+
+USAGE:
+
+    ray symmetric-run --address ADDRESS
+[--min-nodes NUM_NODES] [RAY_START_OPTIONS] -- [ENTRYPOINT_COMMAND]
+
+DESCRIPTION:
+
+    This command (1) starts a Ray cluster across all nodes,
+(2) runs a command on the head node, and (3) stops the Ray cluster.
+
+    The '--' separator is required to distinguish between Ray start arguments
+and the entrypoint command. The --min-nodes option is optional and
+can be used to wait for a specific number of nodes to start.
+
+EXAMPLES:
+
+    # Start Ray with default settings and run a Python script
+
+    ray symmetric-run --address 127.0.0.1:6379 -- python my_script.py
+
+    # Start Ray with specific head node and run a command
+
+    ray symmetric-run --address 127.0.0.1:6379 --min-nodes 4 -- python train_model.py --epochs=100
+
+    # Start Ray and run a multi-word command
+
+    ray symmetric-run --address 127.0.0.1:6379 --min-nodes 4 --num-cpus=4 -- python -m my_module --config=prod
+
+RAY START OPTIONS:
+
+    Most ray start command options are supported. Arguments that are not
+supported are: --head, --node-ip-address, --port, --block.
+
+SEPARATOR REQUIREMENT:
+
+    The '--' separator is mandatory and must appear between Ray start
+    arguments and the entrypoint command. This ensures clear separation
+    between the two sets of arguments.
+""",
+)
+@click.option(
+    "--address", required=True, type=str, help="The address of the Ray cluster."
+)
+@click.option(
+    "--min-nodes",
+    type=int,
+    help="If provided, wait for this number of nodes to start.",
+)
+@click.argument("ray_args_and_entrypoint", nargs=-1, type=click.UNPROCESSED)
+def symmetric_run(address, min_nodes, ray_args_and_entrypoint):
+    all_args = sys.argv[1:]
+
+    if all_args and all_args[0] == "symmetric-run":
+        all_args = all_args[1:]
+
+    try:
+        separator = all_args.index("--")
+    except ValueError:
+        raise click.ClickException(
+            "No separator '--' found in arguments. Please use '--' to "
+            "separate Ray start arguments and the entrypoint command."
+        )
+
+    run_and_start_args, entrypoint_on_head = (
+        all_args[:separator],
+        all_args[separator + 1 :],
+    )
+
+    ray_start_args = curate_and_validate_ray_start_args(run_and_start_args)
+
+    min_nodes = 1 if min_nodes is None else min_nodes
+
+    if not entrypoint_on_head:
+        raise click.ClickException("No entrypoint command provided.")
+
+    if check_ray_already_started():
+        raise click.ClickException("Ray is already started on this node.")
+
+    # 1. Parse address and check if we are on the head node.
+    gcs_host_port = ray._common.network_utils.parse_address(address)
+    if gcs_host_port is None:
+        raise click.ClickException(
+            f"Invalid address format: {address}, should be `host:port`"
+        )
+    gcs_host, gcs_port = gcs_host_port
+
+    try:
+        # AF_UNSPEC allows resolving both IPv4 and IPv6
+        addrinfo = socket.getaddrinfo(
+            gcs_host, gcs_port, socket.AF_UNSPEC, socket.SOCK_STREAM
+        )
+        resolved_gcs_host = addrinfo[0][4][0]
+    except socket.gaierror:
+        raise click.ClickException(f"Could not resolve hostname: {gcs_host}")
+
+    my_ips = []
+    for iface, addrs in psutil.net_if_addrs().items():
+        for addr in addrs:
+            # Look for AF_INET (IPv4) or AF_INET6 (IPv6)
+            if addr.family in [
+                socket.AddressFamily.AF_INET,
+                socket.AddressFamily.AF_INET6,
+            ]:
+                my_ips.append(addr.address)
+
+    if min_nodes > 1:
+        # Ban localhost ips if we are not running on a single node
+        # to avoid starting N head nodes
+        my_ips = [ip for ip in my_ips if not is_localhost(ip)]
+
+    is_head = resolved_gcs_host in my_ips
+
+    result = None
+    # 2. Start Ray and run commands.
+    try:
+        if is_head:
+            # On the head node, start Ray, run the command, then stop Ray.
+            click.echo("On head node. Starting Ray cluster head...")
+
+            # Build the ray start command with all parameters
+            ray_start_cmd = [
+                "ray",
+                "start",
+                "--head",
+                f"--node-ip-address={resolved_gcs_host}",
+                f"--port={gcs_port}",
+                *ray_start_args,
+            ]
+
+            # Start Ray head. This runs in the background and hides output.
+            subprocess.run(ray_start_cmd, check=True, capture_output=True)
+            click.echo("Head node started.")
+            click.echo("=======================")
+            if min_nodes > 1 and not check_cluster_ready(min_nodes):
+                raise click.ClickException(
+                    "Timed out waiting for other nodes to start."
+                )
+
+            click.echo(
+                f"Running command on head node: {entrypoint_on_head}",
+            )
+            click.echo("=======================")
+            result = subprocess.run(entrypoint_on_head)
+            click.echo("=======================")
+        else:
+            # On a worker node, start Ray and connect to the head.
+            click.echo(f"On worker node. Connecting to Ray cluster at {address}...")
+
+            if not check_head_node_ready(address):
+                raise click.ClickException("Timed out waiting for head node to start.")
+
+            # Build the ray start command for worker nodes with all parameters
+            ray_start_cmd = [
+                "ray",
+                "start",
+                "--address",
+                address,
+                "--block",
+                *ray_start_args,
+            ]
+
+            # This command will block until the Ray cluster is stopped.
+            subprocess.run(ray_start_cmd, check=True)
+
+    except subprocess.CalledProcessError as e:
+        click.echo(f"Failed to start Ray: {e}", err=True)
+        if e.stdout:
+            click.echo(f"stdout:\n{e.stdout.decode()}", err=True)
+        if e.stderr:
+            click.echo(f"stderr:\n{e.stderr.decode()}", err=True)
+    except KeyboardInterrupt:
+        # This can be triggered by ctrl-c on the user's side.
+        click.echo("Interrupted by user.", err=True)
+    finally:
+        # Stop Ray cluster.
+        subprocess.run(["ray", "stop"])
+
+        # Propagate the exit code of the user script.
+        if result is not None and result.returncode != 0:
+            click.echo(f"Command failed with return code {result.returncode}", err=True)
+            sys.exit(result.returncode)
+
+
+if __name__ == "__main__":
+    symmetric_run()