chore: import upstream snapshot with attribution

tests/runner/__init__.py

@@ -0,0 +1 @@
+# Copyright (c) Microsoft. All rights reserved.

tests/runner/test_agent_integration.py

@@ -0,0 +1,264 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import os
+from contextlib import asynccontextmanager
+from typing import Any, AsyncGenerator, Awaitable, Dict, Optional, cast
+
+import litellm
+import openai
+import pytest
+
+from agentlightning.litagent import LitAgent
+from agentlightning.llm_proxy import LLMProxy
+from agentlightning.reward import emit_reward
+from agentlightning.runner import LitAgentRunner
+from agentlightning.store.client_server import LightningStoreClient, LightningStoreServer
+from agentlightning.store.memory import InMemoryLightningStore
+from agentlightning.tracer.agentops import AgentOpsTracer
+from agentlightning.types import LLM, AttemptedRollout, NamedResources, Rollout
+
+from ..common.network import get_free_port
+from ..common.tracer import clear_tracer_provider
+from ..common.vllm import VLLM_AVAILABLE, RemoteOpenAIServer
+
+
+class InitRunnerFunction:
+
+    def __call__(
+        self,
+        agent: LitAgent[Any],
+        *,
+        resources: Optional[Dict[str, LLM]] = None,
+    ) -> Awaitable[tuple[LitAgentRunner[Any], InMemoryLightningStore]]: ...
+
+
+@pytest.fixture(
+    params=[
+        pytest.param("agentops", marks=pytest.mark.agentops),
+    ]
+)
+def init_runner(request: pytest.FixtureRequest) -> InitRunnerFunction:
+    async def init_runner_fn(
+        agent: LitAgent[Any],
+        *,
+        resources: Optional[Dict[str, LLM]] = None,
+    ) -> tuple[LitAgentRunner[Any], InMemoryLightningStore]:
+        store = InMemoryLightningStore()
+        llm_resource: NamedResources = resources or {"llm": LLM(endpoint="http://localhost", model="dummy")}  # type: ignore[assignment]
+        await store.update_resources("default", llm_resource)
+
+        # This is always AgentOpsTracer for now
+        runner = LitAgentRunner[Any](tracer=AgentOpsTracer(), poll_interval=0.01)
+        runner.init(agent)
+        runner.init_worker(worker_id=0, store=store)
+        return runner, store
+
+    return init_runner_fn  # type: ignore
+
+
+def teardown_runner(runner: LitAgentRunner[Any]) -> None:
+    runner.teardown_worker(worker_id=0)
+    runner.teardown()
+
+
+@pytest.fixture(scope="module", autouse=True)
+def setup_module():
+    # This must execute only once for this module.
+    # Once agentops tracer is initialized, it cannot be reset,
+    # otherwise it will never be rewired.
+    clear_tracer_provider()
+    yield
+
+
+@pytest.mark.asyncio
+async def test_runner_integration_basic_rollout(init_runner: InitRunnerFunction) -> None:
+    class EchoAgent(LitAgent[str]):
+        async def validation_rollout_async(self, task: str, resources: Dict[str, Any], rollout: Any) -> None:
+            emit_reward(1.0)
+
+    agent = EchoAgent()
+    runner, store = await init_runner(agent)
+    try:
+        await runner.step("hello integration")
+    finally:
+        teardown_runner(runner)
+
+    rollouts = await store.query_rollouts()
+    assert rollouts and rollouts[0].status == "succeeded"
+    attempts = await store.query_attempts(rollouts[0].rollout_id)
+    spans = await store.query_spans(rollouts[0].rollout_id, attempts[-1].attempt_id)
+    assert any(span.attributes.get("agentlightning.reward.0.value") == 1.0 for span in spans)
+
+
+@pytest.mark.asyncio
+@pytest.mark.openai
+async def test_runner_integration_with_openai(init_runner: InitRunnerFunction) -> None:
+    class OpenAIAgent(LitAgent[str]):
+        async def validation_rollout_async(self, task: str, resources: NamedResources, rollout: Rollout) -> float:
+            llm = cast(LLM, resources["llm"])
+            client = openai.AsyncOpenAI(base_url=llm.endpoint, api_key=llm.api_key)
+            response = await client.chat.completions.create(
+                model=llm.model,
+                messages=[{"role": "user", "content": task}],
+            )
+            assert response.choices, "OpenAI response should contain choices"
+            return 0.0
+
+    if not (os.getenv("OPENAI_BASE_URL") and os.getenv("OPENAI_API_KEY")):
+        raise RuntimeError("OpenAI endpoint or key not configured")
+
+    base_url = os.environ["OPENAI_BASE_URL"]
+    api_key = os.environ["OPENAI_API_KEY"]
+    model = os.getenv("OPENAI_MODEL", "gpt-4o-mini")
+
+    agent = OpenAIAgent()
+    resources = {"llm": LLM(endpoint=base_url, model=model, api_key=api_key)}
+    runner, store = await init_runner(agent, resources=resources)
+    try:
+        await runner.step("Say hello in one word")
+    finally:
+        teardown_runner(runner)
+
+    rollouts = await store.query_rollouts()
+    assert rollouts and rollouts[0].status == "succeeded"
+
+
+@pytest.mark.asyncio
+@pytest.mark.openai
+@pytest.mark.skipif(
+    not (os.getenv("OPENAI_BASE_URL") and os.getenv("OPENAI_API_KEY")),
+    reason="OpenAI endpoint or key not configured",
+)
+async def test_runner_integration_with_litellm_proxy(init_runner: InitRunnerFunction) -> None:
+    class LiteLLMAgent(LitAgent[str]):
+        def validation_rollout(self, task: str, resources: NamedResources, rollout: Rollout) -> float:
+            llm = cast(LLM, resources["llm"])
+            response = litellm.completion(  # type: ignore
+                model=llm.model,
+                messages=[{"role": "user", "content": task}],
+            )
+            assert response.get("choices"), "litellm proxy should return choices"  # type: ignore
+            return 0.0
+
+    if not (os.getenv("OPENAI_BASE_URL") and os.getenv("OPENAI_API_KEY")):
+        raise RuntimeError("OpenAI endpoint or key not configured")
+
+    agent = LiteLLMAgent()
+    resources = {"llm": LLM(endpoint="http://dummy", model="openai/gpt-4o-mini")}
+    runner, store = await init_runner(agent, resources=resources)
+    try:
+        await runner.step("Give me a short greeting")
+    finally:
+        teardown_runner(runner)
+
+    rollouts = await store.query_rollouts()
+    assert rollouts and rollouts[0].status == "succeeded"
+
+
+@pytest.fixture
+def server():
+    if not VLLM_AVAILABLE:
+        pytest.skip("vLLM is not available")
+    vllm_port = get_free_port()
+    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(vllm_port),
+        ],
+    ) as server:
+        yield server
+
+
+class LLMProxyWithClearedTracerProvider(LLMProxy):
+    """LLMProxy that clears the tracer provider before serving.
+
+    It will be run in a separate process, so the tracer provider initialized there does not
+    interfere with the main process's tracer provider.
+    """
+
+    @asynccontextmanager
+    async def _serve_context(self) -> AsyncGenerator[None, None]:
+        # This will be run inside the LLM proxy's own process
+        clear_tracer_provider()
+        async with super()._serve_context():
+            yield
+
+
+@pytest.mark.asyncio
+@pytest.mark.gpu
+async def test_runner_integration_with_spawned_litellm_proxy(
+    init_runner: InitRunnerFunction, server: RemoteOpenAIServer
+) -> None:
+    class ProxyAgent(LitAgent[str]):
+        async def validation_rollout_async(self, task: str, resources: NamedResources, rollout: Rollout) -> float:
+            attempted_rollout = cast(AttemptedRollout, rollout)
+            llm_resource = cast(LLM, resources["llm"])
+            client = openai.AsyncOpenAI(
+                base_url=llm_resource.get_base_url(attempted_rollout.rollout_id, attempted_rollout.attempt.attempt_id),
+                api_key="dummy",
+            )
+            response = await client.chat.completions.create(
+                model=llm_resource.model,
+                messages=[{"role": "user", "content": task}],
+            )
+            assert response.choices, "Proxy should return at least one choice"
+            return 0.5
+
+    agent = ProxyAgent()
+    runner, store = await init_runner(agent)
+
+    server_store = LightningStoreServer(store=store, host="127.0.0.1", port=get_free_port())
+    await server_store.start()
+    client_store = LightningStoreClient(server_store.endpoint)
+
+    proxy = LLMProxyWithClearedTracerProvider(
+        port=get_free_port(),
+        model_list=[
+            {
+                "model_name": "gpt-4o-arbitrary",
+                "litellm_params": {
+                    "model": "hosted_vllm/" + server.model,
+                    "api_base": server.url_for("v1"),
+                },
+            }
+        ],
+        store=client_store,
+    )
+
+    await proxy.start()
+
+    try:
+        await runner.step("Say hello to Agent Lightning", resources={"llm": proxy.as_resource()})
+
+        rollouts = await client_store.query_rollouts()
+        assert rollouts and rollouts[0].status == "succeeded"
+
+        spans = await client_store.query_spans(rollouts[0].rollout_id, "latest")
+        assert len(spans) > 1
+        first_spans = [span for span in spans if span.sequence_id == 1]
+        assert len(first_spans) > 1
+        assert any("llm.hosted_vllm.choices" in span.attributes for span in first_spans)
+        assert any("llm.hosted_vllm.prompt_token_ids" in span.attributes for span in first_spans)
+        assert any("gen_ai.prompt.0.content" in span.attributes for span in first_spans)
+
+        second_spans = [span for span in spans if span.sequence_id == 2]
+        assert len(second_spans) == 1
+        assert second_spans[0].name == "openai.chat.completion"
+
+        last_spans = [span for span in spans if span.sequence_id == max(span.sequence_id for span in spans)]
+        assert len(last_spans) == 1
+        assert last_spans[0].name == "agentlightning.annotation"
+        assert (
+            last_spans[0].attributes.get("agentlightning.reward.0.value") == 0.5
+        ), f"Expected reward to be 0.5, found {last_spans[0].attributes}"
+    finally:
+        teardown_runner(runner)
+        await proxy.stop()
+        await client_store.close()
+        await server_store.stop()

tests/runner/test_runner_context.py

@@ -0,0 +1,320 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from contextlib import asynccontextmanager
+from typing import Any, AsyncGenerator, Dict, List, Optional
+
+import pytest
+from opentelemetry import trace as trace_api
+from opentelemetry.sdk.trace import TracerProvider
+
+from agentlightning.litagent import LitAgent
+from agentlightning.runner import LitAgentRunner
+from agentlightning.store.base import LightningStore
+from agentlightning.store.memory import InMemoryLightningStore
+from agentlightning.tracer.base import Tracer
+from agentlightning.types import LLM, Hook, Rollout, Span
+
+from ..common.tracer import clear_tracer_provider
+
+
+@pytest.fixture(scope="module", autouse=True)
+def setup_module():
+    """Setup the tracer provider for the tests."""
+
+    clear_tracer_provider()
+
+    trace_api.set_tracer_provider(TracerProvider())
+
+
+class DummyTracer(Tracer):
+    """Minimal tracer for testing."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self._last_trace: List[Span] = []
+        self.init_called = False
+        self.init_worker_called = False
+        self.teardown_called = False
+        self.teardown_worker_called = False
+
+    def init(self, *args: Any, **kwargs: Any) -> None:
+        self.init_called = True
+        self._last_trace.clear()
+
+    def init_worker(self, worker_id: int, *args: Any, **kwargs: Any) -> None:
+        self.init_worker_called = True
+
+    def teardown(self, *args: Any, **kwargs: Any) -> None:
+        self.teardown_called = True
+        self._last_trace.clear()
+
+    def teardown_worker(self, worker_id: int, *args: Any, **kwargs: Any) -> None:
+        self.teardown_worker_called = True
+
+    def get_last_trace(self) -> List[Span]:
+        return list(self._last_trace)
+
+    @asynccontextmanager
+    async def trace_context(
+        self,
+        name: Optional[str] = None,
+        *,
+        store: Optional[LightningStore] = None,
+        rollout_id: Optional[str] = None,
+        attempt_id: Optional[str] = None,
+    ) -> AsyncGenerator[List[Span], None]:
+        self._last_trace = []
+        try:
+            yield self._last_trace
+        finally:
+            pass
+
+
+class DummyAgent(LitAgent[Dict[str, Any]]):
+    """Minimal agent for testing."""
+
+    def validation_rollout(self, task: Dict[str, Any], resources: Dict[str, Any], rollout: Any) -> None:
+        return None
+
+
+class RecordingHook(Hook):
+    """Hook that records lifecycle events."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.calls: List[str] = []
+
+    async def on_rollout_start(self, *, agent: LitAgent[Any], runner: Any, rollout: Rollout) -> None:
+        self.calls.append("on_rollout_start")
+
+    async def on_trace_start(self, *, agent: LitAgent[Any], runner: Any, tracer: Tracer, rollout: Rollout) -> None:
+        self.calls.append("on_trace_start")
+
+    async def on_trace_end(self, *, agent: LitAgent[Any], runner: Any, tracer: Tracer, rollout: Rollout) -> None:
+        self.calls.append("on_trace_end")
+
+    async def on_rollout_end(self, *, agent: LitAgent[Any], runner: Any, rollout: Rollout, spans: Any) -> None:
+        self.calls.append("on_rollout_end")
+
+
+@pytest.mark.asyncio
+async def test_run_context_basic_lifecycle() -> None:
+    """Test that run_context properly initializes and tears down the runner."""
+    tracer = DummyTracer()
+    agent = DummyAgent()
+    store = InMemoryLightningStore()
+    runner = LitAgentRunner[Dict[str, Any]](tracer=tracer)
+
+    with runner.run_context(agent=agent, store=store):
+        # Verify initialization happened
+        assert tracer.init_called
+        assert tracer.init_worker_called
+        assert runner.worker_id == 0
+        assert runner.get_agent() is agent
+        assert runner.get_store() is store
+
+    # Verify teardown happened
+    assert tracer.teardown_worker_called
+    assert tracer.teardown_called
+    assert runner.worker_id is None
+
+
+@pytest.mark.asyncio
+async def test_run_context_yields_runner() -> None:
+    """Test that run_context yields the runner instance itself."""
+    tracer = DummyTracer()
+    agent = DummyAgent()
+    store = InMemoryLightningStore()
+    runner = LitAgentRunner[Dict[str, Any]](tracer=tracer)
+
+    with runner.run_context(agent=agent, store=store) as yielded_runner:
+        assert yielded_runner is runner
+
+
+@pytest.mark.asyncio
+async def test_run_context_with_hooks() -> None:
+    """Test that run_context properly passes hooks to init()."""
+    tracer = DummyTracer()
+    agent = DummyAgent()
+    store = InMemoryLightningStore()
+    runner = LitAgentRunner[Dict[str, Any]](tracer=tracer)
+    hook = RecordingHook()
+
+    with runner.run_context(agent=agent, store=store, hooks=[hook]):
+        # Verify hooks were registered
+        assert runner._hooks == [hook]  # pyright: ignore[reportPrivateUsage]
+
+
+@pytest.mark.asyncio
+async def test_run_context_teardown_on_exception_in_context() -> None:
+    """Test that run_context properly tears down even when exception occurs in with block."""
+    tracer = DummyTracer()
+    agent = DummyAgent()
+    store = InMemoryLightningStore()
+    runner = LitAgentRunner[Dict[str, Any]](tracer=tracer)
+
+    with pytest.raises(RuntimeError, match="test error"):
+        with runner.run_context(agent=agent, store=store):
+            raise RuntimeError("test error")
+
+    # Verify teardown still happened
+    assert tracer.teardown_worker_called
+    assert tracer.teardown_called
+
+
+@pytest.mark.asyncio
+async def test_run_context_no_teardown_worker_if_init_worker_fails() -> None:
+    """Test that teardown_worker is not called if init_worker fails."""
+    tracer = DummyTracer()
+    agent = DummyAgent()
+    store = InMemoryLightningStore()
+    runner = LitAgentRunner[Dict[str, Any]](tracer=tracer)
+
+    # Mock init_worker to raise an exception
+    original_init_worker = runner.init_worker
+
+    def failing_init_worker(worker_id: int, store: LightningStore, **kwargs: Any) -> None:
+        raise RuntimeError("init_worker failed")
+
+    runner.init_worker = failing_init_worker  # type: ignore[method-assign]
+
+    try:
+        with pytest.raises(RuntimeError, match="init_worker failed"):
+            with runner.run_context(agent=agent, store=store):
+                pass
+    finally:
+        # Restore original method
+        runner.init_worker = original_init_worker  # type: ignore[method-assign]
+
+    # Verify teardown was called but teardown_worker was not
+    assert tracer.init_called
+    assert not tracer.init_worker_called
+    assert not tracer.teardown_worker_called
+    assert tracer.teardown_called
+
+
+@pytest.mark.asyncio
+async def test_run_context_no_teardown_if_init_fails() -> None:
+    """Test that teardown is not called if init fails."""
+    tracer = DummyTracer()
+    agent = DummyAgent()
+    store = InMemoryLightningStore()
+    runner = LitAgentRunner[Dict[str, Any]](tracer=tracer)
+
+    # Mock init to raise an exception
+    original_init = runner.init
+
+    def failing_init(agent: LitAgent[Any], **kwargs: Any) -> None:
+        raise RuntimeError("init failed")
+
+    runner.init = failing_init  # type: ignore[method-assign]
+
+    try:
+        with pytest.raises(RuntimeError, match="init failed"):
+            with runner.run_context(agent=agent, store=store):
+                pass
+    finally:
+        # Restore original method
+        runner.init = original_init  # type: ignore[method-assign]
+
+    # Verify neither init_worker nor teardown methods were called
+    assert not tracer.init_called
+    assert not tracer.init_worker_called
+    assert not tracer.teardown_worker_called
+    assert not tracer.teardown_called
+
+
+@pytest.mark.asyncio
+async def test_run_context_handles_teardown_worker_exception(caplog: pytest.LogCaptureFixture) -> None:
+    """Test that exceptions in teardown_worker are caught and logged."""
+    tracer = DummyTracer()
+    agent = DummyAgent()
+    store = InMemoryLightningStore()
+    runner = LitAgentRunner[Dict[str, Any]](tracer=tracer)
+
+    # Mock teardown_worker to raise an exception
+    original_teardown_worker = runner.teardown_worker
+
+    def failing_teardown_worker(worker_id: int, *args: Any, **kwargs: Any) -> None:
+        original_teardown_worker(worker_id, *args, **kwargs)
+        raise RuntimeError("teardown_worker failed")
+
+    runner.teardown_worker = failing_teardown_worker  # type: ignore[method-assign]
+
+    try:
+        with runner.run_context(agent=agent, store=store):
+            pass
+    finally:
+        # Restore original method
+        runner.teardown_worker = original_teardown_worker  # type: ignore[method-assign]
+
+    # Verify both teardown methods were attempted
+    assert tracer.teardown_worker_called
+    assert tracer.teardown_called
+
+    # Verify error was logged
+    assert any("Error during runner worker teardown" in record.message for record in caplog.records)
+
+
+@pytest.mark.asyncio
+async def test_run_context_handles_teardown_exception(caplog: pytest.LogCaptureFixture) -> None:
+    """Test that exceptions in teardown are caught and logged."""
+    tracer = DummyTracer()
+    agent = DummyAgent()
+    store = InMemoryLightningStore()
+    runner = LitAgentRunner[Dict[str, Any]](tracer=tracer)
+
+    # Mock teardown to raise an exception
+    original_teardown = runner.teardown
+
+    def failing_teardown(*args: Any, **kwargs: Any) -> None:
+        original_teardown(*args, **kwargs)
+        raise RuntimeError("teardown failed")
+
+    runner.teardown = failing_teardown  # type: ignore[method-assign]
+
+    try:
+        with runner.run_context(agent=agent, store=store):
+            pass
+    finally:
+        # Restore original method
+        runner.teardown = original_teardown  # type: ignore[method-assign]
+
+    # Verify both teardown methods were attempted
+    assert tracer.teardown_worker_called
+    assert tracer.teardown_called
+
+    # Verify error was logged
+    assert any("Error during runner teardown" in record.message for record in caplog.records)
+
+
+@pytest.mark.asyncio
+async def test_run_context_can_be_used_for_step() -> None:
+    """Test that run_context can be used to execute runner.step()."""
+
+    class CountingAgent(LitAgent[str]):
+        def __init__(self) -> None:
+            super().__init__()
+            self.call_count = 0
+
+        def validation_rollout(self, task: str, resources: Dict[str, Any], rollout: Any) -> float:
+            self.call_count += 1
+            return 0.5
+
+    tracer = DummyTracer()
+    agent = CountingAgent()
+    store = InMemoryLightningStore()
+    runner = LitAgentRunner[str](tracer=tracer)
+
+    await store.update_resources("default", {"llm": LLM(endpoint="http://localhost", model="dummy")})
+
+    with runner.run_context(agent=agent, store=store):
+        await runner.step("test task")
+
+    # Verify the agent was called
+    assert agent.call_count == 1
+
+    # Verify the rollout was recorded in the store
+    rollouts = await store.query_rollouts()
+    assert len(rollouts) == 1
+    assert rollouts[0].status == "succeeded"