Files
2026-07-13 13:17:40 +08:00

703 lines
21 KiB
Python

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__]))