def _parse_proto_conda_runtime_env(runtime_env: ProtoRuntimeEnv, runtime_env_dict: dict): """ Parse conda runtime env protobuf to runtime env dict. """ if runtime_env.HasField("conda_runtime_env"): runtime_env_dict["conda"] = json.loads( runtime_env.conda_runtime_env.config)
def _parse_proto_pip_runtime_env(runtime_env: ProtoRuntimeEnv, runtime_env_dict: dict): """ Parse pip runtime env protobuf to runtime env dict. """ if runtime_env.HasField("pip_runtime_env"): runtime_env_dict["pip"] = list( runtime_env.pip_runtime_env.config.packages)
def _parse_proto_plugin_runtime_env(runtime_env: ProtoRuntimeEnv, runtime_env_dict: dict): """ Parse plugin runtime env protobuf to runtime env dict. """ if runtime_env.HasField("py_plugin_runtime_env"): for plugin in runtime_env.py_plugin_runtime_env.plugins: runtime_env_dict["plugins"][plugin.class_path] = dict( json.loads(plugin.config))
def _parse_proto_container_runtime_env(runtime_env: ProtoRuntimeEnv, runtime_env_dict: dict): """ Parse container runtime env protobuf to runtime env dict. """ if runtime_env.HasField("py_container_runtime_env"): runtime_env_dict["container"][ "image"] = runtime_env.container_runtime_env.image runtime_env_dict["container"][ "worker_path"] = runtime_env.container_runtime_env.worker_path runtime_env_dict["container"]["run_options"] = list( runtime_env.container_runtime_env.run_options)
class RuntimeEnv: """ A wrap class of runtime env protobuf. """ def __init__(self, serialized_runtime_env=None, proto_runtime_env: ProtoRuntimeEnv = None): if serialized_runtime_env: self._proto_runtime_env = json_format.Parse( serialized_runtime_env, ProtoRuntimeEnv()) elif proto_runtime_env: self._proto_runtime_env = proto_runtime_env else: self._proto_runtime_env = ProtoRuntimeEnv() def to_dict(self) -> Dict: initialize_dict: Dict[str, Any] = {} if self._proto_runtime_env.py_modules: initialize_dict["py_modules"] = list( self._proto_runtime_env.py_modules) if self._proto_runtime_env.working_dir: initialize_dict[ "working_dir"] = self._proto_runtime_env.working_dir if self._proto_runtime_env.env_vars: initialize_dict["env_vars"] = dict( self._proto_runtime_env.env_vars) if self._proto_runtime_env.extensions: initialize_dict.update(dict(self._proto_runtime_env.extensions)) _parse_proto_pip_runtime_env(self._proto_runtime_env, initialize_dict) _parse_proto_conda_runtime_env(self._proto_runtime_env, initialize_dict) _parse_proto_container_runtime_env(self._proto_runtime_env, initialize_dict) _parse_proto_plugin_runtime_env(self._proto_runtime_env, initialize_dict) return initialize_dict def has_uris(self) -> bool: uris = self._proto_runtime_env.uris if uris.working_dir_uri \ or uris.py_modules_uris \ or uris.conda_uri \ or uris.plugin_uris: return True return False def working_dir_uri(self) -> str: return self._proto_runtime_env.uris.working_dir_uri def py_modules_uris(self) -> List[str]: return list(self._proto_runtime_env.uris.py_modules_uris) def conda_uri(self) -> str: return self._proto_runtime_env.uris.conda_uri def plugin_uris(self) -> List[str]: return list(self._proto_runtime_env.uris.plugin_uris) def working_dir(self) -> str: return self._proto_runtime_env.working_dir def py_modules(self) -> List[str]: return list(self._proto_runtime_env.py_modules) def env_vars(self) -> Dict: return dict(self._proto_runtime_env.env_vars) def plugins(self) -> List[Tuple[str, str]]: result = list() for plugin in self._proto_runtime_env.py_plugin_runtime_env.plugins: result.append((plugin.class_path, plugin.config)) return result def has_conda(self) -> str: return self._proto_runtime_env.HasField("conda_runtime_env") def conda_env_name(self) -> str: if not self.has_conda(): return None if not self._proto_runtime_env.conda_runtime_env.HasField( "conda_env_name"): return None return self._proto_runtime_env.conda_runtime_env.conda_env_name def conda_config(self) -> str: if not self.has_conda(): return None if not self._proto_runtime_env.conda_runtime_env.HasField("config"): return None return self._proto_runtime_env.conda_runtime_env.config def has_pip(self) -> bool: return self._proto_runtime_env.HasField("pip_runtime_env") def pip_packages(self) -> List: if not self.has_pip(): return [] return list(self._proto_runtime_env.pip_runtime_env.config.packages) def serialize(self) -> str: # Sort the keys we can compare the serialized string for equality. return json.dumps(json.loads( json_format.MessageToJson(self._proto_runtime_env)), sort_keys=True) def get_extension(self, key) -> str: return self._proto_runtime_env.extensions.get(key) def has_py_container(self) -> bool: return self._proto_runtime_env.HasField("py_container_runtime_env") def py_container_image(self) -> str: if not self.has_py_container(): return None return self._proto_runtime_env.py_container_runtime_env.image def py_container_run_options(self) -> List: if not self.has_py_container(): return None return list( self._proto_runtime_env.py_container_runtime_env.run_options) @classmethod def from_dict(cls, runtime_env_dict: Dict[str, Any], conda_get_uri_fn) -> "RuntimeEnv": proto_runtime_env = ProtoRuntimeEnv() proto_runtime_env.py_modules.extend( runtime_env_dict.get("py_modules", [])) proto_runtime_env.working_dir = runtime_env_dict.get("working_dir", "") if "working_dir" in runtime_env_dict: proto_runtime_env.uris.working_dir_uri = runtime_env_dict[ "working_dir"] if "py_modules" in runtime_env_dict: for uri in runtime_env_dict["py_modules"]: proto_runtime_env.uris.py_modules_uris.append(uri) if "conda" or "pip" in runtime_env_dict: uri = conda_get_uri_fn(runtime_env_dict) if uri is not None: proto_runtime_env.uris.conda_uri = uri env_vars = runtime_env_dict.get("env_vars", {}) proto_runtime_env.env_vars.update(env_vars.items()) if "_ray_release" in runtime_env_dict: proto_runtime_env.extensions["_ray_release"] = str( runtime_env_dict["_ray_release"]) if "_ray_commit" in runtime_env_dict: proto_runtime_env.extensions["_ray_commit"] = str( runtime_env_dict["_ray_commit"]) if "_inject_current_ray" in runtime_env_dict: proto_runtime_env.extensions["_inject_current_ray"] = str( runtime_env_dict["_inject_current_ray"]) _build_proto_pip_runtime_env(runtime_env_dict, proto_runtime_env) _build_proto_conda_runtime_env(runtime_env_dict, proto_runtime_env) _build_proto_container_runtime_env(runtime_env_dict, proto_runtime_env) _build_proto_plugin_runtime_env(runtime_env_dict, proto_runtime_env) return cls(proto_runtime_env=proto_runtime_env)