[core] Introduce get_compatible_keys for ImageURIPlugin

python/ray/_private/runtime_env/agent/runtime_env_agent.py

@@ -14,7 +14,7 @@
 from ray._private.ray_logging import setup_component_logger
 from ray._private.runtime_env.conda import CondaPlugin
 from ray._private.runtime_env.context import RuntimeEnvContext
-from ray._private.runtime_env.default_impl import get_image_uri_plugin
+from ray._private.runtime_env.default_impl import get_image_uri_plugin_cls
 from ray._private.runtime_env.java_jars import JavaJarsPlugin
 from ray._private.runtime_env.image_uri import ContainerPlugin
 from ray._private.runtime_env.pip import PipPlugin
@@ -223,7 +223,7 @@ def __init__(
         # and unify with nsight and other profilers.
         self._nsight_plugin = NsightPlugin(self._runtime_env_dir)
         self._mpi_plugin = MPIPlugin()
-        self._image_uri_plugin = get_image_uri_plugin(temp_dir)
+        self._image_uri_plugin = get_image_uri_plugin_cls()(temp_dir)
 
         # TODO(architkulkarni): "base plugins" and third-party plugins should all go
         # through the same code path.  We should never need to refer to

python/ray/_private/runtime_env/default_impl.py

@@ -1,8 +1,8 @@
 from ray._private.runtime_env.image_uri import ImageURIPlugin
 
 
-def get_image_uri_plugin(ray_tmp_dir: str):
-    return ImageURIPlugin(ray_tmp_dir)
+def get_image_uri_plugin_cls():
+    return ImageURIPlugin
 
 
 def get_protocols_provider():

python/ray/_private/runtime_env/image_uri.py

@@ -108,6 +108,10 @@ class ImageURIPlugin(RuntimeEnvPlugin):
 
     name = "image_uri"
 
+    @staticmethod
+    def get_compatible_keys():
+        return {"image_uri", "config", "env_vars"}
+
     def __init__(self, ray_tmp_dir: str):
         self._ray_tmp_dir = ray_tmp_dir
 

python/ray/runtime_env/runtime_env.py

@@ -8,6 +8,7 @@
 import ray
 from ray._private.ray_constants import DEFAULT_RUNTIME_ENV_TIMEOUT_SECONDS
 from ray._private.runtime_env.conda import get_uri as get_conda_uri
+from ray._private.runtime_env.default_impl import get_image_uri_plugin_cls
 from ray._private.runtime_env.pip import get_uri as get_pip_uri
 from ray._private.runtime_env.plugin_schema_manager import RuntimeEnvPluginSchemaManager
 from ray._private.runtime_env.uv import get_uri as get_uv_uri
@@ -389,7 +390,10 @@ def __init__(
                 )
 
         if self.get("image_uri"):
-            invalid_keys = set(runtime_env.keys()) - {"image_uri", "config", "env_vars"}
+            image_uri_plugin_cls = get_image_uri_plugin_cls()
+            invalid_keys = (
+                set(runtime_env.keys()) - image_uri_plugin_cls.get_compatible_keys()
+            )
             if len(invalid_keys):
                 raise ValueError(
                     "The 'image_uri' field currently cannot be used "