chore: import upstream snapshot with attribution

tests/llm_proxy/test_llm_proxy_cpu.py

@@ -0,0 +1,382 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import asyncio
+import logging
+import multiprocessing
+import random
+from typing import Any, List, cast
+
+import litellm
+import openai
+import opentelemetry.trace as trace_api
+import pytest
+from agentops.sdk.core import BatchSpanProcessor
+from litellm.llms.custom_llm import CustomLLM
+from litellm.types.utils import ModelResponse
+from litellm.utils import custom_llm_setup
+from opentelemetry.sdk.trace import ReadableSpan
+from opentelemetry.sdk.trace.export import SimpleSpanProcessor
+
+from agentlightning.llm_proxy import LightningSpanExporter, LLMProxy
+from agentlightning.store import LightningStoreServer
+from agentlightning.store.memory import InMemoryLightningStore
+from agentlightning.store.threading import LightningStoreThreaded
+from agentlightning.types import Span
+from agentlightning.utils.server_launcher import PythonServerLauncherArgs
+
+from ..common.network import get_free_port
+from ..common.tracer import clear_tracer_provider
+
+pytestmark = pytest.mark.llmproxy
+
+
+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
+
+
+@pytest.mark.asyncio
+async def test_update_model_list():
+    store = InMemoryLightningStore()
+    proxy = LLMProxy(
+        model_list=[
+            {
+                "model_name": "gpt-4o-arbitrary",
+                "litellm_params": {
+                    "model": "openai/gpt-4o",
+                },
+            }
+        ],
+        launch_mode="asyncio",
+        port=get_free_port(),
+        store=store,
+    )
+    await proxy.start()
+    assert proxy.is_running()
+    assert proxy.model_list == [
+        {
+            "model_name": "gpt-4o-arbitrary",
+            "litellm_params": {
+                "model": "openai/gpt-4o",
+            },
+        }
+    ]
+    proxy.update_model_list(
+        [
+            {
+                "model_name": "gpt-4o-arbitrary",
+                "litellm_params": {
+                    "model": "openai/gpt-4o-mini",
+                },
+            }
+        ]
+    )
+    assert proxy.model_list == [
+        {
+            "model_name": "gpt-4o-arbitrary",
+            "litellm_params": {
+                "model": "openai/gpt-4o-mini",
+            },
+        }
+    ]
+    assert proxy.is_running()
+    await proxy.stop()
+
+
+@pytest.mark.asyncio
+async def test_restart_resets_litellm_logging_worker() -> None:
+    """LLMProxy.start() should recreate LiteLLM's logging worker on each run."""
+
+    try:
+        from litellm.litellm_core_utils import logging_worker as litellm_logging_worker
+    except ImportError:
+        pytest.skip("LiteLLM logging worker not available")
+
+    store = InMemoryLightningStore()
+    proxy = LLMProxy(
+        model_list=[
+            {
+                "model_name": "dummy-model",
+                # The backend is never invoked; only the proxy lifecycle matters here.
+                "litellm_params": {"model": "gpt-3.5-turbo"},
+            }
+        ],
+        store=store,
+        launcher_args=PythonServerLauncherArgs(
+            port=get_free_port(),
+            launch_mode="asyncio",
+            healthcheck_url="/health",
+            startup_timeout=10.0,
+            process_join_timeout=10.0,
+        ),
+    )
+
+    try:
+        await proxy.start()
+        first_worker = litellm_logging_worker.GLOBAL_LOGGING_WORKER
+        await proxy.stop()
+
+        await proxy.start()
+        second_worker = litellm_logging_worker.GLOBAL_LOGGING_WORKER
+    finally:
+        await proxy.stop()
+
+    assert first_worker is not second_worker, "LiteLLM logging worker should be refreshed after restart"
+
+
+class TestLLM(CustomLLM):
+    def __init__(self, content: str) -> None:
+        super().__init__()
+        self.content = content
+
+    def completion(self, *args: Any, **kwargs: Any) -> ModelResponse:
+        return litellm.completion(  # type: ignore
+            model="gpt-4o",
+            messages=[{"role": "user", "content": "Hello world"}],
+            mock_response=self.content,
+        )
+
+    async def acompletion(self, *args: Any, **kwargs: Any) -> ModelResponse:
+        return litellm.completion(  # type: ignore
+            model="gpt-4o",
+            messages=[{"role": "user", "content": "Hello world"}],
+            mock_response=self.content,
+        )
+
+
+@pytest.mark.asyncio
+async def test_custom_llm_restarted_multiple_times(caplog: pytest.LogCaptureFixture) -> None:
+    clear_tracer_provider()
+
+    restart_times: int = 30
+
+    store = LightningStoreThreaded(InMemoryLightningStore())
+    caplog.set_level(logging.WARNING)
+
+    port = get_free_port()
+    try:
+        llm_proxy = LLMProxy(
+            model_list=[
+                {
+                    "model_name": "gpt-4o-arbitrary",
+                    "litellm_params": {
+                        # NOTE: The model after "/" cannot be an openai model like gpt-4o
+                        # This might be a bug with litellm
+                        "model": "test-llm/any-llm",
+                    },
+                }
+            ],
+            launcher_args=PythonServerLauncherArgs(
+                launch_mode="thread",
+                healthcheck_url="/health",
+                port=port,
+            ),
+            store=store,
+        )
+        for restart_idx in range(restart_times):
+            llm_instance = TestLLM(f"Hi! {restart_idx}")
+            litellm.custom_provider_map = [{"provider": "test-llm", "custom_handler": llm_instance}]
+            custom_llm_setup()
+            await llm_proxy.restart()
+            assert llm_proxy.is_running()
+
+            openai_client = openai.AsyncOpenAI(
+                base_url=llm_proxy.server_launcher.access_endpoint,
+                api_key="token-abc123",
+                timeout=5,
+                max_retries=0,
+            )
+            response = await openai_client.chat.completions.create(
+                model="gpt-4o-arbitrary",
+                messages=[{"role": "user", "content": "Hello world"}],
+                stream=False,
+            )
+            assert response.choices[0].message.content == f"Hi! {restart_idx}"
+
+            error_logs = [record.message for record in caplog.records if record.levelno >= logging.ERROR]
+            assert not error_logs, f"Found error logs: {error_logs}"
+            assert not any("Cannot add callback" in record.message for record in caplog.records)
+
+        await llm_proxy.stop()
+    finally:
+        litellm.custom_provider_map = []
+        custom_llm_setup()
+
+
+async def llm_proxy_span_exporter_loop(otlp_enabled: bool = False):
+    store = LightningStoreThreaded(InMemoryLightningStore())
+
+    if otlp_enabled:
+        store = LightningStoreServer(store, "127.0.0.1", get_free_port())
+        await store.start()
+
+    llm_instance = TestLLM(f"Hi! I'm a test LLM")
+    litellm.custom_provider_map = [{"provider": "test-llm", "custom_handler": llm_instance}]
+    custom_llm_setup()
+    proxy = LLMProxy(
+        launcher_args=PythonServerLauncherArgs(
+            launch_mode="thread",
+            healthcheck_url="/health",
+            port=get_free_port(),
+        ),
+        store=store,
+        model_list=[
+            {
+                "model_name": "gpt-4o-arbitrary",
+                "litellm_params": {
+                    "model": "test-llm/any-llm",
+                },
+            }
+        ],
+    )
+    await proxy.start()
+
+    rollout = await store.start_rollout(None)
+    resource = proxy.as_resource(rollout.rollout_id, rollout.attempt.attempt_id)
+
+    client = openai.AsyncOpenAI(
+        base_url=resource.endpoint,
+        api_key="token-abc123",
+        timeout=5,
+        max_retries=0,
+    )
+    response = await client.chat.completions.create(
+        model="gpt-4o-arbitrary",
+        messages=[{"role": "user", "content": "Hello world"}],
+        stream=False,
+    )
+    assert response.choices[0].message.content == "Hi! I'm a test LLM"
+
+    spans = await store.query_spans(rollout.rollout_id, rollout.attempt.attempt_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"
+
+    tracer_provider = trace_api.get_tracer_provider()
+
+    have_asserted_loop = False
+    for span_processor in tracer_provider._active_span_processor._span_processors:  # type: ignore
+        if isinstance(span_processor, (SimpleSpanProcessor, BatchSpanProcessor)):
+            if isinstance(span_processor.span_exporter, LightningSpanExporter):
+                if otlp_enabled:
+                    assert span_processor.span_exporter._loop is None  # type: ignore
+                else:
+                    assert span_processor.span_exporter._loop is not None  # type: ignore
+                have_asserted_loop = True
+                break
+    assert have_asserted_loop, f"LightningSpanExporter should be used with otlp_enabled={otlp_enabled}"
+
+    await proxy.stop()
+
+    if isinstance(store, LightningStoreServer):
+        await store.stop()
+
+
+def llm_proxy_span_exporter_loop_sync(otlp_enabled: bool = False):
+    asyncio.run(llm_proxy_span_exporter_loop(otlp_enabled))
+
+
+@pytest.mark.parametrize("otlp_enabled", [True, False])
+def test_llm_proxy_span_exporter_loop(otlp_enabled: bool):
+    context = multiprocessing.get_context("spawn")
+    process = context.Process(target=llm_proxy_span_exporter_loop_sync, args=(otlp_enabled,))
+    process.start()
+    process.join(timeout=30.0)
+    assert process.exitcode == 0