[cli] install nightly wheels via ray install-nightly (#10054)

python/ray/scripts/scripts.py

@@ -9,6 +9,7 @@ import sys
 import time
 import urllib
 import urllib.parse
+import yaml
 
 import ray
 import psutil
@@ -1487,6 +1488,60 @@ def globalgc(address):
     print("Triggered gc.collect() on all workers.")
 
 
+@cli.command()
+@click.option("-v", "--verbose", is_flag=True)
+@click.option(
+    "--dryrun",
+    is_flag=True,
+    help="Identifies the wheel but does not execute the installation.")
+def install_nightly(verbose, dryrun):
+    """Install the latest wheels for Ray.
+
+    This uses the same python environment as the one that Ray is currently
+    installed in. Make sure that there is no Ray processes on this
+    machine (ray stop) when running this command.
+    """
+    raydir = os.path.abspath(os.path.dirname(ray.__file__))
+    all_wheels_path = os.path.join(raydir, "nightly-wheels.yaml")
+
+    wheels = None
+    if os.path.exists(all_wheels_path):
+        with open(all_wheels_path) as f:
+            wheels = yaml.safe_load(f)
+
+    if not wheels:
+        raise click.ClickException(
+            f"Wheels not found in '{all_wheels_path}'! "
+            "Please visit https://docs.ray.io/en/master/installation.html to "
+            "obtain the latest wheels.")
+
+    platform = sys.platform
+    py_version = "{0}.{1}".format(*sys.version_info[:2])
+
+    matching_wheel = None
+    for target_platform, wheel_map in wheels.items():
+        if verbose:
+            print(f"Evaluating os={target_platform}, python={list(wheel_map)}")
+        if platform.startswith(target_platform):
+            if py_version in wheel_map:
+                matching_wheel = wheel_map[py_version]
+                break
+        if verbose:
+            print("Not matched.")
+
+    if matching_wheel is None:
+        raise click.ClickException(
+            "Unable to identify a matching platform. "
+            "Please visit https://docs.ray.io/en/master/installation.html to "
+            "obtain the latest wheels.")
+    if dryrun:
+        print(f"Found wheel: {matching_wheel}")
+    else:
+        cmd = [sys.executable, "-m", "pip", "install", "-U", matching_wheel]
+        print(f"Running: {' '.join(cmd)}.")
+        subprocess.check_call(cmd)
+
+
 def add_command_alias(command, name, hidden):
     new_command = copy.deepcopy(command)
     new_command.hidden = hidden
@@ -1518,6 +1573,7 @@ cli.add_command(globalgc)
 cli.add_command(timeline)
 cli.add_command(project_cli)
 cli.add_command(session_cli)
+cli.add_command(install_nightly)
 
 try:
     from ray.serve.scripts import serve_cli