def _server_init(self,
job_config: JobConfig,
ray_init_kwargs: Optional[Dict[str, Any]] = None):
"""Initialize the server"""
if ray_init_kwargs is None:
ray_init_kwargs = {}
try:
if job_config is None:
serialized_job_config = None
else:
with tempfile.TemporaryDirectory() as tmp_dir:
runtime_env = job_config.runtime_env or {}
runtime_env = upload_py_modules_if_needed(runtime_env,
tmp_dir,
logger=logger)
runtime_env = upload_working_dir_if_needed(runtime_env,
tmp_dir,
logger=logger)
# Remove excludes, it isn't relevant after the upload step.
runtime_env.pop("excludes", None)
job_config.set_runtime_env(runtime_env, validate=True)
serialized_job_config = pickle.dumps(job_config)
response = self.data_client.Init(
ray_client_pb2.InitRequest(
job_config=serialized_job_config,
ray_init_kwargs=json.dumps(ray_init_kwargs),
reconnect_grace_period=self._reconnect_grace_period))
if not response.ok:
raise ConnectionAbortedError(
f"Initialization failure from server:\n{response.msg}")
except grpc.RpcError as e:
raise decode_exception(e)
def test_prepare_runtime_init_req_no_modification():
"""
Check that `prepare_runtime_init_req` properly extracts the JobConfig.
"""
job_config = JobConfig(worker_env={"KEY": "VALUE"}, ray_namespace="abc")
init_req = ray_client_pb2.DataRequest(
init=ray_client_pb2.InitRequest(job_config=pickle.dumps(job_config)))
req, new_config = proxier.prepare_runtime_init_req(init_req)
assert new_config.serialize() == job_config.serialize()
assert isinstance(req, ray_client_pb2.DataRequest)
assert pickle.loads(
req.init.job_config).serialize() == new_config.serialize()
def _server_init(self, job_config: JobConfig):
"""Initialize the server"""
try:
if job_config is None:
init_req = ray_client_pb2.InitRequest()
self._call_init(init_req)
return
import ray._private.runtime_env as runtime_env
import tempfile
with tempfile.TemporaryDirectory() as tmp_dir:
(old_dir, runtime_env.PKG_DIR) = (runtime_env.PKG_DIR, tmp_dir)
# Generate the uri for runtime env
runtime_env.rewrite_runtime_env_uris(job_config)
init_req = ray_client_pb2.InitRequest(
job_config=pickle.dumps(job_config))
self._call_init(init_req)
runtime_env.upload_runtime_env_package_if_needed(job_config)
runtime_env.PKG_DIR = old_dir
prep_req = ray_client_pb2.PrepRuntimeEnvRequest()
self.data_client.PrepRuntimeEnv(prep_req)
except grpc.RpcError as e:
raise decode_exception(e.details())
def test_prepare_runtime_init_req_no_modification():
"""
Check that `prepare_runtime_init_req` properly extracts the JobConfig.
"""
job_config = JobConfig(runtime_env={"env_vars": {
"KEY": "VALUE"
}},
ray_namespace="abc")
init_req = ray_client_pb2.DataRequest(init=ray_client_pb2.InitRequest(
job_config=pickle.dumps(job_config),
ray_init_kwargs=json.dumps({"log_to_driver": False})), )
req, new_config = proxier.prepare_runtime_init_req(init_req)
assert new_config.serialize() == job_config.serialize()
assert isinstance(req, ray_client_pb2.DataRequest)
assert pickle.loads(
req.init.job_config).serialize() == new_config.serialize()
assert json.loads(req.init.ray_init_kwargs) == {"log_to_driver": False}
def test_prepare_runtime_init_req_modified_job():
"""
Check that `prepare_runtime_init_req` properly extracts the JobConfig and
modifies it according to `ray_client_server_env_prep`.
"""
job_config = JobConfig(worker_env={"KEY": "VALUE"}, ray_namespace="abc")
init_req = ray_client_pb2.DataRequest(
init=ray_client_pb2.InitRequest(job_config=pickle.dumps(job_config)))
def modify_namespace(job_config: JobConfig):
job_config.set_ray_namespace("test_value")
return job_config
with patch.object(proxier, "ray_client_server_env_prep", modify_namespace):
req, new_config = proxier.prepare_runtime_init_req(init_req)
assert new_config.ray_namespace == "test_value"
assert pickle.loads(
req.init.job_config).serialize() == new_config.serialize()
def prepare_runtime_init_req(iterator: Iterator[ray_client_pb2.DataRequest]
) -> Tuple[ray_client_pb2.DataRequest, JobConfig]:
"""
Extract JobConfig and possibly mutate InitRequest before it is passed to
the specific RayClient Server.
"""
init_req = next(iterator)
init_type = init_req.WhichOneof("type")
assert init_type == "init", ("Received initial message of type "
f"{init_type}, not 'init'.")
req = init_req.init
job_config = JobConfig()
if req.job_config:
job_config = pickle.loads(req.job_config)
new_job_config = ray_client_server_env_prep(job_config)
modified_init_req = ray_client_pb2.InitRequest(
job_config=pickle.dumps(new_job_config))
init_req.init.CopyFrom(modified_init_req)
return (init_req, new_job_config)
def test_proxy_manager_internal_kv(shutdown_only, with_specific_server):
"""
Test that proxy manager can use internal kv with and without a
SpecificServer and that once a SpecificServer is started up, it
goes through it.
"""
proxier.CHECK_PROCESS_INTERVAL_S = 1
# The timeout has likely been set to 1 in an earlier test. Increase timeout
# to wait for the channel to become ready.
proxier.CHECK_CHANNEL_TIMEOUT_S = 5
os.environ["TIMEOUT_FOR_SPECIFIC_SERVER_S"] = "5"
pm, free_ports = start_ray_and_proxy_manager(n_ports=2)
client = "client1"
task_servicer = proxier.RayletServicerProxy(None, pm)
def make_internal_kv_calls():
response = task_servicer.KVPut(
ray_client_pb2.KVPutRequest(key=b"key", value=b"val")
)
assert isinstance(response, ray_client_pb2.KVPutResponse)
assert not response.already_exists
response = task_servicer.KVPut(
ray_client_pb2.KVPutRequest(key=b"key", value=b"val2")
)
assert isinstance(response, ray_client_pb2.KVPutResponse)
assert response.already_exists
response = task_servicer.KVGet(ray_client_pb2.KVGetRequest(key=b"key"))
assert isinstance(response, ray_client_pb2.KVGetResponse)
assert response.value == b"val"
response = task_servicer.KVPut(
ray_client_pb2.KVPutRequest(key=b"key", value=b"val2", overwrite=True)
)
assert isinstance(response, ray_client_pb2.KVPutResponse)
assert response.already_exists
response = task_servicer.KVGet(ray_client_pb2.KVGetRequest(key=b"key"))
assert isinstance(response, ray_client_pb2.KVGetResponse)
assert response.value == b"val2"
with patch(
"ray.util.client.server.proxier._get_client_id_from_context"
) as mock_get_client_id:
mock_get_client_id.return_value = client
if with_specific_server:
pm.create_specific_server(client)
assert pm.start_specific_server(client, JobConfig())
channel = pm.get_channel(client)
assert channel is not None
task_servicer.Init(
ray_client_pb2.InitRequest(job_config=pickle.dumps(JobConfig()))
)
# Mock out the internal kv calls in this process to raise an
# exception if they're called. This verifies that we are not
# making any calls in the proxier if there is a SpecificServer
# started up.
with patch(
"ray.experimental.internal_kv._internal_kv_put"
) as mock_put, patch(
"ray.experimental.internal_kv._internal_kv_get"
) as mock_get, patch(
"ray.experimental.internal_kv._internal_kv_initialized"
) as mock_initialized:
mock_put.side_effect = Exception("This shouldn't be called!")
mock_get.side_effect = Exception("This shouldn't be called!")
mock_initialized.side_effect = Exception("This shouldn't be called!")
make_internal_kv_calls()
else:
make_internal_kv_calls()