eval-protocol 0.0.3__py3-none-any.whl

eval_protocol/cli_commands/deploy.py

@@ -0,0 +1,486 @@
+"""
+CLI command for creating and deploying an evaluator,
+or registering a pre-deployed remote evaluator.
+"""
+
+import importlib  # For dynamically importing modules
+import json
+import os  # For os.path.join, os.makedirs, os.getcwd (already imported but good to be explicit if used extensively)
+import secrets  # For API key generation (already imported but good to be explicit)
+import sys  # For sys.executable
+import time  # For sleep
+from pathlib import Path  # For path operations
+from typing import Any, Dict
+
+import yaml  # For saving config if save_config helper doesn't exist
+
+# TODO: Consider moving subprocess_manager functions to a more central location if used by core CLI
+try:
+    from development.utils.subprocess_manager import (
+        start_ngrok_and_get_url,  # Added ngrok function
+        start_process,
+        start_serveo_and_get_url,
+        stop_process,
+    )
+except ImportError:
+    # Fallback implementations when development module is not available
+    import subprocess
+    import signal
+    import socket
+
+    def start_process(command, log_path, env=None):
+        """Fallback process starter."""
+        try:
+            with open(log_path, "w") as log_file:
+                process = subprocess.Popen(command, stdout=log_file, stderr=subprocess.STDOUT, env=env)
+                return process
+        except Exception as e:
+            print(f"Error starting process: {e}")
+            return None
+
+    def stop_process(pid):
+        """Fallback process stopper."""
+        try:
+            import os
+
+            os.kill(pid, signal.SIGTERM)
+        except Exception:
+            pass
+
+    def start_serveo_and_get_url(local_port, log_path):
+        """Fallback serveo tunnel - returns None to indicate unavailable."""
+        print("Serveo tunneling not available - development module not found")
+        return None, None
+
+    def start_ngrok_and_get_url(local_port, log_path):
+        """Fallback ngrok tunnel - returns None to indicate unavailable."""
+        print("ngrok tunneling not available - development module not found")
+        return None, None
+
+
+from eval_protocol.auth import get_fireworks_account_id
+from eval_protocol.config import GCPCloudRunConfig, RewardKitConfig
+from eval_protocol.config import _config_file_path as global_loaded_config_path
+from eval_protocol.config import get_config
+from eval_protocol.evaluation import create_evaluation
+from eval_protocol.gcp_tools import (
+    build_and_push_docker_image,
+    deploy_to_cloud_run,
+    ensure_artifact_registry_repo_exists,
+    ensure_gcp_secret,
+)
+from eval_protocol.packaging import generate_dockerfile_content
+from eval_protocol.platform_api import (  # For catching errors from create_evaluation
+    PlatformAPIError,
+    create_or_update_fireworks_secret,
+)
+
+from .common import check_environment
+
+
+def _establish_local_server_and_tunnel(args):
+    """
+    Handles starting the local generic server and establishing a public tunnel
+    using Serveo, with a fallback to ngrok.
+    Returns: (public_url, tunnel_provider_name, local_server_pid, tunnel_process_pid)
+             Returns (None, None, server_pid_or_None, None) if tunneling fails.
+    """
+    if not args.function_ref:
+        print("Error: --function-ref is required for local-serve target.")
+        return None, None, None, None
+
+    evaluator_id = args.id
+    function_ref = args.function_ref
+    local_server_port = args.local_port
+
+    log_dir = os.path.join(os.getcwd(), "logs", "reward-kit-local")
+    os.makedirs(log_dir, exist_ok=True)
+    generic_server_log_path = os.path.join(log_dir, f"generic_server_{evaluator_id}.log")
+
+    server_env = None  # Run local server without API key protection
+    print(f"Note: Local server for '{evaluator_id}' will run without API key protection.")
+
+    print(f"Starting local reward function server for '{function_ref}' on port {local_server_port}...")
+    server_command = [
+        sys.executable,
+        "-m",
+        "eval_protocol.generic_server",
+        function_ref,
+        "--port",
+        str(local_server_port),
+    ]
+
+    local_server_process = start_process(server_command, generic_server_log_path, env=server_env)
+
+    if not local_server_process or local_server_process.poll() is not None:
+        print(f"Error: Failed to start local generic server. Check log: {generic_server_log_path}")
+        return None, None, None, None  # No server, no tunnel
+
+    local_server_pid = local_server_process.pid
+    print(f"Local server started (PID: {local_server_pid}). Log: {generic_server_log_path}")
+    print("Waiting for server to initialize...")
+    time.sleep(5)
+
+    # Attempt Serveo first
+    print(f"Attempting Serveo tunnel for local port {local_server_port}...")
+    serveo_log_path = os.path.join(log_dir, f"serveo_{evaluator_id}.log")
+    serveo_tunnel_process, serveo_url = start_serveo_and_get_url(local_server_port, serveo_log_path)
+
+    if serveo_url and serveo_tunnel_process:
+        print(f"Serveo tunnel established: {serveo_url} (PID: {serveo_tunnel_process.pid}). Log: {serveo_log_path}")
+        return serveo_url, "serveo", local_server_pid, serveo_tunnel_process.pid
+    else:
+        print(f"Serveo tunnel failed. Check log: {serveo_log_path}")
+        print("Attempting fallback to ngrok...")
+
+        ngrok_log_path = os.path.join(log_dir, f"ngrok_{evaluator_id}.log")
+        # Assuming ngrok authtoken is pre-configured by the user or via NGROK_AUTHTOKEN env var
+        ngrok_tunnel_process, ngrok_url = start_ngrok_and_get_url(local_server_port, ngrok_log_path)
+
+        if ngrok_url and ngrok_tunnel_process:
+            print(f"ngrok tunnel established: {ngrok_url} (PID: {ngrok_tunnel_process.pid}). Log: {ngrok_log_path}")
+            return ngrok_url, "ngrok", local_server_pid, ngrok_tunnel_process.pid
+        else:
+            print(f"ngrok tunnel also failed. Check log: {ngrok_log_path}")
+            # Both failed, stop the local server we started
+            if local_server_pid:
+                stop_process(local_server_pid)
+            return (
+                None,
+                None,
+                local_server_pid,
+                None,
+            )  # URL, provider, server_pid, tunnel_pid
+
+
+def _deploy_to_gcp_cloud_run(args, current_config, gcp_config_from_yaml):
+    """Handles the logic for --target gcp-cloud-run up to service deployment."""
+    print(f"Starting GCP Cloud Run deployment for evaluator '{args.id}'...")
+
+    # Resolve function_ref (must be from CLI for GCP)
+    if not args.function_ref:  # This check is also in main, but good for helper too
+        print("Error: --function-ref is required for GCP Cloud Run deployment.")
+        return None
+
+    # Dynamically import the reward function to get its requirements
+    inline_requirements_content = None
+    try:
+        module_name, func_name = args.function_ref.rsplit(".", 1)
+        module = importlib.import_module(module_name)
+        reward_func = getattr(module, func_name)
+        if hasattr(reward_func, "_reward_function_requirements"):
+            inline_requirements_content = reward_func._reward_function_requirements
+            if inline_requirements_content:
+                print(f"Found inline requirements for {args.function_ref}")
+    except Exception as e:
+        print(f"Warning: Could not load reward function {args.function_ref} to check for inline requirements: {e}")
+        # Continue without inline requirements if loading fails
+
+    # Resolve GCP project_id
+    gcp_project_id = args.gcp_project
+    if not gcp_project_id and gcp_config_from_yaml:
+        gcp_project_id = gcp_config_from_yaml.project_id
+    if not gcp_project_id:
+        print("Error: GCP Project ID must be provided via --gcp-project argument or in rewardkit.yaml.")
+        return None
+
+    # Resolve GCP region
+    gcp_region = args.gcp_region
+    if not gcp_region and gcp_config_from_yaml:
+        gcp_region = gcp_config_from_yaml.region
+    if not gcp_region:
+        print("Error: GCP Region must be provided via --gcp-region argument or in rewardkit.yaml.")
+        return None
+
+    # Resolve GCP AR repo name
+    gcp_ar_repo_name = args.gcp_ar_repo
+    if not gcp_ar_repo_name and gcp_config_from_yaml:
+        gcp_ar_repo_name = gcp_config_from_yaml.artifact_registry_repository
+    if not gcp_ar_repo_name:
+        gcp_ar_repo_name = "reward-kit-evaluators"
+
+    print(f"Using GCP Project: {gcp_project_id}, Region: {gcp_region}, AR Repo: {gcp_ar_repo_name}")
+
+    if not ensure_artifact_registry_repo_exists(
+        project_id=gcp_project_id, region=gcp_region, repo_name=gcp_ar_repo_name
+    ):
+        print(f"Failed to ensure Artifact Registry repository '{gcp_ar_repo_name}' exists. Aborting.")
+        return None
+
+    dockerfile_content = generate_dockerfile_content(
+        function_ref=args.function_ref,
+        python_version=(
+            f"{args.runtime[6]}.{args.runtime[7:]}"
+            if args.runtime.startswith("python") and len(args.runtime) > 7
+            else args.runtime.replace("python", "")
+        ),
+        eval_protocol_install_source=".",
+        user_requirements_path=None,  # Explicitly None, inline_requirements_content will be used
+        inline_requirements_content=inline_requirements_content,
+        service_port=8080,
+    )
+    if not dockerfile_content:
+        print("Failed to generate Dockerfile content. Aborting.")
+        return None
+
+    image_tag = "latest"
+    image_name_tag = f"{gcp_region}-docker.pkg.dev/{gcp_project_id}/{gcp_ar_repo_name}/{args.id}:{image_tag}"
+    build_context_dir = os.getcwd()
+
+    if not build_and_push_docker_image(
+        image_name_tag=image_name_tag,
+        dockerfile_content=dockerfile_content,
+        build_context_dir=build_context_dir,
+        gcp_project_id=gcp_project_id,
+    ):
+        print(f"Failed to build and push Docker image {image_name_tag}. Aborting.")
+        return None
+    print(f"Successfully built and pushed Docker image: {image_name_tag}")
+
+    gcp_env_vars: Dict[str, str] = {}
+    parsed_gcp_secrets: Dict[str, Any] = {}
+    allow_unauthenticated_gcp = True
+
+    resolved_auth_mode = "api-key"
+    if gcp_config_from_yaml and gcp_config_from_yaml.default_auth_mode:
+        resolved_auth_mode = gcp_config_from_yaml.default_auth_mode
+    if args.gcp_auth_mode is not None:
+        resolved_auth_mode = args.gcp_auth_mode
+    print(f"Using GCP Auth Mode for service: {resolved_auth_mode}")
+
+    if resolved_auth_mode == "api-key":
+        print("Configuring GCP Cloud Run service for API key authentication (application layer).")
+        evaluator_id = args.id
+        api_key_for_service = None  # This is the key the service itself will use
+        config_path = global_loaded_config_path
+
+        if current_config.evaluator_endpoint_keys and evaluator_id in current_config.evaluator_endpoint_keys:
+            api_key_for_service = current_config.evaluator_endpoint_keys[evaluator_id]
+            print(f"Using existing API key for '{evaluator_id}' from configuration for the service.")
+        else:
+            api_key_for_service = secrets.token_hex(32)
+            print(f"Generated new API key for '{evaluator_id}' for the service.")
+            if not current_config.evaluator_endpoint_keys:
+                current_config.evaluator_endpoint_keys = {}
+            current_config.evaluator_endpoint_keys[evaluator_id] = api_key_for_service
+            if config_path:
+                _save_config(current_config, config_path)
+            else:
+                print(f"Warning: No rewardkit.yaml found to save API key for '{evaluator_id}'.")
+
+        gcp_sanitized_eval_id = "".join(filter(lambda char: char.isalnum() or char in ["-", "_"], args.id))
+        if not gcp_sanitized_eval_id:
+            gcp_sanitized_eval_id = "rewardkit-evaluator"
+        secret_id_for_auth_key = f"rk-eval-{gcp_sanitized_eval_id}-authkey"
+        secret_labels = {"managed-by": "reward-kit", "evaluator-id": evaluator_id}
+
+        api_key_secret_version_id = ensure_gcp_secret(
+            project_id=gcp_project_id,
+            secret_id=secret_id_for_auth_key,
+            secret_value=api_key_for_service,
+            labels=secret_labels,
+        )
+        if not api_key_secret_version_id:
+            print(f"Error: Failed to store API key in GCP Secret Manager for '{evaluator_id}'. Aborting.")
+            return None
+        print(f"API key for service stored in GCP Secret Manager: {secret_id_for_auth_key}")
+        parsed_gcp_secrets["RK_ENDPOINT_API_KEY"] = api_key_secret_version_id
+
+        # Register this key with Fireworks secrets for the shim
+        fireworks_account_id_for_secret = get_fireworks_account_id()
+        if fireworks_account_id_for_secret:
+            fw_eval_id_sanitized = args.id.lower()
+            fw_eval_id_sanitized = "".join(filter(lambda char: char.isalnum() or char == "-", fw_eval_id_sanitized))
+            fw_eval_id_sanitized = "-".join(filter(None, fw_eval_id_sanitized.split("-")))
+            if not fw_eval_id_sanitized:
+                fw_eval_id_sanitized = "evaluator"
+            fw_eval_id_sanitized = fw_eval_id_sanitized[:40]
+            fw_secret_key_name = f"rkeval-{fw_eval_id_sanitized}-shim-key"
+            print(f"Registering API key on Fireworks platform as secret '{fw_secret_key_name}' for shim...")
+            if create_or_update_fireworks_secret(
+                account_id=fireworks_account_id_for_secret,
+                key_name=fw_secret_key_name,
+                secret_value=api_key_for_service,
+            ):
+                print(f"Successfully registered/updated secret '{fw_secret_key_name}' on Fireworks platform.")
+            else:
+                print(f"Warning: Failed to register/update secret '{fw_secret_key_name}' on Fireworks platform.")
+        else:
+            print("Warning: Fireworks Account ID not found, cannot store shim API key on Fireworks platform.")
+
+    cloud_run_service_url = deploy_to_cloud_run(
+        service_name=args.id,
+        image_name_tag=image_name_tag,
+        gcp_project_id=gcp_project_id,
+        gcp_region=gcp_region,
+        allow_unauthenticated=allow_unauthenticated_gcp,  # True if api-key mode, app handles auth
+        env_vars=gcp_env_vars if gcp_env_vars else None,
+        secrets_to_mount=parsed_gcp_secrets,
+    )
+
+    if not cloud_run_service_url:
+        print("Failed to deploy to Cloud Run or retrieve service URL. Aborting.")
+        return None
+
+    print(f"Successfully deployed to Cloud Run. Service URL: {cloud_run_service_url}")
+    return cloud_run_service_url
+
+
+# Helper to save config (can be moved to config.py later)
+def _save_config(config_data: RewardKitConfig, path: str):
+    # Basic save, ideally config.py would provide a robust method
+    try:
+        with open(path, "w") as f:
+            yaml.dump(config_data.model_dump(exclude_none=True), f, sort_keys=False)
+        print(f"Config updated and saved to {path}")
+    except Exception as e:
+        print(f"Warning: Failed to save updated config to {path}: {e}")
+
+
+def deploy_command(args):
+    """Create and deploy an evaluator or register a remote one."""
+
+    # Check environment variables
+    if not check_environment():
+        return 1
+
+    if not args.id:  # ID is always required
+        print("Error: Evaluator ID (--id) is required.")
+        return 1
+
+    # Process HuggingFace key mapping if provided
+    huggingface_message_key_map = None
+    if args.huggingface_key_map:
+        try:
+            huggingface_message_key_map = json.loads(args.huggingface_key_map)
+        except json.JSONDecodeError:
+            print("Error: Invalid JSON format for --huggingface-key-map")
+            return 1
+
+    # Initialize variables for URL registration path
+    service_url_to_register = None
+    # api_key_for_shim = None # Not currently used by create_evaluation for shim auth directly
+
+    # PIDs for cleanup if registration fails for local-serve
+    local_server_pid_to_clean = None
+    # serveo_pid_to_clean = None # This was old, replaced by local_tunnel_pid_to_clean
+    local_tunnel_pid_to_clean = None  # Initialize here
+
+    if args.target == "gcp-cloud-run":
+        current_config = get_config()  # Needed by the helper
+        gcp_config_from_yaml = current_config.gcp_cloud_run if current_config.gcp_cloud_run else None
+
+        cloud_run_service_url = _deploy_to_gcp_cloud_run(args, current_config, gcp_config_from_yaml)
+        if not cloud_run_service_url:
+            return 1  # Error already printed by helper
+        service_url_to_register = cloud_run_service_url
+
+    elif args.target == "local-serve":
+        # Renamed helper and updated return values
+        url, tunnel_provider, server_pid, tunnel_pid = _establish_local_server_and_tunnel(args)
+        if not url:
+            # _establish_local_server_and_tunnel handles cleanup of server if tunnel fails completely
+            return 1  # Error already printed by helper
+        service_url_to_register = url
+        local_server_pid_to_clean = server_pid
+        # serveo_pid_to_clean was specific, now it's generic tunnel_pid
+        # Let's rename it for clarity in the cleanup logic
+        local_tunnel_pid_to_clean = tunnel_pid
+        print(f"Tunnel established using {tunnel_provider}.")
+
+    elif args.remote_url:
+        # This is for --target fireworks (default) but with --remote-url
+        print(f"Registering remote URL: {args.remote_url} for evaluator '{args.id}'")
+        if not (args.remote_url.startswith("http://") or args.remote_url.startswith("https://")):
+            print(f"Error: Invalid --remote-url '{args.remote_url}'. Must start with http:// or https://")
+            return 1
+        if args.metrics_folders:  # This check might be redundant if --target is explicit
+            print("Info: --metrics-folders are ignored when deploying with --remote-url.")
+        service_url_to_register = args.remote_url
+        # No specific shim auth provided by this path.
+
+    # Common registration step for targets that produce a URL
+    if service_url_to_register:
+        try:
+            print(f"Registering URL '{service_url_to_register}' with Fireworks AI for evaluator '{args.id}'...")
+            evaluator = create_evaluation(
+                evaluator_id=args.id,
+                remote_url=service_url_to_register,
+                display_name=args.display_name or args.id,
+                description=args.description or f"Evaluator for {args.id} at {service_url_to_register}",
+                force=args.force,
+                huggingface_dataset=args.huggingface_dataset,
+                huggingface_split=args.huggingface_split,
+                huggingface_message_key_map=huggingface_message_key_map,
+                huggingface_prompt_key=args.huggingface_prompt_key,
+                huggingface_response_key=args.huggingface_response_key,
+                # remote_auth_header_name="X-Api-Key" if api_key_for_shim else None, # No API key for shim for now
+                # remote_auth_header_value=api_key_for_shim # No API key for shim for now
+            )
+            evaluator_name = evaluator.get("name", args.id)
+            print(
+                f"Successfully registered evaluator '{evaluator_name}' on Fireworks AI, pointing to '{service_url_to_register}'."
+            )
+            if args.target == "local-serve":
+                # tunnel_provider is defined in the local-serve block
+                # We need to ensure it's accessible here or pass it through.
+                # For now, let's assume tunnel_provider was defined in the calling scope of this block.
+                # This will require a small adjustment to how tunnel_provider is scoped.
+                # Let's fetch it from args if we store it there, or pass it.
+                # Simpler: just make the message generic or re-fetch from the PIDs.
+                # The variable `tunnel_provider` is set in the `elif args.target == "local-serve":` block.
+                # It needs to be available here.
+                # For now, I'll adjust the print statement to be more generic or rely on the PIDs.
+                # The `tunnel_provider` variable is indeed set in the correct scope.
+                print(
+                    f"Local server (PID: {local_server_pid_to_clean}) and {tunnel_provider} tunnel (PID: {local_tunnel_pid_to_clean}) are running."
+                )
+                print("They will be stopped automatically when this command exits (e.g., Ctrl+C).")
+            return 0
+        except PlatformAPIError as e:
+            print(f"Error registering URL with Fireworks AI: {str(e)}")
+        except Exception as e:
+            print(f"An unexpected error occurred during Fireworks AI registration: {str(e)}")
+        finally:
+            # If registration fails for local-serve, clean up the started processes
+            if args.target == "local-serve" and ("evaluator" not in locals() or not locals().get("evaluator")):
+                print("Registration failed or was interrupted for local-serve. Cleaning up local processes...")
+                if local_tunnel_pid_to_clean:  # Use the new generic tunnel PID variable
+                    stop_process(local_tunnel_pid_to_clean)
+                if local_server_pid_to_clean:
+                    stop_process(local_server_pid_to_clean)
+        return 1
+
+    # Fallback to original behavior: Deploying by packaging local metrics_folders (target=fireworks, no remote_url)
+    # This is when args.target == "fireworks" (default) AND args.remote_url is NOT provided.
+    elif args.target == "fireworks" and not args.remote_url:
+        if not args.metrics_folders:
+            print("Error: --metrics-folders are required for 'fireworks' target if --remote-url is not provided.")
+            return 1
+        for folder_spec in args.metrics_folders:
+            if "=" not in folder_spec:
+                print(f"Error: Metric folder format should be 'name=path', got '{folder_spec}'")
+                return 1
+        try:
+            print(f"Packaging and deploying metrics for evaluator '{args.id}' to Fireworks AI...")
+            evaluator = create_evaluation(
+                evaluator_id=args.id,
+                metric_folders=args.metrics_folders,
+                display_name=args.display_name or args.id,
+                description=args.description or f"Evaluator: {args.id}",
+                force=args.force,
+                huggingface_dataset=args.huggingface_dataset,
+                huggingface_split=args.huggingface_split,
+                huggingface_message_key_map=huggingface_message_key_map,
+                huggingface_prompt_key=args.huggingface_prompt_key,
+                huggingface_response_key=args.huggingface_response_key,
+            )
+            evaluator_name = evaluator.get("name", args.id)
+            print(f"Successfully created/updated evaluator: {evaluator_name}")
+            return 0
+        except PlatformAPIError as e:
+            print(f"Error creating/updating evaluator '{args.id}': {str(e)}")
+            return 1
+        except Exception as e:
+            print(f"Error creating/updating evaluator '{args.id}': {str(e)}")
+            return 1