import asyncio import json import os from unittest.mock import patch import numpy as np import pytest from fastapi.encoders import jsonable_encoder import ray from ray import serve from ray._common.constants import HEAD_NODE_RESOURCE_NAME from ray._common.test_utils import wait_for_condition from ray.serve._private.constants import SERVE_NAMESPACE from ray.serve._private.utils import ( Semaphore, calculate_remaining_timeout, get_active_placement_group_ids, get_all_live_placement_group_names, get_current_actor_id, get_head_node_id, is_running_in_asyncio_loop, merge_dict, msgpack_deserialize, msgpack_serialize, override_runtime_envs_except_env_vars, serve_encoders, snake_to_camel_case, ) from ray.serve.tests.common.remote_uris import ( TEST_DAG_PINNED_URI, TEST_DEPLOY_GROUP_PINNED_URI, ) def test_serialize(): data = msgpack_serialize(5) obj = msgpack_deserialize(data) assert 5 == obj def test_merge_dict(): dict1 = {"pending": 1, "running": 1, "finished": 1} dict2 = {"pending": 4, "finished": 1} merge = merge_dict(dict1, dict2) assert merge["pending"] == 5 assert merge["running"] == 1 assert merge["finished"] == 2 dict1 = None merge = merge_dict(dict1, dict2) assert merge["pending"] == 4 assert merge["finished"] == 1 try: assert merge["running"] == 1 assert False except KeyError: assert True dict2 = None merge = merge_dict(dict1, dict2) assert merge is None def test_bytes_encoder(): data_before = {"inp": {"nest": b"bytes"}} data_after = {"inp": {"nest": "bytes"}} assert json.loads(json.dumps(jsonable_encoder(data_before))) == data_after def test_numpy_encoding(): data = [1, 2] floats = np.array(data).astype(np.float32) ints = floats.astype(np.int32) uints = floats.astype(np.uint32) list_of_uints = [np.int64(1), np.int64(2)] for np_data in [floats, ints, uints, list_of_uints]: assert ( json.loads( json.dumps(jsonable_encoder(np_data, custom_encoder=serve_encoders)) ) == data ) nested = {"a": np.array([1, 2])} assert json.loads( json.dumps(jsonable_encoder(nested, custom_encoder=serve_encoders)) ) == {"a": [1, 2]} @serve.deployment def decorated_f(*args): return "reached decorated_f" @serve.deployment class DecoratedActor: def __call__(self, *args): return "reached decorated_actor" def gen_func(): @serve.deployment def f(): pass return f def gen_class(): @serve.deployment class A: pass return A class TestOverrideRuntimeEnvsExceptEnvVars: def test_merge_empty(self): assert {"env_vars": {}} == override_runtime_envs_except_env_vars({}, {}) def test_merge_empty_parent(self): child = {"env_vars": {"test1": "test_val"}, "working_dir": "."} assert child == override_runtime_envs_except_env_vars({}, child) def test_merge_empty_child(self): parent = {"env_vars": {"test1": "test_val"}, "working_dir": "."} assert parent == override_runtime_envs_except_env_vars(parent, {}) @pytest.mark.parametrize("invalid_env", [None, 0, "runtime_env", set()]) def test_invalid_type(self, invalid_env): with pytest.raises(TypeError): override_runtime_envs_except_env_vars(invalid_env, {}) with pytest.raises(TypeError): override_runtime_envs_except_env_vars({}, invalid_env) with pytest.raises(TypeError): override_runtime_envs_except_env_vars(invalid_env, invalid_env) def test_basic_merge(self): parent = { "py_modules": ["http://test.com/test0.zip", "s3://path/test1.zip"], "working_dir": "gs://path/test2.zip", "env_vars": {"test": "val", "trial": "val2"}, "pip": ["pandas", "numpy"], "excludes": ["my_file.txt"], } original_parent = parent.copy() child = { "py_modules": [], "working_dir": "s3://path/test1.zip", "env_vars": {"test": "val", "trial": "val2"}, "pip": ["numpy"], } original_child = child.copy() merged = override_runtime_envs_except_env_vars(parent, child) assert original_parent == parent assert original_child == child assert merged == { "py_modules": [], "working_dir": "s3://path/test1.zip", "env_vars": {"test": "val", "trial": "val2"}, "pip": ["numpy"], "excludes": ["my_file.txt"], } def test_merge_deep_copy(self): """Check that the env values are actually deep-copied.""" parent_env_vars = {"parent": "pval"} child_env_vars = {"child": "cval"} parent = {"env_vars": parent_env_vars} child = {"env_vars": child_env_vars} original_parent = parent.copy() original_child = child.copy() merged = override_runtime_envs_except_env_vars(parent, child) assert merged["env_vars"] == {"parent": "pval", "child": "cval"} assert original_parent == parent assert original_child == child def test_merge_empty_env_vars(self): env_vars = {"test": "val", "trial": "val2"} non_empty = {"env_vars": {"test": "val", "trial": "val2"}} empty = {} assert ( env_vars == override_runtime_envs_except_env_vars(non_empty, empty)["env_vars"] ) assert ( env_vars == override_runtime_envs_except_env_vars(empty, non_empty)["env_vars"] ) assert {} == override_runtime_envs_except_env_vars(empty, empty)["env_vars"] def test_merge_env_vars(self): parent = { "py_modules": ["http://test.com/test0.zip", "s3://path/test1.zip"], "working_dir": "gs://path/test2.zip", "env_vars": {"parent": "pval", "override": "old"}, "pip": ["pandas", "numpy"], "excludes": ["my_file.txt"], } child = { "py_modules": [], "working_dir": "s3://path/test1.zip", "env_vars": {"child": "cval", "override": "new"}, "pip": ["numpy"], } merged = override_runtime_envs_except_env_vars(parent, child) assert merged == { "py_modules": [], "working_dir": "s3://path/test1.zip", "env_vars": {"parent": "pval", "child": "cval", "override": "new"}, "pip": ["numpy"], "excludes": ["my_file.txt"], } def test_inheritance_regression(self): """Check if the general Ray runtime_env inheritance behavior matches. override_runtime_envs_except_env_vars should match the general Ray runtime_env inheritance behavior. This test checks if that behavior has changed, which would indicate a regression in override_runtime_envs_except_env_vars. If the runtime_env inheritance behavior changes, override_runtime_envs_except_env_vars should also change to match. """ with ray.init( runtime_env={ "py_modules": [TEST_DAG_PINNED_URI], "env_vars": {"var1": "hello"}, } ): @ray.remote def check_module(): # Check that Ray job's py_module loaded correctly from conditional_dag import serve_dag # noqa: F401 return os.getenv("var1") assert ray.get(check_module.remote()) == "hello" @ray.remote( runtime_env={ "py_modules": [TEST_DEPLOY_GROUP_PINNED_URI], "env_vars": {"var2": "world"}, } ) def test_task(): with pytest.raises(ImportError): # Check that Ray job's py_module was overwritten from conditional_dag import serve_dag # noqa: F401 from test_env.shallow_import import ShallowClass if ShallowClass()() == "Hello shallow world!": return os.getenv("var1") + " " + os.getenv("var2") assert ray.get(test_task.remote()) == "hello world" class TestSnakeToCamelCase: def test_empty(self): assert "" == snake_to_camel_case("") def test_single_word(self): assert snake_to_camel_case("oneword") == "oneword" def test_multiple_words(self): assert ( snake_to_camel_case("there_are_multiple_words_in_this_phrase") == "thereAreMultipleWordsInThisPhrase" ) def test_single_char_words(self): assert snake_to_camel_case("this_is_a_test") == "thisIsATest" def test_leading_capitalization(self): """If the leading character is already capitalized, leave it capitalized.""" assert snake_to_camel_case("Leading_cap") == "LeadingCap" def test_leading_alphanumeric(self): assert ( snake_to_camel_case("check_@lphanum3ric_©har_behavior") == "check@lphanum3ric©harBehavior" ) def test_embedded_capitalization(self): assert snake_to_camel_case("check_eMbEDDed_caPs") == "checkEMbEDDedCaPs" def test_mixed_caps_alphanumeric(self): assert ( snake_to_camel_case("check_3Mb3DD*d_©a!s_behAvior_Here_wIth_MIxed_cAPs") == "check3Mb3DD*d©a!sBehAviorHereWIthMIxedCAPs" ) def test_leading_underscore(self): """Should strip leading underscores.""" assert snake_to_camel_case("_leading_underscore") == "leadingUnderscore" def test_trailing_underscore(self): """Should strip trailing underscores.""" assert snake_to_camel_case("trailing_underscore_") == "trailingUnderscore" def test_leading_and_trailing_underscores(self): """Should strip leading and trailing underscores""" assert snake_to_camel_case(f"{'_' * 5}hello__world{'_' * 10}") == "helloWorld" def test_double_underscore(self): """Should treat repeated underscores as single underscore.""" assert snake_to_camel_case("double__underscore") == "doubleUnderscore" assert snake_to_camel_case(f"many{'_' * 30}underscore") == "manyUnderscore" def test_get_head_node_id(): """Test get_head_node_id() returning the correct head node id. When there are woker node, dead head node, and other alive head nodes, get_head_node_id() should return the node id of the first alive head node. When there are no alive head nodes, get_head_node_id() should raise assertion error. """ nodes = [ {"NodeID": "worker_node1", "Alive": True, "Resources": {"CPU": 1}}, { "NodeID": "dead_head_node1", "Alive": False, "Resources": {"CPU": 1, HEAD_NODE_RESOURCE_NAME: 1.0}, }, { "NodeID": "alive_head_node1", "Alive": True, "Resources": {"CPU": 1, HEAD_NODE_RESOURCE_NAME: 1.0}, }, { "NodeID": "alive_head_node2", "Alive": True, "Resources": {"CPU": 1, HEAD_NODE_RESOURCE_NAME: 1.0}, }, ] with patch("ray.nodes", return_value=nodes): assert get_head_node_id() == "alive_head_node1" with patch("ray.nodes", return_value=[]): with pytest.raises(AssertionError): get_head_node_id() def test_calculate_remaining_timeout(): # Always return `None` or negative value. assert ( calculate_remaining_timeout( timeout_s=None, start_time_s=100, curr_time_s=101, ) is None ) assert ( calculate_remaining_timeout( timeout_s=-1, start_time_s=100, curr_time_s=101, ) == -1 ) # Return delta from start. assert ( calculate_remaining_timeout( timeout_s=10, start_time_s=100, curr_time_s=101, ) == 9 ) assert ( calculate_remaining_timeout( timeout_s=100, start_time_s=100, curr_time_s=101.1, ) == 98.9 ) # Never return a negative timeout once it has elapsed. assert ( calculate_remaining_timeout( timeout_s=10, start_time_s=100, curr_time_s=111, ) == 0 ) def test_get_all_live_placement_group_names(ray_instance): """Test the logic to parse the Ray placement group table. The test contains three cases: - A named placement group that was created and is still alive ("pg2"). - A named placement group that's still being scheduled ("pg3"). - An unnamed placement group. Only "pg2" and "pg3" should be returned as live placement group names. """ # Named placement group that's been removed (should not be returned). pg1 = ray.util.placement_group([{"CPU": 0.1}], name="pg1") ray.util.remove_placement_group(pg1) # Named, detached placement group that's been removed (should not be returned). pg2 = ray.util.placement_group([{"CPU": 0.1}], name="pg2", lifetime="detached") ray.util.remove_placement_group(pg2) # Named placement group that's still alive (should be returned). pg3 = ray.util.placement_group([{"CPU": 0.1}], name="pg3") assert pg3.wait() # Named, detached placement group that's still alive (should be returned). pg4 = ray.util.placement_group([{"CPU": 0.1}], name="pg4", lifetime="detached") assert pg4.wait() # Named placement group that's being scheduled (should be returned). pg5 = ray.util.placement_group([{"CPU": 1000}], name="pg5") assert not pg5.wait(timeout_seconds=0.001) # Named, detached placement group that's being scheduled (should be returned). pg6 = ray.util.placement_group([{"CPU": 1000}], name="pg6", lifetime="detached") assert not pg6.wait(timeout_seconds=0.001) # Unnamed placement group (should not be returned). pg7 = ray.util.placement_group([{"CPU": 0.1}]) assert pg7.wait() # Unnamed, detached placement group (should not be returned). pg8 = ray.util.placement_group([{"CPU": 0.1}], lifetime="detached") assert pg8.wait() # Use wait_for_condition to allow GCS state to propagate after # remove_placement_group(). Removal is async — the scheduling state may not # have flipped to REMOVED by the time we query. def check_live_placement_group_names(): assert set(get_all_live_placement_group_names()) == { "pg3", "pg4", "pg5", "pg6", } return True wait_for_condition(check_live_placement_group_names, timeout=10) def test_get_active_placement_group_ids(ray_instance): """Test that get_active_placement_group_ids returns PG IDs for alive actors only. Cases covered: - A PG with a live Serve actor scheduled on it (should be returned). - A PG with no actors (should NOT be returned). - An actor without a PG (should NOT cause its absence to break anything). - A PG whose actor has been killed (should NOT be returned). - A PG with a live non-Serve actor (should NOT be returned). """ # PG with a live actor scheduled on it pg_with_actor = ray.util.placement_group([{"CPU": 0.1}], name="pg_with_actor") assert pg_with_actor.wait() @ray.remote(num_cpus=0.1) class DummyActor: def ready(self): return True actor_on_pg = DummyActor.options( namespace=SERVE_NAMESPACE, placement_group=pg_with_actor, ).remote() ray.get(actor_on_pg.ready.remote()) # PG with no actors scheduled on it pg_no_actor = ray.util.placement_group([{"CPU": 0.1}], name="pg_no_actor") assert pg_no_actor.wait() # Actor without any PG actor_no_pg = DummyActor.options(namespace=SERVE_NAMESPACE).remote() ray.get(actor_no_pg.ready.remote()) # PG whose actor will be killed pg_killed = ray.util.placement_group([{"CPU": 0.1}], name="pg_killed") assert pg_killed.wait() actor_killed = DummyActor.options( namespace=SERVE_NAMESPACE, placement_group=pg_killed, ).remote() ray.get(actor_killed.ready.remote()) ray.kill(actor_killed) # PG with a live actor in a non-Serve namespace (should be excluded) pg_non_serve = ray.util.placement_group([{"CPU": 0.1}], name="pg_non_serve") assert pg_non_serve.wait() actor_non_serve = DummyActor.options( namespace="non_serve", placement_group=pg_non_serve, ).remote() ray.get(actor_non_serve.ready.remote()) pg_with_actor_id = pg_with_actor.id.hex() pg_no_actor_id = pg_no_actor.id.hex() pg_killed_id = pg_killed.id.hex() pg_non_serve_id = pg_non_serve.id.hex() # Use wait_for_condition to allow GCS state to propagate after ray.kill(). # ray.kill() is async — the actor's ALIVE -> DEAD transition in GCS may lag. def check_active_placement_group_ids(): active_ids = get_active_placement_group_ids() assert pg_with_actor_id in active_ids assert pg_no_actor_id not in active_ids assert pg_killed_id not in active_ids assert pg_non_serve_id not in active_ids return True wait_for_condition(check_active_placement_group_ids, timeout=10) def test_is_running_in_asyncio_loop_false(): assert is_running_in_asyncio_loop() is False @pytest.mark.asyncio async def test_is_running_in_asyncio_loop_true(): assert is_running_in_asyncio_loop() is True async def check(): return is_running_in_asyncio_loop() # Verify that it also works in a task. assert await asyncio.ensure_future(check()) is True def test_get_current_actor_id(ray_instance): @ray.remote class A: def call_get_current_actor_id(self): return get_current_actor_id() a = A.remote() actor_id = ray.get(a.call_get_current_actor_id.remote()) assert len(actor_id) > 0 assert actor_id != "DRIVER" assert get_current_actor_id() == "DRIVER" @pytest.mark.asyncio async def test_semaphore(): """Test core Semaphore functionality.""" max_value = 2 sema = Semaphore(get_value_fn=lambda: max_value) # Test get_max_value functionality assert sema.get_max_value() == max_value # Initially, semaphore should not be locked and should allow acquisitions assert not sema.locked() # Acquire one await sema.acquire() assert not sema.locked() assert sema._value == 1 # Acquire one await sema.acquire() assert sema.locked() # Should now be locked (2 out of 2) assert sema._value == 2 # Release one sema.release() assert not sema.locked() # Should not be locked anymore (1 out of 2) assert sema._value == 1 # Acquire one await sema.acquire() assert sema.locked() assert sema._value == 2 @pytest.mark.asyncio async def test_semaphore_waiters_and_single_release(): """Test that release() wakes up exactly one waiter.""" max_value = 1 sema = Semaphore(get_value_fn=lambda: max_value) # Fill the semaphore to capacity await sema.acquire() assert sema.locked() assert sema._value == 1 # Create multiple waiters waiters_completed = [] async def waiter(waiter_id): await sema.acquire() waiters_completed.append(waiter_id) # Start 2 waiters that will all block waiter_tasks = [ asyncio.create_task(waiter(1)), asyncio.create_task(waiter(2)), ] # Yield the event loop await asyncio.sleep(0.01) # Verify they are all waiting assert len(waiters_completed) == 0 assert sema.locked() assert len(sema._waiters) == 2 # Release once - this should wake up exactly ONE waiter sema.release() await asyncio.sleep(0.01) # Verify exactly one waiter was woken up and completed assert len(waiters_completed) == 1 assert sema._value == 1 assert sema.locked() assert len(sema._waiters) == 1 # Release again - should wake up exactly one more waiter sema.release() await asyncio.sleep(0.01) # Verify exactly one more waiter was woken up assert len(waiters_completed) == 2 assert sema._value == 1 assert sema.locked() assert len(sema._waiters) == 0 assert len(await asyncio.gather(*waiter_tasks)) == 2 @pytest.mark.asyncio async def test_semaphore_dynamic_max_value(): """Test that Semaphore respects dynamic changes to max_value.""" current_max = 2 def get_dynamic_max(): return current_max sema = Semaphore(get_value_fn=get_dynamic_max) # Initially max is 2 assert sema.get_max_value() == 2 # Acquire up to the limit await sema.acquire() await sema.acquire() assert sema.locked() # Increase the max value dynamically current_max = 3 assert sema.get_max_value() == 3 assert not sema.locked() # Should be able to acquire one more await sema.acquire() assert sema.locked() assert sema._value == 3 # Decrease the max value current_max = 1 assert sema.get_max_value() == 1 assert sema.locked() # Release to get back within limits sema.release() sema.release() assert sema.locked() sema.release() assert not sema.locked() assert sema._value == 0 if __name__ == "__main__": import sys sys.exit(pytest.main(["-v", "-s", __file__]))