def test_serialization(self):
env1 = RuntimeEnv(pip=["requests"],
env_vars={
"hi1": "hi1",
"hi2": "hi2"
})
env2 = RuntimeEnv(env_vars={
"hi2": "hi2",
"hi1": "hi1"
},
pip=["requests"])
assert env1 == env2
serialized_env1 = env1.serialize()
serialized_env2 = env2.serialize()
# Key ordering shouldn't matter.
assert serialized_env1 == serialized_env2
deserialized_env1 = RuntimeEnv.deserialize(serialized_env1)
deserialized_env2 = RuntimeEnv.deserialize(serialized_env2)
assert env1 == deserialized_env1 == env2 == deserialized_env2
def get_runtime_env_info(
runtime_env: RuntimeEnv,
*,
is_job_runtime_env: bool = False,
serialize: bool = False,
):
"""Create runtime env info from runtime env.
In the user interface, the argument `runtime_env` contains some fields
which not contained in `ProtoRuntimeEnv` but in `ProtoRuntimeEnvInfo`,
such as `eager_install`. This function will extract those fields from
`RuntimeEnv` and create a new `ProtoRuntimeEnvInfo`, and serialize it.
"""
proto_runtime_env_info = ProtoRuntimeEnvInfo()
proto_runtime_env_info.uris[:] = runtime_env.get_uris()
# Normally, `RuntimeEnv` should guarantee the accuracy of field eager_install,
# but so far, the internal code has not completely prohibited direct
# modification of fields in RuntimeEnv, so we should check it for insurance.
# TODO(Catch-Bull): overload `__setitem__` for `RuntimeEnv`, change the
# runtime_env of all internal code from dict to RuntimeEnv.
eager_install = runtime_env.get("eager_install")
if is_job_runtime_env or eager_install is not None:
if eager_install is None:
eager_install = True
elif not isinstance(eager_install, bool):
raise TypeError(
f"eager_install must be a boolean. got {type(eager_install)}")
proto_runtime_env_info.runtime_env_eager_install = eager_install
runtime_env_config = runtime_env.get("config")
if runtime_env_config is None:
runtime_env_config = RuntimeEnvConfig.default_config()
else:
runtime_env_config = RuntimeEnvConfig.parse_and_validate_runtime_env_config(
runtime_env_config)
proto_runtime_env_info.runtime_env_config.CopyFrom(
runtime_env_config.build_proto_runtime_env_config())
proto_runtime_env_info.serialized_runtime_env = runtime_env.serialize()
if not serialize:
return proto_runtime_env_info
return json_format.MessageToJson(proto_runtime_env_info)
def test_convert_from_and_to_dataclass():
runtime_env = RuntimeEnv()
test_plugin = TestPlugin(
field1=[
ValueType(nfield1=["a", "b", "c"], nfield2=False),
ValueType(nfield1=["d", "e"], nfield2=True),
],
field2="abc",
)
runtime_env.set("test_plugin", test_plugin)
serialized_runtime_env = runtime_env.serialize()
assert "test_plugin" in serialized_runtime_env
runtime_env_2 = RuntimeEnv.deserialize(serialized_runtime_env)
test_plugin_2 = runtime_env_2.get("test_plugin", data_class=TestPlugin)
assert len(test_plugin_2.field1) == 2
assert test_plugin_2.field1[0].nfield1 == ["a", "b", "c"]
assert test_plugin_2.field1[0].nfield2 is False
assert test_plugin_2.field1[1].nfield1 == ["d", "e"]
assert test_plugin_2.field1[1].nfield2 is True
assert test_plugin_2.field2 == "abc"
def test_serialize_deserialize(option):
runtime_env = dict()
if option == "pip_list":
runtime_env["pip"] = ["pkg1", "pkg2"]
elif option == "pip_dict":
runtime_env["pip"] = {
"packages": ["pkg1", "pkg2"],
"pip_check": False,
"pip_version": "<22,>20",
}
elif option == "conda_name":
runtime_env["conda"] = "env_name"
elif option == "conda_dict":
runtime_env["conda"] = {"dependencies": ["dep1", "dep2"]}
elif option == "container":
runtime_env["container"] = {
"image": "anyscale/ray-ml:nightly-py38-cpu",
"worker_path":
"/root/python/ray/_private/workers/default_worker.py",
"run_options": ["--cap-drop SYS_ADMIN", "--log-level=debug"],
}
else:
raise ValueError("unexpected option " + str(option))
typed_runtime_env = RuntimeEnv(**runtime_env)
serialized_runtime_env = typed_runtime_env.serialize()
cls_runtime_env = RuntimeEnv.deserialize(serialized_runtime_env)
cls_runtime_env_dict = cls_runtime_env.to_dict()
if "pip" in typed_runtime_env and isinstance(typed_runtime_env["pip"],
list):
pip_config_in_cls_runtime_env = cls_runtime_env_dict.pop("pip")
pip_config_in_runtime_env = typed_runtime_env.pop("pip")
assert {
"packages": pip_config_in_runtime_env,
"pip_check": False,
} == pip_config_in_cls_runtime_env
assert cls_runtime_env_dict == typed_runtime_env
def test_runtime_env_interface():
# Test the interface related to working_dir
default_working_dir = "s3://bucket/key.zip"
modify_working_dir = "s3://bucket/key_A.zip"
runtime_env = RuntimeEnv(working_dir=default_working_dir)
runtime_env_dict = runtime_env.to_dict()
assert runtime_env.working_dir_uri() == default_working_dir
runtime_env["working_dir"] = modify_working_dir
runtime_env_dict["working_dir"] = modify_working_dir
assert runtime_env.working_dir_uri() == modify_working_dir
assert runtime_env.to_dict() == runtime_env_dict
# Test that the modification of working_dir also works on
# proto serialization
assert runtime_env_dict == RuntimeEnv.from_proto(
runtime_env.build_proto_runtime_env())
runtime_env.pop("working_dir")
assert runtime_env.to_dict() == {}
# Test the interface related to py_modules
init_py_modules = ["s3://bucket/key_1.zip", "s3://bucket/key_2.zip"]
addition_py_modules = ["s3://bucket/key_3.zip", "s3://bucket/key_4.zip"]
runtime_env = RuntimeEnv(py_modules=init_py_modules)
runtime_env_dict = runtime_env.to_dict()
assert set(runtime_env.py_modules_uris()) == set(init_py_modules)
runtime_env["py_modules"].extend(addition_py_modules)
runtime_env_dict["py_modules"].extend(addition_py_modules)
assert set(runtime_env.py_modules_uris()) == set(init_py_modules +
addition_py_modules)
assert runtime_env.to_dict() == runtime_env_dict
# Test that the modification of py_modules also works on
# proto serialization
assert runtime_env_dict == RuntimeEnv.from_proto(
runtime_env.build_proto_runtime_env())
runtime_env.pop("py_modules")
assert runtime_env.to_dict() == {}
# Test the interface related to env_vars
init_env_vars = {"A": "a", "B": "b"}
update_env_vars = {"C": "c"}
runtime_env = RuntimeEnv(env_vars=init_env_vars)
runtime_env_dict = runtime_env.to_dict()
runtime_env["env_vars"].update(update_env_vars)
runtime_env_dict["env_vars"].update(update_env_vars)
init_env_vars_copy = init_env_vars.copy()
init_env_vars_copy.update(update_env_vars)
assert runtime_env["env_vars"] == init_env_vars_copy
assert runtime_env_dict == runtime_env.to_dict()
# Test that the modification of env_vars also works on
# proto serialization
assert runtime_env_dict == RuntimeEnv.from_proto(
runtime_env.build_proto_runtime_env())
runtime_env.pop("env_vars")
assert runtime_env.to_dict() == {}
# Test the interface related to conda
conda_name = "conda"
modify_conda_name = "conda_A"
conda_config = {"dependencies": ["dep1", "dep2"]}
runtime_env = RuntimeEnv(conda=conda_name)
runtime_env_dict = runtime_env.to_dict()
assert runtime_env.has_conda()
assert runtime_env.conda_env_name() == conda_name
assert runtime_env.conda_config() is None
runtime_env["conda"] = modify_conda_name
runtime_env_dict["conda"] = modify_conda_name
assert runtime_env_dict == runtime_env.to_dict()
assert runtime_env.has_conda()
assert runtime_env.conda_env_name() == modify_conda_name
assert runtime_env.conda_config() is None
runtime_env["conda"] = conda_config
runtime_env_dict["conda"] = conda_config
assert runtime_env_dict == runtime_env.to_dict()
assert runtime_env.has_conda()
assert runtime_env.conda_env_name() is None
assert runtime_env.conda_config() == json.dumps(conda_config,
sort_keys=True)
# Test that the modification of conda also works on
# proto serialization
assert runtime_env_dict == RuntimeEnv.from_proto(
runtime_env.build_proto_runtime_env())
runtime_env.pop("conda")
assert runtime_env.to_dict() == {"_ray_commit": "{{RAY_COMMIT_SHA}}"}
# Test the interface related to pip
with tempfile.TemporaryDirectory() as tmpdir, chdir(tmpdir):
requirement_file = os.path.join(tmpdir, "requirements.txt")
requirement_packages = ["dep5", "dep6"]
with open(requirement_file, "wt") as f:
for package in requirement_packages:
f.write(package)
f.write("\n")
pip_packages = ["dep1", "dep2"]
addition_pip_packages = ["dep3", "dep4"]
runtime_env = RuntimeEnv(pip=pip_packages)
runtime_env_dict = runtime_env.to_dict()
assert runtime_env.has_pip()
assert set(runtime_env.pip_config()["packages"]) == set(pip_packages)
assert runtime_env.virtualenv_name() is None
runtime_env["pip"]["packages"].extend(addition_pip_packages)
runtime_env_dict["pip"]["packages"].extend(addition_pip_packages)
# The default value of pip_check is False
runtime_env_dict["pip"]["pip_check"] = False
assert runtime_env_dict == runtime_env.to_dict()
assert runtime_env.has_pip()
assert set(
runtime_env.pip_config()["packages"]) == set(pip_packages +
addition_pip_packages)
assert runtime_env.virtualenv_name() is None
runtime_env["pip"] = requirement_file
runtime_env_dict["pip"] = requirement_packages
assert runtime_env.has_pip()
assert set(
runtime_env.pip_config()["packages"]) == set(requirement_packages)
assert runtime_env.virtualenv_name() is None
# The default value of pip_check is False
runtime_env_dict["pip"] = dict(packages=runtime_env_dict["pip"],
pip_check=False)
assert runtime_env_dict == runtime_env.to_dict()
# Test that the modification of pip also works on
# proto serialization
assert runtime_env_dict == RuntimeEnv.from_proto(
runtime_env.build_proto_runtime_env())
runtime_env.pop("pip")
assert runtime_env.to_dict() == {"_ray_commit": "{{RAY_COMMIT_SHA}}"}
# Test conflict
with pytest.raises(ValueError):
RuntimeEnv(pip=pip_packages, conda=conda_name)
runtime_env = RuntimeEnv(pip=pip_packages)
runtime_env["conda"] = conda_name
with pytest.raises(ValueError):
runtime_env.serialize()
# Test the interface related to container
container_init = {
"image": "anyscale/ray-ml:nightly-py38-cpu",
"worker_path": "/root/python/ray/workers/default_worker.py",
"run_options": ["--cap-drop SYS_ADMIN", "--log-level=debug"],
}
update_container = {"image": "test_modify"}
runtime_env = RuntimeEnv(container=container_init)
runtime_env_dict = runtime_env.to_dict()
assert runtime_env.has_py_container()
assert runtime_env.py_container_image() == container_init["image"]
assert runtime_env.py_container_worker_path(
) == container_init["worker_path"]
assert runtime_env.py_container_run_options(
) == container_init["run_options"]
runtime_env["container"].update(update_container)
runtime_env_dict["container"].update(update_container)
container_copy = container_init
container_copy.update(update_container)
assert runtime_env_dict == runtime_env.to_dict()
assert runtime_env.has_py_container()
assert runtime_env.py_container_image() == container_copy["image"]
assert runtime_env.py_container_worker_path(
) == container_copy["worker_path"]
assert runtime_env.py_container_run_options(
) == container_copy["run_options"]
# Test that the modification of container also works on
# proto serialization
assert runtime_env_dict == RuntimeEnv.from_proto(
runtime_env.build_proto_runtime_env())
runtime_env.pop("container")
assert runtime_env.to_dict() == {}
def test_reference_table():
expected_unused_uris = []
expected_unused_runtime_env = str()
def uris_parser(runtime_env) -> Tuple[str, UriType]:
result = list()
result.append((runtime_env.working_dir(), UriType.WORKING_DIR))
py_module_uris = runtime_env.py_modules()
for uri in py_module_uris:
result.append((uri, UriType.PY_MODULES))
return result
def unused_uris_processor(unused_uris: List[Tuple[str, UriType]]) -> None:
nonlocal expected_unused_uris
assert expected_unused_uris
for unused in unused_uris:
assert unused in expected_unused_uris
expected_unused_uris.remove(unused)
assert not expected_unused_uris
def unused_runtime_env_processor(unused_runtime_env: str) -> None:
nonlocal expected_unused_runtime_env
assert expected_unused_runtime_env
assert expected_unused_runtime_env == unused_runtime_env
expected_unused_runtime_env = None
reference_table = ReferenceTable(
uris_parser, unused_uris_processor, unused_runtime_env_processor
)
runtime_env_1 = RuntimeEnv(
working_dir="s3://working_dir_1.zip",
py_modules=["s3://py_module_A.zip", "s3://py_module_B.zip"],
)
runtime_env_2 = RuntimeEnv(
working_dir="s3://working_dir_2.zip",
py_modules=["s3://py_module_A.zip", "s3://py_module_C.zip"],
)
# Add runtime env 1
reference_table.increase_reference(
runtime_env_1, runtime_env_1.serialize(), "raylet"
)
# Add runtime env 2
reference_table.increase_reference(
runtime_env_2, runtime_env_2.serialize(), "raylet"
)
# Add runtime env 1 by `client_server`, this will be skipped by reference table.
reference_table.increase_reference(
runtime_env_1, runtime_env_1.serialize(), "client_server"
)
# Remove runtime env 1
expected_unused_uris.append(("s3://working_dir_1.zip", UriType.WORKING_DIR))
expected_unused_uris.append(("s3://py_module_B.zip", UriType.PY_MODULES))
expected_unused_runtime_env = runtime_env_1.serialize()
reference_table.decrease_reference(
runtime_env_1, runtime_env_1.serialize(), "raylet"
)
assert not expected_unused_uris
assert not expected_unused_runtime_env
# Remove runtime env 2
expected_unused_uris.append(("s3://working_dir_2.zip", UriType.WORKING_DIR))
expected_unused_uris.append(("s3://py_module_A.zip", UriType.PY_MODULES))
expected_unused_uris.append(("s3://py_module_C.zip", UriType.PY_MODULES))
expected_unused_runtime_env = runtime_env_2.serialize()
reference_table.decrease_reference(
runtime_env_2, runtime_env_2.serialize(), "raylet"
)
assert not expected_unused_uris
assert not expected_unused_runtime_env
