miku-discord/bot/utils/container_manager.py

# container_manager.py
"""
Manages Docker containers for STT and TTS services.
Handles startup, shutdown, and warmup detection.
"""

import asyncio
import subprocess
import aiohttp
from utils.logger import get_logger

logger = get_logger('container_manager')

class ContainerManager:
    """Manages STT and TTS Docker containers."""
    
    # Container names from docker-compose.yml
    STT_CONTAINER = "miku-stt"
    TTS_CONTAINER = "miku-rvc-api"
    
    # Warmup check endpoints
    STT_HEALTH_URL = "http://miku-stt:8767/health"  # HTTP health check endpoint
    TTS_HEALTH_URL = "http://miku-rvc-api:8765/health"
    
    # Warmup timeouts
    STT_WARMUP_TIMEOUT = 30  # seconds
    TTS_WARMUP_TIMEOUT = 60  # seconds (RVC takes longer)
    
    @classmethod
    async def start_voice_containers(cls) -> bool:
        """
        Start STT and TTS containers and wait for them to warm up.
        
        Returns:
            bool: True if both containers started and warmed up successfully
        """
        logger.info("🚀 Starting voice chat containers...")
        
        try:
            # Start STT container using docker start (assumes container exists)
            logger.info(f"Starting {cls.STT_CONTAINER}...")
            result = subprocess.run(
                ["docker", "start", cls.STT_CONTAINER],
                capture_output=True,
                text=True
            )
            
            if result.returncode != 0:
                logger.error(f"Failed to start {cls.STT_CONTAINER}: {result.stderr}")
                return False
            
            logger.info(f"✓ {cls.STT_CONTAINER} started")
            
            # Start TTS container
            logger.info(f"Starting {cls.TTS_CONTAINER}...")
            result = subprocess.run(
                ["docker", "start", cls.TTS_CONTAINER],
                capture_output=True,
                text=True
            )
            
            if result.returncode != 0:
                logger.error(f"Failed to start {cls.TTS_CONTAINER}: {result.stderr}")
                return False
            
            logger.info(f"✓ {cls.TTS_CONTAINER} started")
            
            # Wait for warmup
            logger.info("⏳ Waiting for containers to warm up...")
            
            stt_ready = await cls._wait_for_stt_warmup()
            if not stt_ready:
                logger.error("STT failed to warm up")
                return False
            
            tts_ready = await cls._wait_for_tts_warmup()
            if not tts_ready:
                logger.error("TTS failed to warm up")
                return False
            
            logger.info("✅ All voice containers ready!")
            return True
            
        except Exception as e:
            logger.error(f"Error starting voice containers: {e}")
            return False
    
    @classmethod
    async def stop_voice_containers(cls) -> bool:
        """
        Stop STT and TTS containers.
        
        Returns:
            bool: True if containers stopped successfully
        """
        logger.info("🛑 Stopping voice chat containers...")
        
        try:
            # Stop both containers
            result = subprocess.run(
                ["docker", "stop", cls.STT_CONTAINER, cls.TTS_CONTAINER],
                capture_output=True,
                text=True
            )
            
            if result.returncode != 0:
                logger.error(f"Failed to stop containers: {result.stderr}")
                return False
            
            logger.info("✓ Voice containers stopped")
            return True
            
        except Exception as e:
            logger.error(f"Error stopping voice containers: {e}")
            return False
    
    @classmethod
    async def _wait_for_stt_warmup(cls) -> bool:
        """
        Wait for STT container to be ready by checking health endpoint.
        
        Returns:
            bool: True if STT is ready within timeout
        """
        start_time = asyncio.get_event_loop().time()
        
        async with aiohttp.ClientSession() as session:
            while (asyncio.get_event_loop().time() - start_time) < cls.STT_WARMUP_TIMEOUT:
                try:
                    async with session.get(cls.STT_HEALTH_URL, timeout=aiohttp.ClientTimeout(total=2)) as resp:
                        if resp.status == 200:
                            data = await resp.json()
                            if data.get("status") == "ready" and data.get("warmed_up"):
                                logger.info("✓ STT is ready")
                                return True
                except Exception:
                    # Not ready yet, wait and retry
                    pass
                
                await asyncio.sleep(2)
        
        logger.error(f"STT warmup timeout ({cls.STT_WARMUP_TIMEOUT}s)")
        return False
    
    @classmethod
    async def _wait_for_tts_warmup(cls) -> bool:
        """
        Wait for TTS container to be ready by checking health endpoint.
        
        Returns:
            bool: True if TTS is ready within timeout
        """
        start_time = asyncio.get_event_loop().time()
        
        async with aiohttp.ClientSession() as session:
            while (asyncio.get_event_loop().time() - start_time) < cls.TTS_WARMUP_TIMEOUT:
                try:
                    async with session.get(cls.TTS_HEALTH_URL, timeout=aiohttp.ClientTimeout(total=2)) as resp:
                        if resp.status == 200:
                            data = await resp.json()
                            # RVC API returns "status": "healthy", not "ready"
                            status_ok = data.get("status") in ["ready", "healthy"]
                            if status_ok and data.get("warmed_up"):
                                logger.info("✓ TTS is ready")
                                return True
                except Exception:
                    # Not ready yet, wait and retry
                    pass
                
                await asyncio.sleep(2)
        
        logger.error(f"TTS warmup timeout ({cls.TTS_WARMUP_TIMEOUT}s)")
        return False
        return False
    
    @classmethod
    async def are_containers_running(cls) -> tuple[bool, bool]:
        """
        Check if STT and TTS containers are currently running.
        
        Returns:
            tuple[bool, bool]: (stt_running, tts_running)
        """
        try:
            # Check STT
            result = subprocess.run(
                ["docker", "inspect", "-f", "{{.State.Running}}", cls.STT_CONTAINER],
                capture_output=True,
                text=True
            )
            stt_running = result.returncode == 0 and result.stdout.strip() == "true"
            
            # Check TTS
            result = subprocess.run(
                ["docker", "inspect", "-f", "{{.State.Running}}", cls.TTS_CONTAINER],
                capture_output=True,
                text=True
            )
            tts_running = result.returncode == 0 and result.stdout.strip() == "true"
            
            return (stt_running, tts_running)
            
        except Exception as e:
            logger.error(f"Error checking container status: {e}")
            return (False, False)