create_elasticsearch_config_from_env (#271)

Author: dphuang2

eval_protocol/cli.py

@@ -301,6 +301,12 @@ def parse_args(args=None):
     logs_parser = subparsers.add_parser("logs", help="Serve logs with file watching and real-time updates")
     logs_parser.add_argument("--port", type=int, default=8000, help="Port to bind to (default: 8000)")
     logs_parser.add_argument("--debug", action="store_true", help="Enable debug mode")
+    logs_parser.add_argument("--disable-elasticsearch-setup", action="store_true", help="Disable Elasticsearch setup")
+    logs_parser.add_argument(
+        "--use-env-elasticsearch-confi",
+        action="store_true",
+        help="Use env vars for Elasticsearch config (requires ELASTICSEARCH_URL, ELASTICSEARCH_API_KEY, ELASTICSEARCH_INDEX_NAME)",
+    )
 
     # Upload command
     upload_parser = subparsers.add_parser(

eval_protocol/cli_commands/logs.py

@@ -20,10 +20,28 @@ def logs_command(args):
     print("Press Ctrl+C to stop the server")
     print("-" * 50)
 
-    # setup Elasticsearch
-    from eval_protocol.pytest.elasticsearch_setup import ElasticsearchSetup
-
-    elasticsearch_config = ElasticsearchSetup().setup_elasticsearch()
+    # Setup Elasticsearch based on flags
+    elasticsearch_config = None
+    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()
+        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")
+    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)

eval_protocol/pytest/remote_rollout_processor.py

@@ -26,6 +26,25 @@
 logger = logging.getLogger(__name__)
 
 
+def create_elasticsearch_config_from_env() -> ElasticsearchConfig:
+    """Setup Elasticsearch config from environment variables."""
+    url = os.getenv("ELASTICSEARCH_URL")
+    api_key = os.getenv("ELASTICSEARCH_API_KEY")
+    index_name = os.getenv("ELASTICSEARCH_INDEX_NAME")
+
+    if url is None:
+        raise ValueError("ELASTICSEARCH_URL must be set")
+    if api_key is None:
+        raise ValueError("ELASTICSEARCH_API_KEY must be set")
+    if index_name is None:
+        raise ValueError("ELASTICSEARCH_INDEX_NAME must be set")
+    return ElasticsearchConfig(
+        url=url,
+        api_key=api_key,
+        index_name=index_name,
+    )
+
+
 def _build_fireworks_tracing_url(
     base_url: str, metadata: RolloutMetadata, completion_params_base_url: Optional[str] = None
 ) -> str: