docker_manager.py

#!/usr/bin/env python3
"""
Docker 설치 감지, 이미지 빌드, 컨테이너 관리, 로그 스트리밍을 담당.
"""

import os
import subprocess
import threading
import time
import platform
import re
from pathlib import Path
from typing import Optional, Callable, Dict, Any, List

# ==============================================================================
# CONSTANTS
# ==============================================================================

IMAGE_NAME = "linux_node"
IMAGE_TAG = "latest"
FULL_IMAGE_NAME = f"{IMAGE_NAME}:{IMAGE_TAG}"

# Docker Desktop 다운로드 URL
DOCKER_DOWNLOAD_URLS = {
    'Darwin': 'https://www.docker.com/products/docker-desktop/',
    'Windows': 'https://www.docker.com/products/docker-desktop/',
    'Linux': 'https://docs.docker.com/engine/install/',
}


# ==============================================================================
# UTILITY FUNCTIONS
# ==============================================================================

def get_platform() -> str:
    """현재 OS 플랫폼 반환"""
    return platform.system()


def should_use_docker() -> bool:
    """
    Docker 모드를 사용해야 하는지 판단.
    - Windows/Mac: Docker 필수
    - Linux: 네이티브 실행 가능
    """
    if os.environ.get('FORCE_DOCKER', '').strip() == '1': # FORCE_DOCKER==1이면 linux여도 Docker 환경 사용
        return True
    return get_platform() in ('Darwin', 'Windows')


def get_docker_platform() -> str:
    """호스트 CPU 아키텍처에 맞는 Docker 플랫폼 문자열 반환.

    에뮬레이션 없이 네이티브로 실행하기 위해 호스트 아키텍처를 Docker에 전달한다.
    seccomp-BPF는 커널 수준이므로 아키텍처 에뮬레이션 환경에서는 동작하지 않는다.
    """
    machine = platform.machine()
    if machine in ('arm64', 'aarch64'):
        return 'linux/arm64'
    return 'linux/amd64'


def get_project_root() -> Path:
    """프로젝트 루트 디렉토리 반환"""
    return Path(__file__).parent.absolute()


# ==============================================================================
# DOCKER MANAGER CLASS
# ==============================================================================

class DockerManager:
    def __init__(self):
        self.project_root = get_project_root()
        self.running_containers: Dict[str, subprocess.Popen] = {}
        self._log_threads: Dict[str, threading.Thread] = {}
        self._stop_events: Dict[str, threading.Event] = {}

    # --------------------------------------------------------------------------
    # Docker 설치 및 상태 확인
    # --------------------------------------------------------------------------

    def is_docker_installed(self) -> bool:
        try:
            result = subprocess.run(
                ['docker', '--version'],
                capture_output=True,
                text=True,
                timeout=10
            )
            return result.returncode == 0
        except (subprocess.TimeoutExpired, FileNotFoundError):
            return False

    def is_docker_running(self) -> bool:
        try:
            result = subprocess.run(
                ['docker', 'info'],
                capture_output=True,
                text=True,
                timeout=15
            )
            return result.returncode == 0
        except (subprocess.TimeoutExpired, FileNotFoundError):
            return False

    def get_docker_version(self) -> Optional[str]:
        try:
            result = subprocess.run(
                ['docker', '--version'],
                capture_output=True,
                text=True,
                timeout=10
            )
            if result.returncode == 0:
                return result.stdout.strip()
            return None
        except (subprocess.TimeoutExpired, FileNotFoundError):
            return None

    def get_docker_download_url(self) -> str:
        return DOCKER_DOWNLOAD_URLS.get(get_platform(), DOCKER_DOWNLOAD_URLS['Linux'])

    # --------------------------------------------------------------------------
    # 이미지 관리
    # --------------------------------------------------------------------------

    def is_image_available(self, image_name: str = FULL_IMAGE_NAME) -> bool:
        try:
            result = subprocess.run(
                ['docker', 'image', 'inspect', image_name],
                capture_output=True,
                text=True,
                timeout=10
            )
            return result.returncode == 0
        except (subprocess.TimeoutExpired, FileNotFoundError):
            return False

    def build_image(
        self,
        progress_callback: Optional[Callable[[str, int], None]] = None,
        force_rebuild: bool = False # 강제 재빌드 여부
    ) -> bool:
        if not force_rebuild and self.is_image_available():
            if progress_callback:
                progress_callback("이미지가 이미 준비되어 있습니다.", 100)
            return True

        if progress_callback:
            progress_callback("Docker 이미지를 빌드하고 있습니다...", 0)

        try:
            # 빌드 프로세스 실행
            process = subprocess.Popen(
                [
                    'docker', 'build',
                    '--platform', get_docker_platform(),
                    '-t', FULL_IMAGE_NAME,
                    '--progress=plain',
                    str(self.project_root)
                ],
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT,
                text=True,
                bufsize=1
            )

            # 빌드 단계 추적을 위한 패턴
            step_pattern = re.compile(r'Step (\d+)/(\d+)')
            progress = 0

            for line in iter(process.stdout.readline, ''):
                if not line:
                    break

                line = line.strip()

                # 진행률 추정
                match = step_pattern.search(line)
                if match:
                    current, total = int(match.group(1)), int(match.group(2))
                    progress = int((current / total) * 90)  # 90%까지 빌드

                if progress_callback:
                    # 간략한 메시지로 변환
                    if 'FROM' in line:
                        progress_callback("베이스 이미지 다운로드 중...", progress)
                    elif 'RUN pip install' in line:
                        progress_callback("Python 패키지 설치 중...", progress)
                    elif 'COPY' in line:
                        progress_callback("프로젝트 파일 복사 중...", progress)

            process.wait()

            if process.returncode == 0:
                if progress_callback:
                    progress_callback("이미지 빌드 완료!", 100)
                return True
            else:
                if progress_callback:
                    progress_callback("이미지 빌드 실패", 0)
                return False

        except Exception as e:
            if progress_callback:
                progress_callback(f"빌드 오류: {str(e)}", 0)
            return False

    # --------------------------------------------------------------------------
    # 컨테이너 관리
    # --------------------------------------------------------------------------

    def run_node(
        self,
        config: Dict[str, Any],
        log_callback: Optional[Callable[[str, str], None]] = None,
        container_name: Optional[str] = None
    ) -> Optional[str]:
        """
        Node를 Docker 컨테이너로 실행.
        Orchestrator/Helper는 호스트에서 네이티브로 실행되므로 Docker 대상 아님.

        Args:
            config: Node 설정 딕셔너리
            log_callback: 로그 콜백 (level, message)
            container_name: 컨테이너 이름 (None이면 자동 생성)

        Returns:
            컨테이너 이름 (실패 시 None)
        """
        if container_name is None:
            container_name = f"co-op-{int(time.time())}"

        # 기존에 같은 이름의 컨테이너가 있으면 제거
        self._remove_container_if_exists(container_name)

        # Docker 명령어 구성
        cmd = self._build_node_command(config, container_name)

        try:
            # 컨테이너 실행
            process = subprocess.Popen(
                cmd,
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT,
                text=True,
                bufsize=1
            )

            self.running_containers[container_name] = process

            # 로그 스트리밍 스레드 시작
            if log_callback:
                stop_event = threading.Event()
                self._stop_events[container_name] = stop_event

                log_thread = threading.Thread(
                    target=self._stream_logs,
                    args=(process, log_callback, stop_event),
                    daemon=True
                )
                log_thread.start()
                self._log_threads[container_name] = log_thread

            return container_name

        except Exception as e:
            if log_callback:
                log_callback('ERROR', f"컨테이너 실행 실패: {str(e)}")
            return None

    def _get_host_address(self) -> str:
        """
        컨테이너에서 호스트 머신에 접근하기 위한 주소 반환.
        - Mac/Windows: host.docker.internal (Docker Desktop 제공)
        - Linux: 172.17.0.1 (docker0 브릿지 기본 게이트웨이)
        """
        plat = get_platform()
        if plat in ('Darwin', 'Windows'):
            return 'host.docker.internal'
        return '172.17.0.1'

    def _build_node_command(
        self,
        config: Dict[str, Any],
        container_name: str
    ) -> List[str]:
        """Node 컨테이너용 Docker 명령어 구성"""
        host_addr = self._get_host_address()

        # 호스트의 data 디렉토리를 컨테이너에 마운트 (CIFAR-10 재다운로드 방지)
        data_dir = str(self.project_root / 'data')

        cmd = [
            'docker', 'run',
            '--rm',
            '--privileged',
            '--platform', get_docker_platform(),
            '--name', container_name,
            '-v', f'{data_dir}:/app/data',
        ]

        # 이미지 및 역할 지정
        cmd.append(FULL_IMAGE_NAME)

        # Node 인자 추가 — 호스트의 Orchestrator/Helper에 접속
        if config.get('secure', False):
            cmd.extend([
                '--secure',
                '--orch_host', config.get('orch_host', host_addr),
                '--orch_port', str(config.get('orch_port', 8080)),
                '--helper_host', config.get('helper_host', host_addr),
                '--helper_port', str(config.get('helper_port', 8081)),
            ])
        else:
            cmd.extend([
                '--host', config.get('host', host_addr),
                '--port', str(config.get('port', 8080)),
            ])

        if config.get('no_accel', False):
            cmd.append('--no_accel')
        if config.get('isolated', False):
            cmd.append('--isolated')

        return cmd

    def _stream_logs(
        self,
        process: subprocess.Popen,
        callback: Callable[[str, str], None],
        stop_event: threading.Event
    ):
        """컨테이너 로그를 실시간으로 스트리밍"""
        try:
            for line in iter(process.stdout.readline, ''):
                if stop_event.is_set():
                    break
                if not line:
                    break

                line = line.strip()
                if not line:
                    continue

                # 로그 레벨 파싱
                level = 'INFO'
                if 'ERROR' in line.upper() or 'EXCEPTION' in line.upper():
                    level = 'ERROR'
                elif 'WARNING' in line.upper() or 'WARN' in line.upper():
                    level = 'WARNING'
                elif 'DEBUG' in line.upper():
                    level = 'DEBUG'

                callback(level, line)

        except Exception as e:
            callback('ERROR', f"로그 스트리밍 오류: {str(e)}")

    def stop_container(self, container_name: str, timeout: int = 10):
        """컨테이너 중지"""
        # 로그 스트리밍 중지
        if container_name in self._stop_events:
            self._stop_events[container_name].set()

        # 프로세스 종료
        if container_name in self.running_containers:
            process = self.running_containers[container_name]
            try:
                process.terminate()
                process.wait(timeout=timeout)
            except subprocess.TimeoutExpired:
                process.kill()
            del self.running_containers[container_name]

        # docker stop 명령 실행 (컨테이너가 아직 실행 중일 수 있음)
        try:
            subprocess.run(
                ['docker', 'stop', container_name],
                capture_output=True,
                timeout=timeout
            )
        except (subprocess.TimeoutExpired, subprocess.CalledProcessError):
            # 강제 종료
            subprocess.run(
                ['docker', 'kill', container_name],
                capture_output=True,
                timeout=5
            )

    def _remove_container_if_exists(self, container_name: str):
        """동일한 이름의 컨테이너가 있으면 제거"""
        try:
            subprocess.run(
                ['docker', 'rm', '-f', container_name],
                capture_output=True,
                timeout=10
            )
        except (subprocess.TimeoutExpired, subprocess.CalledProcessError):
            pass

    def cleanup(self):
        """모든 실행 중인 컨테이너 정리"""
        for container_name in list(self.running_containers.keys()):
            self.stop_container(container_name)

        self.running_containers.clear()
        self._log_threads.clear()
        self._stop_events.clear()

    def run_node_auto(self, config: Dict[str, Any]) -> bool:
        """
        CLI용: Docker 환경 확인 → 이미지 빌드 → 컨테이너 실행을 자동 수행.

        Args:
            config: Node 설정 딕셔너리

        Returns:
            True if 정상 완료, False if 실패
        """
        import logging

        if not self.is_docker_installed():
            logging.error("Docker가 설치되어 있지 않습니다.")
            logging.error(f"설치: {self.get_docker_download_url()}")
            return False

        if not self.is_docker_running():
            logging.error("Docker가 실행 중이 아닙니다. Docker Desktop을 시작해주세요.")
            return False

        def on_progress(msg, pct):
            logging.info(f"[Docker Build] {msg} ({pct}%)")

        if not self.build_image(progress_callback=on_progress):
            logging.error("Docker 이미지 빌드 실패")
            return False

        def on_log(level, message):
            print(f"[{level}] {message}")

        container = self.run_node(config, log_callback=on_log)
        if not container:
            logging.error("컨테이너 실행 실패")
            return False

        process = self.running_containers.get(container)
        if process:
            try:
                process.wait()
            except KeyboardInterrupt:
                logging.info("Stopping container...")
                self.stop_container(container)

        self.cleanup()
        return True

    def get_container_status(self, container_name: str) -> Optional[str]:
        """컨테이너 상태 조회"""
        try:
            result = subprocess.run(
                ['docker', 'inspect', '--format', '{{.State.Status}}', container_name],
                capture_output=True,
                text=True,
                timeout=10
            )
            if result.returncode == 0:
                return result.stdout.strip()
            return None
        except (subprocess.TimeoutExpired, subprocess.CalledProcessError):
            return None


# ==============================================================================
# DOCKER STATUS CHECKER
# ==============================================================================

class DockerStatusChecker:
    """Docker 환경 상태를 체크하고 문제를 진단하는 클래스"""

    def __init__(self):
        self.manager = DockerManager()

    def check_all(self) -> Dict[str, Any]:
        """모든 Docker 상태 체크"""
        return {
            'platform': get_platform(),
            'docker_required': should_use_docker(),
            'docker_installed': self.manager.is_docker_installed(),
            'docker_running': self.manager.is_docker_running(),
            'docker_version': self.manager.get_docker_version(),
            'image_available': self.manager.is_image_available(),
            'download_url': self.manager.get_docker_download_url(),
        }

    def get_status_message(self) -> tuple[str, str]:
        """
        현재 상태에 대한 메시지 반환.

        Returns:
            (status_type, message) - status_type: 'ok', 'warning', 'error'
        """
        status = self.check_all()

        if not status['docker_required']:
            return ('ok', 'Linux 환경입니다. 네이티브 모드로 실행할 수 있습니다.')

        if not status['docker_installed']:
            return ('error', 'Docker가 설치되어 있지 않습니다. Docker Desktop을 설치해주세요.')

        if not status['docker_running']:
            return ('warning', 'Docker가 실행 중이지 않습니다. Docker Desktop을 시작해주세요.')

        if not status['image_available']:
            return ('warning', 'TL++ 이미지가 없습니다. 첫 실행 시 자동으로 빌드됩니다.')

        return ('ok', 'Docker 환경이 준비되었습니다.')


# ==============================================================================
# MODULE TEST
# ==============================================================================

if __name__ == '__main__':
    # 간단한 테스트
    checker = DockerStatusChecker()
    status = checker.check_all()

    print("=" * 50)
    print("Docker Environment Status")
    print("=" * 50)

    for key, value in status.items():
        print(f"  {key}: {value}")

    print()
    status_type, message = checker.get_status_message()
    print(f"[{status_type.upper()}] {message}")