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