comfygit-deploy 0.3.4__py3-none-any.whl

comfygit_deploy/startup/scripts.py

@@ -0,0 +1,210 @@
+"""Startup script generator for RunPod deployments.
+
+Generates bash scripts that set up ComfyGit environments on RunPod pods
+using the cg CLI commands.
+"""
+
+import re
+import secrets
+from datetime import datetime
+
+
+def generate_deployment_id(env_name: str) -> str:
+    """Generate a unique deployment ID.
+
+    Format: deploy-{sanitized_env_name}-{YYYYMMDD}-{HHMMSS}-{random}
+
+    Args:
+        env_name: Source environment name
+
+    Returns:
+        Unique deployment identifier
+    """
+    # Sanitize env name: replace special chars with dashes, collapse multiple dashes
+    sanitized = re.sub(r"[^a-zA-Z0-9-]", "-", env_name)
+    sanitized = re.sub(r"-+", "-", sanitized).strip("-").lower()
+    if not sanitized:
+        sanitized = "env"
+
+    now = datetime.now()
+    timestamp = now.strftime("%Y%m%d-%H%M%S")
+    random_suffix = secrets.token_hex(2)  # 4 hex chars
+
+    return f"deploy-{sanitized}-{timestamp}-{random_suffix}"
+
+
+def generate_startup_script(
+    deployment_id: str,
+    import_source: str,
+    branch: str | None = None,
+    comfyui_port: int = 8188,
+) -> str:
+    """Generate the pod startup script for ComfyGit environment setup.
+
+    This script runs on the RunPod pod and:
+    1. Sets up COMFYGIT_HOME environment variable
+    2. Initializes workspace with cg init --yes
+    3. Imports environment with cg import --yes --name {deployment_id}
+    4. Starts ComfyUI with cg -e {deployment_id} run --listen 0.0.0.0
+
+    Args:
+        deployment_id: Unique identifier for this deployment
+        import_source: Git URL or local path for cg import
+        branch: Optional git branch/tag for import
+        comfyui_port: Port for ComfyUI (default 8188)
+
+    Returns:
+        Bash script string
+    """
+    # Build branch flag if specified
+    branch_flag = f" -b {branch}" if branch else ""
+
+    return f'''#!/bin/bash
+set -e
+
+# =============================================================================
+# ComfyGit Deployment Startup Script
+# Deployment ID: {deployment_id}
+# =============================================================================
+
+# Status tracking
+STATUS_FILE="/workspace/.comfygit_deploy_status.json"
+START_TIME=$(date -Iseconds)
+
+update_status() {{
+    local phase="$1"
+    local detail="$2"
+    local progress="$3"
+    cat > "$STATUS_FILE" << EOF
+{{
+    "deployment_id": "{deployment_id}",
+    "phase": "$phase",
+    "phase_detail": "$detail",
+    "progress_percent": $progress,
+    "started_at": "$START_TIME",
+    "error": null
+}}
+EOF
+}}
+
+set_error() {{
+    local message="$1"
+    cat > "$STATUS_FILE" << EOF
+{{
+    "deployment_id": "{deployment_id}",
+    "phase": "ERROR",
+    "phase_detail": "$message",
+    "progress_percent": 0,
+    "started_at": "$START_TIME",
+    "error": "$message"
+}}
+EOF
+    echo "ERROR: $message" >&2
+    exit 1
+}}
+
+# Ensure workspace directory exists
+mkdir -p /workspace
+
+# =============================================================================
+# Phase: INITIALIZING
+# =============================================================================
+update_status "INITIALIZING" "Setting up environment..." 5
+
+# Set ComfyGit home directory
+export COMFYGIT_HOME=/workspace/comfygit
+export PATH="$HOME/.local/bin:$PATH"
+
+update_status "INITIALIZING" "Installing uv package manager..." 10
+
+# Install uv if not present (default RunPod template doesn't include it)
+if ! command -v uv &> /dev/null; then
+    curl -LsSf https://astral.sh/uv/install.sh | sh || set_error "Failed to install uv"
+fi
+
+# Source uv's env file to ensure PATH is configured correctly
+source "$HOME/.local/bin/env" 2>/dev/null || true
+export PATH="$HOME/.local/bin:$PATH"
+
+update_status "INITIALIZING" "Installing ComfyGit CLI..." 12
+
+# Install comfygit CLI via uv
+if ! command -v cg &> /dev/null; then
+    uv tool install comfygit || set_error "Failed to install comfygit CLI"
+fi
+
+# Verify cg is available after installation (catch PATH issues early)
+command -v cg &> /dev/null || set_error "cg command not found after installation - check PATH"
+
+# =============================================================================
+# Phase: WORKSPACE_INIT
+# =============================================================================
+update_status "WORKSPACE_INIT" "Initializing ComfyGit workspace..." 15
+
+# Initialize workspace (idempotent - fails gracefully if exists)
+mkdir -p "$COMFYGIT_HOME"
+cd "$COMFYGIT_HOME"
+cg init --yes "$COMFYGIT_HOME" 2>/dev/null || true
+
+# =============================================================================
+# Phase: ENVIRONMENT_CHECK (handles restart vs fresh deploy)
+# =============================================================================
+ENV_PATH="$COMFYGIT_HOME/environments/{deployment_id}"
+
+if [ -d "$ENV_PATH" ] && [ -d "$ENV_PATH/ComfyUI" ]; then
+    # Restart case: environment already exists from previous run
+    update_status "RESTARTING" "Environment exists, skipping import..." 60
+    echo "Detected existing environment: {deployment_id}"
+    echo "   Skipping import, proceeding to start ComfyUI..."
+    cg use {deployment_id} 2>/dev/null || true
+else
+    # Fresh deploy: import the environment
+    update_status "IMPORTING" "Preparing import source..." 25
+
+    IMPORT_SOURCE="{import_source}"
+
+    update_status "IMPORTING" "Importing environment {deployment_id}..." 30
+
+    # Import the environment (--models all downloads all models with sources)
+    cd "$COMFYGIT_HOME"
+    cg import "$IMPORT_SOURCE"{branch_flag} --name {deployment_id} --yes --use --models all || set_error "Failed to import environment"
+
+    update_status "IMPORTING" "Environment imported successfully" 60
+fi
+
+# =============================================================================
+# Phase: STARTING_COMFYUI
+# =============================================================================
+update_status "STARTING_COMFYUI" "Starting ComfyUI server..." 80
+
+# Start ComfyUI with the imported environment
+cd "$COMFYGIT_HOME"
+nohup cg -e {deployment_id} run --listen 0.0.0.0 --port {comfyui_port} > /workspace/comfyui.log 2>&1 &
+COMFYUI_PID=$!
+
+# Wait for ComfyUI to start (check port)
+update_status "STARTING_COMFYUI" "Waiting for ComfyUI to become ready..." 90
+
+for i in {{1..120}}; do
+    if curl -s http://localhost:{comfyui_port} > /dev/null 2>&1; then
+        break
+    fi
+    sleep 2
+done
+
+if ! curl -s http://localhost:{comfyui_port} > /dev/null 2>&1; then
+    set_error "ComfyUI failed to start within 4 minutes. Check /workspace/comfyui.log"
+fi
+
+# =============================================================================
+# Phase: READY
+# =============================================================================
+update_status "READY" "ComfyUI is running" 100
+
+echo "=== Deployment Complete ==="
+echo "Deployment ID: {deployment_id}"
+echo "ComfyUI is running at http://localhost:{comfyui_port}"
+
+# Keep the script running (RunPod terminates if script exits)
+tail -f /workspace/comfyui.log
+'''

comfygit_deploy/worker/__init__.py

@@ -0,0 +1,12 @@
+"""Worker server components for self-hosted deployment."""
+
+from .server import WorkerServer, create_worker_app
+from .state import InstanceState, PortAllocator, WorkerState
+
+__all__ = [
+    "InstanceState",
+    "PortAllocator",
+    "WorkerState",
+    "WorkerServer",
+    "create_worker_app",
+]

comfygit_deploy/worker/mdns.py

@@ -0,0 +1,154 @@
+"""mDNS service broadcasting and scanning for worker discovery.
+
+Registers the worker as a _cg-deploy._tcp.local. service so it can be
+discovered by frontends scanning the local network. Also provides client-side
+scanning to discover workers on the network.
+"""
+
+import socket
+import time
+from dataclasses import dataclass
+
+from zeroconf import ServiceBrowser, ServiceInfo, ServiceStateChange, Zeroconf
+
+from .. import __version__
+
+SERVICE_TYPE = "_cg-deploy._tcp.local."
+
+
+@dataclass
+class DiscoveredWorker:
+    """Worker discovered via mDNS scan."""
+
+    name: str
+    host: str
+    port: int
+    version: str = "unknown"
+    mode: str = "docker"
+
+
+def get_local_ip() -> str:
+    """Get the local IP address of this machine.
+
+    Returns:
+        Local IP address as string, or 127.0.0.1 if detection fails.
+    """
+    try:
+        # Create a socket to determine local IP
+        s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+        s.connect(("8.8.8.8", 80))
+        ip = s.getsockname()[0]
+        s.close()
+        return ip
+    except Exception:
+        return "127.0.0.1"
+
+
+class MDNSBroadcaster:
+    """Broadcasts worker availability via mDNS/Zeroconf."""
+
+    def __init__(self, port: int, worker_name: str | None = None, mode: str = "docker"):
+        """Initialize broadcaster.
+
+        Args:
+            port: Port the worker HTTP server is listening on
+            worker_name: Optional name for the worker (defaults to hostname)
+            mode: Worker mode (docker or native)
+        """
+        self.port = port
+        self.worker_name = worker_name or socket.gethostname()
+        self.mode = mode
+        self.zeroconf: Zeroconf | None = None
+        self.service_info: ServiceInfo | None = None
+
+    def start(self) -> None:
+        """Register the mDNS service."""
+        local_ip = get_local_ip()
+
+        self.service_info = ServiceInfo(
+            SERVICE_TYPE,
+            f"{self.worker_name}.{SERVICE_TYPE}",
+            addresses=[socket.inet_aton(local_ip)],
+            port=self.port,
+            properties={
+                "version": __version__,
+                "name": self.worker_name,
+                "mode": self.mode,
+            },
+        )
+
+        self.zeroconf = Zeroconf()
+        self.zeroconf.register_service(self.service_info)
+        print(f"  mDNS: Broadcasting as {self.worker_name} on {local_ip}:{self.port}")
+
+    def stop(self) -> None:
+        """Unregister the mDNS service."""
+        if self.zeroconf and self.service_info:
+            self.zeroconf.unregister_service(self.service_info)
+            self.zeroconf.close()
+            self.zeroconf = None
+            self.service_info = None
+
+
+class MDNSScanner:
+    """Scans the network for ComfyGit workers via mDNS/Zeroconf."""
+
+    def __init__(self, timeout: float = 5.0):
+        """Initialize scanner.
+
+        Args:
+            timeout: How long to scan (seconds)
+        """
+        self.timeout = timeout
+        self._discovered: list[DiscoveredWorker] = []
+        self._zeroconf: Zeroconf | None = None
+
+    def _on_service_state_change(
+        self,
+        zeroconf: Zeroconf,
+        service_type: str,
+        name: str,
+        state_change: str | ServiceStateChange,
+    ) -> None:
+        """Called when a service is discovered or removed."""
+        # Handle both string and enum state change types
+        state = state_change if isinstance(state_change, str) else state_change.name
+
+        if state in ("Added", "Updated"):
+            info = zeroconf.get_service_info(service_type, name)
+            if info:
+                addresses = info.parsed_addresses()
+                if addresses:
+                    props = info.properties or {}
+                    worker = DiscoveredWorker(
+                        name=props.get(b"name", b"unknown").decode(),
+                        host=addresses[0],
+                        port=info.port,
+                        version=props.get(b"version", b"unknown").decode(),
+                        mode=props.get(b"mode", b"docker").decode(),
+                    )
+                    # Avoid duplicates
+                    if not any(w.host == worker.host and w.port == worker.port for w in self._discovered):
+                        self._discovered.append(worker)
+
+    def scan(self) -> list[DiscoveredWorker]:
+        """Scan for workers on the network.
+
+        Returns:
+            List of discovered workers
+        """
+        self._discovered = []
+        zeroconf = Zeroconf()
+        self._zeroconf = zeroconf
+
+        try:
+            ServiceBrowser(zeroconf, SERVICE_TYPE, handlers=[self._on_service_state_change])
+            time.sleep(self.timeout)
+        finally:
+            zeroconf.close()
+
+        return self._discovered
+
+    def get_discovered(self) -> list[DiscoveredWorker]:
+        """Get workers discovered so far (for manual callback testing)."""
+        return self._discovered