clean up

Author: benjibc

eval_protocol/pytest/remote_rollout_processor.py

@@ -35,19 +35,34 @@ class RemoteRolloutProcessor(RolloutProcessor):
       Returns: {"terminated": bool, "info": {...}?}
     """
 
-    def __init__(self):
-        pass
+    def __init__(
+        self,
+        *,
+        remote_base_url: Optional[str] = None,
+        num_turns: int = 2,
+        poll_interval: float = 1.0,
+        timeout_seconds: float = 120.0,
+    ):
+        # Prefer constructor-provided configuration. These can be overridden via
+        # config.kwargs at call time for backward compatibility.
+        self._remote_base_url = remote_base_url
+        self._num_turns = num_turns
+        self._poll_interval = poll_interval
+        self._timeout_seconds = timeout_seconds
 
     def __call__(self, rows: List[EvaluationRow], config: RolloutProcessorConfig) -> List[asyncio.Task[EvaluationRow]]:
         tasks: List[asyncio.Task[EvaluationRow]] = []
 
-        remote_base_url: Optional[str] = None
-        num_turns: int = 2
-        poll_interval: float = 1.0
-        timeout_seconds: float = 120.0
+        # Start with constructor values
+        remote_base_url: Optional[str] = self._remote_base_url
+        num_turns: int = self._num_turns
+        poll_interval: float = self._poll_interval
+        timeout_seconds: float = self._timeout_seconds
 
+        # Backward compatibility: allow overrides via config.kwargs
         if config.kwargs:
-            remote_base_url = config.kwargs.get("remote_base_url")
+            if remote_base_url is None:
+                remote_base_url = config.kwargs.get("remote_base_url", remote_base_url)
             num_turns = int(config.kwargs.get("num_turns", num_turns))
             poll_interval = float(config.kwargs.get("poll_interval", poll_interval))
             timeout_seconds = float(config.kwargs.get("timeout_seconds", timeout_seconds))

tests/chinook/langfuse/remote_server.py

@@ -45,37 +45,6 @@ def init(req: InitRequest):
     # Kick off worker thread that runs multi-turn chat via LiteLLM proxy
     def _worker():
         try:
-            # Try to set up Langfuse trace to guarantee observability, independent of proxy wiring
-            langfuse = None
-            trace = None
-            try:
-                from langfuse import get_client  # pyright: ignore[reportPrivateImportUsage]
-
-                langfuse = get_client()
-                id_tags = []
-                try:
-                    id_tags = [
-                        f"inv:{req.metadata.get('invocation_id')}",
-                        f"exp:{req.metadata.get('experiment_id')}",
-                        f"rollout:{req.metadata.get('rollout_id')}",
-                    ]
-                except Exception:
-                    id_tags = []
-                trace = langfuse.api.trace.create(
-                    name="remote_chinook_rollout",
-                    metadata=req.metadata,
-                    requester_metadata=req.metadata,
-                    tags=["chinook_remote", "chinook_sql", *[t for t in id_tags if t]],
-                    input={
-                        "messages": _clean_messages_for_api(req.messages),
-                        "tools": req.tools,
-                        "metadata": req.metadata,
-                    },
-                )
-            except Exception:
-                langfuse = None
-                trace = None
-
             base_url = os.getenv(
                 "LITELLM_BASE_URL",
                 "https://litellm-cloud-proxy-prod-644257448872.us-central1.run.app",
@@ -110,48 +79,13 @@ def _worker():
                 r.raise_for_status()
                 data = r.json()
                 assistant = data.get("choices", [{}])[0].get("message", {})
-                # Optionally record a generation on Langfuse
-                try:
-                    if langfuse and trace and getattr(langfuse.api, "generation", None):
-                        langfuse.api.generation.create(
-                            trace_id=trace.id,
-                            name="assistant",
-                            input={"messages": _clean_messages_for_api(messages)},
-                            output=assistant,
-                        )
-                except Exception:
-                    pass
                 # Append assistant for next turn
                 messages = messages + [assistant]
 
-            # Update final trace output for easier adapter extraction
-            try:
-                if langfuse and trace:
-                    langfuse.api.trace.update(
-                        id=trace.id,
-                        output={
-                            "messages": _clean_messages_for_api(messages),
-                            "metadata": req.metadata,
-                        },
-                    )
-            except Exception:
-                pass
-
         except Exception:
             # Best-effort; mark as done even on error to unblock polling
             pass
         finally:
-            try:
-                if "langfuse" in locals() and langfuse is not None:
-                    # Ensure buffered telemetry is sent
-                    flush = getattr(langfuse, "flush", None)
-                    if callable(flush):
-                        flush()
-                    shutdown = getattr(langfuse, "shutdown", None)
-                    if callable(shutdown):
-                        shutdown()
-            except Exception:
-                pass
             _STATE[req.rollout_id]["terminated"] = True
 
     t = threading.Thread(target=_worker, daemon=True)

tests/chinook/langfuse/test_remote_langfuse_chinook.py

@@ -6,6 +6,7 @@
 import atexit
 
 import pytest
+import requests
 
 from eval_protocol.models import EvaluationRow, Message
 from eval_protocol.pytest import evaluation_test
@@ -23,17 +24,36 @@ def _start_remote_server():
 
 
 def _ensure_server_running():
+    host = os.getenv("REMOTE_SERVER_HOST", "127.0.0.1")
+    port = int(os.getenv("REMOTE_SERVER_PORT", "7077"))
+    base_url = f"http://{host}:{port}"
+
+    def _is_up() -> bool:
+        try:
+            r = requests.get(f"{base_url}/status", params={"rollout_id": "ping"}, timeout=1.0)
+            return r.status_code in (200, 404)
+        except Exception:
+            return False
+
+    if _is_up():
+        return None
+
     # Launch in a background process
     proc = multiprocessing.Process(target=_start_remote_server, daemon=True)
     proc.start()
-    # Give it a moment to boot
-    time.sleep(1.5)
+
+    # Poll for readiness up to 10s
+    deadline = time.time() + 10
+    while time.time() < deadline:
+        if _is_up():
+            break
+        time.sleep(0.5)
     return proc
 
 
 # Ensure server is running BEFORE rollouts start (evaluation_test triggers rollouts before test body)
 _SERVER_PROC = _ensure_server_running()
-atexit.register(lambda: (_SERVER_PROC.terminate() if _SERVER_PROC.is_alive() else None))
+atexit.register(lambda: (_SERVER_PROC and _SERVER_PROC.is_alive() and _SERVER_PROC.terminate()))
 
 
 def _make_input_rows() -> List[EvaluationRow]:
@@ -47,12 +67,11 @@ def _make_input_rows() -> List[EvaluationRow]:
 @evaluation_test(
     input_rows=[_make_input_rows()],
     completion_params=[{"model": "fireworks_ai/accounts/fireworks/models/kimi-k2-instruct"}],
-    rollout_processor=RemoteRolloutProcessor(),
-    rollout_processor_kwargs={
-        "remote_base_url": "http://127.0.0.1:7077",
-        "num_turns": 2,
-        "timeout_seconds": 30,
-    },
+    rollout_processor=RemoteRolloutProcessor(
+        remote_base_url="http://127.0.0.1:7077",
+        num_turns=2,
+        timeout_seconds=30,
+    ),
     mode="pointwise",
 )
 async def test_remote_rollout_and_fetch_langfuse(row: EvaluationRow) -> EvaluationRow: