[runtime env] Make env_vars take effect when pip install packages (#2…

…2730)

Previously, for the stability of pip installation, we set env to empty, but when pip installs some gzip package, maybe need env_vars.  like this issue: #22610

python/ray/_private/runtime_env/pip.py

@@ -198,6 +198,7 @@ async def _install_pip_packages(
         path: str,
         pip_packages: List[str],
         cwd: str,
+        env_vars: Dict,
         logger: logging.Logger,
     ):
         virtualenv_path = _PathHelper.get_virtualenv_path(path)
@@ -234,7 +235,9 @@ def _gen_requirements_txt():
             pip_requirements_file,
         ]
         logger.info("Installing python requirements to %s", virtualenv_path)
-        await check_output_cmd(pip_install_cmd, logger=logger, cwd=cwd, env={})
+        pip_env = os.environ.copy()
+        pip_env.update(env_vars)
+        await check_output_cmd(pip_install_cmd, logger=logger, cwd=cwd, env=pip_env)
 
     async def _run(self):
         path = self._target_dir
@@ -249,7 +252,9 @@ async def _run(self):
             await self._create_or_get_virtualenv(path, exec_cwd, logger)
             python = _PathHelper.get_virtualenv_python(path)
             async with self._check_ray(python, exec_cwd, logger):
-                await self._install_pip_packages(path, pip_packages, exec_cwd, logger)
+                await self._install_pip_packages(
+                    path, pip_packages, exec_cwd, self._runtime_env.env_vars(), logger
+                )
             # TODO(fyrestone): pip check.
         except Exception:
             logger.info("Delete incomplete virtualenv: %s", path)

python/ray/tests/test_runtime_env_complicated.py

@@ -24,7 +24,11 @@
     wait_for_condition,
     chdir,
 )
-from ray._private.utils import get_conda_env_dir, get_conda_bin_executable
+from ray._private.utils import (
+    get_conda_env_dir,
+    get_conda_bin_executable,
+    try_to_create_directory,
+)
 
 if not os.environ.get("CI"):
     # This flags turns on the local development that link against current ray
@@ -934,6 +938,76 @@ def spawn_child(self, name, runtime_env):
         ray.shutdown()
 
 
+@pytest.mark.skipif(
+    os.environ.get("CI") and sys.platform != "linux",
+    reason="This test is only run on linux CI machines.",
+)
+def test_pip_with_env_vars(start_cluster):
+
+    with tempfile.TemporaryDirectory() as tmpdir, chdir(tmpdir):
+        TEST_ENV_NAME = "TEST_ENV_VARS"
+        TEST_ENV_VALUE = "TEST"
+        package_name = "test_package"
+        package_dir = os.path.join(tmpdir, package_name)
+        try_to_create_directory(package_dir)
+
+        setup_filename = os.path.join(package_dir, "setup.py")
+        setup_code = """import os
+from setuptools import setup, find_packages
+from setuptools.command.install import install
+
+class InstallTestPackage(install):
+    # this function will be called when pip install this package
+    def run(self):
+        assert os.environ.get('{TEST_ENV_NAME}') == '{TEST_ENV_VALUE}'
+
+setup(
+    name='test_package',
+    version='0.0.1',
+    packages=find_packages(),
+    cmdclass=dict(install=InstallTestPackage),
+    license="MIT",
+    zip_safe=False,
+)
+""".format(
+            TEST_ENV_NAME=TEST_ENV_NAME, TEST_ENV_VALUE=TEST_ENV_VALUE
+        )
+        with open(setup_filename, "wt") as f:
+            f.writelines(setup_code)
+
+        python_filename = os.path.join(package_dir, "test.py")
+        python_code = "import os; print(os.environ)"
+        with open(python_filename, "wt") as f:
+            f.writelines(python_code)
+
+        gz_filename = os.path.join(tmpdir, package_name + ".tar.gz")
+        subprocess.check_call(["tar", "-zcvf", gz_filename, package_name])
+
+        with pytest.raises(ray.exceptions.RuntimeEnvSetupError):
+
+            @ray.remote(
+                runtime_env={
+                    "env_vars": {TEST_ENV_NAME: "failed"},
+                    "pip": [gz_filename],
+                }
+            )
+            def f1():
+                return True
+
+            ray.get(f1.remote())
+
+        @ray.remote(
+            runtime_env={
+                "env_vars": {TEST_ENV_NAME: TEST_ENV_VALUE},
+                "pip": [gz_filename],
+            }
+        )
+        def f2():
+            return True
+
+        assert ray.get(f2.remote())
+
+
 if __name__ == "__main__":
     import sys