Consoleモジュール: WebApiHandlerの実装

[Geson-anko]

タスク内容

AMI Systemの WebApiHandlerクラスをほぼそのまま移植します。
https://github.com/MLShukai/ami/blob/main/ami/threads/web_api_handler.py

スケッチ

SystemStatus に関しては、どのように実装するか検討中です。

"""Web API interface for system control.

This module provides a simple web API for controlling the system, allowing
external applications to pause, resume, shutdown the system and save checkpoints.
"""

import json
import logging
import threading
from enum import Enum, auto
from queue import Empty, Queue
from typing import Any

import bottle
from bottle import Bottle, request, response

from pamiq_core.threads import ThreadTypes

# Constants for commonly used responses
_RESULT_OK = {"result": "ok"}
_ERROR_INVALID_ENDPOINT = {"error": "Invalid API endpoint"}
_ERROR_INVALID_METHOD = {"error": "Invalid API method"}
_ERROR_INTERNAL_SERVER = {"error": "Internal server error"}

# Type aliases
PayloadType = dict[str, str]


class ControlCommands(Enum):
    """Enumerates the commands for the system control."""

    SHUTDOWN = auto()
    PAUSE = auto()
    RESUME = auto()
    SAVE_CHECKPOINT = auto()

class WebApiHandler:
    """Web API handler for controlling the system.

    This class provides a simple Web API for controlling the thread controller,
    allowing external applications to pause, resume, and shutdown the system.
    """

    def __init__(
        self,
        system_status: SystemStatus,
        host: str = "localhost",
        port: int = 8391,
    ) -> None:
        """Initialize the WebApiHandler.

        Args:
            controller_status: To watch the thread controller status.
            host: Hostname to run the API server on.
            port: Port to run the API server on.
        """
        self._logger = logging.getLogger(f"{__name__}.{self.__class__.__name__}")
        self._ system_status = system_status
        self._host = host
        self._port = port
        
        # Create a dedicated Bottle app instead of using the default app
        self._app = Bottle()
        self._register_handlers()
        
        self._handler_thread = threading.Thread(
            target=self.run, 
            daemon=True,
            name=f"{ThreadTypes.CONTROL.thread_name}_webapi"
        )

        self._received_commands_queue: Queue[ControlCommands] = Queue()

    def run(self) -> None:
        """Run the API server."""
        failed = True
        while failed:
            try:
                self._logger.info(f"Serving system command at '{self._host}:{self._port}'")
                self._app.run(host=self._host, port=self._port, quiet=True)
                failed = False
            except OSError:
                self._logger.info(f"Address '{self._host}:{self._port}' is already used, increment port number...")
                self._port += 1
            except Exception:
                self._logger.exception("Error starting API server")
                raise

    def run_in_background(self) -> None:
        """Run the API server in a background thread."""
        self._handler_thread.start()

    def has_commands(self) -> bool:
        """Check if there are commands in the queue.
        
        Returns:
            True if there are commands, False otherwise.
        """
        return not self._received_commands_queue.empty()

    def receive_command(self) -> ControlCommands:
        """Receive a command from the queue without blocking.
        
        Returns:
            The command from the queue.
            
        Raises:
            Empty: If there are no commands in the queue.
        """
        return self._received_commands_queue.get_nowait()

    def _register_handlers(self) -> None:
        """Register API handlers with bottle."""
        # Use the app's route decorator instead of the global one
        self._app.route("/api/status", method="GET", callback=self._get_status)
        self._app.route("/api/pause", method="POST", callback=self._post_pause)
        self._app.route("/api/resume", method="POST", callback=self._post_resume)
        self._app.route("/api/shutdown", method="POST", callback=self._post_shutdown)
        self._app.route("/api/save-checkpoint", method="POST", callback=self._post_save_checkpoint)

        # Register error handlers
        self._app.error(404)(self._error_404)
        self._app.error(405)(self._error_405)
        self._app.error(500)(self._error_500)

    def _json_response(self, data: dict[str, Any]) -> str:
        """Helper method to return JSON responses.
        
        Args:
            data: The data to convert to JSON.
            
        Returns:
            JSON string representation of the data.
        """
        response.content_type = "application/json"
        return json.dumps(data)

    def _get_status(self) -> str:
        """Handle GET /api/status request.
        
        Returns:
            JSON payload with the current system status.
        """
        status: str
        if self._system_status.is_shutdown():
            status = "stopped"
        elif self._system_status.is_paused():
            status = "paused"
        else:
            status = "active"

        self._logger.info(f"Status request: returning {status}")
        return self._json_response({"status": status})

    def _post_pause(self) -> str:
        """Handle POST /api/pause request.
        
        Returns:
            JSON payload with the result.
        """
        self._logger.info("Pause command received")
        self._received_commands_queue.put(ControlCommands.PAUSE)
        return self._json_response(_RESULT_OK)

    def _post_resume(self) -> str:
        """Handle POST /api/resume request.
        
        Returns:
            JSON payload with the result.
        """
        self._logger.info("Resume command received")
        self._received_commands_queue.put(ControlCommands.RESUME)
        return self._json_response(_RESULT_OK)

    def _post_shutdown(self) -> str:
        """Handle POST /api/shutdown request.
        
        Returns:
            JSON payload with the result.
        """
        self._logger.info("Shutdown command received")
        self._received_commands_queue.put(ControlCommands.SHUTDOWN)
        return self._json_response(_RESULT_OK)

    def _post_save_checkpoint(self) -> str:
        """Handle POST /api/save-checkpoint request.
        
        Returns:
            JSON payload with the result.
        """
        self._logger.info("Save checkpoint command received")
        self._received_commands_queue.put(ControlCommands.SAVE_CHECKPOINT)
        return self._json_response(_RESULT_OK)

    def _error_404(self, error: bottle.HTTPError) -> str:
        """Handle 404 errors.
        
        Args:
            error: The HTTP error.
            
        Returns:
            JSON payload with the error message.
        """
        self._logger.error(f"404: {request.method} {request.path} is invalid API endpoint")
        return self._json_response(_ERROR_INVALID_ENDPOINT)

    def _error_405(self, error: bottle.HTTPError) -> str:
        """Handle 405 errors.
        
        Args:
            error: The HTTP error.
            
        Returns:
            JSON payload with the error message.
        """
        self._logger.error(f"405: {request.method} {request.path} is invalid API method")
        return self._json_response(_ERROR_INVALID_METHOD)

    def _error_500(self, error: bottle.HTTPError) -> str:
        """Handle 500 errors.
        
        Args:
            error: The HTTP error.
            
        Returns:
            JSON payload with the error message.
        """
        self._logger.error(f"500: {request.method} {request.path} caused an error")
        return self._json_response(_ERROR_INTERNAL_SERVER)

達成条件

• WebApiHandlerが実装された
• テスト