pipelines/deployment_pipeline.py

"""Model deployment pipeline for Myanmar Ghost."""

import argparse
import logging
import shutil
import sys
from pathlib import Path
from datetime import datetime

sys.path.insert(0, str(Path(__file__).parent.parent.parent))

from src.utils.logger import setup_logger

logger = setup_logger("deployment_pipeline")


def export_model(model_path: str, output_dir: str, format: str = "pytorch") -> str:
    """Export model for deployment."""
    logger.info(f"Exporting model from {model_path}")
    
    import torch
    
    # Load model
    checkpoint = torch.load(model_path, map_location="cpu")
    
    # Create output directory
    export_path = Path(output_dir) / "exported_model"
    export_path.mkdir(parents=True, exist_ok=True)
    
    if format == "pytorch":
        # Save as PyTorch model
        torch.save(checkpoint, export_path / "model.pt")
        logger.info(f"Exported to {export_path}")
    
    elif format == "onnx":
        # Export to ONNX (requires model forward pass)
        logger.warning("ONNX export not implemented yet")
    
    elif format == "safetensors":
        # Save as safetensors
        try:
            from safetensors.torch import save_file
            if "model_state_dict" in checkpoint:
                state_dict = checkpoint["model_state_dict"]
            else:
                state_dict = checkpoint
            
            save_file(state_dict, export_path / "model.safetensors")
            logger.info(f"Exported safetensors to {export_path}")
        except ImportError:
            logger.warning("safetensors not installed, using PyTorch format")
            torch.save(checkpoint, export_path / "model.pt")
    
    # Save metadata
    metadata = {
        "exported_at": datetime.now().isoformat(),
        "format": format,
        "original_path": model_path,
    }
    
    import json
    with open(export_path / "metadata.json", "w") as f:
        json.dump(metadata, f, indent=2)
    
    return str(export_path)


def create_docker_image(model_path: str, output_dir: str, tag: str = "myanmar-ghost") -> str:
    """Create Docker image for deployment."""
    logger.info(f"Creating Docker image: {tag}")
    
    import subprocess
    
    # Create deployment directory
    deploy_dir = Path(output_dir) / "docker"
    deploy_dir.mkdir(parents=True, exist_ok=True)
    
    # Copy model
    shutil.copytree(model_path, deploy_dir / "model", dirs_exist_ok=True)
    
    # Create Dockerfile
    dockerfile = f"""
FROM python:3.10-slim

WORKDIR /app

COPY model/ /app/model/
COPY deployment/api/ /app/api/

RUN pip install --no-cache-dir \\
    torch \\
    transformers \\
    fastapi \\
    uvicorn

EXPOSE 8000

CMD ["uvicorn", "api.app:app", "--host", "0.0.0.0", "--port", "8000"]
"""
    
    with open(deploy_dir / "Dockerfile", "w") as f:
        f.write(dockerfile)
    
    # Build image
    result = subprocess.run(
        ["docker", "build", "-t", tag, "."],
        cwd=deploy_dir,
        capture_output=True,
        text=True,
    )
    
    if result.returncode != 0:
        logger.error(f"Docker build failed: {result.stderr}")
        raise RuntimeError("Docker build failed")
    
    logger.info(f"Docker image created: {tag}")
    return tag


def push_to_huggingface(model_path: str, repo_id: str) -> str:
    """Push model to HuggingFace Hub."""
    logger.info(f"Pushing model to HuggingFace: {repo_id}")
    
    import subprocess
    
    try:
        # Use huggingface_hub
        from huggingface_hub import HfApi, create_repo
        
        api = HfApi()
        
        # Create repo if doesn't exist
        try:
            create_repo(repo_id, repo_type="model", exist_ok=True)
        except Exception:
            pass
        
        # Upload folder
        api.upload_folder(
            folder_path=model_path,
            repo_id=repo_id,
            repo_type="model",
        )
        
        logger.info(f"Pushed to https://huggingface.co/{repo_id}")
        return f"https://huggingface.co/{repo_id}"
    
    except ImportError:
        logger.warning("huggingface_hub not installed, using CLI")
        
        result = subprocess.run([
            "huggingface-cli", "upload",
            repo_id,
            model_path,
        ], capture_output=True, text=True)
        
        if result.returncode != 0:
            logger.error(f"HuggingFace upload failed: {result.stderr}")
            raise RuntimeError("HuggingFace upload failed")
        
        return f"https://huggingface.co/{repo_id}"


def deploy_to_render(api_token: str, service_name: str, model_path: str) -> str:
    """Deploy to Render using blueprint."""
    logger.info(f"Deploying to Render: {service_name}")
    
    import subprocess
    
    # Create render.yaml
    render_yaml = f"""
services:
  - type: web
    name: {service_name}
    env: docker
    repo: {model_path}
    envVars:
      - key: MODEL_PATH
        value: /app/model
"""
    
    with open("render.yaml", "w") as f:
        f.write(render_yaml)
    
    logger.info("Created render.yaml - deploy manually via Render dashboard")
    
    return "render.yaml created"


def run_deployment_pipeline(
    model_path: str,
    output_dir: str = "outputs/deployment",
    format: str = "safetensors",
    push_hub: bool = False,
    repo_id: str = None,
) -> dict:
    """Run full deployment pipeline."""
    logger.info("Starting deployment pipeline...")
    
    Path(output_dir).mkdir(parents=True, exist_ok=True)
    
    # Export model
    export_path = export_model(model_path, output_dir, format)
    
    results = {
        "export_path": export_path,
    }
    
    # Push to HuggingFace if requested
    if push_hub and repo_id:
        hub_url = push_to_huggingface(export_path, repo_id)
        results["hub_url"] = hub_url
    
    logger.info("Deployment pipeline complete!")
    
    return results


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Model deployment pipeline")
    parser.add_argument("--model_path", type=str, required=True)
    parser.add_argument("--output_dir", type=str, default="outputs/deployment")
    parser.add_argument("--format", type=str, default="safetensors")
    parser.add_argument("--push_hub", action="store_true")
    parser.add_argument("--repo_id", type=str, default=None)
    
    args = parser.parse_args()
    run_deployment_pipeline(
        args.model_path,
        args.output_dir,
        args.format,
        args.push_hub,
        args.repo_id,
    )