use fireworks tracing and logging directly

benjibc · benjibc · commit 35aef2b69ea3 · 2025-10-16T16:07:19.000-07:00
diff --git a/eval_protocol/adapters/fireworks_tracing.py b/eval_protocol/adapters/fireworks_tracing.py
@@ -265,6 +265,40 @@ def __init__(
         self.base_url = base_url.rstrip("/")
         self.timeout = timeout
 
+    def search_logs(self, tags: List[str], limit: int = 100, hours_back: int = 24) -> List[Dict[str, Any]]:
+        """Fetch logs from Fireworks tracing gateway /logs endpoint.
+
+        Returns entries with keys: timestamp, message, severity, tags.
+        """
+        if not tags:
+            raise ValueError("At least one tag is required to fetch logs")
+
+        url = f"{self.base_url}/logs"
+        headers = {"Authorization": f"Bearer {os.environ.get('FIREWORKS_API_KEY')}"}
+        params: Dict[str, Any] = {"tags": tags, "limit": limit, "hours_back": hours_back, "program": "eval_protocol"}
+
+        try:
+            response = requests.get(url, params=params, timeout=self.timeout, headers=headers)
+            response.raise_for_status()
+            data = response.json() or {}
+        except requests.exceptions.RequestException as e:
+            logger.error("Failed to fetch logs from Fireworks /logs: %s", str(e))
+            return []
+
+        entries: List[Dict[str, Any]] = data.get("entries", []) or []
+        # Normalize minimal shape
+        results: List[Dict[str, Any]] = []
+        for e in entries:
+            results.append(
+                {
+                    "timestamp": e.get("timestamp"),
+                    "message": e.get("message"),
+                    "severity": e.get("severity", "INFO"),
+                    "tags": e.get("tags", []),
+                }
+            )
+        return results
+
     def get_evaluation_rows(
         self,
         tags: List[str],
diff --git a/eval_protocol/cli.py b/eval_protocol/cli.py
@@ -307,6 +307,16 @@ def parse_args(args=None):
         action="store_true",
         help="Use env vars for Elasticsearch config (requires ELASTICSEARCH_URL, ELASTICSEARCH_API_KEY, ELASTICSEARCH_INDEX_NAME)",
     )
+    logs_parser.add_argument(
+        "--use-fireworks",
+        action="store_true",
+        help="Force Fireworks tracing backend for logs UI (overrides env auto-detection)",
+    )
+    logs_parser.add_argument(
+        "--use-elasticsearch",
+        action="store_true",
+        help="Force Elasticsearch backend for logs UI (overrides env auto-detection)",
+    )
 
     # Upload command
     upload_parser = subparsers.add_parser(
diff --git a/eval_protocol/cli_commands/logs.py b/eval_protocol/cli_commands/logs.py
@@ -5,6 +5,7 @@
 import sys
 from pathlib import Path
 
+import os
 from ..utils.logs_server import serve_logs
 
 
@@ -20,53 +21,70 @@ def logs_command(args):
     print("Press Ctrl+C to stop the server")
     print("-" * 50)
 
-    # Setup Elasticsearch based on flags
+    # Backend selection: Fireworks first when API key present, unless overridden
+    use_fireworks = False
+    if getattr(args, "use_fireworks", False):
+        use_fireworks = True
+    elif getattr(args, "use_elasticsearch", False):
+        use_fireworks = False
+    else:
+        use_fireworks = bool(os.environ.get("FIREWORKS_API_KEY"))
+
+    # Setup backend configs
     elasticsearch_config = None
+    fireworks_base_url = os.environ.get("FW_TRACING_GATEWAY_BASE_URL") or "https://tracing.fireworks.ai"
     try:
-        if getattr(args, "use_env_elasticsearch_config", False):
-            # Use environment variables for configuration
-            print("⚙️ Using environment variables for Elasticsearch config")
-            from eval_protocol.pytest.remote_rollout_processor import (
-                create_elasticsearch_config_from_env,
-            )
-
-            elasticsearch_config = create_elasticsearch_config_from_env()
-            # Ensure index exists with correct mapping, mirroring Docker setup path
-            try:
-                from eval_protocol.log_utils.elasticsearch_index_manager import (
-                    ElasticsearchIndexManager,
+        if not use_fireworks:
+            if getattr(args, "use_env_elasticsearch_config", False):
+                # Use environment variables for configuration
+                print("⚙️ Using environment variables for Elasticsearch config")
+                from eval_protocol.pytest.remote_rollout_processor import (
+                    create_elasticsearch_config_from_env,
                 )
 
-                index_manager = ElasticsearchIndexManager(
-                    elasticsearch_config.url,
-                    elasticsearch_config.index_name,
-                    elasticsearch_config.api_key,
-                )
-                created = index_manager.create_logging_index_mapping()
-                if created:
-                    print(
-                        f"🧭 Verified Elasticsearch index '{elasticsearch_config.index_name}' mapping (created or already correct)"
+                elasticsearch_config = create_elasticsearch_config_from_env()
+                # Ensure index exists with correct mapping, mirroring Docker setup path
+                try:
+                    from eval_protocol.log_utils.elasticsearch_index_manager import (
+                        ElasticsearchIndexManager,
                     )
-                else:
-                    print(
-                        f"⚠️ Could not verify/create mapping for index '{elasticsearch_config.index_name}'. Searches may behave unexpectedly."
+
+                    index_manager = ElasticsearchIndexManager(
+                        elasticsearch_config.url,
+                        elasticsearch_config.index_name,
+                        elasticsearch_config.api_key,
                     )
-            except Exception as e:
-                print(f"⚠️ Failed to ensure index mapping via IndexManager: {e}")
-        elif not getattr(args, "disable_elasticsearch_setup", False):
-            # Default behavior: start or connect to local Elasticsearch via Docker helper
-            from eval_protocol.pytest.elasticsearch_setup import ElasticsearchSetup
+                    created = index_manager.create_logging_index_mapping()
+                    if created:
+                        print(
+                            f"🧭 Verified Elasticsearch index '{elasticsearch_config.index_name}' mapping (created or already correct)"
+                        )
+                    else:
+                        print(
+                            f"⚠️ Could not verify/create mapping for index '{elasticsearch_config.index_name}'. Searches may behave unexpectedly."
+                        )
+                except Exception as e:
+                    print(f"⚠️ Failed to ensure index mapping via IndexManager: {e}")
+            elif not getattr(args, "disable_elasticsearch_setup", False):
+                # Default behavior: start or connect to local Elasticsearch via Docker helper
+                from eval_protocol.pytest.elasticsearch_setup import ElasticsearchSetup
 
-            print("🧰 Auto-configuring local Elasticsearch (Docker)")
-            elasticsearch_config = ElasticsearchSetup().setup_elasticsearch()
-        else:
-            print("🚫 Elasticsearch setup disabled; running without Elasticsearch integration")
+                print("🧰 Auto-configuring local Elasticsearch (Docker)")
+                elasticsearch_config = ElasticsearchSetup().setup_elasticsearch()
+            else:
+                print("🚫 Elasticsearch setup disabled; running without Elasticsearch integration")
     except Exception as e:
         print(f"❌ Failed to configure Elasticsearch: {e}")
         return 1
 
     try:
-        serve_logs(port=args.port, elasticsearch_config=elasticsearch_config, debug=args.debug)
+        serve_logs(
+            port=args.port,
+            elasticsearch_config=elasticsearch_config,
+            debug=args.debug,
+            backend="fireworks" if use_fireworks else "elasticsearch",
+            fireworks_base_url=fireworks_base_url if use_fireworks else None,
+        )
         return 0
     except KeyboardInterrupt:
         print("\n🛑 Server stopped by user")
diff --git a/eval_protocol/log_utils/fireworks_tracing_http_handler.py b/eval_protocol/log_utils/fireworks_tracing_http_handler.py
@@ -12,10 +12,19 @@ class FireworksTracingHttpHandler(logging.Handler):
 
     def __init__(self, gateway_base_url: Optional[str] = None, rollout_id_env: str = "EP_ROLLOUT_ID") -> None:
         super().__init__()
-        self.gateway_base_url = gateway_base_url or os.getenv("FW_TRACING_GATEWAY_BASE_URL")
+        self.gateway_base_url = (
+            gateway_base_url or os.getenv("FW_TRACING_GATEWAY_BASE_URL") or "https://tracing.fireworks.ai"
+        )
         self.rollout_id_env = rollout_id_env
         self._session = requests.Session()
         self._lock = threading.Lock()
+        # Include Authorization header if FIREWORKS_API_KEY is available
+        api_key = os.environ.get("FIREWORKS_API_KEY")
+        if api_key:
+            try:
+                self._session.headers.update({"Authorization": f"Bearer {api_key}"})
+            except Exception:
+                pass
 
     def emit(self, record: logging.LogRecord) -> None:
         try:
diff --git a/eval_protocol/log_utils/init.py b/eval_protocol/log_utils/init.py
@@ -39,9 +39,10 @@ def init_external_logging_from_env() -> None:
     # Ensure we do not add duplicate handlers if already present
     existing_handler_types = {type(h).__name__ for h in root_logger.handlers}
 
-    # Fireworks tracing
-    fw_url = _get_env("FW_TRACING_GATEWAY_BASE_URL")
-    if fw_url and "FireworksTracingHttpHandler" not in existing_handler_types:
+    # Fireworks tracing: prefer if FIREWORKS_API_KEY is present; default base URL if not provided
+    fw_key = _get_env("FIREWORKS_API_KEY")
+    fw_url = _get_env("FW_TRACING_GATEWAY_BASE_URL") or "https://tracing.fireworks.ai"
+    if fw_key and "FireworksTracingHttpHandler" not in existing_handler_types:
         fw_handler = FireworksTracingHttpHandler(gateway_base_url=fw_url)
         fw_handler.setLevel(logging.INFO)
         fw_handler.addFilter(ContextRolloutIdFilter())
@@ -51,7 +52,13 @@ def init_external_logging_from_env() -> None:
     es_url = _get_env("EP_ELASTICSEARCH_URL")
     es_api_key = _get_env("EP_ELASTICSEARCH_API_KEY")
     es_index = _get_env("EP_ELASTICSEARCH_INDEX")
-    if es_url and es_api_key and es_index and "ElasticsearchDirectHttpHandler" not in existing_handler_types:
+    if (
+        not fw_key
+        and es_url
+        and es_api_key
+        and es_index
+        and "ElasticsearchDirectHttpHandler" not in existing_handler_types
+    ):
         es_config = ElasticsearchConfig(url=es_url, api_key=es_api_key, index_name=es_index)
         es_handler = ElasticsearchDirectHttpHandler(elasticsearch_config=es_config)
         es_handler.setLevel(logging.INFO)
diff --git a/eval_protocol/utils/logs_server.py b/eval_protocol/utils/logs_server.py
@@ -8,7 +8,7 @@
 from contextlib import asynccontextmanager
 from pathlib import Path
 from queue import Queue
-from typing import TYPE_CHECKING, Any, Dict, List, Optional
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Literal
 
 import psutil
 import uvicorn
@@ -327,6 +327,8 @@ def __init__(
         port: Optional[int] = 8000,
         index_file: str = "index.html",
         elasticsearch_config: Optional[ElasticsearchConfig] = None,
+        backend: Literal["fireworks", "elasticsearch"] = "elasticsearch",
+        fireworks_base_url: Optional[str] = None,
         debug: bool = False,
     ):
         # Enable debug mode if requested
@@ -336,6 +338,10 @@ def __init__(
         # Initialize WebSocket manager
         self.websocket_manager = WebSocketManager()
 
+        # Backend selection
+        self.backend: Literal["fireworks", "elasticsearch"] = backend
+        self.fireworks_base_url = fireworks_base_url
+
         # Initialize Elasticsearch client if config is provided
         self.elasticsearch_client: Optional[ElasticsearchClient] = None
         if elasticsearch_config:
@@ -414,6 +420,8 @@ async def status():
                 # Expose an empty list to satisfy consumers and type checker
                 "watch_paths": [],
                 "elasticsearch_enabled": self.elasticsearch_client is not None,
+                "backend": self.backend,
+                "fireworks_enabled": self.backend == "fireworks",
             }
 
         @self.app.get("/api/logs/{rollout_id}", response_model=LogsResponse, response_model_exclude_none=True)
@@ -422,7 +430,47 @@ async def get_logs(
             level: Optional[str] = Query(None, description="Filter by log level (DEBUG, INFO, WARNING, ERROR)"),
             limit: int = Query(100, description="Maximum number of log entries to return"),
         ) -> LogsResponse:
-            """Get logs for a specific rollout ID from Elasticsearch."""
+            """Get logs for a specific rollout ID from the configured backend."""
+            # Fireworks backend
+            if self.backend == "fireworks":
+                try:
+                    from eval_protocol.adapters.fireworks_tracing import FireworksTracingAdapter
+
+                    base_url = self.fireworks_base_url or "https://tracing.fireworks.ai"
+                    adapter = FireworksTracingAdapter(base_url=base_url)
+                    # Fetch lightweight log entries filtered by rollout_id tag
+                    tags = [f"rollout_id:{rollout_id}"]
+                    entries = adapter.search_logs(tags=tags, limit=limit)
+                    # Map to LogEntry responses
+                    log_entries: List[LogEntry] = []
+                    for e in entries:
+                        ts = e.get("timestamp") or datetime.utcnow().isoformat() + "Z"
+                        msg = e.get("message") or "trace"
+                        sev = e.get("severity") or "INFO"
+                        entry = LogEntry(
+                            **{
+                                "@timestamp": ts,
+                                "level": sev,
+                                "message": str(msg),
+                                "logger_name": "fireworks",
+                                "rollout_id": rollout_id,
+                            }
+                        )
+                        log_entries.append(entry)
+
+                    return LogsResponse(
+                        logs=log_entries,
+                        total=len(log_entries),
+                        rollout_id=rollout_id,
+                        filtered_by_level=level,
+                    )
+                except HTTPException:
+                    raise
+                except Exception as e:
+                    logger.error(f"Error retrieving Fireworks logs for rollout {rollout_id}: {e}")
+                    raise HTTPException(status_code=500, detail=f"Failed to retrieve Fireworks logs: {str(e)}")
+
+            # Elasticsearch backend
             if not self.elasticsearch_client:
                 raise HTTPException(status_code=503, detail="Elasticsearch is not configured for this logs server")
 
@@ -574,6 +622,8 @@ def create_app(
     port: int = 8000,
     build_dir: Optional[str] = None,
     elasticsearch_config: Optional[ElasticsearchConfig] = None,
+    backend: Literal["fireworks", "elasticsearch"] = "elasticsearch",
+    fireworks_base_url: Optional[str] = None,
     debug: bool = False,
 ) -> FastAPI:
     """
@@ -597,20 +647,36 @@ def create_app(
         )
 
     server = LogsServer(
-        host=host, port=port, build_dir=build_dir, elasticsearch_config=elasticsearch_config, debug=debug
+        host=host,
+        port=port,
+        build_dir=build_dir,
+        elasticsearch_config=elasticsearch_config,
+        backend=backend,
+        fireworks_base_url=fireworks_base_url,
+        debug=debug,
     )
     server.start_loops()
     return server.app
 
 
 # For backward compatibility and direct usage
 def serve_logs(
-    port: Optional[int] = None, elasticsearch_config: Optional[ElasticsearchConfig] = None, debug: bool = False
+    port: Optional[int] = None,
+    elasticsearch_config: Optional[ElasticsearchConfig] = None,
+    debug: bool = False,
+    backend: Literal["fireworks", "elasticsearch"] = "elasticsearch",
+    fireworks_base_url: Optional[str] = None,
 ):
     """
     Convenience function to create and run a LogsServer.
     """
-    server = LogsServer(port=port, elasticsearch_config=elasticsearch_config, debug=debug)
+    server = LogsServer(
+        port=port,
+        elasticsearch_config=elasticsearch_config,
+        debug=debug,
+        backend=backend,
+        fireworks_base_url=fireworks_base_url,
+    )
     server.run()
 
 
diff --git a/scripts/verify_logging_locally.py b/scripts/verify_logging_locally.py
