nemo-evaluator-launcher 0.1.28__py3-none-any.whl

nemo_evaluator_launcher/executors/lepton/deployment_helpers.py

@@ -0,0 +1,609 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+"""Lepton deployment helper functions for nemo-evaluator-launcher.
+
+Handles Lepton endpoint creation, management, and health checks.
+"""
+
+import json
+import subprocess
+import time
+from pathlib import Path
+from typing import Any, Dict, Optional
+
+# Import lepton dependencies
+from omegaconf import DictConfig
+
+from nemo_evaluator_launcher.common.logging_utils import logger
+
+
+def deep_merge(base: Dict[Any, Any], override: Dict[Any, Any]) -> Dict[Any, Any]:
+    """Deep merge two dictionaries, with override taking precedence."""
+    result = base.copy()
+
+    for key, value in override.items():
+        if key in result and isinstance(result[key], dict) and isinstance(value, dict):
+            result[key] = deep_merge(result[key], value)
+        else:
+            result[key] = value
+
+    return result
+
+
+def replace_placeholders(data: Any, replacements: Dict[str, str]) -> Any:
+    """Replace placeholders in the data structure."""
+
+    def replace_in_obj(obj: Any) -> Any:
+        if isinstance(obj, dict):
+            return {k: replace_in_obj(v) for k, v in obj.items()}
+        elif isinstance(obj, list):
+            return [replace_in_obj(item) for item in obj]
+        elif isinstance(obj, str):
+            result = obj
+            for placeholder, value in replacements.items():
+                result = result.replace(f"{{{{{placeholder}}}}}", value)
+            return result
+        else:
+            return obj
+
+    return replace_in_obj(data)
+
+
+def generate_lepton_spec(cfg: DictConfig) -> Dict[str, Any]:
+    """Generate a Lepton endpoint specification from nemo-evaluator-launcher configuration.
+
+    This function creates a layered configuration by merging:
+    1. Platform defaults (from execution.lepton_platform.platform_defaults)
+    2. Environment settings (from execution.lepton_platform)
+    3. Inference engine config (from deployment.* - vllm/sglang settings)
+    4. Lepton platform config (from deployment.lepton_config - Lepton-specific settings)
+
+    Args:
+        cfg: The nemo-evaluator-launcher configuration object containing all settings.
+
+    Returns:
+        Dict containing the Lepton endpoint specification.
+    """
+
+    # Step 1: Start with platform defaults from execution config
+    platform_defaults = {}
+    if hasattr(cfg, "execution") and hasattr(cfg.execution, "lepton_platform"):
+        deployment_config = cfg.execution.lepton_platform.get("deployment", {})
+        platform_defaults = deployment_config.get("platform_defaults", {})
+
+    base_config = deep_merge({}, platform_defaults)
+
+    # Step 2: Apply deployment-specific settings from execution config
+    if hasattr(cfg, "execution") and hasattr(cfg.execution, "lepton_platform"):
+        lepton_platform = cfg.execution.lepton_platform
+        deployment_config = lepton_platform.get("deployment", {})
+
+        # Add deployment node group as affinity constraint
+        deployment_node_group = deployment_config.get("node_group")
+        if deployment_node_group:
+            if not base_config.get("resource_requirement"):
+                base_config["resource_requirement"] = {}
+            base_config["resource_requirement"]["affinity"] = {
+                "allowed_dedicated_node_groups": [deployment_node_group]
+            }
+
+        # Add queue config from platform defaults
+        platform_defaults = deployment_config.get("platform_defaults", {})
+        if platform_defaults.get("queue_config"):
+            base_config["queue_config"] = platform_defaults.get("queue_config")
+
+    # Step 3: Get Lepton-specific config from deployment.lepton_config
+    if not hasattr(cfg.deployment, "lepton_config"):
+        raise ValueError(
+            "deployment.lepton_config is required when using Lepton executor"
+        )
+
+    lepton_config = cfg.deployment.lepton_config
+
+    # Step 4: Convert inference engine config to container spec
+    container_spec = _create_inference_container_spec(cfg.deployment)
+
+    # Step 5: Apply Lepton platform deployment configurations
+    deployment_config = {
+        "resource_requirement": {
+            **base_config.get("resource_requirement", {}),
+            "resource_shape": lepton_config.resource_shape,
+            "min_replicas": lepton_config.min_replicas,
+            "max_replicas": lepton_config.max_replicas,
+        },
+        "auto_scaler": lepton_config.auto_scaler,
+        "container": container_spec,
+        "envs": [],
+    }
+
+    # Add health check configuration if provided
+    if hasattr(lepton_config, "health") and lepton_config.health:
+        deployment_config["health"] = lepton_config.health
+    # Merge deployment config into base config
+    final_config = deep_merge(base_config, deployment_config)
+
+    # Step 6: Add environment variables from lepton_config
+    if hasattr(lepton_config, "envs") and lepton_config.envs:
+        from omegaconf import DictConfig
+
+        for key, value in lepton_config.envs.items():
+            env_var: Dict[str, Any] = {"name": key}
+
+            # Support both direct values and secret references
+            if isinstance(value, (dict, DictConfig)) and "value_from" in value:
+                # Secret reference: {value_from: {secret_name_ref: "secret_name"}}
+                env_var["value_from"] = dict(value["value_from"])
+            else:
+                # Direct value: "direct_value"
+                env_var["value"] = str(value)
+
+            final_config["envs"].append(env_var)
+
+    # Step 6b: Auto-populate environment variables from deployment parameters
+    _add_deployment_derived_envs(final_config["envs"], cfg.deployment)
+
+    # Step 7: Add mounts with intelligent path construction
+    if hasattr(lepton_config, "mounts") and lepton_config.mounts.enabled:
+        # Get storage source from task config mounts (since mounts are shared between tasks and deployments)
+        storage_source = "node-nfs:lepton-shared-fs"  # default
+        if hasattr(cfg, "execution") and hasattr(cfg.execution, "lepton_platform"):
+            task_config = cfg.execution.lepton_platform.get("tasks", {})
+            task_mounts = task_config.get("mounts", [])
+            if task_mounts:
+                storage_source = task_mounts[0].get("from", storage_source)
+
+        final_config["mounts"] = [
+            {
+                "path": lepton_config.mounts.cache_path,
+                "from": storage_source,
+                "mount_path": lepton_config.mounts.mount_path,
+                "mount_options": {},
+            }
+        ]
+
+    # Step 8: Extract image_pull_secrets to top level (required by Lepton API)
+    if "image_pull_secrets" in final_config:
+        image_pull_secrets = final_config["image_pull_secrets"]
+        # Convert OmegaConf ListConfig to regular Python list
+        from omegaconf import ListConfig
+
+        if isinstance(image_pull_secrets, (list, ListConfig)):
+            final_config["image_pull_secrets"] = list(image_pull_secrets)
+        else:
+            # Remove invalid image_pull_secrets
+            final_config.pop("image_pull_secrets", None)
+
+    # Step 9: Add API tokens if provided (supports both single and multiple tokens)
+    if hasattr(lepton_config, "api_tokens") and lepton_config.api_tokens:
+        from omegaconf import DictConfig
+
+        api_tokens_list = []
+
+        for token_config in lepton_config.api_tokens:
+            token_var: Dict[str, Any] = {}
+
+            # Support both direct values and secret references
+            if isinstance(token_config, (dict, DictConfig)):
+                if "value" in token_config:
+                    # Direct value: {value: "token_string"}
+                    token_var["value"] = str(token_config["value"])
+                elif "value_from" in token_config:
+                    # Secret reference: {value_from: {secret_name_ref: "secret_name"}}
+                    token_var["value_from"] = dict(token_config["value_from"])
+            else:
+                # Simple string value
+                token_var["value"] = str(token_config)
+
+            api_tokens_list.append(token_var)
+
+        final_config["api_tokens"] = api_tokens_list
+
+    # Backward compatibility: support legacy single api_token
+    elif hasattr(lepton_config, "api_token") and lepton_config.api_token:
+        final_config["api_tokens"] = [{"value": lepton_config.api_token}]
+
+    # Step 10: Replace placeholders
+    replacements = {
+        "MODEL_CACHE_NAME": _generate_model_cache_name(cfg.deployment.image)
+    }
+    final_config_with_replacements: Dict[str, Any] = replace_placeholders(
+        final_config, replacements
+    )
+
+    return final_config_with_replacements
+
+
+def _create_inference_container_spec(deployment_cfg: DictConfig) -> Dict[str, Any]:
+    """Create container specification from inference engine config (vLLM/SGLang/NIM).
+
+    Args:
+        deployment_cfg: Deployment configuration containing vLLM/SGLang/NIM settings.
+
+    Returns:
+        Container specification for Lepton.
+    """
+    container_spec = {
+        "image": deployment_cfg.image,
+        "ports": [{"container_port": deployment_cfg.port}],
+    }
+
+    # Generate command based on deployment type
+    if deployment_cfg.type == "vllm":
+        # Convert vLLM command template to actual command
+        command_parts = [
+            "vllm",
+            "serve",
+            deployment_cfg.checkpoint_path,
+            f"--tensor-parallel-size={deployment_cfg.tensor_parallel_size}",
+            f"--pipeline-parallel-size={deployment_cfg.pipeline_parallel_size}",
+            f"--data-parallel-size={deployment_cfg.data_parallel_size}",
+            f"--port={deployment_cfg.port}",
+            f"--served-model-name={deployment_cfg.served_model_name}",
+        ]
+
+        # Add extra args if provided
+        if hasattr(deployment_cfg, "extra_args") and deployment_cfg.extra_args:
+            command_parts.extend(deployment_cfg.extra_args.split())
+
+        container_spec["command"] = command_parts
+
+    elif deployment_cfg.type == "sglang":
+        # Convert SGLang command template to actual command
+        command_parts = [
+            "python3",
+            "-m",
+            "sglang.launch_server",
+            f"--model-path={deployment_cfg.checkpoint_path}",
+            "--host=0.0.0.0",
+            f"--port={deployment_cfg.port}",
+            f"--served-model-name={deployment_cfg.served_model_name}",
+            f"--tp={deployment_cfg.tensor_parallel_size}",
+            f"--dp={deployment_cfg.data_parallel_size}",
+        ]
+
+        # Add extra args if provided
+        if hasattr(deployment_cfg, "extra_args") and deployment_cfg.extra_args:
+            command_parts.extend(deployment_cfg.extra_args.split())
+
+        container_spec["command"] = command_parts
+
+    elif deployment_cfg.type == "nim":
+        # NIM containers use their default entrypoint - no custom command needed
+        # Configuration is handled via environment variables
+        pass
+
+    return container_spec
+
+
+def _add_deployment_derived_envs(envs_list: list, deployment_cfg: DictConfig) -> None:
+    """Add environment variables derived from deployment configuration.
+
+    Args:
+        envs_list: List to append environment variables to.
+        deployment_cfg: Deployment configuration to derive from.
+    """
+    deployment_type = deployment_cfg.type
+
+    # Common environment variables for all deployment types
+    if (
+        hasattr(deployment_cfg, "served_model_name")
+        and deployment_cfg.served_model_name
+    ):
+        envs_list.append(
+            {"name": "SERVED_MODEL_NAME", "value": deployment_cfg.served_model_name}
+        )
+
+    if hasattr(deployment_cfg, "port") and deployment_cfg.port:
+        envs_list.append({"name": "MODEL_PORT", "value": str(deployment_cfg.port)})
+
+    # Deployment-specific environment variables
+    if deployment_type == "vllm":
+        if (
+            hasattr(deployment_cfg, "checkpoint_path")
+            and deployment_cfg.checkpoint_path
+        ):
+            envs_list.append(
+                {"name": "MODEL_PATH", "value": deployment_cfg.checkpoint_path}
+            )
+        if (
+            hasattr(deployment_cfg, "tensor_parallel_size")
+            and deployment_cfg.tensor_parallel_size
+        ):
+            envs_list.append(
+                {
+                    "name": "TENSOR_PARALLEL_SIZE",
+                    "value": str(deployment_cfg.tensor_parallel_size),
+                }
+            )
+
+    elif deployment_type == "sglang":
+        if (
+            hasattr(deployment_cfg, "checkpoint_path")
+            and deployment_cfg.checkpoint_path
+        ):
+            envs_list.append(
+                {"name": "MODEL_PATH", "value": deployment_cfg.checkpoint_path}
+            )
+        if (
+            hasattr(deployment_cfg, "tensor_parallel_size")
+            and deployment_cfg.tensor_parallel_size
+        ):
+            envs_list.append(
+                {
+                    "name": "TENSOR_PARALLEL_SIZE",
+                    "value": str(deployment_cfg.tensor_parallel_size),
+                }
+            )
+
+    elif deployment_type == "nim":
+        # NIM-specific derived environment variables
+        if (
+            hasattr(deployment_cfg, "served_model_name")
+            and deployment_cfg.served_model_name
+        ):
+            envs_list.append(
+                {"name": "NIM_MODEL_NAME", "value": deployment_cfg.served_model_name}
+            )
+
+
+def _generate_model_cache_name(image: str) -> str:
+    """Generate a cache directory name from the container image.
+
+    Args:
+        image: Container image string like 'nvcr.io/nim/meta/llama-3.1-8b-instruct:1.8.6'
+
+    Returns:
+        Clean cache name like 'llama-3-1-8b-instruct'
+    """
+    # Extract model name from image path
+    if "/" in image:
+        model_part = image.split("/")[-1]  # Get 'llama-3.1-8b-instruct:1.8.6'
+    else:
+        model_part = image
+
+    # Remove version tag
+    if ":" in model_part:
+        model_part = model_part.split(":")[0]  # Get 'llama-3.1-8b-instruct'
+
+    # Replace dots with dashes for filesystem compatibility
+    return model_part.replace(".", "-")
+
+
+def create_lepton_endpoint(cfg: DictConfig, endpoint_name: str) -> bool:
+    """Create a Lepton endpoint using the lep CLI.
+
+    Args:
+        cfg: The nemo-evaluator-launcher configuration object.
+        endpoint_name: Name for the endpoint.
+
+    Returns:
+        True if endpoint creation succeeded, False otherwise.
+    """
+    spec = generate_lepton_spec(cfg)
+
+    # Convert OmegaConf objects to regular Python objects for JSON serialization
+    from omegaconf import DictConfig, ListConfig
+
+    def convert_to_json_serializable(obj: Any) -> Any:
+        """Recursively convert OmegaConf objects to regular Python objects."""
+        if isinstance(obj, (DictConfig, dict)):
+            return {k: convert_to_json_serializable(v) for k, v in obj.items()}
+        elif isinstance(obj, (ListConfig, list)):
+            return [convert_to_json_serializable(item) for item in obj]
+        else:
+            return obj
+
+    json_spec = convert_to_json_serializable(spec)
+
+    # Write spec to temporary file
+    import tempfile
+
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f:
+        json.dump(json_spec, f, indent=2)
+        spec_file = f.name
+
+    try:
+        # Create endpoint using lep CLI
+        result = subprocess.run(
+            ["lep", "endpoint", "create", "--file", spec_file, "--name", endpoint_name],
+            capture_output=True,
+            text=True,
+            timeout=300,
+        )
+
+        if result.returncode == 0:
+            print(f"✅ Successfully created Lepton endpoint: {endpoint_name}")
+            return True
+        else:
+            error_msg = result.stderr.strip() if result.stderr else ""
+            output_msg = result.stdout.strip() if result.stdout else ""
+            print(
+                f"✗ Failed to create Lepton endpoint | Endpoint: {endpoint_name} | Return code: {result.returncode}"
+            )
+            if error_msg:
+                print(f"   stderr: {error_msg}")
+            if output_msg:
+                print(f"   stdout: {output_msg}")
+            return False
+
+    except subprocess.TimeoutExpired as e:
+        print(
+            f"✗ Timeout creating Lepton endpoint | Endpoint: {endpoint_name} | Timeout: 300s"
+        )
+        if hasattr(e, "stderr") and e.stderr:
+            print(f"   stderr: {e.stderr}")
+        if hasattr(e, "stdout") and e.stdout:
+            print(f"   stdout: {e.stdout}")
+        return False
+    except subprocess.CalledProcessError as e:
+        print(
+            f"✗ Error creating Lepton endpoint | Endpoint: {endpoint_name} | Error: {e}"
+        )
+        if hasattr(e, "stderr") and e.stderr:
+            print(f"   stderr: {e.stderr}")
+        if hasattr(e, "stdout") and e.stdout:
+            print(f"   stdout: {e.stdout}")
+        return False
+    finally:
+        # Clean up temporary file
+        Path(spec_file).unlink(missing_ok=True)
+
+
+def delete_lepton_endpoint(endpoint_name: str) -> bool:
+    """Delete a Lepton endpoint.
+
+    Args:
+        endpoint_name: Name of the endpoint to delete.
+
+    Returns:
+        True if deletion succeeded, False otherwise.
+    """
+    try:
+        result = subprocess.run(
+            ["lep", "endpoint", "remove", "--name", endpoint_name],
+            capture_output=True,
+            text=True,
+            timeout=60,
+        )
+
+        return result.returncode == 0
+    except (subprocess.TimeoutExpired, subprocess.CalledProcessError):
+        return False
+
+
+def get_lepton_endpoint_status(endpoint_name: str) -> Optional[dict[str, Any]]:
+    """Get the status of a Lepton endpoint.
+
+    Args:
+        endpoint_name: Name of the endpoint.
+
+    Returns:
+        Status dict if endpoint exists, None otherwise. See
+        https://github.com/leptonai/leptonai/blob/7de93b95357126da1e86fa99f54f9a769d5d2646/leptonai/api/v1/types/deployment.py#L338
+        for the definition.
+    """
+    try:
+        # TODO(agronskiy): why not use Python API?
+        cmd = ["lep", "endpoint", "get", "--name", endpoint_name]
+        result = subprocess.run(
+            cmd,
+            capture_output=True,
+            text=True,
+            timeout=30,
+        )
+
+        if result.returncode != 0:
+            return None
+
+        endpoint_info = json.loads(result.stdout)
+        status = endpoint_info.get("status", {})
+        if isinstance(status, dict):
+            return status
+        logger.error(
+            "Result of running lep command returne non-dict status",
+            cmd=cmd,
+            status=status,
+        )
+        return None
+
+    except (
+        subprocess.TimeoutExpired,
+        subprocess.CalledProcessError,
+        json.JSONDecodeError,
+    ):
+        return None
+
+
+def wait_for_lepton_endpoint_ready(endpoint_name: str, timeout: int = 600) -> bool:
+    """Wait for a Lepton endpoint to become ready.
+
+    Args:
+        endpoint_name: Name of the endpoint.
+        timeout: Maximum time to wait in seconds.
+
+    Returns:
+        True if endpoint becomes ready, False if timeout.
+    """
+    start_time = time.time()
+    while time.time() - start_time < timeout:
+        status = get_lepton_endpoint_status(endpoint_name)
+
+        # `get_lepton_endpoint_status` might return `None` if
+        # e.g. there was a network error, see definition.
+        if status is not None:
+            state = status.get("state", "").lower()
+            if state == "ready":
+                logger.info(
+                    "Lepton endpoint is ready",
+                    endpoint_name=endpoint_name,
+                )
+
+                return True
+            elif state in ["failed", "error"]:
+                return False
+
+        logger.debug(
+            "Waiting for lepton endpoint",
+            endpoint_name=endpoint_name,
+            timeout=timeout,
+            time_delta=time.time() - start_time,
+            curr_status=status,
+        )
+        time.sleep(10)
+
+    logger.error(
+        "Timeout waiting for lepton endpoint",
+        endpoint_name=endpoint_name,
+        timeout=timeout,
+    )
+    return False
+
+
+def get_lepton_endpoint_url(endpoint_name: str) -> Optional[str]:
+    """Get the URL of a Lepton endpoint.
+
+    Args:
+        endpoint_name: Name of the endpoint.
+
+    Returns:
+        Endpoint URL if available, None otherwise.
+    """
+    try:
+        result = subprocess.run(
+            ["lep", "endpoint", "get", "--name", endpoint_name],
+            capture_output=True,
+            text=True,
+            timeout=30,
+        )
+
+        if result.returncode == 0:
+            endpoint_info = json.loads(result.stdout)
+            status = endpoint_info.get("status", {})
+            endpoint = status.get("endpoint", {})
+            external_endpoint = endpoint.get("external_endpoint")
+            # Ensure we return a proper string type or None
+            if isinstance(external_endpoint, str):
+                return external_endpoint
+            else:
+                return None
+        else:
+            return None
+    except (
+        subprocess.TimeoutExpired,
+        subprocess.CalledProcessError,
+        json.JSONDecodeError,
+    ):
+        return None