chore: import upstream snapshot with attribution

python/ray/serve/tests/conftest.py

@@ -0,0 +1,458 @@
+import os
+import random
+import socket
+import subprocess
+import tempfile
+from contextlib import contextmanager
+from copy import deepcopy
+from typing import Any, Dict, Generator
+
+import httpx
+import pytest
+import pytest_asyncio
+
+import ray
+from ray import serve
+from ray._common.test_utils import SignalActor, wait_for_condition
+from ray._common.usage import usage_lib
+from ray._common.utils import reset_ray_address
+from ray.cluster_utils import AutoscalingCluster, Cluster
+from ray.serve._private.test_utils import (
+    TELEMETRY_ROUTE_PREFIX,
+    TEST_METRICS_EXPORT_PORT,
+    check_ray_started,
+    check_ray_stopped,
+    start_telemetry_app,
+)
+from ray.serve.config import HTTPOptions, ProxyLocation, gRPCOptions
+from ray.serve.context import _get_global_client
+from ray.tests.conftest import (  # noqa
+    external_redis,
+    propagate_logs,
+    pytest_runtest_makereport,
+)
+
+# https://tools.ietf.org/html/rfc6335#section-6
+MIN_DYNAMIC_PORT = 49152
+MAX_DYNAMIC_PORT = 65535
+
+TEST_GRPC_SERVICER_FUNCTIONS = [
+    "ray.serve.generated.serve_pb2_grpc.add_UserDefinedServiceServicer_to_server",
+    "ray.serve.generated.serve_pb2_grpc.add_FruitServiceServicer_to_server",
+]
+
+if os.environ.get("RAY_SERVE_INTENTIONALLY_CRASH", False) == 1:
+    serve.controller._CRASH_AFTER_CHECKPOINT_PROBABILITY = 0.5
+
+
+@pytest.fixture(autouse=True)
+def _clear_stale_ray_address():
+    # Serve CI runs several test targets per container sharing /tmp/ray; a target
+    # killed mid-run can leave a ray_current_cluster pointing at a dead cluster.
+    # Drop it before each test so an address-less ray.init() starts fresh.
+    reset_ray_address()
+    yield
+
+
+@pytest.fixture
+def ray_shutdown():
+    serve.shutdown()
+    if ray.is_initialized():
+        ray.shutdown()
+    yield
+    serve.shutdown()
+    if ray.is_initialized():
+        ray.shutdown()
+
+
+@pytest.fixture
+def ray_cluster():
+    cluster = Cluster()
+    yield cluster
+    serve.shutdown()
+    ray.shutdown()
+    cluster.shutdown()
+
+
+@pytest.fixture
+def ray_autoscaling_cluster(request):
+    # NOTE(zcin): We have to make a deepcopy here because AutoscalingCluster
+    # modifies the dictionary that's passed in.
+    params = deepcopy(request.param)
+    cluster = AutoscalingCluster(**params)
+    cluster.start()
+    yield
+    serve.shutdown()
+    ray.shutdown()
+    cluster.shutdown()
+
+
+@pytest.fixture
+def ray_start(scope="module"):
+    port = random.randint(MIN_DYNAMIC_PORT, MAX_DYNAMIC_PORT)
+    subprocess.check_output(
+        [
+            "ray",
+            "start",
+            "--head",
+            "--num-cpus",
+            "16",
+            "--ray-client-server-port",
+            f"{port}",
+        ]
+    )
+    try:
+        yield f"localhost:{port}"
+    finally:
+        subprocess.check_output(["ray", "stop", "--force"])
+
+
+def _check_ray_stop():
+    try:
+        httpx.get("http://localhost:8265/api/ray/version")
+        return False
+    except Exception:
+        return True
+
+
+@contextmanager
+def start_and_shutdown_ray_cli():
+    subprocess.check_output(["ray", "stop", "--force"])
+    wait_for_condition(_check_ray_stop, timeout=15)
+    subprocess.check_output(["ray", "start", "--head"])
+
+    yield
+
+    subprocess.check_output(["ray", "stop", "--force"])
+    wait_for_condition(_check_ray_stop, timeout=15)
+
+
+@pytest.fixture(scope="module")
+def start_and_shutdown_ray_cli_module():
+    with start_and_shutdown_ray_cli():
+        yield
+
+
+@pytest.fixture
+def tmp_dir():
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        old_dir = os.getcwd()
+        os.chdir(tmp_dir)
+        yield tmp_dir
+        os.chdir(old_dir)
+
+
+@pytest.fixture(scope="session")
+def _shared_serve_instance():
+    # Note(simon):
+    # This line should be not turned on on master because it leads to very
+    # spammy and not useful log in case of a failure in CI.
+    # To run locally, please use this instead.
+    # SERVE_DEBUG_LOG=1 pytest -v -s test_api.py
+    # os.environ["SERVE_DEBUG_LOG"] = "1" <- Do not uncomment this.
+
+    # Overriding task_retry_delay_ms to relaunch actors more quickly
+    ray.init(
+        address="local",
+        num_cpus=36,
+        namespace="default_test_namespace",
+        _metrics_export_port=9999,
+        _system_config={"metrics_report_interval_ms": 1000, "task_retry_delay_ms": 50},
+    )
+    serve.start(
+        proxy_location=ProxyLocation.HeadOnly,
+        http_options={"host": "0.0.0.0"},
+        grpc_options={
+            "port": 9000,
+            "grpc_servicer_functions": TEST_GRPC_SERVICER_FUNCTIONS,
+        },
+    )
+    yield _get_global_client()
+    # Shutdown Serve and Ray when the session ends so that proxy actors
+    # (e.g. HAProxyManager) run their shutdown logic and stop subprocesses.
+    serve.shutdown()
+
+
+@pytest_asyncio.fixture
+async def serve_instance_async(_shared_serve_instance):
+    yield _shared_serve_instance
+    # Clear all state for 2.x applications and deployments.
+    _shared_serve_instance.delete_all_apps()
+    # Clear the ServeHandle cache between tests to avoid them piling up.
+    await _shared_serve_instance.shutdown_cached_handles_async()
+
+
+@pytest.fixture
+def serve_instance(_shared_serve_instance):
+    yield _shared_serve_instance
+    # Clear all state for 2.x applications and deployments.
+    _shared_serve_instance.delete_all_apps()
+    # Clear the ServeHandle cache between tests to avoid them piling up.
+    _shared_serve_instance.shutdown_cached_handles()
+
+
+@pytest.fixture
+def serve_instance_with_signal(serve_instance):
+    client = serve_instance
+
+    signal = SignalActor.options(name="signal123").remote()
+    yield client, signal
+
+    # Delete signal actor so there is no conflict between tests
+    ray.kill(signal)
+
+
+def check_ray_stop():
+    try:
+        httpx.get("http://localhost:8265/api/ray/version")
+        return False
+    except Exception:
+        return True
+
+
+@pytest.fixture(scope="function")
+def ray_start_stop():
+    subprocess.check_output(["ray", "stop", "--force"])
+    ray.shutdown()
+    wait_for_condition(
+        check_ray_stop,
+        timeout=15,
+    )
+    subprocess.check_output(["ray", "start", "--head"])
+    wait_for_condition(
+        lambda: httpx.get("http://localhost:8265/api/ray/version").status_code == 200,
+        timeout=15,
+    )
+    ray.init("auto")
+    yield
+    serve.shutdown()
+    ray.shutdown()
+    subprocess.check_output(["ray", "stop", "--force"])
+    wait_for_condition(
+        check_ray_stop,
+        timeout=15,
+    )
+
+
+@pytest.fixture(scope="function")
+def ray_start_stop_in_specific_directory(request):
+    original_working_dir = os.getcwd()
+
+    # Change working directory so Ray will start in the requested directory.
+    new_working_dir = request.param
+    os.chdir(new_working_dir)
+    print(f"\nChanged working directory to {new_working_dir}\n")
+
+    subprocess.check_output(["ray", "start", "--head"])
+    wait_for_condition(
+        lambda: httpx.get("http://localhost:8265/api/ray/version").status_code == 200,
+        timeout=15,
+    )
+    try:
+        yield
+    finally:
+        # Change the directory back to the original one.
+        os.chdir(original_working_dir)
+        print(f"\nChanged working directory back to {original_working_dir}\n")
+
+        subprocess.check_output(["ray", "stop", "--force"])
+        wait_for_condition(
+            check_ray_stop,
+            timeout=15,
+        )
+
+
+@pytest.fixture
+def ray_instance(
+    request: pytest.FixtureRequest,
+) -> Generator[Dict[str, Any], None, None]:
+    """Starts and stops a Ray instance for this test.
+
+    Args:
+        request: request.param should contain a dictionary of env vars and
+            their values. The Ray instance will be started with these env vars.
+
+    Yields:
+        Dict[str, Any]: The dict returned by ``ray.init`` for the started cluster.
+    """
+
+    original_env_vars = os.environ.copy()
+
+    try:
+        requested_env_vars = request.param
+    except AttributeError:
+        requested_env_vars = {}
+
+    os.environ.update(requested_env_vars)
+    yield ray.init(
+        address="local",
+        _metrics_export_port=9999,
+        _system_config={
+            "metrics_report_interval_ms": 1000,
+            "task_retry_delay_ms": 50,
+        },
+    )
+
+    serve.shutdown()
+    ray.shutdown()
+
+    os.environ.clear()
+    os.environ.update(original_env_vars)
+
+
+@pytest.fixture
+def manage_ray_with_telemetry(monkeypatch):
+    with monkeypatch.context() as m:
+        m.setenv("RAY_USAGE_STATS_ENABLED", "1")
+        m.setenv(
+            "RAY_USAGE_STATS_REPORT_URL",
+            f"http://127.0.0.1:8000{TELEMETRY_ROUTE_PREFIX}",
+        )
+        m.setenv("RAY_USAGE_STATS_REPORT_INTERVAL_S", "1")
+        subprocess.check_output(["ray", "stop", "--force"])
+        wait_for_condition(check_ray_stopped, timeout=5)
+
+        subprocess.check_output(["ray", "start", "--head"])
+        wait_for_condition(check_ray_started, timeout=5)
+
+        storage = start_telemetry_app()
+        wait_for_condition(
+            lambda: ray.get(storage.get_reports_received.remote()) > 1, timeout=15
+        )
+
+        yield storage
+
+        # Call Python API shutdown() methods to clear global variable state
+        serve.shutdown()
+        ray.shutdown()
+
+        # Reset global state (any keys that may have been set and cached while the
+        # workload was running).
+        usage_lib.reset_global_state()
+
+        # Shut down Ray cluster with CLI
+        subprocess.check_output(["ray", "stop", "--force"])
+        wait_for_condition(check_ray_stopped, timeout=5)
+
+
+def wait_for_metrics_port_free(port=TEST_METRICS_EXPORT_PORT, timeout=30):
+    """
+    Ensures the metrics export port is freed.
+    """
+
+    def port_free():
+        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+        try:
+            s.bind(("", port))
+            return True
+        except OSError:
+            return False
+        finally:
+            s.close()
+
+    wait_for_condition(port_free, timeout=timeout, retry_interval_ms=200)
+
+
+def wait_for_metrics_endpoint(session_name, port=TEST_METRICS_EXPORT_PORT, timeout=30):
+    """
+    Ensures the current dashboard agent is serving the metrics endpoint. A
+    timeout indicates another agent is still running and holding the port.
+    """
+
+    def ready():
+        try:
+            resp = httpx.get(f"http://localhost:{port}/metrics", timeout=1.0)
+        except Exception:
+            return False
+        return resp.status_code == 200 and f'SessionName="{session_name}"' in resp.text
+
+    wait_for_condition(ready, timeout=timeout, retry_interval_ms=500)
+
+
+@pytest.fixture
+def metrics_start_shutdown(request):
+    param = request.param if hasattr(request, "param") else None
+    request_timeout_s = param if param else None
+    """Fixture provides a fresh Ray cluster to prevent metrics state sharing."""
+    wait_for_metrics_port_free()
+    ray.init(
+        address="local",
+        _metrics_export_port=TEST_METRICS_EXPORT_PORT,
+        _system_config={
+            "metrics_report_interval_ms": 100,
+            "task_retry_delay_ms": 50,
+        },
+    )
+
+    try:
+        session_name = ray._private.worker._global_node.session_name
+        wait_for_metrics_endpoint(session_name)
+
+        grpc_port = 9000
+        grpc_servicer_functions = [
+            "ray.serve.generated.serve_pb2_grpc.add_UserDefinedServiceServicer_to_server",
+            "ray.serve.generated.serve_pb2_grpc.add_FruitServiceServicer_to_server",
+        ]
+        yield serve.start(
+            grpc_options=gRPCOptions(
+                port=grpc_port,
+                grpc_servicer_functions=grpc_servicer_functions,
+                request_timeout_s=request_timeout_s,
+            ),
+            http_options=HTTPOptions(
+                host="0.0.0.0",
+                request_timeout_s=request_timeout_s,
+            ),
+        )
+    finally:
+        serve.shutdown()
+        ray.shutdown()
+        reset_ray_address()
+
+
+# Helper function to return the node ID of a remote worker.
+@ray.remote(num_cpus=0)
+def _get_node_id():
+    return ray.get_runtime_context().get_node_id()
+
+
+# Test fixture to start a Serve instance in a RayCluster with two labeled nodes
+@pytest.fixture(scope="module")
+def serve_instance_with_labeled_nodes():
+    cluster = Cluster()
+
+    # Unlabeled default node.
+    cluster.add_node(num_cpus=3, resources={"worker0": 1})
+
+    # Node 1 - labeled A100 node in us-west.
+    cluster.add_node(
+        num_cpus=3,
+        resources={"worker1": 1},
+        labels={"region": "us-west", "gpu-type": "A100"},
+    )
+
+    # Node 2 - labeled H100 node in us-east.
+    cluster.add_node(
+        num_cpus=3,
+        resources={"worker2": 1},
+        labels={"region": "us-east", "gpu-type": "H100"},
+    )
+
+    cluster.wait_for_nodes()
+
+    if ray.is_initialized():
+        ray.shutdown()
+
+    ray.init(address=cluster.address)
+
+    node_1_id = ray.get(_get_node_id.options(resources={"worker1": 1}).remote())
+    node_2_id = ray.get(_get_node_id.options(resources={"worker2": 1}).remote())
+
+    serve.start()
+
+    yield _get_global_client(), node_1_id, node_2_id, cluster
+
+    serve.shutdown()
+    ray.shutdown()
+    cluster.shutdown()