[RuntimeEnv] Support custom podman command (#693)

* [RuntimeEnv] Support custom podman command

* Fix empty env keys

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

* fix

* fix

* Add test for output filter

---------

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

Author: liying919

python/ray/_private/runtime_env/image_uri.py

@@ -1,6 +1,7 @@
 import logging
 import os
 import json
+import shlex
 from typing import List, Optional
 import ray
 import ray._private.runtime_env.constants as runtime_env_constants
@@ -26,22 +27,38 @@
 async def _create_impl(image_uri: str, logger: logging.Logger):
     # Pull image if it doesn't exist
     # Also get path to `default_worker.py` inside the image.
-    pull_image_cmd = [
-        "podman",
-        "run",
-        "--rm",
-        image_uri,
-        "python",
-        "-c",
-        (
-            "import ray._private.workers.default_worker as default_worker; "
-            "print(default_worker.__file__)"
-        ),
-    ]
-    logger.info("Pulling image %s", image_uri)
-    worker_path = await check_output_cmd(pull_image_cmd, logger=logger)
+    custom_pull_cmd = os.getenv("RAY_PODMAN_PULL_CMD", "")
+    if custom_pull_cmd:
+        logger.info("Using custom pull command: %s", custom_pull_cmd)
+        shell_cmd = ["sh", "-c", custom_pull_cmd]
+        worker_path = await check_output_cmd(shell_cmd, logger=logger)
+    else:
+        pull_image_cmd = [
+            "podman",
+            "run",
+            "--rm",
+            image_uri,
+            "python",
+            "-c",
+            (
+                "import ray._private.workers.default_worker as default_worker; "
+                "print(default_worker.__file__)"
+            ),
+        ]
+        logger.info("Pulling image %s", image_uri)
+        worker_path = await check_output_cmd(pull_image_cmd, logger=logger)
+
+    output_filter = os.getenv("RAY_PODMAN_OUTPUT_FILTER", "")
+    if output_filter:
+        worker_path = await _apply_output_filter(logger, worker_path, output_filter)
     return worker_path.strip()
 
+async def _apply_output_filter(logger, worker_path, output_filter):
+    safe_worker_path = shlex.quote(worker_path)
+    filter_cmd = ["sh", "-c", f"printf '%s' {safe_worker_path} | {output_filter}"]
+    filtered_path = await check_output_cmd(filter_cmd, logger=logger)
+    worker_path = filtered_path
+    return worker_path
 
 def _modify_container_context_impl(
     runtime_env: "RuntimeEnv",  # noqa: F821
@@ -157,9 +174,9 @@ def _modify_container_context_impl(
 
     redirected_pyenv_folder = None
     if container_install_ray or container_pip_packages:
-        container_to_host_mount_dict[
-            container_dependencies_installer_path
-        ] = get_dependencies_installer_path()
+        container_to_host_mount_dict[container_dependencies_installer_path] = (
+            get_dependencies_installer_path()
+        )
         if runtime_env_constants.RAY_PODMAN_UES_WHL_PACKAGE:
             container_to_host_mount_dict[get_ray_whl_dir()] = get_ray_whl_dir()
 
@@ -253,6 +270,16 @@ def _modify_context_impl(
     ray_tmp_dir: str,
 ):
     context.override_worker_entrypoint = worker_path
+    custom_container_cmd = os.getenv("RAY_PODMAN_CONTAINER_CMD", "")
+    if custom_container_cmd:
+        custom_container_cmd_str = custom_container_cmd.format(
+            ray_tmp_dir=ray_tmp_dir, image_uri=image_uri
+        )
+        logger.info(
+            f"Starting worker in container with prefix {custom_container_cmd_str}"
+        )
+        context.py_executable = custom_container_cmd_str
+        return
 
     container_driver = "podman"
     container_command = [
@@ -285,6 +312,12 @@ def _modify_context_impl(
         if env_var_name.startswith("RAY_"):
             env_vars[env_var_name] = env_var_value
 
+    extra_env_keys = os.getenv("RAY_PODMAN_EXTRA_ENV_KEYS", "")
+    if extra_env_keys:
+        for key in (k.strip() for k in extra_env_keys.split(",")):
+            if key and key in os.environ:
+                env_vars[key] = os.environ[key]
+
     # Support for runtime_env['env_vars']
     env_vars.update(context.env_vars)
 

python/ray/tests/runtime_env_container/test_image_uri.py

@@ -0,0 +1,16 @@
+import asyncio
+import logging
+from ray._private.runtime_env.image_uri import _apply_output_filter
+
+
+def test_apply_output_filter_with_grep_v_head():    
+    logger = logging.getLogger(__name__)
+    worker_path = """time=1234567890
+                    /path/to/worker.py
+                    time=1234567891
+                    /other/path/file.py
+                    time=1234567892
+                    /final/path/worker.py"""
+    
+    result = asyncio.run(_apply_output_filter(logger, worker_path, "grep -v '^time=' | head -1")) 
+    assert result.strip() == "/path/to/worker.py"