revamp for bridge - add structure and get sky launch and down working

Author: deep1401

pyproject.toml

@@ -11,9 +11,11 @@ dependencies = [
     "fastapi==0.115.7",
     "pydantic>=2.10.4",
     "pyjwt==2.10.1",
-    "python-dotenv==1.0.0",
+    "python-dotenv>=1.0.0",
     "python-multipart>=0.0.18",
     "pyyaml>=6.0",
+    "requests>=2.31.0",
+    "paramiko>=3.0.0",
     "runpod>=1.6.0",
     "sqlalchemy>=2.0.42",
     "uvicorn[standard]==0.35.0",
@@ -50,7 +52,8 @@ dependencies = [
     "msrestazure==0.6.4.post1",
     "openapi-client>=1.1.7",
     "nanoid==2.0.0",
-    "skypilot[all] @ git+https://github.com/transformerlab/skypilot.git@prod-2",
+    # "skypilot[all] @ git+https://github.com/transformerlab/skypilot.git@prod-2",
+    "skypilot[all]==0.10.5",
     "fsspec>=2025.9.0",
     "adlfs>=2022.11.1",
     "s3fs>=2025.9.0",

src/lattice/providers/__init__.py

@@ -0,0 +1,14 @@
+"""Provider bridge system for abstracting GPU orchestration providers."""
+
+from .base import Provider
+from .router import ProviderRouter, get_provider
+from .config import load_providers_config, ProviderConfig
+
+__all__ = [
+    "Provider",
+    "ProviderRouter",
+    "get_provider",
+    "load_providers_config",
+    "ProviderConfig",
+]
+

src/lattice/providers/base.py

@@ -0,0 +1,142 @@
+"""Abstract base class for provider implementations."""
+
+from abc import ABC, abstractmethod
+from typing import Dict, List, Any, Optional, Union
+from .models import (
+    ClusterConfig,
+    JobConfig,
+    ClusterStatus,
+    JobInfo,
+    ResourceInfo,
+)
+
+
+class Provider(ABC):
+    """Abstract base class for all provider implementations."""
+
+    @abstractmethod
+    def launch_cluster(
+        self, cluster_name: str, config: ClusterConfig
+    ) -> Dict[str, Any]:
+        """
+        Launch/provision a new cluster.
+
+        Args:
+            cluster_name: Name of the cluster to launch
+            config: Cluster configuration
+
+        Returns:
+            Dictionary with launch result (e.g., request_id, cluster_name)
+        """
+        pass
+
+    @abstractmethod
+    def stop_cluster(self, cluster_name: str) -> Dict[str, Any]:
+        """
+        Stop a running cluster (but don't tear it down).
+
+        Args:
+            cluster_name: Name of the cluster to stop
+
+        Returns:
+            Dictionary with stop result
+        """
+        pass
+
+    @abstractmethod
+    def get_cluster_status(
+        self, cluster_name: str
+    ) -> ClusterStatus:
+        """
+        Get the status of a cluster.
+
+        Args:
+            cluster_name: Name of the cluster
+
+        Returns:
+            ClusterStatus object with cluster information
+        """
+        pass
+
+    @abstractmethod
+    def get_cluster_resources(
+        self, cluster_name: str
+    ) -> ResourceInfo:
+        """
+        Get resource information for a cluster (GPUs, CPUs, memory, etc.).
+
+        Args:
+            cluster_name: Name of the cluster
+
+        Returns:
+            ResourceInfo object with resource details
+        """
+        pass
+
+    @abstractmethod
+    def submit_job(
+        self, cluster_name: str, job_config: JobConfig
+    ) -> Dict[str, Any]:
+        """
+        Submit a job to an existing cluster.
+
+        Args:
+            cluster_name: Name of the cluster
+            job_config: Job configuration
+
+        Returns:
+            Dictionary with job submission result (e.g., job_id)
+        """
+        pass
+
+    @abstractmethod
+    def get_job_logs(
+        self,
+        cluster_name: str,
+        job_id: Union[str, int],
+        tail_lines: Optional[int] = None,
+        follow: bool = False,
+    ) -> Union[str, Any]:
+        """
+        Get logs for a job.
+
+        Args:
+            cluster_name: Name of the cluster
+            job_id: Job identifier
+            tail_lines: Number of lines to retrieve from the end (None for all)
+            follow: Whether to stream/follow logs (returns stream if True)
+
+        Returns:
+            Log content as string, or stream object if follow=True
+        """
+        pass
+
+    @abstractmethod
+    def cancel_job(
+        self, cluster_name: str, job_id: Union[str, int]
+    ) -> Dict[str, Any]:
+        """
+        Cancel a running or queued job.
+
+        Args:
+            cluster_name: Name of the cluster
+            job_id: Job identifier
+
+        Returns:
+            Dictionary with cancellation result
+        """
+        pass
+
+    @abstractmethod
+    def list_jobs(self, cluster_name: str) -> List[JobInfo]:
+        """
+        List all jobs for a cluster.
+
+        Args:
+            cluster_name: Name of the cluster
+
+        Returns:
+            List of JobInfo objects
+        """
+        pass
+

src/lattice/providers/config.py

@@ -0,0 +1,170 @@
+"""Configuration loading and provider factory."""
+
+import os
+import yaml
+import json
+from typing import Dict, Any, Optional
+from pathlib import Path
+from pydantic import BaseModel, Field
+
+
+class ProviderConfig(BaseModel):
+    """Configuration for a single provider."""
+
+    type: str  # "skypilot" or "slurm"
+    name: str  # Provider name/identifier
+
+    # SkyPilot-specific config
+    server_url: Optional[str] = None
+    api_token: Optional[str] = None
+    default_env_vars: Dict[str, str] = Field(default_factory=dict)
+    default_entrypoint_command: Optional[str] = None
+
+    # SLURM-specific config
+    mode: Optional[str] = None  # "rest" or "ssh"
+    rest_url: Optional[str] = None
+    ssh_host: Optional[str] = None
+    ssh_user: Optional[str] = None
+    ssh_key_path: Optional[str] = None
+    ssh_port: int = 22
+
+    # Additional provider-specific config
+    extra_config: Dict[str, Any] = Field(default_factory=dict)
+
+
+def load_providers_config(
+    config_path: Optional[str] = None,
+) -> Dict[str, ProviderConfig]:
+    """
+    Load provider configurations from YAML or JSON file.
+
+    Args:
+        config_path: Path to config file. If None, uses default location
+                    or PROVIDERS_CONFIG_PATH env var.
+
+    Returns:
+        Dictionary mapping provider names to ProviderConfig objects
+    """
+    if config_path is None:
+        # Check environment variable first
+        env_path = os.getenv("PROVIDERS_CONFIG_PATH")
+        if env_path:
+            config_path = env_path
+        else:
+            # Try to find the config file in multiple locations
+            current_file = Path(__file__).resolve()
+
+            # 1. Check in the same directory as this file (installed package)
+            package_config = current_file.parent / "providers.yaml"
+
+            # 2. Check in source directory (when running from repo)
+            # Go up from src/lattice/providers/config.py to find repo root
+            # Then look for src/lattice/providers/providers.yaml
+            source_config = None
+            for parent in [
+                current_file.parent.parent.parent.parent,
+                current_file.parent.parent.parent.parent.parent,
+            ]:
+                potential = parent / "src" / "lattice" / "providers" / "providers.yaml"
+                if potential.exists():
+                    source_config = potential
+                    break
+
+            # Prefer source config if it exists (for development)
+            if source_config and source_config.exists():
+                config_path = str(source_config)
+            elif package_config.exists():
+                config_path = str(package_config)
+            else:
+                # Default to package directory location
+                config_path = str(package_config)
+
+    config_path = Path(config_path).expanduser().resolve()
+
+    if not config_path.exists():
+        # Provide helpful error message with suggestions
+        suggested_paths = [
+            os.path.join(os.path.dirname(__file__), "providers.yaml"),
+            os.path.join(os.getcwd(), "providers.yaml"),
+        ]
+        suggestions = "\n".join(f"  - {p}" for p in suggested_paths if Path(p).exists())
+        error_msg = (
+            f"Provider config file not found: {config_path}\n"
+            f"Please create a providers.yaml file or set PROVIDERS_CONFIG_PATH environment variable.\n"
+        )
+        if suggestions:
+            error_msg += f"Found config files at:\n{suggestions}\n"
+        raise FileNotFoundError(error_msg)
+
+    with open(config_path, "r") as f:
+        if config_path.suffix in [".yaml", ".yml"]:
+            config_data = yaml.safe_load(f)
+        elif config_path.suffix == ".json":
+            config_data = json.load(f)
+        else:
+            raise ValueError(f"Unsupported config file format: {config_path.suffix}")
+
+    providers = {}
+    providers_data = config_data.get("providers", {})
+
+    for name, provider_data in providers_data.items():
+        provider_data["name"] = name
+        providers[name] = ProviderConfig(**provider_data)
+
+    return providers
+
+
+def create_provider(config: ProviderConfig):
+    """
+    Factory function to create a provider instance from config.
+
+    Args:
+        config: ProviderConfig object
+
+    Returns:
+        Provider instance
+    """
+    from .skypilot import SkyPilotProvider
+    from .slurm import SLURMProvider
+
+    if config.type == "skypilot":
+        if not config.server_url:
+            raise ValueError("SkyPilot provider requires server_url in config")
+        return SkyPilotProvider(
+            server_url=config.server_url,
+            api_token=config.api_token,
+            default_env_vars=config.default_env_vars,
+            default_entrypoint_command=config.default_entrypoint_command,
+            extra_config=config.extra_config,
+        )
+    elif config.type == "slurm":
+        if config.mode == "rest":
+            if not config.rest_url:
+                raise ValueError(
+                    "SLURM provider in REST mode requires rest_url in config"
+                )
+            return SLURMProvider(
+                mode="rest",
+                rest_url=config.rest_url,
+                api_token=config.api_token,
+                extra_config=config.extra_config,
+            )
+        elif config.mode == "ssh":
+            if not config.ssh_host:
+                raise ValueError(
+                    "SLURM provider in SSH mode requires ssh_host in config"
+                )
+            return SLURMProvider(
+                mode="ssh",
+                ssh_host=config.ssh_host,
+                ssh_user=config.ssh_user or os.getenv("USER", "root"),
+                ssh_key_path=config.ssh_key_path,
+                ssh_port=config.ssh_port,
+                extra_config=config.extra_config,
+            )
+        else:
+            raise ValueError(
+                f"SLURM provider mode must be 'rest' or 'ssh', got: {config.mode}"
+            )
+    else:
+        raise ValueError(f"Unknown provider type: {config.type}")