add dataloaderconfig

Author: xzrderek

eval_protocol/pytest/remote_rollout_processor.py

@@ -7,7 +7,12 @@
 from eval_protocol.log_utils.elasticsearch_client import ElasticsearchClient
 from eval_protocol.models import EvaluationRow, Status
 from eval_protocol.data_loader.dynamic_data_loader import DynamicDataLoader
-from eval_protocol.types.remote_rollout_processor import ElasticsearchConfig, InitRequest, RolloutMetadata
+from eval_protocol.types.remote_rollout_processor import (
+    DataLoaderConfig,
+    ElasticsearchConfig,
+    InitRequest,
+    RolloutMetadata,
+)
 from eval_protocol.adapters.fireworks_tracing import create_fireworks_tracing_adapter
 from eval_protocol.quickstart.utils import filter_longest_conversation
 from .rollout_processor import RolloutProcessor
@@ -20,19 +25,20 @@
 logger = logging.getLogger(__name__)
 
 
-def _default_output_data_loader(rollout_id: str, base_url: str) -> DynamicDataLoader:
+def _default_output_data_loader(config: DataLoaderConfig) -> DynamicDataLoader:
     """Default output data loader that fetches traces from Fireworks tracing proxy.
 
     Args:
-        rollout_id: The rollout ID to filter traces by
+        config: Configuration containing rollout_id and optional model_base_url
 
     Returns:
         DynamicDataLoader configured to fetch and process traces
     """
 
     def fetch_traces() -> List[EvaluationRow]:
+        base_url = config.model_base_url or "https://tracing.fireworks.ai"
         adapter = create_fireworks_tracing_adapter(base_url=base_url)
-        return adapter.get_evaluation_rows(tags=[f"rollout_id:{rollout_id}"], max_retries=5)
+        return adapter.get_evaluation_rows(tags=[f"rollout_id:{config.rollout_id}"], max_retries=5)
 
     return DynamicDataLoader(generators=[fetch_traces], preprocess_fn=filter_longest_conversation)
 
@@ -54,7 +60,7 @@ def __init__(
         model_base_url: str = "https://tracing.fireworks.ai",
         poll_interval: float = 1.0,
         timeout_seconds: float = 120.0,
-        output_data_loader: Optional[Callable[[str, str], DynamicDataLoader]] = None,
+        output_data_loader: Optional[Callable[[DataLoaderConfig], DynamicDataLoader]] = None,
         disable_elastic_search: bool = False,
         elastic_search_config: Optional[ElasticsearchConfig] = None,
     ):
@@ -64,7 +70,6 @@ def __init__(
         self._model_base_url = model_base_url
         if os.getenv("EP_REMOTE_ROLLOUT_PROCESSOR_BASE_URL"):
             self._remote_base_url = os.getenv("EP_REMOTE_ROLLOUT_PROCESSOR_BASE_URL")
-        self._model_base_url = model_base_url
         _ep_model_base_url = os.getenv("EP_MODEL_BASE_URL")
         if _ep_model_base_url:
             self._model_base_url = _ep_model_base_url
@@ -268,7 +273,10 @@ def _get_status() -> Dict[str, Any]:
             if row.execution_metadata.rollout_id is None:
                 raise ValueError("Rollout ID is required in RemoteRolloutProcessor")
 
-            data_loader = self._output_data_loader(row.execution_metadata.rollout_id, model_base_url)
+            loader_config = DataLoaderConfig(
+                rollout_id=row.execution_metadata.rollout_id, model_base_url=model_base_url
+            )
+            data_loader = self._output_data_loader(loader_config)
 
             def _load_data():
                 return data_loader.load()

eval_protocol/types/remote_rollout_processor.py

@@ -34,6 +34,13 @@ class RolloutMetadata(BaseModel):
     row_id: str
 
 
+class DataLoaderConfig(BaseModel):
+    """Configuration passed to output_data_loader functions."""
+
+    rollout_id: str
+    model_base_url: Optional[str] = None
+
+
 class InitRequest(BaseModel):
     """Request model for POST /init endpoint."""
 

tests/remote_server/test_remote_fireworks.py

@@ -22,6 +22,7 @@
 from eval_protocol.pytest.remote_rollout_processor import RemoteRolloutProcessor
 from eval_protocol.adapters.fireworks_tracing import create_fireworks_tracing_adapter
 from eval_protocol.quickstart.utils import filter_longest_conversation
+from eval_protocol.types.remote_rollout_processor import DataLoaderConfig
 
 ROLLOUT_IDS = set()
 
@@ -36,17 +37,18 @@ def check_rollout_coverage():
     assert len(ROLLOUT_IDS) == 3, f"Expected to see 3 rollout_ids, but only saw {ROLLOUT_IDS}"
 
 
-def fetch_fireworks_traces(rollout_id: str, base_url: str) -> List[EvaluationRow]:
+def fetch_fireworks_traces(config: DataLoaderConfig) -> List[EvaluationRow]:
     global ROLLOUT_IDS  # Track all rollout_ids we've seen
-    ROLLOUT_IDS.add(rollout_id)
+    ROLLOUT_IDS.add(config.rollout_id)
 
+    base_url = config.model_base_url or "https://tracing.fireworks.ai"
     adapter = create_fireworks_tracing_adapter(base_url=base_url)
-    return adapter.get_evaluation_rows(tags=[f"rollout_id:{rollout_id}"], max_retries=5)
+    return adapter.get_evaluation_rows(tags=[f"rollout_id:{config.rollout_id}"], max_retries=5)
 
 
-def fireworks_output_data_loader(rollout_id: str, base_url: str) -> DynamicDataLoader:
+def fireworks_output_data_loader(config: DataLoaderConfig) -> DynamicDataLoader:
     return DynamicDataLoader(
-        generators=[lambda: fetch_fireworks_traces(rollout_id, base_url)], preprocess_fn=filter_longest_conversation
+        generators=[lambda: fetch_fireworks_traces(config)], preprocess_fn=filter_longest_conversation
     )
 
 

tests/remote_server/test_remote_langfuse.py

@@ -22,6 +22,7 @@
 from eval_protocol.pytest.remote_rollout_processor import RemoteRolloutProcessor
 from eval_protocol.adapters.langfuse import create_langfuse_adapter
 from eval_protocol.quickstart.utils import filter_longest_conversation
+from eval_protocol.types.remote_rollout_processor import DataLoaderConfig
 
 ROLLOUT_IDS = set()
 
@@ -36,17 +37,17 @@ def check_rollout_coverage():
     assert len(ROLLOUT_IDS) == 3, f"Expected to see {ROLLOUT_IDS} rollout_ids, but only saw {ROLLOUT_IDS}"
 
 
-def fetch_langfuse_traces(rollout_id: str) -> List[EvaluationRow]:
+def fetch_langfuse_traces(config: DataLoaderConfig) -> List[EvaluationRow]:
     global ROLLOUT_IDS  # Track all rollout_ids we've seen
-    ROLLOUT_IDS.add(rollout_id)
+    ROLLOUT_IDS.add(config.rollout_id)
 
     adapter = create_langfuse_adapter()
-    return adapter.get_evaluation_rows(tags=[f"rollout_id:{rollout_id}"], max_retries=5)
+    return adapter.get_evaluation_rows(tags=[f"rollout_id:{config.rollout_id}"], max_retries=5)
 
 
-def langfuse_output_data_loader(rollout_id: str) -> DynamicDataLoader:
+def langfuse_output_data_loader(config: DataLoaderConfig) -> DynamicDataLoader:
     return DynamicDataLoader(
-        generators=[lambda: fetch_langfuse_traces(rollout_id)], preprocess_fn=filter_longest_conversation
+        generators=[lambda: fetch_langfuse_traces(config)], preprocess_fn=filter_longest_conversation
     )