chore: import upstream snapshot with attribution

tests/common/__init__.py

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

tests/common/network.py

@@ -0,0 +1,10 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import socket
+from contextlib import closing
+
+
+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])

tests/common/prometheus_stub.py

@@ -0,0 +1,125 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import types
+from typing import Any, Dict, List, Optional, Sequence, Tuple
+
+
+class _CounterChild:
+    def __init__(self) -> None:
+        self.value = 0.0
+
+    def inc(self, amount: float = 1.0) -> None:
+        self.value += amount
+
+
+class _HistogramChild:
+    def __init__(self) -> None:
+        self.values: List[float] = []
+
+    def observe(self, value: float) -> None:
+        self.values.append(value)
+
+
+class PrometheusStub(types.ModuleType):
+    """Minimal prometheus_client replacement for unit tests."""
+
+    def __init__(self, real_client: Any) -> None:
+        super().__init__("prometheus_client")
+        self.counter_instances: List[_PromCounter] = []
+        self.histogram_instances: List[_PromHistogram] = []
+        self.real_client = real_client
+
+        class CollectorRegistry:
+            pass
+
+        class _Multiprocess:
+            def __init__(self) -> None:
+                self.registry: Optional[CollectorRegistry] = None
+
+            def MultiProcessCollector(self, registry: CollectorRegistry) -> None:
+                self.registry = registry
+
+        self.CollectorRegistry = CollectorRegistry
+        self.REGISTRY = CollectorRegistry()
+        self.multiprocess = _Multiprocess()
+
+        self.Counter = _PromCounterFactory(self)
+        self.Histogram = _PromHistogramFactory(self)
+        self.make_asgi_app = self._make_asgi_app
+
+    def _make_asgi_app(self, *args: Any, **kwargs: Any) -> Any:
+        return self.real_client.make_asgi_app(*args, **kwargs)
+
+
+class _PromCounterFactory:
+    def __init__(self, owner: PrometheusStub) -> None:
+        self._owner = owner
+
+    def __call__(self, name: str, doc: str, labelnames: Sequence[str]) -> _PromCounter:
+        counter = _PromCounter(name, doc, labelnames)
+        counter._register(self._owner.counter_instances)  # pyright: ignore[reportPrivateUsage]
+        return counter
+
+
+class _PromHistogramFactory:
+    def __init__(self, owner: PrometheusStub) -> None:
+        self._owner = owner
+
+    def __call__(
+        self,
+        name: str,
+        doc: str,
+        labelnames: Sequence[str],
+        buckets: Sequence[float] | None = None,
+    ) -> _PromHistogram:
+        histogram = _PromHistogram(name, doc, labelnames, buckets or ())
+        histogram._register(self._owner.histogram_instances)  # pyright: ignore[reportPrivateUsage]
+        return histogram
+
+
+class _PromCounter:
+    def __init__(self, name: str, doc: str, labelnames: Sequence[str]) -> None:
+        self.name = name
+        self.doc = doc
+        self.labelnames = tuple(labelnames)
+        self.default = _CounterChild()
+        self.children: Dict[Tuple[Tuple[str, str], ...], _CounterChild] = {}
+
+    def _register(self, sink: List["_PromCounter"]) -> None:
+        sink.append(self)
+
+    def labels(self, **kwargs: str) -> _CounterChild:
+        key = tuple(sorted(kwargs.items()))
+        return self.children.setdefault(key, _CounterChild())
+
+    def inc(self, amount: float = 1.0) -> None:
+        self.default.inc(amount)
+
+
+class _PromHistogram:
+    def __init__(self, name: str, doc: str, labelnames: Sequence[str], buckets: Sequence[float]) -> None:
+        self.name = name
+        self.doc = doc
+        self.labelnames = tuple(labelnames)
+        self.buckets = tuple(buckets)
+        self.default = _HistogramChild()
+        self.children: Dict[Tuple[Tuple[str, str], ...], _HistogramChild] = {}
+
+    def _register(self, sink: List["_PromHistogram"]) -> None:
+        sink.append(self)
+
+    def labels(self, **kwargs: str) -> _HistogramChild:
+        key = tuple(sorted(kwargs.items()))
+        return self.children.setdefault(key, _HistogramChild())
+
+    def observe(self, value: float) -> None:
+        self.default.observe(value)
+
+
+def make_prometheus_stub() -> PrometheusStub:
+    """Factory helper for tests."""
+    import prometheus_client
+
+    return PrometheusStub(prometheus_client)

tests/common/tracer.py

@@ -0,0 +1,61 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+from typing import Optional
+
+import agentops
+import agentops.sdk.core as agentops_core
+import opentelemetry.trace as trace_api
+
+from agentlightning.tracer.dummy import DummyTracer
+from agentlightning.types import Attributes, SpanCoreFields, TraceStatus
+
+
+# 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
+
+
+def clear_agentops_init() -> None:
+    """Make agentops.init() runnable again."""
+    agentops.get_client().initialized = False
+    agentops_core.tracer._initialized = False
+
+
+class RecordingDummyTracer(DummyTracer):
+    """Dummy tracer that captures the most recent span request for assertions."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.last_span: Optional[SpanCoreFields] = None
+
+    def create_span(
+        self,
+        name: str,
+        attributes: Optional[Attributes] = None,
+        timestamp: Optional[float] = None,
+        status: Optional[TraceStatus] = None,
+    ) -> SpanCoreFields:
+        span = super().create_span(name, attributes, timestamp, status)
+        self.last_span = span
+        return span

tests/common/vllm.py

@@ -0,0 +1,181 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import os
+import subprocess
+import sys
+import time
+from typing import Any, Optional
+
+import httpx
+import openai
+
+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
+
+
+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,
+        )