diff --git a/eval_protocol/cli.py b/eval_protocol/cli.py index 1d54cd65..34ad31d6 100644 --- a/eval_protocol/cli.py +++ b/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( diff --git a/eval_protocol/cli_commands/logs.py b/eval_protocol/cli_commands/logs.py index 9f26f373..16d6d6de 100644 --- a/eval_protocol/cli_commands/logs.py +++ b/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) diff --git a/eval_protocol/pytest/remote_rollout_processor.py b/eval_protocol/pytest/remote_rollout_processor.py index 4dce27af..d5bce2ee 100644 --- a/eval_protocol/pytest/remote_rollout_processor.py +++ b/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: