Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions eval_protocol/cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
26 changes: 22 additions & 4 deletions eval_protocol/cli_commands/logs.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
19 changes: 19 additions & 0 deletions eval_protocol/pytest/remote_rollout_processor.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
Loading