comfygit-deploy 0.3.4__py3-none-any.whl

comfygit_deploy/commands/worker.py

@@ -0,0 +1,266 @@
+"""Worker CLI command handlers.
+
+Commands for setting up and managing the worker server on GPU machines.
+"""
+
+import argparse
+import json
+import secrets
+from pathlib import Path
+
+WORKER_CONFIG_PATH = Path.home() / ".config" / "comfygit" / "deploy" / "worker.json"
+
+
+def generate_api_key() -> str:
+    """Generate a new worker API key."""
+    return f"cg_wk_{secrets.token_hex(16)}"
+
+
+def load_worker_config() -> dict | None:
+    """Load worker config from disk."""
+    if not WORKER_CONFIG_PATH.exists():
+        return None
+    try:
+        return json.loads(WORKER_CONFIG_PATH.read_text())
+    except (json.JSONDecodeError, OSError):
+        return None
+
+
+def get_validated_workspace() -> Path | None:
+    """Get workspace path from env or config.
+
+    Checks COMFYGIT_HOME env first, then falls back to worker config.
+
+    Returns:
+        Workspace Path if configured and exists, None otherwise.
+    """
+    import os
+
+    # Check environment variable first (takes precedence)
+    env_home = os.environ.get("COMFYGIT_HOME")
+    if env_home:
+        workspace = Path(env_home)
+        if workspace.exists():
+            return workspace
+
+    # Fall back to worker config
+    config = load_worker_config()
+    if config and config.get("workspace_path"):
+        workspace = Path(config["workspace_path"])
+        if workspace.exists():
+            return workspace
+
+    return None
+
+
+def save_worker_config(config: dict) -> None:
+    """Save worker config to disk."""
+    WORKER_CONFIG_PATH.parent.mkdir(parents=True, exist_ok=True)
+    WORKER_CONFIG_PATH.write_text(json.dumps(config, indent=2))
+
+
+def is_worker_running() -> bool:
+    """Check if worker server is running."""
+    # Simple check - look for PID file or try to connect
+    pid_file = WORKER_CONFIG_PATH.parent / "worker.pid"
+    if not pid_file.exists():
+        return False
+
+    try:
+        import os
+
+        pid = int(pid_file.read_text().strip())
+        os.kill(pid, 0)  # Check if process exists
+        return True
+    except (OSError, ValueError):
+        return False
+
+
+def handle_setup(args: argparse.Namespace) -> int:
+    """Handle 'worker setup' command."""
+    workspace = args.workspace or str(Path.home() / "comfygit")
+    api_key = args.api_key or generate_api_key()
+
+    config = {
+        "version": "1",
+        "api_key": api_key,
+        "workspace_path": workspace,
+        "default_mode": "docker",
+        "server_port": 9090,
+        "instance_port_range": {"start": 8200, "end": 8210},
+    }
+
+    save_worker_config(config)
+
+    print("Worker setup complete!")
+    print(f"  Workspace: {workspace}")
+    print(f"  API Key: {api_key}")
+    print(f"  Config: {WORKER_CONFIG_PATH}")
+    print()
+    print("To start the worker server:")
+    print("  cg-deploy worker up")
+
+    return 0
+
+
+def handle_up(args: argparse.Namespace) -> int:
+    """Handle 'worker up' command."""
+    import os
+
+    config = load_worker_config()
+    if not config:
+        print("Worker not configured. Run 'cg-deploy worker setup' first.")
+        return 1
+
+    # Parse port range
+    port_range = args.port_range.split(":")
+    port_start = int(port_range[0])
+    port_end = int(port_range[1]) if len(port_range) > 1 else port_start + 10
+
+    # Dev mode: explicit paths override saved config
+    dev_core = getattr(args, "dev_core", None)
+    dev_manager = getattr(args, "dev_manager", None)
+
+    # --dev flag loads saved config for any missing paths
+    if getattr(args, "dev", False):
+        from .dev import load_dev_config
+        dev_config = load_dev_config()
+        dev_core = dev_core or dev_config.get("core_path")
+        dev_manager = dev_manager or dev_config.get("manager_path")
+
+    if dev_core:
+        dev_core = str(Path(dev_core).resolve())
+        os.environ["COMFYGIT_DEV_CORE_PATH"] = dev_core
+        print(f"Dev mode: core -> {dev_core}")
+
+    if dev_manager:
+        dev_manager = str(Path(dev_manager).resolve())
+        # Symlink manager to system_nodes
+        workspace = Path(config["workspace_path"])
+        system_nodes = workspace / ".metadata" / "system_nodes"
+        system_nodes.mkdir(parents=True, exist_ok=True)
+        manager_link = system_nodes / "comfygit-manager"
+
+        if manager_link.is_symlink():
+            manager_link.unlink()
+        elif manager_link.is_dir():
+            import shutil
+            shutil.rmtree(manager_link)
+
+        manager_link.symlink_to(dev_manager)
+        print(f"Dev mode: manager -> {dev_manager}")
+
+    print(f"Starting worker server on {args.host}:{args.port}...")
+    print(f"  Mode: {args.mode}")
+    print(f"  Instance ports: {port_start}-{port_end}")
+    print(f"  Broadcast: {args.broadcast}")
+    print()
+    print("Press Ctrl+C to stop.")
+
+    from aiohttp import web
+
+    from ..worker.server import create_worker_app
+
+    app = create_worker_app(
+        api_key=config["api_key"],
+        workspace_path=Path(config["workspace_path"]),
+        default_mode=args.mode,
+        port_range_start=port_start,
+        port_range_end=port_end,
+    )
+
+    # Save PID file
+    pid_file = WORKER_CONFIG_PATH.parent / "worker.pid"
+    import os
+
+    pid_file.write_text(str(os.getpid()))
+
+    # Start mDNS broadcast if requested
+    broadcaster = None
+    if args.broadcast:
+        from ..worker.mdns import MDNSBroadcaster
+
+        broadcaster = MDNSBroadcaster(port=args.port, mode=args.mode)
+        broadcaster.start()
+
+    try:
+        web.run_app(app, host=args.host, port=args.port, print=lambda _: None)
+    finally:
+        if broadcaster:
+            broadcaster.stop()
+        pid_file.unlink(missing_ok=True)
+
+    return 0
+
+
+def handle_down(args: argparse.Namespace) -> int:
+    """Handle 'worker down' command."""
+    pid_file = WORKER_CONFIG_PATH.parent / "worker.pid"
+
+    if not pid_file.exists():
+        print("Worker server is not running.")
+        return 0
+
+    try:
+        import os
+        import signal
+
+        pid = int(pid_file.read_text().strip())
+        os.kill(pid, signal.SIGTERM)
+        pid_file.unlink(missing_ok=True)
+        print("Worker server stopped.")
+    except (OSError, ValueError) as e:
+        print(f"Failed to stop worker: {e}")
+        pid_file.unlink(missing_ok=True)
+        return 1
+
+    return 0
+
+
+def handle_status(args: argparse.Namespace) -> int:
+    """Handle 'worker status' command."""
+    config = load_worker_config()
+
+    if not config:
+        print("Worker not configured. Run 'cg-deploy worker setup' first.")
+        return 0
+
+    running = is_worker_running()
+    status = "RUNNING" if running else "NOT RUNNING"
+
+    print(f"Worker Status: {status}")
+    print(f"  Workspace: {config.get('workspace_path', 'N/A')}")
+    print(f"  Default Mode: {config.get('default_mode', 'docker')}")
+    print(f"  Server Port: {config.get('server_port', 9090)}")
+
+    port_range = config.get("instance_port_range", {})
+    print(
+        f"  Instance Ports: {port_range.get('start', 8200)}-{port_range.get('end', 8210)}"
+    )
+
+    return 0
+
+
+def handle_regenerate_key(args: argparse.Namespace) -> int:
+    """Handle 'worker regenerate-key' command."""
+    config = load_worker_config()
+    if not config:
+        print("Worker not configured. Run 'cg-deploy worker setup' first.")
+        return 1
+
+    old_key = config.get("api_key", "")[:20] + "..."
+    new_key = generate_api_key()
+
+    config["api_key"] = new_key
+    save_worker_config(config)
+
+    print("API key regenerated!")
+    print(f"  Old: {old_key}")
+    print(f"  New: {new_key}")
+    print()
+    print("Note: Update any clients using the old key.")
+
+    if is_worker_running():
+        print("Warning: Restart the worker server for the new key to take effect.")
+
+    return 0

comfygit_deploy/config.py

@@ -0,0 +1,122 @@
+"""Configuration storage for comfygit-deploy.
+
+Stores RunPod API keys and custom worker registry in ~/.config/comfygit/deploy/config.json
+"""
+
+import json
+import os
+from pathlib import Path
+from typing import Any
+
+
+def _get_default_config_path() -> Path:
+    """Get default config path, respecting HOME environment variable."""
+    home = os.environ.get("HOME") or str(Path.home())
+    return Path(home) / ".config" / "comfygit" / "deploy" / "config.json"
+
+
+class DeployConfig:
+    """Configuration storage for deploy CLI.
+
+    Stores:
+    - RunPod API key
+    - Custom worker registry (name -> {host, port, api_key, ...})
+    """
+
+    def __init__(self, path: Path | None = None):
+        """Initialize config.
+
+        Args:
+            path: Config file path. Defaults to ~/.config/comfygit/deploy/config.json
+        """
+        self.path = path or _get_default_config_path()
+        self._data: dict[str, Any] = {"version": "1", "providers": {}, "workers": {}}
+        self._load()
+
+    def _load(self) -> None:
+        """Load config from disk if it exists."""
+        if self.path.exists():
+            try:
+                self._data = json.loads(self.path.read_text())
+            except (json.JSONDecodeError, OSError):
+                pass  # Use defaults
+
+    def save(self) -> None:
+        """Save config to disk."""
+        self.path.parent.mkdir(parents=True, exist_ok=True)
+        self.path.write_text(json.dumps(self._data, indent=2))
+
+    @property
+    def runpod_api_key(self) -> str | None:
+        """Get RunPod API key."""
+        return self._data.get("providers", {}).get("runpod", {}).get("api_key")
+
+    @runpod_api_key.setter
+    def runpod_api_key(self, value: str | None) -> None:
+        """Set RunPod API key."""
+        if "providers" not in self._data:
+            self._data["providers"] = {}
+        if "runpod" not in self._data["providers"]:
+            self._data["providers"]["runpod"] = {}
+
+        if value is None:
+            self._data["providers"]["runpod"].pop("api_key", None)
+        else:
+            self._data["providers"]["runpod"]["api_key"] = value
+
+    @property
+    def workers(self) -> dict[str, dict[str, Any]]:
+        """Get custom workers registry."""
+        return self._data.get("workers", {})
+
+    def add_worker(
+        self,
+        name: str,
+        host: str,
+        port: int,
+        api_key: str,
+        mode: str = "docker",
+    ) -> None:
+        """Add a custom worker to the registry.
+
+        Args:
+            name: Worker name (unique identifier)
+            host: Worker host/IP
+            port: Worker API port
+            api_key: Worker API key
+            mode: Worker mode (docker or native)
+        """
+        if "workers" not in self._data:
+            self._data["workers"] = {}
+
+        self._data["workers"][name] = {
+            "host": host,
+            "port": port,
+            "api_key": api_key,
+            "mode": mode,
+        }
+
+    def remove_worker(self, name: str) -> bool:
+        """Remove a worker from the registry.
+
+        Args:
+            name: Worker name
+
+        Returns:
+            True if worker was removed, False if not found
+        """
+        if name in self._data.get("workers", {}):
+            del self._data["workers"][name]
+            return True
+        return False
+
+    def get_worker(self, name: str) -> dict[str, Any] | None:
+        """Get a worker by name.
+
+        Args:
+            name: Worker name
+
+        Returns:
+            Worker config dict or None if not found
+        """
+        return self._data.get("workers", {}).get(name)

comfygit_deploy/providers/__init__.py

@@ -0,0 +1,11 @@
+"""Provider clients for deployment backends."""
+
+from .custom import CustomWorkerClient, CustomWorkerError
+from .runpod import RunPodAPIError, RunPodClient
+
+__all__ = [
+    "RunPodClient",
+    "RunPodAPIError",
+    "CustomWorkerClient",
+    "CustomWorkerError",
+]

comfygit_deploy/providers/custom.py

@@ -0,0 +1,238 @@
+"""Custom worker HTTP client for connecting to self-hosted workers.
+
+Provides async interface for worker server REST API.
+"""
+
+import json
+from collections.abc import AsyncIterator
+from dataclasses import dataclass
+from typing import Any
+
+import aiohttp
+
+
+@dataclass
+class CustomWorkerError(Exception):
+    """Error from custom worker API."""
+
+    message: str
+    status_code: int
+
+    def __str__(self) -> str:
+        return f"Worker Error ({self.status_code}): {self.message}"
+
+
+@dataclass
+class LogEntry:
+    """A single log entry from streaming."""
+
+    timestamp: str
+    level: str
+    message: str
+
+
+class CustomWorkerClient:
+    """Async client for custom worker REST API."""
+
+    def __init__(self, host: str, port: int, api_key: str):
+        """Initialize client.
+
+        Args:
+            host: Worker host/IP
+            port: Worker API port
+            api_key: Worker API key
+        """
+        self.host = host
+        self.port = port
+        self.api_key = api_key
+        self.base_url = f"http://{host}:{port}"
+
+    def _headers(self) -> dict[str, str]:
+        """Get request headers with authorization."""
+        return {
+            "Authorization": f"Bearer {self.api_key}",
+            "Content-Type": "application/json",
+        }
+
+    async def _get(self, path: str) -> Any:
+        """Make GET request and return JSON response."""
+        async with aiohttp.ClientSession() as session:
+            async with session.get(
+                f"{self.base_url}{path}",
+                headers=self._headers(),
+            ) as response:
+                if response.status >= 400:
+                    await self._handle_error(response)
+                return await response.json()
+
+    async def _post(self, path: str, data: dict | None = None) -> Any:
+        """Make POST request and return JSON response."""
+        async with aiohttp.ClientSession() as session:
+            async with session.post(
+                f"{self.base_url}{path}",
+                json=data,
+                headers=self._headers(),
+            ) as response:
+                if response.status >= 400:
+                    await self._handle_error(response)
+                return await response.json()
+
+    async def _delete(self, path: str) -> Any:
+        """Make DELETE request and return JSON response."""
+        async with aiohttp.ClientSession() as session:
+            async with session.delete(
+                f"{self.base_url}{path}",
+                headers=self._headers(),
+            ) as response:
+                if response.status >= 400:
+                    await self._handle_error(response)
+                return await response.json()
+
+    async def _handle_error(self, response: aiohttp.ClientResponse) -> None:
+        """Handle error response."""
+        try:
+            error_body = await response.json()
+            message = error_body.get("error", str(error_body))
+        except Exception:
+            message = await response.text() or f"HTTP {response.status}"
+        raise CustomWorkerError(message, response.status)
+
+    async def test_connection(self) -> dict[str, Any]:
+        """Test connection to worker.
+
+        Returns:
+            {"success": True, "worker_version": "..."} on success
+            {"success": False, "error": "..."} on failure
+        """
+        try:
+            health = await self._get("/api/v1/health")
+            return {
+                "success": True,
+                "worker_version": health.get("worker_version"),
+            }
+        except CustomWorkerError as e:
+            return {"success": False, "error": e.message}
+        except Exception as e:
+            return {"success": False, "error": str(e)}
+
+    async def get_system_info(self) -> dict[str, Any]:
+        """Get worker system information."""
+        return await self._get("/api/v1/system/info")
+
+    async def list_instances(self) -> list[dict]:
+        """List all instances on worker."""
+        result = await self._get("/api/v1/instances")
+        return result.get("instances", [])
+
+    async def create_instance(
+        self,
+        import_source: str,
+        name: str | None = None,
+        branch: str | None = None,
+        mode: str | None = None,
+    ) -> dict:
+        """Create new instance.
+
+        Args:
+            import_source: Git URL or local path
+            name: Optional instance name
+            branch: Optional git branch
+            mode: docker or native
+
+        Returns:
+            Instance data
+        """
+        data = {"import_source": import_source}
+        if name:
+            data["name"] = name
+        if branch:
+            data["branch"] = branch
+        if mode:
+            data["mode"] = mode
+
+        return await self._post("/api/v1/instances", data)
+
+    async def get_instance(self, instance_id: str) -> dict:
+        """Get instance details."""
+        return await self._get(f"/api/v1/instances/{instance_id}")
+
+    async def stop_instance(self, instance_id: str) -> dict:
+        """Stop a running instance."""
+        return await self._post(f"/api/v1/instances/{instance_id}/stop")
+
+    async def start_instance(self, instance_id: str) -> dict:
+        """Start a stopped instance."""
+        return await self._post(f"/api/v1/instances/{instance_id}/start")
+
+    async def terminate_instance(self, instance_id: str, keep_env: bool = False) -> dict:
+        """Terminate and remove instance."""
+        path = f"/api/v1/instances/{instance_id}"
+        if keep_env:
+            path += "?keep_env=true"
+        return await self._delete(path)
+
+    async def get_logs(self, instance_id: str, lines: int = 100) -> list[dict]:
+        """Get recent logs for an instance.
+
+        Args:
+            instance_id: Instance ID
+            lines: Number of lines to fetch
+
+        Returns:
+            List of log entries
+        """
+        result = await self._get(f"/api/v1/instances/{instance_id}/logs?lines={lines}")
+        return result.get("logs", [])
+
+    def _connect_ws(self, url: str):
+        """Create a WebSocket connection context manager.
+
+        Returns an async context manager that yields the WebSocket connection.
+        """
+        return _WSConnectionManager(url, self.api_key)
+
+    async def stream_logs(self, instance_id: str) -> AsyncIterator[LogEntry]:
+        """Stream logs from an instance via WebSocket.
+
+        Args:
+            instance_id: Instance ID
+
+        Yields:
+            LogEntry objects as they arrive
+        """
+        url = f"ws://{self.host}:{self.port}/api/v1/instances/{instance_id}/logs"
+        async with self._connect_ws(url) as ws:
+            async for msg in ws:
+                if msg.type == aiohttp.WSMsgType.TEXT:
+                    data = json.loads(msg.data)
+                    if data.get("type") == "log":
+                        yield LogEntry(
+                            timestamp=data.get("timestamp", ""),
+                            level=data.get("level", "INFO"),
+                            message=data.get("message", ""),
+                        )
+                elif msg.type == aiohttp.WSMsgType.ERROR:
+                    break
+
+
+class _WSConnectionManager:
+    """Async context manager for WebSocket connections."""
+
+    def __init__(self, url: str, api_key: str):
+        self.url = url
+        self.api_key = api_key
+        self._session: aiohttp.ClientSession | None = None
+        self._ws = None
+
+    async def __aenter__(self):
+        self._session = aiohttp.ClientSession()
+        self._ws = await self._session.ws_connect(
+            self.url, headers={"Authorization": f"Bearer {self.api_key}"}
+        )
+        return self._ws
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        if self._ws:
+            await self._ws.close()
+        if self._session:
+            await self._session.close()