436 lines
13 KiB
Python
436 lines
13 KiB
Python
# flake8: noqa
|
|
"""
|
|
Code examples for the asyncio best practices guide.
|
|
All examples are structured to be runnable and demonstrate key concepts.
|
|
"""
|
|
|
|
# __imports_begin__
|
|
from ray import serve
|
|
import asyncio
|
|
# __imports_end__
|
|
|
|
|
|
# __echo_async_begin__
|
|
@serve.deployment
|
|
class Echo:
|
|
async def __call__(self, request):
|
|
await asyncio.sleep(0.1)
|
|
return "ok"
|
|
# __echo_async_end__
|
|
|
|
|
|
# __blocking_echo_begin__
|
|
@serve.deployment
|
|
class BlockingEcho:
|
|
def __call__(self, request):
|
|
# Blocking.
|
|
import time
|
|
time.sleep(1)
|
|
return "ok"
|
|
# __blocking_echo_end__
|
|
|
|
|
|
# __fastapi_deployment_begin__
|
|
from fastapi import FastAPI
|
|
|
|
app = FastAPI()
|
|
|
|
@serve.deployment
|
|
@serve.ingress(app)
|
|
class FastAPIDeployment:
|
|
@app.get("/sync")
|
|
def sync_endpoint(self):
|
|
# FastAPI runs this in a threadpool.
|
|
import time
|
|
time.sleep(1)
|
|
return "ok"
|
|
|
|
@app.get("/async")
|
|
async def async_endpoint(self):
|
|
# Runs directly on FastAPI's asyncio loop.
|
|
await asyncio.sleep(1)
|
|
return "ok"
|
|
# __fastapi_deployment_end__
|
|
|
|
|
|
# __blocking_http_begin__
|
|
@serve.deployment
|
|
class BlockingHTTP:
|
|
async def __call__(self, request):
|
|
# ❌ This blocks the event loop until the HTTP call finishes.
|
|
import requests
|
|
resp = requests.get("https://example.com/")
|
|
return resp.text
|
|
# __blocking_http_end__
|
|
|
|
|
|
# __async_http_begin__
|
|
@serve.deployment
|
|
class AsyncHTTP:
|
|
async def __call__(self, request):
|
|
import httpx
|
|
|
|
async with httpx.AsyncClient() as client:
|
|
resp = await client.get("https://example.com/")
|
|
return resp.text
|
|
# __async_http_end__
|
|
|
|
|
|
# __threaded_http_begin__
|
|
@serve.deployment
|
|
class ThreadedHTTP:
|
|
async def __call__(self, request):
|
|
import requests
|
|
|
|
def fetch():
|
|
return requests.get("https://example.com/").text
|
|
|
|
# ✅ Offload blocking I/O to a worker thread.
|
|
return await asyncio.to_thread(fetch)
|
|
# __threaded_http_end__
|
|
|
|
# __threadpool_override_begin__
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
|
|
@serve.deployment
|
|
class CustomThreadPool:
|
|
def __init__(self):
|
|
loop = asyncio.get_running_loop()
|
|
loop.set_default_executor(ThreadPoolExecutor(max_workers=16))
|
|
|
|
async def __call__(self, request):
|
|
return await asyncio.to_thread(lambda: "ok")
|
|
# __threadpool_override_end__
|
|
|
|
|
|
# __numpy_deployment_begin__
|
|
@serve.deployment
|
|
class NumpyDeployment:
|
|
def _heavy_numpy(self, array):
|
|
import numpy as np
|
|
# Many NumPy ops release the GIL while executing C/Fortran code.
|
|
return np.linalg.svd(array)[0]
|
|
|
|
async def __call__(self, request):
|
|
import numpy as np
|
|
# Create a sample array from request data
|
|
array = np.random.rand(100, 100)
|
|
# ✅ Multiple threads can run _heavy_numpy in parallel if
|
|
# the underlying implementation releases the GIL.
|
|
return await asyncio.to_thread(self._heavy_numpy, array)
|
|
# __numpy_deployment_end__
|
|
|
|
|
|
# __max_ongoing_requests_begin__
|
|
@serve.deployment(max_ongoing_requests=32)
|
|
class MyService:
|
|
async def __call__(self, request):
|
|
await asyncio.sleep(1)
|
|
return "ok"
|
|
# __max_ongoing_requests_end__
|
|
|
|
|
|
# __async_io_bound_begin__
|
|
@serve.deployment(max_ongoing_requests=100)
|
|
class AsyncIOBound:
|
|
async def __call__(self, request):
|
|
# Mostly waiting on an external system.
|
|
await asyncio.sleep(0.1)
|
|
return "ok"
|
|
# __async_io_bound_end__
|
|
|
|
|
|
# __blocking_cpu_begin__
|
|
@serve.deployment(max_ongoing_requests=100)
|
|
class BlockingCPU:
|
|
def __call__(self, request):
|
|
# ❌ Blocks the user event loop.
|
|
import time
|
|
time.sleep(1)
|
|
return "ok"
|
|
# __blocking_cpu_end__
|
|
|
|
|
|
# __cpu_with_threadpool_begin__
|
|
@serve.deployment(max_ongoing_requests=100)
|
|
class CPUWithThreadpool:
|
|
def __call__(self, request):
|
|
# With RAY_SERVE_RUN_SYNC_IN_THREADPOOL=1, each call runs in a thread.
|
|
import time
|
|
time.sleep(1)
|
|
return "ok"
|
|
# __cpu_with_threadpool_end__
|
|
|
|
|
|
# __batched_model_begin__
|
|
|
|
@serve.deployment(max_ongoing_requests=64)
|
|
class BatchedModel:
|
|
@serve.batch(max_batch_size=32)
|
|
async def __call__(self, requests):
|
|
# requests is a list of request objects.
|
|
inputs = [r for r in requests]
|
|
outputs = await self._run_model(inputs)
|
|
return outputs
|
|
|
|
async def _run_model(self, inputs):
|
|
# Placeholder model function
|
|
return [f"result_{i}" for i in inputs]
|
|
# __batched_model_end__
|
|
|
|
|
|
# __batched_model_offload_begin__
|
|
@serve.deployment(max_ongoing_requests=64)
|
|
class BatchedModelOffload:
|
|
@serve.batch(max_batch_size=32)
|
|
async def __call__(self, requests):
|
|
# requests is a list of request objects.
|
|
inputs = [r for r in requests]
|
|
outputs = await self._run_model(inputs)
|
|
return outputs
|
|
|
|
async def _run_model(self, inputs):
|
|
def run_sync():
|
|
# Heavy CPU or GIL-releasing native code here.
|
|
# Placeholder model function
|
|
return [f"result_{i}" for i in inputs]
|
|
|
|
loop = asyncio.get_running_loop()
|
|
return await loop.run_in_executor(None, run_sync)
|
|
# __batched_model_offload_end__
|
|
|
|
|
|
# __blocking_stream_begin__
|
|
@serve.deployment
|
|
class BlockingStream:
|
|
def __call__(self, request):
|
|
# ❌ Blocks the event loop between yields.
|
|
import time
|
|
for i in range(10):
|
|
time.sleep(1)
|
|
yield f"{i}\n"
|
|
# __blocking_stream_end__
|
|
|
|
|
|
# __async_stream_begin__
|
|
@serve.deployment
|
|
class AsyncStream:
|
|
async def __call__(self, request):
|
|
# ✅ Yields items without blocking the loop.
|
|
async def generator():
|
|
for i in range(10):
|
|
await asyncio.sleep(1)
|
|
yield f"{i}\n"
|
|
|
|
return generator()
|
|
# __async_stream_end__
|
|
|
|
|
|
# __offload_io_begin__
|
|
@serve.deployment
|
|
class OffloadIO:
|
|
async def __call__(self, request):
|
|
import requests
|
|
|
|
def fetch():
|
|
return requests.get("https://example.com/").text
|
|
|
|
# Offload to a thread, free the event loop.
|
|
body = await asyncio.to_thread(fetch)
|
|
return body
|
|
# __offload_io_end__
|
|
|
|
|
|
# __offload_cpu_begin__
|
|
@serve.deployment
|
|
class OffloadCPU:
|
|
def _compute(self, x):
|
|
# CPU-intensive work.
|
|
total = 0
|
|
for i in range(10_000_000):
|
|
total += (i * x) % 7
|
|
return total
|
|
|
|
async def __call__(self, request):
|
|
x = 123
|
|
loop = asyncio.get_running_loop()
|
|
result = await loop.run_in_executor(None, self._compute, x)
|
|
return str(result)
|
|
# __offload_cpu_end__
|
|
|
|
|
|
# __ray_parallel_begin__
|
|
import ray
|
|
|
|
@ray.remote
|
|
def heavy_task(x):
|
|
# Heavy compute runs in its own worker process.
|
|
return x * x
|
|
|
|
|
|
@serve.deployment
|
|
class RayParallel:
|
|
async def __call__(self, request):
|
|
values = [1, 2, 3, 4]
|
|
refs = [heavy_task.remote(v) for v in values]
|
|
results = await asyncio.gather(*[r for r in refs])
|
|
return {"results": results}
|
|
# __ray_parallel_end__
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import ray
|
|
|
|
# Initialize Ray if not already running
|
|
if not ray.is_initialized():
|
|
ray.init()
|
|
|
|
print("Testing Echo deployment...")
|
|
# Test Echo
|
|
echo_handle = serve.run(Echo.bind())
|
|
result = echo_handle.remote(None).result()
|
|
print(f"Echo result: {result}")
|
|
assert result == "ok"
|
|
|
|
print("\nTesting BlockingEcho deployment...")
|
|
# Test BlockingEcho
|
|
blocking_handle = serve.run(BlockingEcho.bind())
|
|
result = blocking_handle.remote(None).result()
|
|
print(f"BlockingEcho result: {result}")
|
|
assert result == "ok"
|
|
|
|
print("\nTesting MyService deployment...")
|
|
# Test MyService
|
|
service_handle = serve.run(MyService.bind())
|
|
result = service_handle.remote(None).result()
|
|
print(f"MyService result: {result}")
|
|
assert result == "ok"
|
|
|
|
print("\nTesting AsyncIOBound deployment...")
|
|
# Test AsyncIOBound
|
|
io_bound_handle = serve.run(AsyncIOBound.bind())
|
|
result = io_bound_handle.remote(None).result()
|
|
print(f"AsyncIOBound result: {result}")
|
|
assert result == "ok"
|
|
|
|
print("\nTesting AsyncStream deployment...")
|
|
# Test AsyncStream (just create it, don't fully consume)
|
|
stream_handle = serve.run(AsyncStream.bind())
|
|
print("AsyncStream deployment created successfully")
|
|
|
|
print("\nTesting OffloadCPU deployment...")
|
|
# Test OffloadCPU
|
|
cpu_handle = serve.run(OffloadCPU.bind())
|
|
result = cpu_handle.remote(None).result()
|
|
print(f"OffloadCPU result: {result}")
|
|
|
|
print("\nTesting NumpyDeployment...")
|
|
# Test NumpyDeployment
|
|
numpy_handle = serve.run(NumpyDeployment.bind())
|
|
result = numpy_handle.remote(None).result()
|
|
print(f"NumpyDeployment result shape: {result.shape}")
|
|
assert result.shape == (100, 100)
|
|
|
|
print("\nTesting BlockingCPU deployment...")
|
|
# Test BlockingCPU
|
|
blocking_cpu_handle = serve.run(BlockingCPU.bind())
|
|
result = blocking_cpu_handle.remote(None).result()
|
|
print(f"BlockingCPU result: {result}")
|
|
assert result == "ok"
|
|
|
|
print("\nTesting CPUWithThreadpool deployment...")
|
|
# Test CPUWithThreadpool
|
|
cpu_threadpool_handle = serve.run(CPUWithThreadpool.bind())
|
|
result = cpu_threadpool_handle.remote(None).result()
|
|
print(f"CPUWithThreadpool result: {result}")
|
|
assert result == "ok"
|
|
|
|
print("\nTesting CustomThreadPool deployment...")
|
|
custom_threadpool_handle = serve.run(CustomThreadPool.bind())
|
|
result = custom_threadpool_handle.remote(None).result()
|
|
print(f"CustomThreadPool result: {result}")
|
|
assert result == "ok"
|
|
|
|
print("\nTesting BlockingStream deployment...")
|
|
# Test BlockingStream - just verify it can be created and called
|
|
blocking_stream_handle = serve.run(BlockingStream.bind())
|
|
# For generator responses, we need to handle them differently
|
|
# Just verify deployment works
|
|
print("BlockingStream deployment created successfully")
|
|
|
|
print("\nTesting RayParallel deployment...")
|
|
# Test RayParallel
|
|
ray_parallel_handle = serve.run(RayParallel.bind())
|
|
result = ray_parallel_handle.remote(None).result()
|
|
print(f"RayParallel result: {result}")
|
|
assert result == {"results": [1, 4, 9, 16]}
|
|
|
|
print("\nTesting BatchedModel deployment...")
|
|
# Test BatchedModel
|
|
batched_model_handle = serve.run(BatchedModel.bind())
|
|
result = batched_model_handle.remote(1).result()
|
|
print(f"BatchedModel result: {result}")
|
|
assert result == "result_1"
|
|
|
|
print("\nTesting BatchedModelOffload deployment...")
|
|
# Test BatchedModelOffload
|
|
batched_model_offload_handle = serve.run(BatchedModelOffload.bind())
|
|
result = batched_model_offload_handle.remote(1).result()
|
|
print(f"BatchedModelOffload result: {result}")
|
|
assert result == "result_1"
|
|
|
|
# Test HTTP-related deployments with try-except
|
|
print("\n--- Testing HTTP-related deployments (may fail due to network) ---")
|
|
|
|
print("\nTesting BlockingHTTP deployment...")
|
|
try:
|
|
blocking_http_handle = serve.run(BlockingHTTP.bind())
|
|
result = blocking_http_handle.remote(None).result()
|
|
print(f"BlockingHTTP result (first 50 chars): {result[:50]}...")
|
|
print("✅ BlockingHTTP test passed")
|
|
except Exception as e:
|
|
print(f"⚠️ BlockingHTTP test failed (expected): {type(e).__name__}: {e}")
|
|
|
|
print("\nTesting AsyncHTTP deployment...")
|
|
try:
|
|
async_http_handle = serve.run(AsyncHTTP.bind())
|
|
result = async_http_handle.remote(None).result()
|
|
print(f"AsyncHTTP result (first 50 chars): {result[:50]}...")
|
|
print("✅ AsyncHTTP test passed")
|
|
except Exception as e:
|
|
print(f"⚠️ AsyncHTTP test failed (expected): {type(e).__name__}: {e}")
|
|
|
|
print("\nTesting ThreadedHTTP deployment...")
|
|
try:
|
|
threaded_http_handle = serve.run(ThreadedHTTP.bind())
|
|
result = threaded_http_handle.remote(None).result()
|
|
print(f"ThreadedHTTP result (first 50 chars): {result[:50]}...")
|
|
print("✅ ThreadedHTTP test passed")
|
|
except Exception as e:
|
|
print(f"⚠️ ThreadedHTTP test failed (expected): {type(e).__name__}: {e}")
|
|
|
|
print("\nTesting OffloadIO deployment...")
|
|
try:
|
|
offload_io_handle = serve.run(OffloadIO.bind())
|
|
result = offload_io_handle.remote(None).result()
|
|
print(f"OffloadIO result (first 50 chars): {result[:50]}...")
|
|
print("✅ OffloadIO test passed")
|
|
except Exception as e:
|
|
print(f"⚠️ OffloadIO test failed (expected): {type(e).__name__}: {e}")
|
|
|
|
print("\nTesting FastAPIDeployment...")
|
|
fastapi_handle = serve.run(FastAPIDeployment.bind())
|
|
# Give it a moment to start
|
|
import time
|
|
import requests
|
|
time.sleep(2)
|
|
# Test the sync endpoint
|
|
response = requests.get("http://127.0.0.1:8000/sync", timeout=5)
|
|
print(f"FastAPIDeployment /sync result: {response.json()}")
|
|
# Test the async endpoint
|
|
response = requests.get("http://127.0.0.1:8000/async", timeout=5)
|
|
print(f"FastAPIDeployment /async result: {response.json()}")
|
|
print("✅ FastAPIDeployment test passed")
|
|
|
|
print("\n✅ All core tests passed!")
|