Add LLM proxy (#122)
This commit is contained in:
Yuge Zhang
@@ -41,6 +41,8 @@ jobs:
|
||||
compression-level: 0
|
||||
- name: Run tests
|
||||
run: |
|
||||
set -ex
|
||||
. .venv/bin/activate
|
||||
pytest -v tests
|
||||
env:
|
||||
PYTEST_ADDOPTS: "--color=yes"
|
||||
|
||||
@@ -0,0 +1,661 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import ast
|
||||
import asyncio
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
import socket
|
||||
import tempfile
|
||||
import threading
|
||||
import time
|
||||
from typing import Any, Awaitable, Callable, Dict, Iterable, List, Optional, Sequence, TypedDict, Union, cast
|
||||
|
||||
import litellm
|
||||
import uvicorn
|
||||
import yaml
|
||||
from fastapi import Request, Response
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.integrations.opentelemetry import OpenTelemetry, OpenTelemetryConfig
|
||||
from litellm.proxy.proxy_server import app, save_worker_config # pyright: ignore[reportUnknownVariableType]
|
||||
from opentelemetry.sdk.trace import ReadableSpan
|
||||
from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult
|
||||
|
||||
from agentlightning.types import LLM
|
||||
|
||||
from .store.base import LightningStore
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ModelConfig(TypedDict):
|
||||
"""LiteLLM model registration entry.
|
||||
|
||||
This mirrors the items in LiteLLM's ``model_list`` section.
|
||||
|
||||
Attributes:
|
||||
model_name: Logical model name exposed by the proxy.
|
||||
litellm_params: Parameters passed to LiteLLM for this model
|
||||
(e.g., backend model id, api_base, additional options).
|
||||
""" # Google style kept concise.
|
||||
|
||||
model_name: str
|
||||
litellm_params: Dict[str, Any]
|
||||
|
||||
|
||||
def _get_pre_call_data(args: Any, kwargs: Any) -> Dict[str, Any]:
|
||||
"""Extract LiteLLM request payload from hook args.
|
||||
|
||||
The LiteLLM logger hooks receive ``(*args, **kwargs)`` whose third positional
|
||||
argument or ``data=`` kwarg contains the request payload.
|
||||
|
||||
Args:
|
||||
args: Positional arguments from the hook.
|
||||
kwargs: Keyword arguments from the hook.
|
||||
|
||||
Returns:
|
||||
The request payload dict.
|
||||
|
||||
Raises:
|
||||
ValueError: If the payload cannot be located or is not a dict.
|
||||
"""
|
||||
if kwargs.get("data"):
|
||||
data = kwargs["data"]
|
||||
elif len(args) >= 3:
|
||||
data = args[2]
|
||||
else:
|
||||
raise ValueError(f"Unable to get request data from args or kwargs: {args}, {kwargs}")
|
||||
if not isinstance(data, dict):
|
||||
raise ValueError(f"Request data is not a dictionary: {data}")
|
||||
return cast(Dict[str, Any], data)
|
||||
|
||||
|
||||
# We need global state because litellm is based on a global app.
|
||||
# Repeatedly initializing the app with different stores will cause errors.
|
||||
_initialized: bool = False
|
||||
_global_store: LightningStore | None = None
|
||||
|
||||
|
||||
def get_global_store() -> LightningStore:
|
||||
"""Return the globally registered LightningStore.
|
||||
|
||||
Used by components that are initialized without an explicit store
|
||||
(e.g., exporter created inside OpenTelemetry).
|
||||
|
||||
Returns:
|
||||
LightningStore: The active global store.
|
||||
|
||||
Raises:
|
||||
ValueError: If the global store has not been set by ``LLMProxy.start()``.
|
||||
"""
|
||||
if _global_store is None:
|
||||
raise ValueError("Global store is not initialized. Please start a LLMProxy first.")
|
||||
return _global_store
|
||||
|
||||
|
||||
def initialize() -> None:
|
||||
"""Initialize global middleware and LiteLLM callbacks once.
|
||||
|
||||
Idempotent. Installs:
|
||||
|
||||
* A FastAPI middleware that rewrites /rollout/{rid}/attempt/{aid}/... paths,
|
||||
injects rollout/attempt/sequence headers, and forwards downstream.
|
||||
* LiteLLM callbacks for token ids and OpenTelemetry export.
|
||||
|
||||
This function does not start any server. It only wires global hooks.
|
||||
"""
|
||||
global _initialized
|
||||
if _initialized:
|
||||
return
|
||||
|
||||
# Add middleware here because it relies on the global store.
|
||||
@app.middleware("http")
|
||||
async def rollout_attempt_middleware( # pyright: ignore[reportUnusedFunction]
|
||||
request: Request, call_next: Callable[[Request], Awaitable[Response]]
|
||||
) -> Response:
|
||||
# Decode rollout and attempt from the URL prefix. Example:
|
||||
# /rollout/r123/attempt/a456/v1/chat/completions
|
||||
# becomes
|
||||
# /v1/chat/completions
|
||||
# while adding request-scoped headers for trace attribution.
|
||||
path = request.url.path
|
||||
|
||||
match = re.match(r"^/rollout/([^/]+)/attempt/([^/]+)(/.*)?$", path)
|
||||
if match:
|
||||
rollout_id = match.group(1)
|
||||
attempt_id = match.group(2)
|
||||
new_path = match.group(3) if match.group(3) is not None else "/"
|
||||
|
||||
# Rewrite the ASGI scope path so downstream sees a clean OpenAI path.
|
||||
request.scope["path"] = new_path
|
||||
request.scope["raw_path"] = new_path.encode()
|
||||
|
||||
# Allocate a monotonic sequence id per (rollout, attempt).
|
||||
sequence_id = await get_global_store().get_next_span_sequence_id(rollout_id, attempt_id)
|
||||
|
||||
# Inject headers so downstream components and exporters can retrieve them.
|
||||
request.scope["headers"] = list(request.scope["headers"]) + [
|
||||
(b"x-rollout-id", rollout_id.encode()),
|
||||
(b"x-attempt-id", attempt_id.encode()),
|
||||
(b"x-sequence-id", str(sequence_id).encode()),
|
||||
]
|
||||
|
||||
response = await call_next(request)
|
||||
return response
|
||||
|
||||
# Register callbacks once on the global LiteLLM callback list.
|
||||
litellm.callbacks.extend( # pyright: ignore[reportUnknownMemberType]
|
||||
[
|
||||
AddReturnTokenIds(),
|
||||
LightningOpenTelemetry(),
|
||||
]
|
||||
)
|
||||
|
||||
_initialized = True
|
||||
|
||||
|
||||
class AddReturnTokenIds(CustomLogger):
|
||||
"""LiteLLM logger hook to request token ids from vLLM.
|
||||
|
||||
This mutates the outgoing request payload to include ``return_token_ids=True``
|
||||
for backends that support token id return (e.g., vLLM).
|
||||
|
||||
See:
|
||||
https://github.com/vllm-project/vllm/pull/22587
|
||||
"""
|
||||
|
||||
async def async_pre_call_hook(self, *args: Any, **kwargs: Any) -> Optional[Union[Exception, str, Dict[str, Any]]]:
|
||||
"""Async pre-call hook to adjust request payload.
|
||||
|
||||
Args:
|
||||
args: Positional args from LiteLLM.
|
||||
kwargs: Keyword args from LiteLLM.
|
||||
|
||||
Returns:
|
||||
Either an updated payload dict or an Exception to short-circuit.
|
||||
"""
|
||||
try:
|
||||
data = _get_pre_call_data(args, kwargs)
|
||||
except Exception as e:
|
||||
return e
|
||||
|
||||
# Ensure token ids are requested from the backend when supported.
|
||||
return {**data, "return_token_ids": True}
|
||||
|
||||
|
||||
class LightningSpanExporter(SpanExporter):
|
||||
"""Buffered OTEL span exporter with subtree flushing and training-store sink.
|
||||
|
||||
Design:
|
||||
|
||||
* Spans are buffered until a root span's entire subtree is available.
|
||||
* A private event loop on a daemon thread runs async flush logic.
|
||||
* Rollout/attempt/sequence metadata is reconstructed by merging headers
|
||||
from any span within a subtree.
|
||||
|
||||
Thread-safety:
|
||||
|
||||
* Buffer access is protected by a re-entrant lock.
|
||||
* Export is synchronous to the caller yet schedules an async flush on the
|
||||
internal loop, then waits for completion.
|
||||
|
||||
Args:
|
||||
store: Optional explicit LightningStore. If None, uses ``get_global_store()``.
|
||||
"""
|
||||
|
||||
def __init__(self, store: Optional[LightningStore] = None):
|
||||
self._store = store
|
||||
self._buffer: List[ReadableSpan] = []
|
||||
self._lock = threading.RLock()
|
||||
|
||||
# Single dedicated event loop running in a daemon thread.
|
||||
# This decouples OTEL SDK threads from our async store I/O.
|
||||
self._loop = asyncio.new_event_loop()
|
||||
self._loop_thread = threading.Thread(target=self._run_loop, name="LightningSpanExporterLoop", daemon=True)
|
||||
self._loop_thread.start()
|
||||
|
||||
def _get_store(self) -> LightningStore:
|
||||
"""Return the LightningStore to use.
|
||||
|
||||
Returns:
|
||||
LightningStore: Explicit store if provided, else the global store.
|
||||
|
||||
Raises:
|
||||
ValueError: If no global store is configured and no explicit store was given.
|
||||
"""
|
||||
if self._store is None:
|
||||
return get_global_store()
|
||||
return self._store
|
||||
|
||||
def _run_loop(self) -> None:
|
||||
"""Run the private asyncio loop forever on the exporter thread."""
|
||||
asyncio.set_event_loop(self._loop)
|
||||
self._loop.run_forever()
|
||||
|
||||
def shutdown(self) -> None:
|
||||
"""Shut down the exporter event loop.
|
||||
|
||||
Safe to call at process exit.
|
||||
|
||||
"""
|
||||
try:
|
||||
|
||||
def _stop():
|
||||
self._loop.stop()
|
||||
|
||||
self._loop.call_soon_threadsafe(_stop)
|
||||
self._loop_thread.join(timeout=2.0)
|
||||
self._loop.close()
|
||||
except Exception:
|
||||
logger.exception("Error during exporter shutdown")
|
||||
|
||||
def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
|
||||
"""Export spans via buffered subtree flush.
|
||||
|
||||
Appends spans to the internal buffer, then triggers an async flush on the
|
||||
private event loop. Blocks until that flush completes.
|
||||
|
||||
Args:
|
||||
spans: Sequence of spans to export.
|
||||
|
||||
Returns:
|
||||
SpanExportResult: SUCCESS on flush success, else FAILURE.
|
||||
"""
|
||||
# Buffer append under lock to protect against concurrent exporters.
|
||||
with self._lock:
|
||||
for span in spans:
|
||||
self._buffer.append(span)
|
||||
|
||||
# Run the async flush on our private loop, synchronously from caller’s POV.
|
||||
async def _locked_flush():
|
||||
# Take the lock inside the coroutine to serialize with other flushes.
|
||||
with self._lock:
|
||||
return await self._maybe_flush()
|
||||
|
||||
try:
|
||||
fut = asyncio.run_coroutine_threadsafe(_locked_flush(), self._loop)
|
||||
fut.result() # Bubble up any exceptions from the coroutine.
|
||||
except Exception as e:
|
||||
logger.exception("Export flush failed: %s", e)
|
||||
return SpanExportResult.FAILURE
|
||||
|
||||
return SpanExportResult.SUCCESS
|
||||
|
||||
async def _maybe_flush(self):
|
||||
"""Flush ready subtrees from the buffer.
|
||||
|
||||
Strategy:
|
||||
We consider a subtree "ready" if we can identify a root span. We
|
||||
then take that root and all its descendants out of the buffer and
|
||||
try to reconstruct rollout/attempt/sequence headers by merging any
|
||||
span's ``metadata.requester_custom_headers`` within the subtree.
|
||||
|
||||
Required headers:
|
||||
``x-rollout-id`` (str), ``x-attempt-id`` (str), ``x-sequence-id`` (str of int)
|
||||
|
||||
Raises:
|
||||
None directly. Logs and skips malformed spans.
|
||||
|
||||
"""
|
||||
# Iterate over current roots. Each iteration pops a whole subtree.
|
||||
for root_span_id in self._get_root_span_ids():
|
||||
subtree_spans = self._pop_subtrees(root_span_id)
|
||||
if not subtree_spans:
|
||||
continue
|
||||
|
||||
# Merge all custom headers found in the subtree.
|
||||
headers_merged: Dict[str, Any] = {}
|
||||
|
||||
for span in subtree_spans:
|
||||
if span.attributes is None:
|
||||
continue
|
||||
headers_str = span.attributes.get("metadata.requester_custom_headers")
|
||||
if headers_str is None:
|
||||
continue
|
||||
if not isinstance(headers_str, str):
|
||||
logger.error(
|
||||
f"metadata.requester_custom_headers is not stored as a string: {headers_str}. Skipping the span."
|
||||
)
|
||||
continue
|
||||
try:
|
||||
# Use literal_eval to parse the stringified dict safely.
|
||||
headers = ast.literal_eval(headers_str)
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Failed to parse metadata.requester_custom_headers: {headers_str}, error: {e}. Skipping the span."
|
||||
)
|
||||
continue
|
||||
if not isinstance(headers, dict):
|
||||
logger.error(
|
||||
f"metadata.requester_custom_headers is not parsed as a dict: {headers}. Skipping the span."
|
||||
)
|
||||
continue
|
||||
headers_merged.update(cast(Dict[str, Any], headers))
|
||||
|
||||
if not headers_merged:
|
||||
logger.warning(f"No headers found in {len(subtree_spans)} subtree spans. Cannot log to store.")
|
||||
continue
|
||||
|
||||
# Validate and normalize required header fields.
|
||||
rollout_id = headers_merged.get("x-rollout-id")
|
||||
attempt_id = headers_merged.get("x-attempt-id")
|
||||
sequence_id = headers_merged.get("x-sequence-id")
|
||||
if not rollout_id or not attempt_id or not sequence_id or not sequence_id.isdigit():
|
||||
logger.warning(
|
||||
f"Missing or invalid rollout_id, attempt_id, or sequence_id in headers: {headers_merged}. Cannot log to store."
|
||||
)
|
||||
continue
|
||||
if not isinstance(rollout_id, str) or not isinstance(attempt_id, str):
|
||||
logger.warning(
|
||||
f"rollout_id or attempt_id is not a string: {rollout_id}, {attempt_id}. Cannot log to store."
|
||||
)
|
||||
continue
|
||||
sequence_id_decimal = int(sequence_id)
|
||||
|
||||
# Persist each span in the subtree with the resolved identifiers.
|
||||
for span in subtree_spans:
|
||||
await self._get_store().add_otel_span(
|
||||
rollout_id=rollout_id, attempt_id=attempt_id, sequence_id=sequence_id_decimal, readable_span=span
|
||||
)
|
||||
|
||||
def _get_root_span_ids(self) -> Iterable[int]:
|
||||
"""Yield span_ids for root spans currently in the buffer.
|
||||
|
||||
A root span is defined as one with ``parent is None``.
|
||||
|
||||
Yields:
|
||||
int: Span id for each root span found.
|
||||
"""
|
||||
for span in self._buffer:
|
||||
if span.parent is None:
|
||||
span_context = span.get_span_context()
|
||||
if span_context is not None:
|
||||
yield span_context.span_id
|
||||
|
||||
def _get_subtrees(self, root_span_id: int) -> Iterable[int]:
|
||||
"""Yield span_ids in the subtree rooted at ``root_span_id``.
|
||||
|
||||
Depth-first traversal over the current buffer.
|
||||
|
||||
Args:
|
||||
root_span_id: The span id of the root.
|
||||
|
||||
Yields:
|
||||
int: Span ids including the root and all descendants found.
|
||||
"""
|
||||
# Yield the root span id first.
|
||||
yield root_span_id
|
||||
for span in self._buffer:
|
||||
# Check whether the span's parent is the root_span_id.
|
||||
if span.parent is not None and span.parent.span_id == root_span_id:
|
||||
span_context = span.get_span_context()
|
||||
if span_context is not None:
|
||||
# Recursively get child spans.
|
||||
yield from self._get_subtrees(span_context.span_id)
|
||||
|
||||
def _pop_subtrees(self, root_span_id: int) -> List[ReadableSpan]:
|
||||
"""Remove and return the subtree for a particular root from the buffer.
|
||||
|
||||
Args:
|
||||
root_span_id: Root span id identifying the subtree.
|
||||
|
||||
Returns:
|
||||
list[ReadableSpan]: Spans that were part of the subtree. Order follows buffer order.
|
||||
"""
|
||||
subtree_span_ids = set(self._get_subtrees(root_span_id))
|
||||
subtree_spans: List[ReadableSpan] = []
|
||||
new_buffer: List[ReadableSpan] = []
|
||||
for span in self._buffer:
|
||||
span_context = span.get_span_context()
|
||||
if span_context is not None and span_context.span_id in subtree_span_ids:
|
||||
subtree_spans.append(span)
|
||||
else:
|
||||
new_buffer.append(span)
|
||||
# Replace buffer with remaining spans to avoid re-processing.
|
||||
self._buffer = new_buffer
|
||||
return subtree_spans
|
||||
|
||||
|
||||
class LightningOpenTelemetry(OpenTelemetry):
|
||||
"""OpenTelemetry integration that exports spans to the Lightning store.
|
||||
|
||||
Responsibilities:
|
||||
|
||||
* Ensures each request is annotated with a per-attempt sequence id so spans
|
||||
are ordered deterministically even with clock skew across nodes.
|
||||
* Uses ``LightningSpanExporter`` to persist spans for analytics and training.
|
||||
|
||||
Args:
|
||||
store: Optional explicit LightningStore for the exporter.
|
||||
"""
|
||||
|
||||
def __init__(self, store: LightningStore | None = None):
|
||||
config = OpenTelemetryConfig(exporter=LightningSpanExporter(store))
|
||||
super().__init__(config=config) # pyright: ignore[reportUnknownMemberType]
|
||||
|
||||
|
||||
class LLMProxy:
|
||||
"""Host a LiteLLM OpenAI-compatible proxy bound to a LightningStore.
|
||||
|
||||
The proxy:
|
||||
|
||||
* Serves an OpenAI-compatible API via uvicorn.
|
||||
* Adds rollout/attempt routing and headers via middleware.
|
||||
* Registers OTEL export and token-id callbacks.
|
||||
* Writes a LiteLLM worker config file with ``model_list`` and settings.
|
||||
|
||||
Lifecycle:
|
||||
|
||||
* ``start()`` writes config, starts uvicorn server in a thread, and waits until ready.
|
||||
* ``stop()`` tears down the server and removes the temp config file.
|
||||
* ``restart()`` convenience wrapper to stop then start.
|
||||
|
||||
Usage Note:
|
||||
As the LLM Proxy sets up an OpenTelemetry tracer, it's recommended to run it in a different
|
||||
process from the main runner (i.e., tracer from agents).
|
||||
|
||||
Args:
|
||||
port: TCP port to bind.
|
||||
model_list: LiteLLM ``model_list`` entries.
|
||||
store: LightningStore used for span sequence and persistence.
|
||||
host: Publicly reachable host used in resource endpoints. Defaults to best-guess IPv4.
|
||||
litellm_config: Extra LiteLLM proxy config merged with ``model_list``.
|
||||
num_retries: Default LiteLLM retry count injected into ``litellm_settings``.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
port: int,
|
||||
model_list: List[ModelConfig],
|
||||
store: LightningStore,
|
||||
host: str | None = None,
|
||||
litellm_config: Dict[str, Any] | None = None,
|
||||
num_retries: int = 0,
|
||||
):
|
||||
self.store = store
|
||||
self.host = host or _get_default_ipv4_address()
|
||||
self.port = port
|
||||
self.model_list = model_list
|
||||
self.litellm_config = litellm_config or {}
|
||||
|
||||
# Ensure num_retries is present inside the litellm_settings block.
|
||||
self.litellm_config.setdefault("litellm_settings", {})
|
||||
self.litellm_config["litellm_settings"].setdefault("num_retries", num_retries)
|
||||
|
||||
self._server_thread = None
|
||||
self._config_file = None
|
||||
self._uvicorn_server = None
|
||||
self._ready_event = threading.Event()
|
||||
|
||||
def update_model_list(self, model_list: List[ModelConfig]) -> None:
|
||||
"""Replace the in-memory model list and hot-restart if running.
|
||||
|
||||
Args:
|
||||
model_list: New list of model entries.
|
||||
"""
|
||||
self.model_list = model_list
|
||||
if self.is_running():
|
||||
self.restart()
|
||||
# Do nothing if the server is not running.
|
||||
|
||||
def _wait_until_started(self, startup_timeout: float = 20.0):
|
||||
"""Block until the uvicorn server reports started or timeout.
|
||||
|
||||
Args:
|
||||
startup_timeout: Maximum seconds to wait.
|
||||
"""
|
||||
start = time.time()
|
||||
while True:
|
||||
if self._uvicorn_server is None:
|
||||
break
|
||||
if self._uvicorn_server.started:
|
||||
self._ready_event.set()
|
||||
break
|
||||
if self._uvicorn_server.should_exit:
|
||||
break
|
||||
if time.time() - start > startup_timeout:
|
||||
break
|
||||
time.sleep(0.01)
|
||||
|
||||
def start(self):
|
||||
"""Start the proxy server thread and initialize global wiring.
|
||||
|
||||
Side effects:
|
||||
|
||||
* Sets the module-level global store for middleware/exporter access.
|
||||
* Calls ``initialize()`` once to register middleware and callbacks.
|
||||
* Writes a temporary YAML config consumed by LiteLLM worker.
|
||||
* Launches uvicorn in a daemon thread and waits for readiness.
|
||||
"""
|
||||
global _global_store
|
||||
|
||||
_global_store = self.store
|
||||
|
||||
# Initialize global middleware and callbacks once.
|
||||
initialize()
|
||||
|
||||
# Persist a temp worker config for LiteLLM and point the proxy at it.
|
||||
self._config_file = tempfile.NamedTemporaryFile(suffix=".yaml", delete=False).name
|
||||
with open(self._config_file, "w") as fp:
|
||||
yaml.safe_dump(
|
||||
{
|
||||
"model_list": self.model_list,
|
||||
**self.litellm_config,
|
||||
},
|
||||
fp,
|
||||
)
|
||||
|
||||
save_worker_config(config=self._config_file)
|
||||
|
||||
# Bind to all interfaces to allow other hosts to reach it if needed.
|
||||
self._uvicorn_server = uvicorn.Server(uvicorn.Config(app, host="0.0.0.0", port=self.port))
|
||||
|
||||
def run_server():
|
||||
# Serve uvicorn in this background thread with its own event loop.
|
||||
assert self._uvicorn_server is not None
|
||||
asyncio.run(self._uvicorn_server.serve())
|
||||
|
||||
self._ready_event.clear()
|
||||
self._server_thread = threading.Thread(target=run_server, daemon=True)
|
||||
self._server_thread.start()
|
||||
self._wait_until_started()
|
||||
|
||||
def stop(self):
|
||||
"""Stop the proxy server and clean up temporary artifacts.
|
||||
|
||||
This is a best-effort graceful shutdown with a bounded join timeout.
|
||||
"""
|
||||
if not self.is_running():
|
||||
logger.warning("LLMProxy is not running. Nothing to stop.")
|
||||
return
|
||||
|
||||
# Remove worker config to avoid stale references.
|
||||
if self._config_file and os.path.exists(self._config_file):
|
||||
os.unlink(self._config_file)
|
||||
|
||||
if self._server_thread is not None and self._uvicorn_server is not None and self._uvicorn_server.started:
|
||||
self._uvicorn_server.should_exit = True
|
||||
self._server_thread.join(timeout=10.0) # Allow time for graceful shutdown.
|
||||
self._server_thread = None
|
||||
self._uvicorn_server = None
|
||||
self._config_file = None
|
||||
self._ready_event.clear()
|
||||
|
||||
def restart(self) -> None:
|
||||
"""Restart the proxy if running, else start it.
|
||||
|
||||
Convenience wrapper calling ``stop()`` followed by ``start()``.
|
||||
"""
|
||||
if self.is_running():
|
||||
self.stop()
|
||||
self.start()
|
||||
|
||||
def is_running(self) -> bool:
|
||||
"""Return whether the uvicorn server is active.
|
||||
|
||||
Returns:
|
||||
bool: True if server was started and did not signal exit.
|
||||
"""
|
||||
return self._uvicorn_server is not None and self._uvicorn_server.started
|
||||
|
||||
def as_resource(
|
||||
self,
|
||||
rollout_id: str,
|
||||
attempt_id: str,
|
||||
model: str | None = None,
|
||||
sampling_parameters: Dict[str, Any] | None = None,
|
||||
) -> LLM:
|
||||
"""Create an ``LLM`` resource pointing at this proxy with rollout context.
|
||||
|
||||
The returned endpoint is:
|
||||
``http://{host}:{port}/rollout/{rollout_id}/attempt/{attempt_id}``
|
||||
|
||||
Args:
|
||||
rollout_id: Rollout identifier used for span attribution.
|
||||
attempt_id: Attempt identifier used for span attribution.
|
||||
model: Logical model name to use. If omitted and exactly one model
|
||||
is configured, that model is used.
|
||||
sampling_parameters: Optional default sampling parameters.
|
||||
|
||||
Returns:
|
||||
LLM: Configured resource ready for OpenAI-compatible calls.
|
||||
|
||||
Raises:
|
||||
ValueError: If ``model`` is omitted and zero or multiple models are configured.
|
||||
"""
|
||||
if model is None:
|
||||
if len(self.model_list) == 1:
|
||||
model = self.model_list[0]["model_name"]
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Multiple or zero models found in model_list: {self.model_list}. Please specify the model."
|
||||
)
|
||||
|
||||
return LLM(
|
||||
endpoint=f"http://{self.host}:{self.port}/rollout/{rollout_id}/attempt/{attempt_id}",
|
||||
model=model,
|
||||
sampling_parameters=dict(sampling_parameters or {}),
|
||||
)
|
||||
|
||||
|
||||
def _get_default_ipv4_address() -> str:
|
||||
"""Determine the default outbound IPv4 address for this machine.
|
||||
|
||||
Implementation:
|
||||
Opens a UDP socket and "connects" to a public address to force route
|
||||
selection, then inspects the socket's local address. No packets are sent.
|
||||
|
||||
Returns:
|
||||
str: Best-guess IPv4 like ``192.168.x.y``. Falls back to ``127.0.0.1``.
|
||||
"""
|
||||
s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
|
||||
try:
|
||||
# Doesn't actually contact 8.8.8.8; just forces the OS to pick a route.
|
||||
s.connect(("8.8.8.8", 80))
|
||||
return s.getsockname()[0]
|
||||
except Exception:
|
||||
return "127.0.0.1"
|
||||
finally:
|
||||
s.close()
|
||||
@@ -16,6 +16,8 @@ dependencies = [
|
||||
"aiohttp",
|
||||
"opentelemetry-api>=1.35",
|
||||
"opentelemetry-sdk>=1.35",
|
||||
"opentelemetry-exporter-otlp>=1.35",
|
||||
"litellm[proxy]",
|
||||
"pydantic>=2.11",
|
||||
"openai",
|
||||
]
|
||||
@@ -55,6 +57,7 @@ agent = [
|
||||
"sqlparse",
|
||||
"nltk",
|
||||
"uv",
|
||||
"anthropic",
|
||||
]
|
||||
|
||||
[build-system]
|
||||
|
||||
@@ -0,0 +1,669 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Test the LLMProxy class. Still under development.
|
||||
|
||||
General TODOs:
|
||||
|
||||
1. Add tests for update model list and server restart
|
||||
2. Add tests for retries
|
||||
3. Add tests for timeout
|
||||
4. Add tests for multiple models in model list
|
||||
5. Add tests for multi-modal models
|
||||
|
||||
There are some specific TODOs for each test function.
|
||||
"""
|
||||
|
||||
import ast
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import random
|
||||
import socket
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
from contextlib import closing
|
||||
from typing import Any, List, Optional, cast
|
||||
|
||||
import anthropic
|
||||
import httpx
|
||||
import openai
|
||||
import pytest
|
||||
from opentelemetry.sdk.trace import ReadableSpan
|
||||
|
||||
from agentlightning.llm_proxy import LightningSpanExporter, LLMProxy
|
||||
from agentlightning.store.memory import InMemoryLightningStore
|
||||
from agentlightning.tracer.types import Span
|
||||
from agentlightning.types import LLM
|
||||
from tests.tracer.utils import clear_tracer_provider
|
||||
|
||||
try:
|
||||
import torch # type: ignore
|
||||
|
||||
GPU_AVAILABLE = torch.cuda.is_available()
|
||||
except Exception:
|
||||
GPU_AVAILABLE = False # type: ignore
|
||||
pytest.skip(reason="GPU not available", allow_module_level=True)
|
||||
|
||||
VLLM_AVAILABLE = False
|
||||
VLLM_UNAVAILABLE_REASON = ""
|
||||
|
||||
try:
|
||||
import vllm
|
||||
from vllm.engine.arg_utils import AsyncEngineArgs
|
||||
from vllm.entrypoints.cli.serve import ServeSubcommand
|
||||
from vllm.model_executor.model_loader import get_model_loader
|
||||
from vllm.utils import FlexibleArgumentParser
|
||||
|
||||
VLLM_AVAILABLE = True # type: ignore
|
||||
VLLM_VERSION = tuple(int(v) for v in vllm.__version__.split("."))
|
||||
except ImportError as e:
|
||||
AsyncEngineArgs = None
|
||||
get_model_loader = None
|
||||
FlexibleArgumentParser = None
|
||||
ServeSubcommand = None
|
||||
VLLM_VERSION = (0, 0, 0) # type: ignore
|
||||
VLLM_UNAVAILABLE_REASON = str(e) # type: ignore
|
||||
|
||||
|
||||
def _get_free_port() -> int:
|
||||
with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as s:
|
||||
s.bind(("127.0.0.1", 0))
|
||||
return int(s.getsockname()[1])
|
||||
|
||||
|
||||
class RemoteOpenAIServer:
|
||||
"""
|
||||
A context manager for launching and interacting with a remote vLLM-based
|
||||
OpenAI-compatible server instance.
|
||||
|
||||
This class handles:
|
||||
- Preparing the environment and spawning the vLLM server process
|
||||
- Ensuring that the requested model is downloaded before server startup
|
||||
- Polling and health-checking the server until it is ready
|
||||
- Providing helper methods to construct URLs for API calls
|
||||
- Returning configured synchronous and asynchronous OpenAI clients
|
||||
that can communicate with the launched server
|
||||
|
||||
Typical usage:
|
||||
with RemoteOpenAIServer(vllm_serve_args, port, model) as server:
|
||||
client = server.get_client()
|
||||
response = client.chat.completions.create(...)
|
||||
|
||||
Attributes:
|
||||
DUMMY_API_KEY (str): A placeholder API key for compatibility
|
||||
(vLLM does not require authentication).
|
||||
host (str): Host address of the server (default: "localhost").
|
||||
port (int): TCP port number for the server.
|
||||
proc (subprocess.Popen): Handle to the launched server process.
|
||||
"""
|
||||
|
||||
DUMMY_API_KEY = "token-abc123" # vLLM's OpenAI server does not need API key
|
||||
|
||||
def _start_server(self, model: str, vllm_serve_args: list[str], env_dict: Optional[dict[str, str]]) -> None:
|
||||
"""Subclasses override this method to customize server process launch"""
|
||||
env = os.environ.copy()
|
||||
env["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" # safer CUDA init
|
||||
if env_dict is not None:
|
||||
env.update(env_dict)
|
||||
|
||||
if VLLM_VERSION >= (0, 10, 2):
|
||||
# Supports return_token_ids
|
||||
self.proc: subprocess.Popen[bytes] = subprocess.Popen(
|
||||
["vllm", "serve", model, *vllm_serve_args],
|
||||
env=env,
|
||||
stdout=sys.stdout,
|
||||
stderr=sys.stderr,
|
||||
)
|
||||
else:
|
||||
# Does not support return_token_ids
|
||||
self.proc = subprocess.Popen(
|
||||
["python", "-m", "agentlightning.cli.vllm", "serve", model, *vllm_serve_args],
|
||||
env=env,
|
||||
stdout=sys.stdout,
|
||||
stderr=sys.stderr,
|
||||
)
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model: str,
|
||||
vllm_serve_args: list[str], # should not include the model name
|
||||
env_dict: Optional[dict[str, str]] = None,
|
||||
seed: Optional[int] = 0,
|
||||
max_wait_seconds: Optional[float] = None,
|
||||
) -> None:
|
||||
if (
|
||||
not VLLM_AVAILABLE
|
||||
or AsyncEngineArgs is None
|
||||
or get_model_loader is None
|
||||
or FlexibleArgumentParser is None
|
||||
or ServeSubcommand is None
|
||||
):
|
||||
raise ImportError("vLLM is not available: " + VLLM_UNAVAILABLE_REASON)
|
||||
|
||||
self.model = model
|
||||
|
||||
parser = FlexibleArgumentParser(description="vLLM's remote OpenAI server.")
|
||||
subparsers = parser.add_subparsers(required=False, dest="subparser")
|
||||
parser = ServeSubcommand().subparser_init(subparsers) # pyright: ignore[reportUnknownMemberType]
|
||||
args = parser.parse_args(["--model", model, *vllm_serve_args])
|
||||
assert args is not None
|
||||
self.host = str(args.host or "localhost")
|
||||
self.port = int(args.port)
|
||||
|
||||
# download the model before starting the server to avoid timeout
|
||||
is_local = os.path.isdir(model)
|
||||
if not is_local:
|
||||
engine_args = AsyncEngineArgs.from_cli_args(args)
|
||||
model_config = engine_args.create_model_config()
|
||||
load_config = engine_args.create_load_config()
|
||||
|
||||
model_loader = get_model_loader(load_config)
|
||||
model_loader.download_model(model_config)
|
||||
|
||||
self._start_server(model, vllm_serve_args, env_dict)
|
||||
max_wait_seconds = max_wait_seconds or 240
|
||||
self._wait_for_server(url=self.url_for("health"), timeout=max_wait_seconds)
|
||||
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any):
|
||||
self.proc.terminate()
|
||||
try:
|
||||
self.proc.wait(8)
|
||||
except subprocess.TimeoutExpired:
|
||||
self.proc.kill()
|
||||
|
||||
def _poll(self) -> Optional[int]:
|
||||
"""Subclasses override this method to customize process polling"""
|
||||
return self.proc.poll()
|
||||
|
||||
def _wait_for_server(self, *, url: str, timeout: float):
|
||||
start = time.time()
|
||||
client = httpx.Client()
|
||||
|
||||
while True:
|
||||
try:
|
||||
if client.get(url).status_code == 200:
|
||||
break
|
||||
except Exception:
|
||||
result = self._poll()
|
||||
if result is not None and result != 0:
|
||||
raise RuntimeError("Server exited unexpectedly.") from None
|
||||
time.sleep(0.5)
|
||||
if time.time() - start > timeout:
|
||||
raise RuntimeError("Server failed to start in time.") from None
|
||||
|
||||
@property
|
||||
def url_root(self) -> str:
|
||||
return f"http://{self.host}:{self.port}"
|
||||
|
||||
def url_for(self, *parts: str) -> str:
|
||||
return self.url_root + "/" + "/".join(parts)
|
||||
|
||||
def get_client(self, **kwargs: Any):
|
||||
if "timeout" not in kwargs:
|
||||
kwargs["timeout"] = 600
|
||||
return openai.OpenAI(
|
||||
base_url=self.url_for("v1"),
|
||||
api_key=self.DUMMY_API_KEY,
|
||||
max_retries=0,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
def get_async_client(self, **kwargs: Any):
|
||||
if "timeout" not in kwargs:
|
||||
kwargs["timeout"] = 600
|
||||
return openai.AsyncOpenAI(
|
||||
base_url=self.url_for("v1"),
|
||||
api_key=self.DUMMY_API_KEY,
|
||||
max_retries=0,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def qwen25_model():
|
||||
with RemoteOpenAIServer(
|
||||
model="Qwen/Qwen2.5-0.5B-Instruct",
|
||||
vllm_serve_args=[
|
||||
"--gpu-memory-utilization",
|
||||
"0.7",
|
||||
"--enable-auto-tool-choice",
|
||||
"--tool-call-parser",
|
||||
"hermes",
|
||||
"--port",
|
||||
str(_get_free_port()),
|
||||
],
|
||||
) as server:
|
||||
yield server
|
||||
|
||||
|
||||
def test_qwen25_model_sanity(qwen25_model: RemoteOpenAIServer):
|
||||
client = qwen25_model.get_client()
|
||||
response = client.chat.completions.create(
|
||||
model="Qwen/Qwen2.5-0.5B-Instruct",
|
||||
messages=[{"role": "user", "content": "Hello, world!"}],
|
||||
stream=False,
|
||||
)
|
||||
assert response.choices[0].message.content is not None
|
||||
|
||||
|
||||
@pytest.fixture(scope="module", autouse=True)
|
||||
def setup_module():
|
||||
clear_tracer_provider()
|
||||
yield
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_basic_integration(qwen25_model: RemoteOpenAIServer):
|
||||
store = InMemoryLightningStore()
|
||||
proxy = LLMProxy(
|
||||
port=_get_free_port(),
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "gpt-4o-arbitrary",
|
||||
"litellm_params": {
|
||||
"model": "hosted_vllm/" + qwen25_model.model,
|
||||
"api_base": qwen25_model.url_for("v1"),
|
||||
},
|
||||
}
|
||||
],
|
||||
store=store,
|
||||
)
|
||||
|
||||
rollout = await store.start_rollout(None)
|
||||
|
||||
proxy.start()
|
||||
|
||||
resource = proxy.as_resource(rollout.rollout_id, rollout.attempt.attempt_id)
|
||||
|
||||
import openai
|
||||
|
||||
client = openai.OpenAI(base_url=resource.endpoint, api_key="token-abc123")
|
||||
response = client.chat.completions.create(
|
||||
model="gpt-4o-arbitrary",
|
||||
messages=[{"role": "user", "content": "Repeat after me: Hello, world!"}],
|
||||
stream=False,
|
||||
)
|
||||
assert response.choices[0].message.content is not None
|
||||
assert "hello, world" in response.choices[0].message.content.lower()
|
||||
|
||||
proxy.stop()
|
||||
|
||||
spans = await store.query_spans(rollout.rollout_id, rollout.attempt.attempt_id)
|
||||
|
||||
# Verify all spans have correct rollout_id, attempt_id, and sequence_id
|
||||
assert len(spans) > 0, "Should have captured spans"
|
||||
for span in spans:
|
||||
assert span.rollout_id == rollout.rollout_id, f"Span {span.name} has incorrect rollout_id"
|
||||
assert span.attempt_id == rollout.attempt.attempt_id, f"Span {span.name} has incorrect attempt_id"
|
||||
assert span.sequence_id == 1, f"Span {span.name} has incorrect sequence_id"
|
||||
|
||||
# Find the raw_gen_ai_request span and verify token IDs
|
||||
raw_gen_ai_spans = [s for s in spans if s.name == "raw_gen_ai_request"]
|
||||
assert len(raw_gen_ai_spans) == 1, f"Expected 1 raw_gen_ai_request span, found {len(raw_gen_ai_spans)}"
|
||||
raw_span = raw_gen_ai_spans[0]
|
||||
|
||||
# Verify prompt_token_ids is present and non-empty
|
||||
assert (
|
||||
"llm.hosted_vllm.prompt_token_ids" in raw_span.attributes
|
||||
), "prompt_token_ids not found in raw_gen_ai_request span"
|
||||
prompt_token_ids: list[int] = ast.literal_eval(raw_span.attributes["llm.hosted_vllm.prompt_token_ids"]) # type: ignore
|
||||
assert isinstance(prompt_token_ids, list), "prompt_token_ids should be a list"
|
||||
assert len(prompt_token_ids) > 0, "prompt_token_ids should not be empty"
|
||||
assert all(isinstance(tid, int) for tid in prompt_token_ids), "All prompt token IDs should be integers"
|
||||
|
||||
# Verify response token_ids is present in choices
|
||||
assert "llm.hosted_vllm.choices" in raw_span.attributes, "choices not found in raw_gen_ai_request span"
|
||||
choices: list[dict[str, Any]] = ast.literal_eval(raw_span.attributes["llm.hosted_vllm.choices"]) # type: ignore
|
||||
assert len(choices) > 0, "Should have at least one choice"
|
||||
if VLLM_VERSION >= (0, 10, 2):
|
||||
assert "token_ids" in choices[0], "token_ids not found in choice"
|
||||
response_token_ids: list[int] = choices[0]["token_ids"]
|
||||
else:
|
||||
assert (
|
||||
"llm.hosted_vllm.response_token_ids" in raw_span.attributes
|
||||
), "response_token_ids not found in raw_gen_ai_request span"
|
||||
response_token_ids_list: list[list[int]] = ast.literal_eval(raw_span.attributes["llm.hosted_vllm.response_token_ids"]) # type: ignore
|
||||
assert isinstance(response_token_ids_list, list), "response_token_ids_list should be a list"
|
||||
assert len(response_token_ids_list) > 0, "response_token_ids_list should not be empty"
|
||||
assert all(
|
||||
isinstance(tid_list, list) for tid_list in response_token_ids_list
|
||||
), "All response token IDs should be lists"
|
||||
assert all(
|
||||
isinstance(tid, int) for tid_list in response_token_ids_list for tid in tid_list
|
||||
), "All response token IDs should be integers"
|
||||
response_token_ids = response_token_ids_list[0]
|
||||
assert isinstance(response_token_ids, list), "response token_ids should be a list"
|
||||
assert len(response_token_ids) > 0, "response token_ids should not be empty"
|
||||
assert all(isinstance(tid, int) for tid in response_token_ids), "All response token IDs should be integers"
|
||||
|
||||
# Find the litellm_request span and verify gen_ai prompts/completions
|
||||
litellm_spans = [s for s in spans if s.name == "litellm_request"]
|
||||
assert len(litellm_spans) == 1, f"Expected 1 litellm_request span, found {len(litellm_spans)}"
|
||||
litellm_span = litellm_spans[0]
|
||||
|
||||
# Verify gen_ai.prompt attributes
|
||||
assert "gen_ai.prompt.0.role" in litellm_span.attributes, "gen_ai.prompt.0.role not found"
|
||||
assert litellm_span.attributes["gen_ai.prompt.0.role"] == "user", "Expected user role in prompt"
|
||||
assert "gen_ai.prompt.0.content" in litellm_span.attributes, "gen_ai.prompt.0.content not found"
|
||||
assert litellm_span.attributes["gen_ai.prompt.0.content"] == "Repeat after me: Hello, world!"
|
||||
|
||||
# Verify gen_ai.completion attributes
|
||||
assert "gen_ai.completion.0.role" in litellm_span.attributes, "gen_ai.completion.0.role not found"
|
||||
assert litellm_span.attributes["gen_ai.completion.0.role"] == "assistant", "Expected assistant role in completion"
|
||||
assert "gen_ai.completion.0.content" in litellm_span.attributes, "gen_ai.completion.0.content not found"
|
||||
assert "gen_ai.completion.0.finish_reason" in litellm_span.attributes, "gen_ai.completion.0.finish_reason not found"
|
||||
|
||||
|
||||
def _make_proxy_and_store(qwen25_model: RemoteOpenAIServer, *, retries: int = 0):
|
||||
store = InMemoryLightningStore()
|
||||
proxy = LLMProxy(
|
||||
port=_get_free_port(),
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "gpt-4o-arbitrary",
|
||||
"litellm_params": {
|
||||
"model": "hosted_vllm/" + qwen25_model.model,
|
||||
"api_base": qwen25_model.url_for("v1"),
|
||||
},
|
||||
}
|
||||
],
|
||||
store=store,
|
||||
num_retries=retries,
|
||||
)
|
||||
proxy.start()
|
||||
return proxy, store
|
||||
|
||||
|
||||
async def _new_resource(proxy: LLMProxy, store: InMemoryLightningStore):
|
||||
rollout = await store.start_rollout(None)
|
||||
return proxy.as_resource(rollout.rollout_id, rollout.attempt.attempt_id), rollout
|
||||
|
||||
|
||||
def _get_client_for_resource(resource: LLM):
|
||||
return openai.OpenAI(base_url=resource.endpoint, api_key="token-abc123", timeout=120, max_retries=0)
|
||||
|
||||
|
||||
def _get_async_client_for_resource(resource: LLM):
|
||||
return openai.AsyncOpenAI(base_url=resource.endpoint, api_key="token-abc123", timeout=120, max_retries=0)
|
||||
|
||||
|
||||
def _find_span(spans: list[Span], name: str):
|
||||
return [s for s in spans if s.name == name]
|
||||
|
||||
|
||||
def _attr(s: Span, key: str, default: Any = None): # type: ignore
|
||||
return s.attributes.get(key, default)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_multiple_requests_one_attempt(qwen25_model: RemoteOpenAIServer):
|
||||
proxy, store = _make_proxy_and_store(qwen25_model)
|
||||
try:
|
||||
resource, rollout = await _new_resource(proxy, store)
|
||||
client = _get_client_for_resource(resource)
|
||||
|
||||
for i in range(3):
|
||||
r = client.chat.completions.create(
|
||||
model="gpt-4o-arbitrary",
|
||||
messages=[{"role": "user", "content": f"Say ping {i}"}],
|
||||
stream=False,
|
||||
)
|
||||
assert r.choices[0].message.content
|
||||
|
||||
spans = await store.query_spans(rollout.rollout_id, rollout.attempt.attempt_id)
|
||||
assert len(spans) > 0
|
||||
# Different requests have different sequence_ids
|
||||
assert {s.sequence_id for s in spans} == {1, 2, 3}
|
||||
# At least 3 requests recorded
|
||||
assert len(_find_span(spans, "raw_gen_ai_request")) == 3
|
||||
# TODO: Check response contents and token ids for the 3 requests respectively
|
||||
finally:
|
||||
proxy.stop()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ten_concurrent_requests(qwen25_model: RemoteOpenAIServer):
|
||||
proxy, store = _make_proxy_and_store(qwen25_model)
|
||||
try:
|
||||
resource, rollout = await _new_resource(proxy, store)
|
||||
aclient = _get_async_client_for_resource(resource)
|
||||
|
||||
async def _one(i: int):
|
||||
r = await aclient.chat.completions.create(
|
||||
model="gpt-4o-arbitrary",
|
||||
messages=[{"role": "user", "content": f"Return #{i}"}],
|
||||
stream=False,
|
||||
)
|
||||
return r.choices[0].message.content
|
||||
|
||||
outs = await asyncio.gather(*[_one(i) for i in range(10)])
|
||||
assert len([o for o in outs if o]) == 10
|
||||
|
||||
spans = await store.query_spans(rollout.rollout_id, rollout.attempt.attempt_id)
|
||||
assert len(_find_span(spans, "raw_gen_ai_request")) == 10
|
||||
assert {s.sequence_id for s in spans} == {1, 2, 3, 4, 5, 6, 7, 8, 9, 10}
|
||||
# TODO: Check whether the sequence ids get mixed up or not
|
||||
finally:
|
||||
proxy.stop()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_anthropic_client_compat(qwen25_model: RemoteOpenAIServer):
|
||||
# litellm proxy accepts Anthropic schema and forwards to OpenAI backend
|
||||
proxy, store = _make_proxy_and_store(qwen25_model)
|
||||
try:
|
||||
resource, rollout = await _new_resource(proxy, store)
|
||||
|
||||
a = anthropic.Anthropic(base_url=resource.endpoint, api_key="token-abc123", timeout=120)
|
||||
msg = a.messages.create(
|
||||
model="gpt-4o-arbitrary",
|
||||
max_tokens=64,
|
||||
messages=[{"role": "user", "content": "Respond with the word: OK"}],
|
||||
)
|
||||
# Anthropic SDK returns content list
|
||||
txt = "".join([b.text for b in msg.content if b.type == "text"])
|
||||
assert "OK" in txt.upper()
|
||||
|
||||
spans = await store.query_spans(rollout.rollout_id, rollout.attempt.attempt_id)
|
||||
assert len(spans) > 0
|
||||
finally:
|
||||
proxy.stop()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_tool_call_roundtrip(qwen25_model: RemoteOpenAIServer):
|
||||
proxy, store = _make_proxy_and_store(qwen25_model)
|
||||
try:
|
||||
resource, rollout = await _new_resource(proxy, store)
|
||||
client = _get_client_for_resource(resource)
|
||||
|
||||
tools = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "echo",
|
||||
"description": "Echo a string",
|
||||
"parameters": {"type": "object", "properties": {"text": {"type": "string"}}, "required": ["text"]},
|
||||
},
|
||||
}
|
||||
]
|
||||
|
||||
r1 = client.chat.completions.create(
|
||||
model="gpt-4o-arbitrary",
|
||||
messages=[{"role": "user", "content": "Call the echo tool with text=hello"}],
|
||||
tools=cast(Any, tools),
|
||||
tool_choice="auto",
|
||||
stream=False,
|
||||
)
|
||||
# If the small model does not tool-call, skip gracefully
|
||||
tool_calls = r1.choices[0].message.tool_calls or []
|
||||
if not tool_calls:
|
||||
pytest.skip("model did not emit tool calls in this environment")
|
||||
|
||||
call = tool_calls[0]
|
||||
assert call.type == "function"
|
||||
assert call.function and call.function.name == "echo"
|
||||
args = json.loads(call.function.arguments)
|
||||
assert "text" in args
|
||||
|
||||
r2 = client.chat.completions.create(
|
||||
model="gpt-4o-arbitrary",
|
||||
messages=cast(
|
||||
Any,
|
||||
[
|
||||
{"role": "user", "content": "Call the echo tool with text=hello"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": call.id,
|
||||
"type": "function",
|
||||
"function": {"name": "echo", "arguments": call.function.arguments},
|
||||
}
|
||||
],
|
||||
},
|
||||
{"role": "tool", "tool_call_id": call.id, "name": "echo", "content": args["text"]},
|
||||
],
|
||||
),
|
||||
stream=False,
|
||||
)
|
||||
assert args["text"] in (r2.choices[0].message.content or "")
|
||||
|
||||
spans = await store.query_spans(rollout.rollout_id, rollout.attempt.attempt_id)
|
||||
assert len(_find_span(spans, "litellm_request")) == 2
|
||||
assert len(_find_span(spans, "raw_gen_ai_request")) == 2
|
||||
|
||||
# TODO: Check response contents and token ids for the 2 requests respectively
|
||||
finally:
|
||||
proxy.stop()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_chunks(qwen25_model: RemoteOpenAIServer):
|
||||
proxy, store = _make_proxy_and_store(qwen25_model)
|
||||
try:
|
||||
resource, rollout = await _new_resource(proxy, store)
|
||||
client = _get_client_for_resource(resource)
|
||||
|
||||
stream = client.chat.completions.create(
|
||||
model="gpt-4o-arbitrary",
|
||||
messages=[{"role": "user", "content": "Say the word 'apple'"}],
|
||||
stream=True,
|
||||
)
|
||||
collected: list[str] = []
|
||||
for evt in stream:
|
||||
for c in evt.choices:
|
||||
if c.delta and getattr(c.delta, "content", None):
|
||||
assert isinstance(c.delta.content, str)
|
||||
collected.append(c.delta.content)
|
||||
assert "apple" in "".join(collected).lower()
|
||||
|
||||
spans = await store.query_spans(rollout.rollout_id, rollout.attempt.attempt_id)
|
||||
assert len(spans) > 0
|
||||
# TODO: didn't test the token ids in streaming chunks here
|
||||
finally:
|
||||
proxy.stop()
|
||||
|
||||
|
||||
class _FakeSpanContext:
|
||||
def __init__(self, span_id: int):
|
||||
self.span_id = span_id
|
||||
|
||||
|
||||
class _FakeParent:
|
||||
def __init__(self, span_id: int):
|
||||
self.span_id = span_id
|
||||
|
||||
|
||||
class _FakeReadableSpan:
|
||||
def __init__(self, span_id: int, parent_id: int | None, attrs: dict[str, str]):
|
||||
self._ctx = _FakeSpanContext(span_id)
|
||||
self.parent = None if parent_id is None else _FakeParent(parent_id)
|
||||
self.attributes = attrs
|
||||
self.name = f"span-{span_id}"
|
||||
|
||||
def get_span_context(self):
|
||||
return self._ctx
|
||||
|
||||
|
||||
class _FakeStore(InMemoryLightningStore):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.added: list[tuple[str, str, int, _FakeReadableSpan]] = []
|
||||
|
||||
async def add_otel_span(
|
||||
self, rollout_id: str, attempt_id: str, readable_span: ReadableSpan, sequence_id: int | None = None
|
||||
) -> Span:
|
||||
assert isinstance(sequence_id, int)
|
||||
assert isinstance(readable_span, _FakeReadableSpan)
|
||||
self.added.append((rollout_id, attempt_id, sequence_id, readable_span))
|
||||
return cast(Span, None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_exporter_tree_and_flush_headers_parsing():
|
||||
store = _FakeStore()
|
||||
exporter = LightningSpanExporter(store)
|
||||
|
||||
# Build a root and two children. Headers distributed across spans.
|
||||
root = _FakeReadableSpan(1, None, {"metadata.requester_custom_headers": "{'x-rollout-id': 'r1'}"})
|
||||
child_a = _FakeReadableSpan(2, 1, {"metadata.requester_custom_headers": "{'x-attempt-id': 'a9'}"})
|
||||
child_b = _FakeReadableSpan(3, 1, {"metadata.requester_custom_headers": "{'x-sequence-id': '7'}"})
|
||||
|
||||
# Push to buffer and export
|
||||
res = exporter.export(cast(List[ReadableSpan], [root, child_a, child_b]))
|
||||
assert res.name == "SUCCESS"
|
||||
|
||||
# Give event loop a moment to run exporter coroutine
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
# Should have flushed all three with merged headers
|
||||
assert len(store.added) == 3
|
||||
for rid, aid, sid, sp in store.added:
|
||||
assert rid == "r1"
|
||||
assert aid == "a9"
|
||||
assert sid == 7
|
||||
assert isinstance(sp, _FakeReadableSpan)
|
||||
|
||||
exporter.shutdown()
|
||||
|
||||
|
||||
def test_exporter_helpers():
|
||||
store = _FakeStore()
|
||||
exporter = LightningSpanExporter(store)
|
||||
|
||||
# Tree: 10(root) -> 11(child) -> 12(grandchild); 20(root2)
|
||||
s10 = _FakeReadableSpan(10, None, {})
|
||||
s11 = _FakeReadableSpan(11, 10, {})
|
||||
s12 = _FakeReadableSpan(12, 11, {})
|
||||
s20 = _FakeReadableSpan(20, None, {})
|
||||
|
||||
for _ in range(10):
|
||||
exporter._buffer = cast(List[ReadableSpan], [s10, s11, s12, s20]) # pyright: ignore[reportPrivateUsage]
|
||||
random.shuffle(exporter._buffer) # pyright: ignore[reportPrivateUsage]
|
||||
|
||||
roots = list(exporter._get_root_span_ids()) # pyright: ignore[reportPrivateUsage]
|
||||
assert set(roots) == {10, 20}
|
||||
|
||||
subtree_ids = set(exporter._get_subtrees(10)) # pyright: ignore[reportPrivateUsage]
|
||||
assert subtree_ids == {10, 11, 12}
|
||||
|
||||
popped = exporter._pop_subtrees(10) # pyright: ignore[reportPrivateUsage]
|
||||
assert {sp.get_span_context().span_id for sp in popped} == { # pyright: ignore[reportOptionalMemberAccess]
|
||||
10,
|
||||
11,
|
||||
12,
|
||||
}
|
||||
# Remaining buffer has only s20
|
||||
assert {
|
||||
sp.get_span_context().span_id # pyright: ignore[reportOptionalMemberAccess]
|
||||
for sp in exporter._buffer # pyright: ignore[reportPrivateUsage]
|
||||
} == {20}
|
||||
|
||||
exporter.shutdown()
|
||||
|
||||
# TODO: add more complex tests for the exporter helper
|
||||
@@ -68,6 +68,8 @@ from agentlightning.tracer.agentops import AgentOpsTracer, LightningSpanProcesso
|
||||
from agentlightning.tracer.http import HttpTracer
|
||||
from agentlightning.types import Triplet
|
||||
|
||||
from .tracer.utils import clear_tracer_provider
|
||||
|
||||
USE_OPENAI = os.environ.get("USE_OPENAI", "false").lower() == "true"
|
||||
if USE_OPENAI:
|
||||
OPENAI_BASE_URL = os.environ.get("OPENAI_BASE_URL", os.environ["OPENAI_API_BASE"])
|
||||
@@ -783,6 +785,12 @@ def create_prompt_caches() -> None:
|
||||
run_all()
|
||||
|
||||
|
||||
@pytest.fixture(scope="module", autouse=True)
|
||||
def setup_module():
|
||||
clear_tracer_provider()
|
||||
yield
|
||||
|
||||
|
||||
@pytest.mark.parametrize("agent_func", list(iterate_over_agents()), ids=lambda f: f.__name__)
|
||||
def test_run_with_agentops_tracer(agent_func):
|
||||
tracer = AgentOpsTracer()
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
@@ -0,0 +1,27 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import opentelemetry.trace as trace_api
|
||||
|
||||
|
||||
# pyright: reportPrivateUsage=false
|
||||
def clear_tracer_provider() -> None:
|
||||
"""OpenTelemetry tracer provider does not allow set twice.
|
||||
|
||||
Reset the tracer provider to allow setting it again in tests.
|
||||
This is a hack to the internal state of OpenTelemetry.
|
||||
Not a good idea to use in production.
|
||||
Always remember to setup two processes if you need two tracers.
|
||||
"""
|
||||
|
||||
if hasattr(trace_api, "_TRACER_PROVIDER"):
|
||||
if trace_api._TRACER_PROVIDER is not None:
|
||||
trace_api._TRACER_PROVIDER = None
|
||||
|
||||
if hasattr(trace_api, "_TRACER_PROVIDER_SET_ONCE"):
|
||||
if hasattr(trace_api._TRACER_PROVIDER_SET_ONCE, "_done"):
|
||||
if trace_api._TRACER_PROVIDER_SET_ONCE._done:
|
||||
trace_api._TRACER_PROVIDER_SET_ONCE._done = False
|
||||
|
||||
if hasattr(trace_api._TRACER_PROVIDER_SET_ONCE, "_flag"):
|
||||
if trace_api._TRACER_PROVIDER_SET_ONCE._flag: # type: ignore
|
||||
trace_api._TRACER_PROVIDER_SET_ONCE._flag = False # type: ignore
|
||||
Reference in New Issue
Block a user