From 03c2efbd06011343ae5a86668dbc35ded9eacb2f Mon Sep 17 00:00:00 2001
From: Dylan Huang <dhuang@fireworks.ai>
Date: Wed, 8 Oct 2025 14:47:50 -0700
Subject: [PATCH 1/2] add util

---
 eval_protocol/__init__.py       |  2 ++
 eval_protocol/log_utils/util.py | 22 ++++++++++++++++++++++
 2 files changed, 24 insertions(+)
 create mode 100644 eval_protocol/log_utils/util.py

diff --git a/eval_protocol/__init__.py b/eval_protocol/__init__.py
index 9f17f8ac..74a2cb6b 100644
--- a/eval_protocol/__init__.py
+++ b/eval_protocol/__init__.py
@@ -34,6 +34,7 @@
 from .pytest.parameterize import DefaultParameterIdGenerator
 from .log_utils.elasticsearch_direct_http_handler import ElasticsearchDirectHttpHandler
 from .log_utils.rollout_id_filter import RolloutIdFilter
+from .log_utils.setup import setup_rollout_logging_for_elasticsearch_handler
 
 from .types.remote_rollout_processor import (
     InitRequest,
@@ -68,6 +69,7 @@
 __all__ = [
     "ElasticsearchDirectHttpHandler",
     "RolloutIdFilter",
+    "setup_rollout_logging_for_elasticsearch_handler",
     "DataLoaderConfig",
     "Status",
     "RemoteRolloutProcessor",
diff --git a/eval_protocol/log_utils/util.py b/eval_protocol/log_utils/util.py
new file mode 100644
index 00000000..ad55ab46
--- /dev/null
+++ b/eval_protocol/log_utils/util.py
@@ -0,0 +1,22 @@
+import os
+from eval_protocol.types.remote_rollout_processor import ElasticsearchConfig
+from tests.logging.test_elasticsearch_direct_http_handler import ElasticsearchDirectHttpHandler
+
+
+def setup_rollout_logging_for_elasticsearch_handler(
+    handler: ElasticsearchDirectHttpHandler, rollout_id: str, elastic_search_config: ElasticsearchConfig
+) -> None:
+    """
+    Whenever a new subprocess is created, we need to setup the rollout context
+    for the subprocess. This is useful when implementing your own remote server
+    for rollout processing.
+
+    1. Set the EP_ROLLOUT_ID environment variable
+    2. Configure the Elasticsearch handler with the Elasticsearch config
+    """
+
+    # this should only affect this subprocess so logs from this subprocess can
+    # be correlated to the rollout
+    os.environ["EP_ROLLOUT_ID"] = rollout_id
+
+    handler.configure(elasticsearch_config=elastic_search_config)

From 32e7eacee8cba75e24cb4e3e648c1fb380dc9e95 Mon Sep 17 00:00:00 2001
From: Dylan Huang <dhuang@fireworks.ai>
Date: Wed, 8 Oct 2025 14:53:09 -0700
Subject: [PATCH 2/2] fix

---
 eval_protocol/__init__.py       | 2 +-
 eval_protocol/log_utils/util.py | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/eval_protocol/__init__.py b/eval_protocol/__init__.py
index 74a2cb6b..c18ee329 100644
--- a/eval_protocol/__init__.py
+++ b/eval_protocol/__init__.py
@@ -34,7 +34,7 @@
 from .pytest.parameterize import DefaultParameterIdGenerator
 from .log_utils.elasticsearch_direct_http_handler import ElasticsearchDirectHttpHandler
 from .log_utils.rollout_id_filter import RolloutIdFilter
-from .log_utils.setup import setup_rollout_logging_for_elasticsearch_handler
+from .log_utils.util import setup_rollout_logging_for_elasticsearch_handler
 
 from .types.remote_rollout_processor import (
     InitRequest,
diff --git a/eval_protocol/log_utils/util.py b/eval_protocol/log_utils/util.py
index ad55ab46..a72907b5 100644
--- a/eval_protocol/log_utils/util.py
+++ b/eval_protocol/log_utils/util.py
@@ -1,6 +1,6 @@
 import os
 from eval_protocol.types.remote_rollout_processor import ElasticsearchConfig
-from tests.logging.test_elasticsearch_direct_http_handler import ElasticsearchDirectHttpHandler
+from .elasticsearch_direct_http_handler import ElasticsearchDirectHttpHandler
 
 
 def setup_rollout_logging_for_elasticsearch_handler(