handler_generator.py

"""Generator for handler_<name>.py files."""

import ast
import logging
from pathlib import Path
from typing import Any, Dict, List, Union

from runpod_flash.runtime.models import Manifest

logger = logging.getLogger(__name__)

DEPLOYED_HANDLER_TEMPLATE = '''"""
Auto-generated deployed handler for resource: {resource_name}
Generated at: {timestamp}

Deployed endpoint handler: accepts plain JSON, no cloudpickle.
One function per endpoint, identified by FLASH_RESOURCE_NAME.

This file is generated by the Flash build process. Do not edit manually.
"""

import asyncio
import importlib
import inspect
import logging
import traceback

_logger = logging.getLogger(__name__)

# Import the function for this endpoint
{import_statement}


def handler(job):
    """Handler for deployed QB endpoint. Accepts plain JSON kwargs."""
    job_input = job.get("input", {{}})
    try:
        result = {function_name}(**job_input)
        if inspect.iscoroutine(result):
            try:
                loop = asyncio.get_running_loop()
            except RuntimeError:
                loop = None
            if loop and loop.is_running():
                import concurrent.futures

                with concurrent.futures.ThreadPoolExecutor() as pool:
                    result = pool.submit(asyncio.run, result).result()
            else:
                result = asyncio.run(result)
        return result
    except Exception as e:
        _logger.error(
            "Deployed handler error for {function_name}: %s",
            e,
            exc_info=True,
        )
        return {{"error": str(e), "traceback": traceback.format_exc()}}


if __name__ == "__main__":
    import runpod
    runpod.serverless.start({{"handler": handler}})
'''

DEPLOYED_ASYNC_HANDLER_TEMPLATE = '''"""
Auto-generated deployed handler for resource: {resource_name}
Generated at: {timestamp}

Concurrent async endpoint handler: accepts plain JSON, no cloudpickle.
max_concurrency={max_concurrency}

This file is generated by the Flash build process. Do not edit manually.
"""

import importlib
import logging
import traceback

_logger = logging.getLogger(__name__)

# Import the function for this endpoint
{import_statement}


async def handler(job):
    """Async handler for concurrent QB endpoint. Accepts plain JSON kwargs."""
    job_input = job.get("input", {{}})
    try:
        result = await {function_name}(**job_input)
        return result
    except Exception as e:
        _logger.error(
            "Deployed handler error for {function_name}: %s",
            e,
            exc_info=True,
        )
        return {{"error": str(e), "traceback": traceback.format_exc()}}


if __name__ == "__main__":
    import runpod
    runpod.serverless.start({{
        "handler": handler,
        "concurrency_modifier": lambda current: {max_concurrency},
    }})
'''

DEPLOYED_CLASS_HANDLER_TEMPLATE = '''"""
Auto-generated deployed handler for class-based resource: {resource_name}
Generated at: {timestamp}

Deployed endpoint handler for a @remote class. Instantiates the class
once at cold start, then dispatches to methods per request.

This file is generated by the Flash build process. Do not edit manually.
"""

import asyncio
import importlib
import inspect
import logging
import traceback

_logger = logging.getLogger(__name__)

# import the class for this endpoint
{import_statement}

# instantiate once at cold start
_instance = {class_name}()

# public methods available for dispatch
_METHODS = {methods_dict}


def _run_maybe_async(result):
    """resolve a possibly-async result to a plain value."""
    if inspect.iscoroutine(result):
        try:
            loop = asyncio.get_running_loop()
        except RuntimeError:
            loop = None
        if loop and loop.is_running():
            import concurrent.futures
            with concurrent.futures.ThreadPoolExecutor() as pool:
                return pool.submit(asyncio.run, result).result()
        return asyncio.run(result)
    return result


def handler(job):
    """Handler for deployed class-based QB endpoint.

    Dispatches to a method on the singleton class instance.
    If the class has exactly one public method, input is passed
    directly as kwargs. If multiple methods exist, input must
    include a "method" key to select the target.
    """
    job_input = job.get("input", {{}})
    try:
        if len(_METHODS) == 1:
            method_name = next(iter(_METHODS))
        else:
            method_name = job_input.pop("method", None)
            if method_name is None:
                return {{
                    "error": (
                        "class {class_name} has multiple methods: "
                        + ", ".join(sorted(_METHODS))
                        + ". include a \\"method\\" key in the input."
                    )
                }}
            if method_name not in _METHODS:
                return {{
                    "error": (
                        f"unknown method '{{method_name}}' on {class_name}. "
                        f"available: {{', '.join(sorted(_METHODS))}}"
                    )
                }}

        method = getattr(_instance, method_name)
        result = _run_maybe_async(method(**job_input))
        return result
    except Exception as e:
        _logger.error(
            "Deployed handler error for {class_name}.%s: %s",
            job_input.get("method", "<default>"),
            e,
            exc_info=True,
        )
        return {{"error": str(e), "traceback": traceback.format_exc()}}


if __name__ == "__main__":
    import runpod
    runpod.serverless.start({{"handler": handler}})
'''


DEPLOYED_ASYNC_CLASS_HANDLER_TEMPLATE = '''"""
Auto-generated deployed handler for class-based resource: {resource_name}
Generated at: {timestamp}

Concurrent async class endpoint handler: accepts plain JSON, no cloudpickle.
max_concurrency={max_concurrency}

This file is generated by the Flash build process. Do not edit manually.
"""

import importlib
import logging
import traceback

_logger = logging.getLogger(__name__)

# import the class for this endpoint
{import_statement}

# instantiate once at cold start
_instance = {class_name}()

# public methods available for dispatch
_METHODS = {methods_dict}


async def handler(job):
    """Async handler for concurrent class-based QB endpoint.

    Dispatches to a method on the singleton class instance.
    If the class has exactly one public method, input is passed
    directly as kwargs. If multiple methods exist, input must
    include a "method" key to select the target.
    """
    job_input = job.get("input", {{}})
    try:
        if len(_METHODS) == 1:
            method_name = next(iter(_METHODS))
        else:
            method_name = job_input.pop("method", None)
            if method_name is None:
                return {{
                    "error": (
                        "class {class_name} has multiple methods: "
                        + ", ".join(sorted(_METHODS))
                        + ". include a \\"method\\" key in the input."
                    )
                }}
            if method_name not in _METHODS:
                return {{
                    "error": (
                        f"unknown method '{{method_name}}' on {class_name}. "
                        f"available: {{', '.join(sorted(_METHODS))}}"
                    )
                }}

        method = getattr(_instance, method_name)
        result = await method(**job_input)
        return result
    except Exception as e:
        _logger.error(
            "Deployed handler error for {class_name}.%s: %s",
            job_input.get("method", "<default>"),
            e,
            exc_info=True,
        )
        return {{"error": str(e), "traceback": traceback.format_exc()}}


if __name__ == "__main__":
    import runpod
    runpod.serverless.start({{
        "handler": handler,
        "concurrency_modifier": lambda current: {max_concurrency},
    }})
'''


class HandlerGenerator:
    """Generates handler_<name>.py files for each resource config."""

    def __init__(self, manifest: Union[Dict[str, Any], Manifest], build_dir: Path):
        self.manifest = manifest
        self.build_dir = build_dir

    def generate_handlers(self) -> List[Path]:
        """Generate all handler files for queue-based (non-LB) resources."""
        handler_paths = []

        # Handle both dict and Manifest types
        resources = (
            self.manifest.resources
            if isinstance(self.manifest, Manifest)
            else self.manifest.get("resources", {})
        )

        for resource_name, resource_data in resources.items():
            # Skip load-balanced resources (handled by LBHandlerGenerator)
            # Use flag determined by isinstance() at scan time
            is_load_balanced = (
                resource_data.is_load_balanced
                if hasattr(resource_data, "is_load_balanced")
                else resource_data.get("is_load_balanced", False)
            )
            if is_load_balanced:
                continue

            handler_path = self._generate_handler(resource_name, resource_data)
            handler_paths.append(handler_path)

        return handler_paths

    def _generate_handler(self, resource_name: str, resource_data: Any) -> Path:
        """Generate a handler file for a deployed QB endpoint."""
        handler_filename = f"handler_{resource_name}.py"
        handler_path = self.build_dir / handler_filename

        timestamp = (
            self.manifest.generated_at
            if isinstance(self.manifest, Manifest)
            else self.manifest.get("generated_at", "")
        )

        functions = (
            resource_data.functions
            if hasattr(resource_data, "functions")
            else resource_data.get("functions", [])
        )

        max_concurrency = (
            resource_data.max_concurrency
            if hasattr(resource_data, "max_concurrency")
            else resource_data.get("max_concurrency", 1)
        )

        handler_code = self._generate_deployed_handler_code(
            resource_name, timestamp, functions, max_concurrency
        )

        handler_path.write_text(handler_code)
        self._validate_handler_imports(handler_path)

        return handler_path

    def _generate_deployed_handler_code(
        self,
        resource_name: str,
        timestamp: str,
        functions: List[Any],
        max_concurrency: int = 1,
    ) -> str:
        """Generate deployed handler code for a QB endpoint.

        Selects template based on max_concurrency and is_async:
        - max_concurrency=1: current sync templates (unchanged behavior)
        - max_concurrency>1 + async: new async templates with concurrency_modifier
        - max_concurrency>1 + sync: sync templates with concurrency_modifier injected
        """
        if not functions:
            raise ValueError(
                f"Resource '{resource_name}' has no functions. "
                f"Cannot generate a deployed handler without at least one function."
            )

        func = functions[0]
        module = func.module if hasattr(func, "module") else func.get("module")
        name = func.name if hasattr(func, "name") else func.get("name")
        is_class = (
            func.is_class if hasattr(func, "is_class") else func.get("is_class", False)
        )
        is_async = (
            func.is_async if hasattr(func, "is_async") else func.get("is_async", False)
        )

        import_statement = (
            f"{name} = importlib.import_module('{module}').{name}"
            if module and name
            else "# No function to import"
        )

        if max_concurrency > 100:
            logger.warning(
                "High max_concurrency=%d for resource '%s'. Ensure your handler "
                "and GPU can support this level of concurrent execution.",
                max_concurrency,
                resource_name,
            )

        if max_concurrency > 1 and not is_async:
            logger.warning(
                "max_concurrency=%d set on sync handler '%s'. "
                "Only async handlers benefit from concurrent execution. "
                "Consider making the handler async.",
                max_concurrency,
                resource_name,
            )

        if is_class:
            class_methods = (
                func.class_methods
                if hasattr(func, "class_methods")
                else func.get("class_methods", [])
            )
            methods_dict = {m: m for m in class_methods} if class_methods else {}

            if max_concurrency > 1 and is_async:
                return DEPLOYED_ASYNC_CLASS_HANDLER_TEMPLATE.format(
                    resource_name=resource_name,
                    timestamp=timestamp,
                    import_statement=import_statement,
                    class_name=name or "None",
                    methods_dict=repr(methods_dict),
                    max_concurrency=max_concurrency,
                )

            code = DEPLOYED_CLASS_HANDLER_TEMPLATE.format(
                resource_name=resource_name,
                timestamp=timestamp,
                import_statement=import_statement,
                class_name=name or "None",
                methods_dict=repr(methods_dict),
            )
            if max_concurrency > 1:
                code = self._inject_concurrency_modifier(code, max_concurrency)
            return code

        # Function-based handler
        if max_concurrency > 1 and is_async:
            return DEPLOYED_ASYNC_HANDLER_TEMPLATE.format(
                resource_name=resource_name,
                timestamp=timestamp,
                import_statement=import_statement,
                function_name=name or "None",
                max_concurrency=max_concurrency,
            )

        code = DEPLOYED_HANDLER_TEMPLATE.format(
            resource_name=resource_name,
            timestamp=timestamp,
            import_statement=import_statement,
            function_name=name or "None",
        )
        if max_concurrency > 1:
            code = self._inject_concurrency_modifier(code, max_concurrency)
        return code

    @staticmethod
    def _inject_concurrency_modifier(code: str, max_concurrency: int) -> str:
        """Replace the default runpod.serverless.start call with one including concurrency_modifier."""
        return code.replace(
            'runpod.serverless.start({"handler": handler})',
            "runpod.serverless.start({\n"
            '        "handler": handler,\n'
            f'        "concurrency_modifier": lambda current: {max_concurrency},\n'
            "    })",
        )

    def _validate_handler_imports(self, handler_path: Path) -> None:
        """Validate that generated handler has valid Python syntax.

        Uses ast.parse() to check syntax without executing the module.
        This avoids ImportErrors from modules that only resolve at runtime
        inside Docker (e.g., numeric-prefixed module paths).

        Args:
            handler_path: Path to generated handler file

        Raises:
            ValueError: If handler has syntax errors
        """
        try:
            source = handler_path.read_text(encoding="utf-8")
            ast.parse(source, filename=str(handler_path))
        except SyntaxError as e:
            raise ValueError(f"Handler has syntax errors: {e}") from e