chore: import upstream snapshot with attribution

python/ray/serve/tests/test_util.py

@@ -0,0 +1,702 @@
+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__]))