chore: import upstream snapshot with attribution

python/ray/serve/tests/test_task_consumer_autoscaling.py

@@ -0,0 +1,222 @@
+import os
+import sys
+
+import pytest
+
+import ray
+from ray import serve
+from ray._common.test_utils import SignalActor, wait_for_condition
+from ray.serve._private.common import DeploymentID, ReplicaState
+from ray.serve.config import AutoscalingConfig, AutoscalingPolicy
+from ray.serve.schema import CeleryAdapterConfig, TaskProcessorConfig
+from ray.serve.task_consumer import (
+    instantiate_adapter_from_config,
+    task_consumer,
+    task_handler,
+)
+from ray.tests.conftest import external_redis  # noqa: F401
+
+
+@ray.remote
+def enqueue_task(processor_config: TaskProcessorConfig, data, task_name="process"):
+    adapter = instantiate_adapter_from_config(task_processor_config=processor_config)
+    result = adapter.enqueue_task_sync(task_name, args=[data])
+    assert result.id is not None
+    return result.id
+
+
+def get_num_running_replicas(controller, deployment_name, app_name):
+    """Get the number of running replicas for a deployment."""
+    deployment_id = DeploymentID(name=deployment_name, app_name=app_name)
+    replicas = ray.get(
+        controller._dump_replica_states_for_testing.remote(deployment_id)
+    )
+    return len(replicas.get([ReplicaState.RUNNING]))
+
+
+@pytest.mark.skipif(sys.platform == "win32", reason="Flaky on Windows.")
+class TestTaskConsumerQueueAutoscaling:
+    """Test queue-based autoscaling for TaskConsumer deployments."""
+
+    def test_task_consumer_queue_autoscaling(
+        self, external_redis, serve_instance  # noqa: F811
+    ):
+        """Test that TaskConsumer deployments autoscale based on queue length.
+
+        Verifies the full e2e flow:
+        1. Replicas scale up when messages pile up in the queue
+        2. Replicas scale down when the queue drains
+        """
+        redis_address = os.environ.get("RAY_REDIS_ADDRESS")
+        app_name = "autoscaling_app"
+        deployment_name = "AutoscalingConsumer"
+
+        processor_config = TaskProcessorConfig(
+            queue_name="autoscaling_test_queue",
+            adapter_config=CeleryAdapterConfig(
+                broker_url=f"redis://{redis_address}/0",
+                backend_url=f"redis://{redis_address}/1",
+            ),
+        )
+
+        signal = SignalActor.remote()
+
+        @serve.deployment(
+            name=deployment_name,
+            max_ongoing_requests=1,
+            autoscaling_config=AutoscalingConfig(
+                min_replicas=1,
+                max_replicas=5,
+                target_ongoing_requests=1,
+                upscale_delay_s=0,
+                downscale_delay_s=0,
+                metrics_interval_s=0.1,
+                look_back_period_s=0.5,
+                policy=AutoscalingPolicy(
+                    policy_function="ray.serve.async_inference_autoscaling_policy:AsyncInferenceAutoscalingPolicy",
+                    policy_kwargs={
+                        "broker_url": f"redis://{redis_address}/0",
+                        "queue_name": "autoscaling_test_queue",
+                    },
+                ),
+            ),
+        )
+        @task_consumer(task_processor_config=processor_config)
+        class AutoscalingConsumer:
+            def __init__(self, signal_actor):
+                self._signal = signal_actor
+
+            @task_handler(name="process")
+            def process(self, data):
+                ray.get(self._signal.wait.remote())
+
+        _ = serve.run(
+            AutoscalingConsumer.bind(signal),
+            name=app_name,
+            route_prefix="/autoscaling",
+        )
+
+        controller = serve_instance._controller
+
+        # Wait for initial replica to be running
+        wait_for_condition(
+            lambda: get_num_running_replicas(controller, deployment_name, app_name)
+            == 1,
+            timeout=30,
+        )
+
+        # Enqueue tasks to build up the queue (signal blocks processing)
+        num_tasks = 10
+        for i in range(num_tasks):
+            enqueue_task.remote(processor_config, f"data_{i}")
+
+        # Wait for replicas to scale up to max_replicas
+        wait_for_condition(
+            lambda: get_num_running_replicas(controller, deployment_name, app_name)
+            == 5,
+            timeout=60,
+        )
+
+        # Release the signal to let all tasks drain
+        ray.get(signal.send.remote())
+
+        # Wait for replicas to scale back down to min_replicas
+        wait_for_condition(
+            lambda: get_num_running_replicas(controller, deployment_name, app_name)
+            == 1,
+            timeout=60,
+        )
+
+        serve.delete(app_name)
+
+    def test_task_consumer_scale_from_and_to_zero(
+        self, external_redis, serve_instance  # noqa: F811
+    ):
+        """Test that TaskConsumer deployments can scale down to zero.
+
+        Verifies:
+        1. Replicas scale up when messages pile up in the queue
+        2. Replicas scale down to 0 when the queue drains
+        """
+        redis_address = os.environ.get("RAY_REDIS_ADDRESS")
+        app_name = "scale_to_zero_app"
+        deployment_name = "ScaleToZeroConsumer"
+
+        processor_config = TaskProcessorConfig(
+            queue_name="scale_to_zero_queue",
+            adapter_config=CeleryAdapterConfig(
+                broker_url=f"redis://{redis_address}/0",
+                backend_url=f"redis://{redis_address}/1",
+            ),
+        )
+
+        signal = SignalActor.remote()
+
+        @serve.deployment(
+            name=deployment_name,
+            max_ongoing_requests=1,
+            autoscaling_config=AutoscalingConfig(
+                min_replicas=0,
+                max_replicas=5,
+                target_ongoing_requests=1,
+                upscale_delay_s=0,
+                downscale_delay_s=0,
+                downscale_to_zero_delay_s=5,
+                metrics_interval_s=0.1,
+                look_back_period_s=0.5,
+                policy=AutoscalingPolicy(
+                    policy_function="ray.serve.async_inference_autoscaling_policy:AsyncInferenceAutoscalingPolicy",
+                    policy_kwargs={
+                        "broker_url": f"redis://{redis_address}/0",
+                        "queue_name": "scale_to_zero_queue",
+                        "poll_interval_s": 1,
+                    },
+                ),
+            ),
+        )
+        @task_consumer(task_processor_config=processor_config)
+        class ScaleToZeroConsumer:
+            def __init__(self, signal_actor):
+                self._signal = signal_actor
+
+            @task_handler(name="process")
+            def process(self, data):
+                ray.get(self._signal.wait.remote())
+
+        _ = serve.run(
+            ScaleToZeroConsumer.bind(signal),
+            name=app_name,
+            route_prefix="/scale_to_zero",
+        )
+
+        controller = serve_instance._controller
+
+        wait_for_condition(
+            lambda: get_num_running_replicas(controller, deployment_name, app_name)
+            == 0,
+            timeout=60,
+        )
+
+        enqueue_task.remote(processor_config, "data_0")
+
+        wait_for_condition(
+            lambda: get_num_running_replicas(controller, deployment_name, app_name)
+            == 1,
+            timeout=60,
+        )
+
+        # Release the signal to let all tasks drain
+        ray.get(signal.send.remote())
+
+        # Wait for replicas to scale down to 0
+        wait_for_condition(
+            lambda: get_num_running_replicas(controller, deployment_name, app_name)
+            == 0,
+            timeout=60,
+        )
+
+        serve.delete(app_name)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", "-s", __file__]))