github-featured/microsoft--agent-lightning
2
0
Fork 0
Code Issues Pull Requests Deployments Models Agent Notes Packages Projects Releases Wiki Activity

Replace AgentOps mock server with bypassable client (#202)

Browse Source 
---------

Co-authored-by: Hao Ni (CSI Interfusion Co Ltd) <v-nhao@microsoft.com>
This commit is contained in:
Ni Hao
2025-10-31 23:05:39 +08:00
committed by GitHub
parent c453c41fd2
commit 3f372ff7b3
6 changed files with 292 additions and 202 deletions
Download Patch File Download Diff File Expand all files Collapse all files
agentlightning/adapter/triplet.py
+9
View File
@@ -429,7 +429,16 @@ class TraceTree:
If we don't, when we want to select the LLM completion span with agent as filter.
We will never get the correct span underneath.
"""
# If the current node has only one child, recursively repair its hierarchy directly.
# This special-case handling is needed because when a trace is manually ended
# (via agentops.end_trace), the AgentOps provider automatically wraps all spans
# under an extra synthetic root node (e.g., "run_one.session").
if len(self.children) == 1:
self.children[0].repair_hierarchy()
return
nodes_to_repair = list(self.children)
for repair_node in nodes_to_repair:
if len(self.children) == 1:
# If there is only one child, we don't need to repair the hierarchy.
agentlightning/cli/agentops_server.py
-30
View File
@@ -1,30 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
from __future__ import annotations
import argparse
import time
from typing import Iterable
from agentlightning.instrumentation.agentops import AgentOpsServerManager
def main(argv: Iterable[str] | None = None) -> int:
parser = argparse.ArgumentParser(description="Start AgentOps server")
parser.add_argument("--daemon", action="store_true", help="Run server as a daemon")
parser.add_argument("--port", type=int, default=8002, help="Port to run the server on")
args = parser.parse_args(list(argv) if argv is not None else None)
manager = AgentOpsServerManager(daemon=args.daemon, port=args.port)
try:
manager.start()
# Wait forever
while True:
time.sleep(1)
except KeyboardInterrupt:
manager.stop()
return 0
if __name__ == "__main__":
raise SystemExit(main())
agentlightning/instrumentation/agentops.py
+115 -110
View File
@@ -4,28 +4,68 @@ from __future__ import annotations
import json
import logging
import multiprocessing
import signal
import socket
import time
from typing import Any, Callable, no_type_check
import flask
import requests
import setproctitle
from agentops.client.api import V3Client, V4Client
from agentops.client.api.types import AuthTokenResponse
from agentops.sdk.exporters import AuthenticatedOTLPExporter
from opentelemetry.exporter.otlp.proto.http.metric_exporter import OTLPMetricExporter
from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
from opentelemetry.sdk.metrics.export import MetricExportResult
from opentelemetry.sdk.trace.export import SpanExportResult
logger = logging.getLogger(__name__)
__all__ = [
"instrument_agentops",
"uninstrument_agentops",
"agentops_local_server",
"AgentOpsServerManager",
]
# Module-level storage for originals
_original_handle_chat_attributes: Callable[..., Any] | None = None
_original_handle_response: Callable[..., Any] | None = None
_agentops_service_enabled = False
def enable_agentops_service(enabled: bool = True) -> None:
"""
Enable or disable communication with the AgentOps service.
False (default): AgentOps exporters and clients will run in local mode
and will not attempt to communicate with the remote AgentOps service.
True: all exporters and clients will operate in normal mode and send data
to the AgentOps service as expected.
"""
global _agentops_service_enabled
_agentops_service_enabled = enabled
logger.info(f"Switch set to {enabled} for exporters and clients.")
def _patch_exporters():
import agentops.client.api
import agentops.sdk.core
import opentelemetry.exporter.otlp.proto.http.metric_exporter
import opentelemetry.exporter.otlp.proto.http.trace_exporter
agentops.sdk.core.AuthenticatedOTLPExporter = BypassableAuthenticatedOTLPExporter # type: ignore
opentelemetry.exporter.otlp.proto.http.metric_exporter.OTLPMetricExporter = BypassableOTLPMetricExporter
opentelemetry.exporter.otlp.proto.http.trace_exporter.OTLPSpanExporter = BypassableOTLPSpanExporter
agentops.client.api.V3Client = BypassableV3Client
agentops.client.api.V4Client = BypassableV4Client
def _unpatch_exporters():
import agentops.client.api
import agentops.sdk.core
import opentelemetry.exporter.otlp.proto.http.metric_exporter
import opentelemetry.exporter.otlp.proto.http.trace_exporter
agentops.sdk.core.AuthenticatedOTLPExporter = AuthenticatedOTLPExporter # type: ignore
opentelemetry.exporter.otlp.proto.http.metric_exporter.OTLPMetricExporter = OTLPMetricExporter
opentelemetry.exporter.otlp.proto.http.trace_exporter.OTLPSpanExporter = OTLPSpanExporter
agentops.client.api.V3Client = V3Client
agentops.client.api.V4Client = V4Client
def _unwrap_legacy_response(response: Any) -> Any:
@@ -170,6 +210,8 @@ def instrument_agentops():
Instrument agentops to capture token IDs.
Automatically detects and uses the appropriate patching method based on the installed agentops version.
"""
_patch_exporters()
# Try newest version first (tested for 0.4.16)
try:
return _patch_new_agentops()
@@ -189,6 +231,8 @@ def instrument_agentops():
def uninstrument_agentops():
"""Uninstrument agentops to stop capturing token IDs."""
_unpatch_exporters()
try:
_unpatch_new_agentops()
except Exception:
@@ -199,114 +243,75 @@ def uninstrument_agentops():
pass
def agentops_local_server():
class BypassableAuthenticatedOTLPExporter(AuthenticatedOTLPExporter):
"""
Returns a Flask app that can be used to test agentops integration.
This server provides endpoints for token fetching and a catch-all endpoint.
AuthenticatedOTLPExporter with switchable service control.
When `_agentops_service_enabled` is False, skip export and return success.
"""
app = flask.Flask(__name__)
@app.route("/v3/auth/token", methods=["POST"])
def fetch_token(): # type: ignore
return {"token": "dummy", "project_id": "dummy"}
@app.route("/", defaults={"path": ""}, methods=["GET", "POST"])
@app.route("/<path:path>", methods=["GET", "POST"])
def catch_all(path: str): # type: ignore
return {"path": path}
return app
def _run_server(**kwargs: Any): # type: ignore
"""
Internal function to run the Flask server.
This is used to avoid issues with multiprocessing and Flask's reloader.
"""
signal.signal(signal.SIGINT, signal.SIG_IGN) # Ignore SIGINT in worker processes
setproctitle.setproctitle(multiprocessing.current_process().name)
app = agentops_local_server()
app.run(**kwargs)
class AgentOpsServerManager:
"""Manages a AgentOps local server to bypass the online service of AgentOps."""
def __init__(self, daemon: bool = True, port: int | None = None):
self.server_process: multiprocessing.Process | None = None
self.server_port = port
self.daemon = daemon
logger.info("AgentOpsServerManager initialized.")
def _find_available_port(self) -> int:
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
s.bind(("", 0))
return s.getsockname()[1]
def start(self):
if self.server_process and self.server_process.is_alive():
logger.warning("AgentOps server process appears to be already running.")
return
if self.server_port is None:
self.server_port = self._find_available_port()
logger.info(f"Starting AgentOps local server on port {self.server_port}...")
self.server_process = multiprocessing.Process(
target=_run_server,
kwargs={"host": "127.0.0.1", "port": self.server_port, "use_reloader": False, "debug": False},
daemon=self.daemon,
name="AgentLightning-AgentOpsServer",
)
self.server_process.start()
logger.info(
f"AgentOps local server process (PID: {self.server_process.pid}) started, targeting port {self.server_port}."
)
for attempt in range(20): # 10 seconds total
time.sleep(0.5) # Brief wait for server to start up
try:
result = requests.get(f"http://127.0.0.1:{self.server_port}/")
if result.status_code == 200:
break
except Exception as e:
logger.debug(f"Error checking AgentOps server: {e}")
logger.warning(f"AgentOps still not ready after {attempt} attempts. Retrying...")
def export(self, *args: Any, **kwargs: Any) -> SpanExportResult:
if _agentops_service_enabled:
return super().export(*args, **kwargs)
else:
logger.error(f"AgentOps local server failed to start or exited prematurely.")
return
logger.debug("SwitchableAuthenticatedOTLPExporter is switched off, skipping export.")
return SpanExportResult.SUCCESS
if not self.server_process.is_alive():
logger.error(f"AgentOps local server failed to start or exited prematurely.")
def is_alive(self) -> bool:
if self.server_process and self.server_process.is_alive():
return True
return False
class BypassableOTLPMetricExporter(OTLPMetricExporter):
"""
OTLPMetricExporter with switchable service control.
When `_agentops_service_enabled` is False, skip export and return success.
"""
def stop(self):
if self.server_process is not None and self.server_process.is_alive():
logger.info(f"Stopping AgentOps local server (PID: {self.server_process.pid})...")
self.server_process.terminate() # Send SIGTERM
self.server_process.join(timeout=5) # Wait for clean exit
if self.server_process.is_alive():
logger.warning(
f"AgentOps server (PID: {self.server_process.pid}) did not terminate gracefully, killing..."
)
self.server_process.kill() # Force kill
self.server_process.join(timeout=10) # Wait for kill
self.server_process = None
logger.info(f"AgentOps local server stopped.")
def export(self, *args: Any, **kwargs: Any) -> MetricExportResult:
if _agentops_service_enabled:
return super().export(*args, **kwargs) # type: ignore[reportUnknownMemberType]
else:
logger.info("AgentOps local server was not running or already stopped.")
logger.debug("SwitchableOTLPMetricExporter is switched off, skipping export.")
return MetricExportResult.SUCCESS
def get_port(self) -> int | None:
# Check liveness again in case it died since start()
if self.is_alive() and self.server_port is not None:
return self.server_port
# If called after server stopped or failed, port might be stale or None
if self.server_port is not None and (self.server_process is None or not self.server_process.is_alive()):
logger.warning(
f"AgentOps server port {self.server_port} is stored, but server process is not alive. Returning stored port."
)
return self.server_port
class BypassableOTLPSpanExporter(OTLPSpanExporter):
"""
OTLPSpanExporter with switchable service control.
When `_agentops_service_enabled` is False, skip export and return success.
"""
def export(self, *args: Any, **kwargs: Any) -> SpanExportResult:
if _agentops_service_enabled:
return super().export(*args, **kwargs)
else:
logger.debug("SwitchableOTLPSpanExporter is switched off, skipping export.")
return SpanExportResult.SUCCESS
class BypassableV3Client(V3Client):
"""
V3Client with toggleable authentication calls.
Returns dummy auth response when `_agentops_service_enabled` is False.
"""
# Temporary synchronous override of fetch_auth_token for mock purposes.
def fetch_auth_token(self, *args: Any, **kwargs: Any) -> AuthTokenResponse: # type: ignore[override]
if _agentops_service_enabled:
return super().fetch_auth_token(*args, **kwargs) # type: ignore[override]
else:
logger.debug("SwitchableV3Client is switched off, skipping fetch_auth_token request.")
return AuthTokenResponse(token="dummy", project_id="dummy")
class BypassableV4Client(V4Client):
"""
V4Client with toggleable post requests.
Returns dummy response when `_agentops_service_enabled` is False.
"""
def post(self, *args: Any, **kwargs: Any) -> requests.Response:
if _agentops_service_enabled:
return super().post(*args, **kwargs)
else:
logger.debug("SwitchableV4Client is switched off, skipping post request.")
response = requests.Response()
response.status_code = 200
response._content = b"{}"
return response
agentlightning/tracer/agentops.py
+27 -62
View File
@@ -15,9 +15,9 @@ from agentops.sdk.core import TracingCore
from agentops.sdk.processors import SpanProcessor
from opentelemetry.instrumentation.utils import suppress_instrumentation
from opentelemetry.sdk.trace import ReadableSpan
from opentelemetry.trace.status import StatusCode
from agentlightning.instrumentation import instrument_all, uninstrument_all
from agentlightning.instrumentation.agentops import AgentOpsServerManager
from agentlightning.store.base import LightningStore
from .base import Tracer
@@ -56,46 +56,11 @@ class AgentOpsTracer(Tracer):
self.instrument_managed = instrument_managed
self.daemon = daemon
self._agentops_server_manager = AgentOpsServerManager(self.daemon)
self._agentops_server_port_val: Optional[int] = None
if not self.agentops_managed:
logger.warning("agentops_managed=False. You are responsible for AgentOps setup.")
if not self.instrument_managed:
logger.warning("instrument_managed=False. You are responsible for all instrumentation.")
def __getstate__(self):
state = self.__dict__.copy()
state["_agentops_server_manager"] = None # Exclude the unpicklable server manager
# _agentops_server_port_val (int) is inherently picklable and will be included.
logger.debug(f"Getting state for pickling Trainer (PID {os.getpid()}). _agentops_server_manager excluded.")
return state
def __setstate__(self, state: Any):
self.__dict__.update(state)
# In child process, self._agentops_server_manager will be None.
logger.debug(f"Setting state for unpickled Trainer (PID {os.getpid()}). _agentops_server_manager is None.")
def init(self, *args: Any, **kwargs: Any):
if self.agentops_managed and self._agentops_server_manager:
self._agentops_server_manager.start()
self._agentops_server_port_val = self._agentops_server_manager.get_port()
if self._agentops_server_port_val is None:
if (
self._agentops_server_manager.server_process is not None
and self._agentops_server_manager.server_process.is_alive()
):
raise RuntimeError("AgentOps server started but port is None. Check server manager logic.")
elif (
self._agentops_server_port_val is None and self._agentops_server_manager.server_process is None
): # Server failed to start
raise RuntimeError("AgentOps server manager indicates server is not running and port is None.")
def teardown(self):
if self.agentops_managed:
self._agentops_server_manager.stop()
logger.info("AgentOps server stopped.")
def instrument(self, worker_id: int):
instrument_all()
@@ -111,24 +76,9 @@ class AgentOpsTracer(Tracer):
logger.info(f"[Worker {worker_id}] Instrumentation applied.")
if self.agentops_managed:
if self._agentops_server_port_val: # Use the stored, picklable port value
base_url = f"http://localhost:{self._agentops_server_port_val}"
env_vars_to_set = {
"AGENTOPS_API_KEY": "dummy",
"AGENTOPS_API_ENDPOINT": base_url,
"AGENTOPS_APP_URL": f"{base_url}/notavailable",
"AGENTOPS_EXPORTER_ENDPOINT": f"{base_url}/traces",
}
for key, value in env_vars_to_set.items():
os.environ[key] = value
logger.info(f"[Worker {worker_id}] Env var set: {key}={value}")
else:
logger.warning(
f"[Worker {worker_id}] AgentOps managed, but local server port is not available. Client may not connect as expected."
)
os.environ.setdefault("AGENTOPS_API_KEY", "dummy")
if not agentops.get_client().initialized:
agentops.init() # type: ignore
agentops.init(auto_start_session=False) # type: ignore
logger.info(f"[Worker {worker_id}] AgentOps client initialized.")
else:
logger.warning(f"[Worker {worker_id}] AgentOps client was already initialized.")
@@ -192,15 +142,30 @@ class AgentOpsTracer(Tracer):
if not self._lightning_span_processor:
raise RuntimeError("LightningSpanProcessor is not initialized. Call init_worker() first.")
if store is not None and rollout_id is not None and attempt_id is not None:
ctx = self._lightning_span_processor.with_context(store=store, rollout_id=rollout_id, attempt_id=attempt_id)
with ctx as processor:
yield processor
elif store is None and rollout_id is None and attempt_id is None:
with self._lightning_span_processor:
yield self._lightning_span_processor
else:
raise ValueError("store, rollout_id, and attempt_id must be either all provided or all None")
kwargs: dict[str, Any] = {}
if name is not None:
kwargs["trace_name"] = name
elif rollout_id is not None:
kwargs["trace_name"] = rollout_id
trace = agentops.start_trace(**kwargs)
status = StatusCode.OK # type: ignore
try:
if store is not None and rollout_id is not None and attempt_id is not None:
ctx = self._lightning_span_processor.with_context(
store=store, rollout_id=rollout_id, attempt_id=attempt_id
)
with ctx as processor:
yield processor
elif store is None and rollout_id is None and attempt_id is None:
with self._lightning_span_processor:
yield self._lightning_span_processor
else:
raise ValueError("store, rollout_id, and attempt_id must be either all provided or all None")
except Exception as e:
status = StatusCode.ERROR # type: ignore
logger.error(f"Trace failed for rollout_id={rollout_id}, attempt_id={attempt_id}, error={e}")
finally:
agentops.end_trace(trace, end_state=status) # type: ignore
def get_last_trace(self) -> List[ReadableSpan]:
"""
tests/instrumentation/test_agentops.py
+64
View File
@@ -0,0 +1,64 @@
# Copyright (c) Microsoft. All rights reserved.
from unittest.mock import MagicMock, patch
from opentelemetry.sdk.metrics.export import MetricExportResult
from opentelemetry.sdk.trace.export import SpanExportResult
from agentlightning.instrumentation.agentops import (
BypassableAuthenticatedOTLPExporter,
BypassableOTLPMetricExporter,
BypassableOTLPSpanExporter,
enable_agentops_service,
)
def test_switchable_authenticated_exporter():
switchable_authenticated_exporter = BypassableAuthenticatedOTLPExporter(endpoint="http://dummy", jwt="dummy")
with patch.object(
switchable_authenticated_exporter.__class__.__bases__[0], "export", return_value=SpanExportResult.SUCCESS
) as mock_export:
enable_agentops_service()
result = switchable_authenticated_exporter.export([])
assert result == SpanExportResult.SUCCESS
mock_export.assert_called_once()
enable_agentops_service(False)
result = switchable_authenticated_exporter.export([])
assert result == SpanExportResult.SUCCESS
assert mock_export.call_count == 1
def test_switchable_otlp_metric_exporter():
switchable_otlp_metric_exporter = BypassableOTLPMetricExporter()
with patch.object(
switchable_otlp_metric_exporter.__class__.__bases__[0], "export", return_value=MetricExportResult.SUCCESS
) as mock_export:
enable_agentops_service()
result = switchable_otlp_metric_exporter.export(metrics_data=MagicMock())
assert result == MetricExportResult.SUCCESS
mock_export.assert_called_once()
enable_agentops_service(False)
result = switchable_otlp_metric_exporter.export(metrics_data=MagicMock())
assert result == MetricExportResult.SUCCESS
assert mock_export.call_count == 1
def test_switchable_otlp_span_exporter():
switchable_otlp_span_exporter = BypassableOTLPSpanExporter()
with patch.object(
switchable_otlp_span_exporter.__class__.__bases__[0], "export", return_value=SpanExportResult.SUCCESS
) as mock_export:
enable_agentops_service()
result = switchable_otlp_span_exporter.export([])
assert result == SpanExportResult.SUCCESS
mock_export.assert_called_once()
enable_agentops_service(False)
result = switchable_otlp_span_exporter.export([])
assert result == SpanExportResult.SUCCESS
assert mock_export.call_count == 1
tests/tracer/test_agentops.py
+77
View File
@@ -0,0 +1,77 @@
# Copyright (c) Microsoft. All rights reserved.
import multiprocessing
from typing import Any, Optional, Union
import agentops
from agentops.sdk.core import TraceContext
from opentelemetry.trace.status import StatusCode
from agentlightning.tracer.agentops import AgentOpsTracer
def _func_with_exception():
"""Function that always raises an exception to test error tracing."""
raise ValueError("This is a test exception")
def _func_without_exception():
"""Function that always executed successfully to test success tracing."""
pass
def test_trace_error_status_from_instance():
"""
Test that AgentOpsTracer correctly sets trace end state based on execution result.
This test replaces `agentops.end_trace` with a custom function to capture
the `end_state` passed in. It verifies that traces ending after a raised
exception have `StatusCode.ERROR`, while normal runs have `StatusCode.OK`.
"""
ctx = multiprocessing.get_context("spawn")
proc = ctx.Process(target=_test_trace_error_status_from_instance_imp)
proc.start()
proc.join(30.0) # On GPU server, the time is around 10 seconds.
if proc.is_alive():
proc.terminate()
proc.join(5)
if proc.is_alive():
proc.kill()
assert False, "Child process hung. Check test output for details."
assert proc.exitcode == 0, (
f"Child process for test_trace_error_status_from_instance failed with exit code {proc.exitcode}. "
"Check child traceback in test output."
)
def _test_trace_error_status_from_instance_imp():
captured_state = {}
old_end_trace = agentops.end_trace
def custom_end_trace(
trace_context: Optional[TraceContext] = None, end_state: Union[Any, StatusCode, str] = None
) -> None:
captured_state["state"] = end_state
return old_end_trace(trace_context, end_state=end_state)
agentops.end_trace = custom_end_trace
tracer = AgentOpsTracer()
tracer.init()
tracer.init_worker(0)
try:
tracer.trace_run(_func_with_exception)
assert captured_state["state"] == StatusCode.ERROR
tracer.trace_run(_func_without_exception)
assert captured_state["state"] == StatusCode.OK
finally:
agentops.end_trace = old_end_trace
tracer.teardown_worker(0)
tracer.teardown()