chore: import upstream snapshot with attribution

python/ray/tests/chaos/streaming_llm.py

@@ -0,0 +1,98 @@
+import argparse
+import asyncio
+import logging
+
+import requests
+from fastapi import FastAPI
+from starlette.responses import StreamingResponse
+
+import ray
+from ray import serve
+
+logger = logging.getLogger("ray.serve")
+
+fastapi_app = FastAPI()
+
+
+# Input: a prompt of words
+# Output: each word reversed and produced N times.
+@serve.deployment(
+    num_replicas=6, ray_actor_options={"num_cpus": 0.01, "memory": 10 * 1024 * 1024}
+)
+class ReverseAndDupEachWord:
+    def __init__(self, dup_times: int):
+        self.dup_times = dup_times
+
+    async def __call__(self, prompt: str):
+        for word in prompt.split():
+            rev = word[::-1]
+            for _ in range(self.dup_times):
+                await asyncio.sleep(0.001)
+                # Ideally we want to do " ".join(words), but for the sake of
+                # simplicity we also have an extra trailing space.
+                yield rev + " "
+
+
+@serve.deployment(
+    num_replicas=6, ray_actor_options={"num_cpus": 0.01, "memory": 10 * 1024 * 1024}
+)
+@serve.ingress(fastapi_app)
+class Textbot:
+    def __init__(self, llm):
+        self.llm = llm.options(stream=True)
+
+    @fastapi_app.post("/")
+    async def handle_request(self, prompt: str) -> StreamingResponse:
+        logger.info(f'Got prompt with size "{len(prompt)}"')
+        return StreamingResponse(self.llm.remote(prompt), media_type="text/plain")
+
+
+@ray.remote(num_cpus=0.1, memory=10 * 1024 * 1024)
+def make_http_query(num_words, num_queries):
+    for _ in range(num_queries):
+        words = "Lorem ipsum dolor sit amet".split()
+        prompt_words = [words[i % len(words)] for i in range(num_words)]
+        prompt = " ".join(prompt_words)
+        expected_words = [word[::-1] for word in prompt_words for _ in range(2)]
+
+        response = requests.post(f"http://localhost:8000/?prompt={prompt}", stream=True)
+        response.raise_for_status()
+        content = response.content.decode()
+        assert content == " ".join(expected_words) + " ", content
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Generates HTTP workloads with Ray.")
+
+    parser.add_argument("--num_tasks", type=int, required=True, help="Number of tasks.")
+    parser.add_argument(
+        "--num_queries_per_task",
+        type=int,
+        required=True,
+        help="Number of queries per task.",
+    )
+    parser.add_argument(
+        "--num_words_per_query",
+        type=int,
+        required=True,
+        help="Number of words per query",
+    )
+
+    args = parser.parse_args()
+
+    # Run the serve, run the client, then showdown serve.
+    llm = ReverseAndDupEachWord.bind(2)
+    app = Textbot.bind(llm)
+
+    serve.run(app)
+
+    objs = [
+        make_http_query.remote(args.num_words_per_query, args.num_queries_per_task)
+        for _ in range(args.num_tasks)
+    ]
+    ray.get(objs)
+
+    serve.shutdown()
+
+
+main()