[Core] Automatically inject ray and python version in runtime_env (#15958)

.buildkite/pipeline.yml

@@ -29,6 +29,7 @@
       --test_env=CONDA_PREFIX
       --test_env=CONDA_DEFAULT_ENV
       --test_env=CI
+      --test_env=RAY_CI_POST_WHEEL_TESTS=True
       python/ray/tests/... python/ray/serve/... python/ray/tune/... rllib/...
 
 - label: ":docker: Build Images"

python/ray/test_utils.py

@@ -544,19 +544,52 @@ def set_setup_func():
     runtime_env.VAR = "hello world"
 
 
-def get_wheel_filename() -> str:
-    """Returns the filename used for the Ray wheel of the current build."""
-    ray_version = ray.__version__
-    python_version = f"{sys.version_info.major}{sys.version_info.minor}"
+def get_wheel_filename(
+        sys_platform: str = sys.platform,
+        ray_version: str = ray.__version__,
+        py_version: str = f"{sys.version_info.major}{sys.version_info.minor}"
+) -> str:
+    """Returns the filename used for the nightly Ray wheel.
+
+    Args:
+        sys_platform (str): The platform as returned by sys.platform. Examples:
+            "darwin", "linux", "win32"
+        ray_version (str): The Ray version as returned by ray.__version__ or
+            `ray --version`.  Examples: "2.0.0.dev0"
+        py_version (str):
+            The major and minor Python versions concatenated.  Examples: "36",
+            "37", "38"
+    Returns:
+        The wheel file name.  Examples:
+            ray-2.0.0.dev0-cp38-cp38-manylinux2014_x86_64.whl
+    """
+    assert py_version in ["36", "37", "38"], ("py_version must be one of '36',"
+                                              " '37', or '38'")
+
     os_strings = {
         "darwin": "macosx_10_13_x86_64"
-        if python_version == "38" else "macosx_10_13_intel",
+        if py_version == "38" else "macosx_10_13_intel",
         "linux": "manylinux2014_x86_64",
         "win32": "win_amd64"
     }
 
-    wheel_filename = (
-        f"ray-{ray_version}-cp{python_version}-"
-        f"cp{python_version}{'m' if python_version != '38' else ''}"
-        f"-{os_strings[sys.platform]}.whl")
+    assert sys_platform in os_strings, ("sys_platform must be one of 'darwin',"
+                                        " 'linux', or 'win32'")
+
+    wheel_filename = (f"ray-{ray_version}-cp{py_version}-"
+                      f"cp{py_version}{'m' if py_version != '38' else ''}"
+                      f"-{os_strings[sys_platform]}.whl")
+
     return wheel_filename
+
+
+def get_master_wheel_url(
+        ray_commit: str = ray.__commit__,
+        sys_platform: str = sys.platform,
+        ray_version: str = ray.__version__,
+        py_version: str = f"{sys.version_info.major}{sys.version_info.minor}"
+) -> str:
+    """Return the URL for the wheel from a specific commit."""
+
+    return (f"https://s3-us-west-2.amazonaws.com/ray-wheels/master/"
+            f"{ray_commit}/{get_wheel_filename()}")

python/ray/tests/test_runtime_env.py

@@ -4,10 +4,12 @@ import sys
 import unittest
 import random
 import tempfile
+import requests
 from pathlib import Path
 import ray
 from ray.test_utils import (run_string_as_driver,
-                            run_string_as_driver_nonblocking)
+                            run_string_as_driver_nonblocking,
+                            get_wheel_filename, get_master_wheel_url)
 import ray.experimental.internal_kv as kv
 from time import sleep
 driver_script = """
@@ -655,6 +657,27 @@ def test_init(shutdown_only):
         os.chdir(old_dir)
 
 
+def test_get_wheel_filename():
+    ray_version = "2.0.0.dev0"
+    for sys_platform in ["darwin", "linux", "win32"]:
+        for py_version in ["36", "37", "38"]:
+            filename = get_wheel_filename(sys_platform, ray_version,
+                                          py_version)
+            prefix = "https://s3-us-west-2.amazonaws.com/ray-wheels/latest/"
+            url = f"{prefix}{filename}"
+            assert requests.head(url).status_code == 200
+
+
+def test_get_master_wheel_url():
+    ray_version = "2.0.0.dev0"
+    test_commit = "ba6cebe30fab6925e5b2d9e859ad064d53015246"
+    for sys_platform in ["darwin", "linux", "win32"]:
+        for py_version in ["36", "37", "38"]:
+            url = get_master_wheel_url(test_commit, sys_platform, ray_version,
+                                       py_version)
+            assert requests.head(url).status_code == 200
+
+
 if __name__ == "__main__":
     import sys
     sys.exit(pytest.main(["-sv", __file__]))

python/ray/tests/test_runtime_env_complicated.py

@@ -1,4 +1,5 @@
 import os
+from ray.workers.setup_runtime_env import inject_ray_and_python
 import pytest
 import sys
 import unittest
@@ -9,7 +10,7 @@ from unittest import mock
 import ray
 from ray._private.utils import get_conda_env_dir, get_conda_bin_executable
 from ray.job_config import JobConfig
-from ray.test_utils import get_wheel_filename, run_string_as_driver
+from ray.test_utils import run_string_as_driver
 
 
 @pytest.fixture(scope="session")
@@ -258,17 +259,12 @@ in the runtime env.  Buildkite only supports Linux for now.
 def test_conda_create_task(shutdown_only):
     """Tests dynamic creation of a conda env in a task's runtime env."""
     ray.init()
-    ray_wheel_filename = get_wheel_filename()
-    # E.g. 3.6.13
-    python_micro_version_dots = ".".join(map(str, sys.version_info[:3]))
-    ray_wheel_path = os.path.join("/ray/.whl", ray_wheel_filename)
     runtime_env = {
         "conda": {
             "dependencies": [
-                f"python={python_micro_version_dots}", "pip", {
+                "pip", {
                     "pip": [
-                        ray_wheel_path, "pip-install-test==0.5",
-                        "opentelemetry-api==1.0.0rc1",
+                        "pip-install-test==0.5", "opentelemetry-api==1.0.0rc1",
                         "opentelemetry-sdk==1.0.0rc1"
                     ]
                 }
@@ -298,18 +294,12 @@ def test_conda_create_task(shutdown_only):
 def test_conda_create_job_config(shutdown_only):
     """Tests dynamic conda env creation in a runtime env in the JobConfig."""
 
-    ray_wheel_filename = get_wheel_filename()
-    # E.g. 3.6.13
-    python_micro_version_dots = ".".join(map(str, sys.version_info[:3]))
-    ray_wheel_path = os.path.join("/ray/.whl", ray_wheel_filename)
-
     runtime_env = {
         "conda": {
             "dependencies": [
-                f"python={python_micro_version_dots}", "pip", {
+                "pip", {
                     "pip": [
-                        ray_wheel_path, "pip-install-test==0.5",
-                        "opentelemetry-api==1.0.0rc1",
+                        "pip-install-test==0.5", "opentelemetry-api==1.0.0rc1",
                         "opentelemetry-sdk==1.0.0rc1"
                     ]
                 }
@@ -329,6 +319,50 @@ def test_conda_create_job_config(shutdown_only):
     assert ray.get(f.remote())
 
 
+def test_inject_ray_and_python():
+    num_tests = 4
+    conda_dicts = [None] * num_tests
+    outputs = [None] * num_tests
+
+    conda_dicts[0] = {}
+    outputs[0] = {
+        "dependencies": ["python=7.8", "pip", {
+            "pip": ["ray==1.2.3"]
+        }]
+    }
+
+    conda_dicts[1] = {"dependencies": ["blah"]}
+    outputs[1] = {
+        "dependencies": ["blah", "python=7.8", "pip", {
+            "pip": ["ray==1.2.3"]
+        }]
+    }
+
+    conda_dicts[2] = {"dependencies": ["blah", "pip"]}
+    outputs[2] = {
+        "dependencies": ["blah", "pip", "python=7.8", {
+            "pip": ["ray==1.2.3"]
+        }]
+    }
+
+    conda_dicts[3] = {"dependencies": ["blah", "pip", {"pip": ["some_pkg"]}]}
+    outputs[3] = {
+        "dependencies": [
+            "blah", "pip", {
+                "pip": ["some_pkg", "ray==1.2.3"]
+            }, "python=7.8"
+        ]
+    }
+
+    for i in range(num_tests):
+        output = inject_ray_and_python(conda_dicts[i], "ray==1.2.3", "7.8")
+        error_msg = (f"failed on input {i}."
+                     f"Input: {conda_dicts[i]} \n"
+                     f"Output: {output} \n"
+                     f"Expected output: {outputs[i]}")
+        assert (output == outputs[i]), error_msg
+
+
 @pytest.mark.skipif(
     os.environ.get("CI") is None, reason="This test is only run on CI.")
 @pytest.mark.skipif(
@@ -340,18 +374,12 @@ def test_conda_create_job_config(shutdown_only):
 def test_conda_create_ray_client(call_ray_start):
     """Tests dynamic conda env creation in RayClient."""
 
-    ray_wheel_filename = get_wheel_filename()
-    # E.g. 3.6.13
-    python_micro_version_dots = ".".join(map(str, sys.version_info[:3]))
-    ray_wheel_path = os.path.join("/ray/.whl", ray_wheel_filename)
-
     runtime_env = {
         "conda": {
             "dependencies": [
-                f"python={python_micro_version_dots}", "pip", {
+                "pip", {
                     "pip": [
-                        ray_wheel_path, "pip-install-test==0.5",
-                        "opentelemetry-api==1.0.0rc1",
+                        "pip-install-test==0.5", "opentelemetry-api==1.0.0rc1",
                         "opentelemetry-sdk==1.0.0rc1"
                     ]
                 }
@@ -392,10 +420,8 @@ def test_pip_task(shutdown_only, pip_as_str):
     """Tests pip installs in the runtime env specified in the job config."""
 
     ray.init()
-    ray_wheel_path = os.path.join("/ray/.whl", get_wheel_filename())
     if pip_as_str:
-        requirements_txt = f"""
-        {ray_wheel_path}
+        requirements_txt = """
         pip-install-test==0.5
         opentelemetry-api==1.0.0rc1
         opentelemetry-sdk==1.0.0rc1
@@ -404,8 +430,8 @@ def test_pip_task(shutdown_only, pip_as_str):
     else:
         runtime_env = {
             "pip": [
-                ray_wheel_path, "pip-install-test==0.5",
-                "opentelemetry-api==1.0.0rc1", "opentelemetry-sdk==1.0.0rc1"
+                "pip-install-test==0.5", "opentelemetry-api==1.0.0rc1",
+                "opentelemetry-sdk==1.0.0rc1"
             ]
         }
 
@@ -432,10 +458,8 @@ def test_pip_task(shutdown_only, pip_as_str):
 def test_pip_job_config(shutdown_only, pip_as_str):
     """Tests dynamic installation of pip packages in a task's runtime env."""
 
-    ray_wheel_path = os.path.join("/ray/.whl", get_wheel_filename())
     if pip_as_str:
-        requirements_txt = f"""
-        {ray_wheel_path}
+        requirements_txt = """
         pip-install-test==0.5
         opentelemetry-api==1.0.0rc1
         opentelemetry-sdk==1.0.0rc1
@@ -444,8 +468,8 @@ def test_pip_job_config(shutdown_only, pip_as_str):
     else:
         runtime_env = {
             "pip": [
-                ray_wheel_path, "pip-install-test==0.5",
-                "opentelemetry-api==1.0.0rc1", "opentelemetry-sdk==1.0.0rc1"
+                "pip-install-test==0.5", "opentelemetry-api==1.0.0rc1",
+                "opentelemetry-sdk==1.0.0rc1"
             ]
         }
 

python/ray/workers/setup_runtime_env.py

@@ -7,10 +7,14 @@ import yaml
 import hashlib
 
 from filelock import FileLock
+from typing import Optional
+from pathlib import Path
 
+import ray
 from ray._private.conda import (get_conda_activate_commands,
                                 get_or_create_conda_env)
 from ray._private.utils import try_to_create_directory
+from ray.test_utils import get_wheel_filename, get_master_wheel_url
 logger = logging.getLogger(__name__)
 
 parser = argparse.ArgumentParser()
@@ -40,6 +44,11 @@ def setup(input_args):
         if isinstance(runtime_env["conda"], str):
             commands += get_conda_activate_commands(runtime_env["conda"])
         elif isinstance(runtime_env["conda"], dict):
+            py_version = ".".join(map(str,
+                                      sys.version_info[:3]))  # like 3.6.10
+            conda_dict = inject_ray_and_python(runtime_env["conda"],
+                                               current_ray_pip_specifier(),
+                                               py_version)
             # Locking to avoid multiple processes installing concurrently
             conda_hash = hashlib.sha1(
                 json.dumps(runtime_env["conda"],
@@ -55,7 +64,7 @@ def setup(input_args):
                     # Sort keys because we hash based on the file contents,
                     # and we don't want the hash to depend on the order
                     # of the dependencies.
-                    yaml.dump(runtime_env["conda"], file, sort_keys=True)
+                    yaml.dump(conda_dict, file, sort_keys=True)
                 conda_env_name = get_or_create_conda_env(
                     conda_yaml_path, conda_dir)
             commands += get_conda_activate_commands(conda_env_name)
@@ -74,13 +83,15 @@ def setup(input_args):
         py_version = ".".join(map(str, sys.version_info[:3]))  # E.g. 3.6.13
         conda_dict = {
             "name": pip_hash_str,
-            "dependencies": [
-                f"python={py_version}", "pip", {
-                    "pip": [f"-r {requirements_txt_path}"]
-                }
-            ]
+            "dependencies": ["pip", {
+                "pip": [f"-r {requirements_txt_path}"]
+            }]
         }
 
+        conda_dict = inject_ray_and_python(conda_dict,
+                                           current_ray_pip_specifier(),
+                                           py_version)
+
         file_lock_name = f"ray-{pip_hash_str}.lock"
         with FileLock(os.path.join(args.session_dir, file_lock_name)):
             try_to_create_directory(conda_dir)
@@ -101,3 +112,73 @@ def setup(input_args):
     command_separator = " && "
     command_str = command_separator.join(commands)
     os.execvp("bash", ["bash", "-c", command_str])
+
+
+def current_ray_pip_specifier() -> Optional[str]:
+    """The pip requirement specifier for the running version of Ray.
+
+    Returns:
+        A string which can be passed to `pip install` to install the
+        currently running Ray version, or None if running on a version
+        built from source locally (likely if you are developing Ray).
+
+    Examples:
+        Returns "ray[all]==1.4.0" if running the stable release
+        Returns "https://s3-us-west-2.amazonaws.com/ray-wheels/master/[..].whl"
+            if running the nightly or a specific commit
+    """
+    if os.environ.get("RAY_CI_POST_WHEEL_TESTS"):
+        # Running in Buildkite CI after the wheel has been built.
+        # Wheels are at in the ray/.whl directory, and the present file is
+        # at ray/python/ray/workers.  Use relative paths to allow for
+        # testing locally if needed.
+        return os.path.join(
+            Path(__file__).resolve().parents[3], ".whl", get_wheel_filename())
+    elif ray.__commit__ == "{{RAY_COMMIT_SHA}}":
+        # Running on a version built from source locally.
+        return None
+    elif "dev" in ray.__version__:
+        # Running on a nightly wheel.
+        return get_master_wheel_url()
+    else:
+        return f"ray[all]=={ray.__version__}"
+
+
+def inject_ray_and_python(conda_dict, ray_pip_specifier: Optional[str],
+                          py_version: str) -> None:
+    if conda_dict.get("dependencies") is None:
+        conda_dict["dependencies"] = []
+
+    # Inject Python dependency.
+    deps = conda_dict["dependencies"]
+
+    # Add current python dependency.  If the user has already included a
+    # python version dependency, conda will raise a readable error if the two
+    # are incompatible, e.g:
+    #   ResolvePackageNotFound: - python[version='3.5.*,>=3.6']
+    deps.append(f"python={py_version}")
+
+    if "pip" not in deps:
+        deps.append("pip")
+
+    # Insert Ray dependency. If the user has already included Ray, conda
+    # will raise an error only if the two are incompatible.
+
+    if ray_pip_specifier is not None:
+        found_pip_dict = False
+        for dep in deps:
+            if isinstance(dep, dict) and dep.get("pip"):
+                dep["pip"].append(ray_pip_specifier)
+                found_pip_dict = True
+                break
+        if not found_pip_dict:
+            deps.append({"pip": [ray_pip_specifier]})
+    else:
+        logger.warning("Current Ray version could not be inserted "
+                       "into conda's pip dependencies, most likely "
+                       "because you are using a version of Ray "
+                       "built from source.  If so, you can try "
+                       "building a wheel and including the wheel "
+                       "as a dependency.")
+
+    return conda_dict