isa-model 0.2.0__py3-none-any.whl → 0.2.8__py3-none-any.whl

isa_model/training/cloud/storage_manager.py

@@ -0,0 +1,482 @@
+"""
+Cloud Storage Manager for Training Assets
+
+This module handles storage of datasets, models, and training artifacts
+across different cloud storage providers (S3, GCS, Azure, etc.).
+"""
+
+import os
+import json
+import logging
+from typing import Dict, List, Optional, Any, Union
+from dataclasses import dataclass
+from pathlib import Path
+from urllib.parse import urlparse
+
+try:
+    import boto3
+    from botocore.exceptions import ClientError
+    S3_AVAILABLE = True
+except ImportError:
+    S3_AVAILABLE = False
+    boto3 = None
+
+try:
+    from google.cloud import storage as gcs
+    GCS_AVAILABLE = True
+except ImportError:
+    GCS_AVAILABLE = False
+    gcs = None
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class StorageConfig:
+    """Configuration for cloud storage."""
+    
+    provider: str  # "s3", "gcs", "azure"
+    bucket_name: str
+    region: Optional[str] = None
+    
+    # Authentication
+    access_key: Optional[str] = None
+    secret_key: Optional[str] = None
+    service_account_path: Optional[str] = None
+    
+    # Paths
+    datasets_prefix: str = "datasets/"
+    models_prefix: str = "models/"
+    logs_prefix: str = "logs/"
+    
+    def __post_init__(self):
+        """Validate configuration."""
+        if self.provider not in ["s3", "gcs", "azure"]:
+            raise ValueError(f"Unsupported storage provider: {self.provider}")
+        
+        if not self.bucket_name:
+            raise ValueError("Bucket name is required")
+
+
+class CloudStorageManager:
+    """
+    Cloud storage manager for training assets.
+    
+    Handles upload/download of datasets, models, and training artifacts
+    across different cloud storage providers.
+    
+    Example:
+        ```python
+        # Configure S3 storage
+        storage_config = StorageConfig(
+            provider="s3",
+            bucket_name="my-training-bucket",
+            region="us-west-2",
+            access_key="your-access-key",
+            secret_key="your-secret-key"
+        )
+        
+        # Initialize storage manager
+        storage = CloudStorageManager(storage_config)
+        
+        # Upload dataset
+        dataset_url = storage.upload_dataset("local_data.json", "my-dataset")
+        
+        # Upload trained model
+        model_url = storage.upload_model("./trained_model/", "gemma-finetuned-v1")
+        
+        # Download model
+        local_path = storage.download_model("gemma-finetuned-v1", "./downloaded_model/")
+        ```
+    """
+    
+    def __init__(self, config: StorageConfig):
+        """
+        Initialize cloud storage manager.
+        
+        Args:
+            config: Storage configuration
+        """
+        self.config = config
+        self._client = None
+        self._initialize_client()
+        
+        logger.info(f"Storage manager initialized for {config.provider}://{config.bucket_name}")
+    
+    def _initialize_client(self) -> None:
+        """Initialize storage client based on provider."""
+        if self.config.provider == "s3":
+            if not S3_AVAILABLE:
+                raise ImportError("boto3 is required for S3 storage. Install with: pip install boto3")
+            
+            self._client = boto3.client(
+                's3',
+                aws_access_key_id=self.config.access_key,
+                aws_secret_access_key=self.config.secret_key,
+                region_name=self.config.region
+            )
+            
+        elif self.config.provider == "gcs":
+            if not GCS_AVAILABLE:
+                raise ImportError("google-cloud-storage is required for GCS. Install with: pip install google-cloud-storage")
+            
+            if self.config.service_account_path:
+                os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = self.config.service_account_path
+            
+            self._client = gcs.Client()
+            
+        else:
+            raise NotImplementedError(f"Provider {self.config.provider} not implemented yet")
+    
+    def upload_dataset(self, local_path: str, dataset_name: str, 
+                      metadata: Optional[Dict[str, Any]] = None) -> str:
+        """
+        Upload dataset to cloud storage.
+        
+        Args:
+            local_path: Local path to dataset file
+            dataset_name: Name for the dataset
+            metadata: Optional metadata to store with dataset
+            
+        Returns:
+            Cloud storage URL for the dataset
+        """
+        remote_path = f"{self.config.datasets_prefix}{dataset_name}.json"
+        
+        try:
+            # Upload dataset file
+            self._upload_file(local_path, remote_path)
+            
+            # Upload metadata if provided
+            if metadata:
+                metadata_path = f"{self.config.datasets_prefix}{dataset_name}_metadata.json"
+                metadata_content = json.dumps(metadata, indent=2)
+                self._upload_content(metadata_content, metadata_path)
+            
+            dataset_url = self._get_public_url(remote_path)
+            logger.info(f"Dataset uploaded: {dataset_url}")
+            return dataset_url
+            
+        except Exception as e:
+            logger.error(f"Failed to upload dataset {dataset_name}: {e}")
+            raise
+    
+    def download_dataset(self, dataset_name: str, local_path: str) -> str:
+        """
+        Download dataset from cloud storage.
+        
+        Args:
+            dataset_name: Name of the dataset
+            local_path: Local path to save dataset
+            
+        Returns:
+            Local path to downloaded dataset
+        """
+        remote_path = f"{self.config.datasets_prefix}{dataset_name}.json"
+        
+        try:
+            os.makedirs(os.path.dirname(local_path), exist_ok=True)
+            self._download_file(remote_path, local_path)
+            
+            logger.info(f"Dataset downloaded to: {local_path}")
+            return local_path
+            
+        except Exception as e:
+            logger.error(f"Failed to download dataset {dataset_name}: {e}")
+            raise
+    
+    def upload_model(self, local_model_dir: str, model_name: str,
+                    metadata: Optional[Dict[str, Any]] = None) -> str:
+        """
+        Upload trained model to cloud storage.
+        
+        Args:
+            local_model_dir: Local directory containing model files
+            model_name: Name for the model
+            metadata: Optional model metadata
+            
+        Returns:
+            Cloud storage URL for the model
+        """
+        model_prefix = f"{self.config.models_prefix}{model_name}/"
+        
+        try:
+            # Upload all model files
+            model_files = []
+            for root, dirs, files in os.walk(local_model_dir):
+                for file in files:
+                    local_file_path = os.path.join(root, file)
+                    relative_path = os.path.relpath(local_file_path, local_model_dir)
+                    remote_path = f"{model_prefix}{relative_path}"
+                    
+                    self._upload_file(local_file_path, remote_path)
+                    model_files.append(relative_path)
+            
+            # Upload model metadata
+            if metadata is None:
+                metadata = {}
+            
+            metadata.update({
+                "model_name": model_name,
+                "files": model_files,
+                "upload_timestamp": self._get_timestamp()
+            })
+            
+            metadata_path = f"{model_prefix}model_metadata.json"
+            metadata_content = json.dumps(metadata, indent=2)
+            self._upload_content(metadata_content, metadata_path)
+            
+            model_url = self._get_public_url(model_prefix)
+            logger.info(f"Model uploaded: {model_url}")
+            return model_url
+            
+        except Exception as e:
+            logger.error(f"Failed to upload model {model_name}: {e}")
+            raise
+    
+    def download_model(self, model_name: str, local_dir: str) -> str:
+        """
+        Download model from cloud storage.
+        
+        Args:
+            model_name: Name of the model
+            local_dir: Local directory to save model
+            
+        Returns:
+            Local path to downloaded model
+        """
+        model_prefix = f"{self.config.models_prefix}{model_name}/"
+        
+        try:
+            os.makedirs(local_dir, exist_ok=True)
+            
+            # First, get model metadata to know which files to download
+            metadata_path = f"{model_prefix}model_metadata.json"
+            metadata_content = self._download_content(metadata_path)
+            metadata = json.loads(metadata_content)
+            
+            # Download all model files
+            for file_path in metadata.get("files", []):
+                remote_path = f"{model_prefix}{file_path}"
+                local_file_path = os.path.join(local_dir, file_path)
+                
+                os.makedirs(os.path.dirname(local_file_path), exist_ok=True)
+                self._download_file(remote_path, local_file_path)
+            
+            logger.info(f"Model downloaded to: {local_dir}")
+            return local_dir
+            
+        except Exception as e:
+            logger.error(f"Failed to download model {model_name}: {e}")
+            raise
+    
+    def upload_training_logs(self, local_log_dir: str, job_id: str) -> str:
+        """Upload training logs to cloud storage."""
+        logs_prefix = f"{self.config.logs_prefix}{job_id}/"
+        
+        try:
+            for root, dirs, files in os.walk(local_log_dir):
+                for file in files:
+                    local_file_path = os.path.join(root, file)
+                    relative_path = os.path.relpath(local_file_path, local_log_dir)
+                    remote_path = f"{logs_prefix}{relative_path}"
+                    
+                    self._upload_file(local_file_path, remote_path)
+            
+            logs_url = self._get_public_url(logs_prefix)
+            logger.info(f"Training logs uploaded: {logs_url}")
+            return logs_url
+            
+        except Exception as e:
+            logger.error(f"Failed to upload training logs for job {job_id}: {e}")
+            raise
+    
+    def list_datasets(self) -> List[Dict[str, Any]]:
+        """List all datasets in storage."""
+        try:
+            datasets = []
+            objects = self._list_objects(self.config.datasets_prefix)
+            
+            for obj in objects:
+                if obj.endswith('.json') and not obj.endswith('_metadata.json'):
+                    dataset_name = os.path.basename(obj).replace('.json', '')
+                    datasets.append({
+                        "name": dataset_name,
+                        "path": obj,
+                        "url": self._get_public_url(obj),
+                        "size": self._get_object_size(obj),
+                        "modified": self._get_object_modified_time(obj)
+                    })
+            
+            return datasets
+            
+        except Exception as e:
+            logger.error(f"Failed to list datasets: {e}")
+            return []
+    
+    def list_models(self) -> List[Dict[str, Any]]:
+        """List all models in storage."""
+        try:
+            models = []
+            prefixes = self._list_prefixes(self.config.models_prefix)
+            
+            for prefix in prefixes:
+                model_name = prefix.rstrip('/').split('/')[-1]
+                metadata_path = f"{prefix}model_metadata.json"
+                
+                try:
+                    metadata_content = self._download_content(metadata_path)
+                    metadata = json.loads(metadata_content)
+                    
+                    models.append({
+                        "name": model_name,
+                        "path": prefix,
+                        "url": self._get_public_url(prefix),
+                        "metadata": metadata,
+                        "files_count": len(metadata.get("files", [])),
+                        "upload_time": metadata.get("upload_timestamp", "")
+                    })
+                except:
+                    # If metadata doesn't exist, add basic info
+                    models.append({
+                        "name": model_name,
+                        "path": prefix,
+                        "url": self._get_public_url(prefix),
+                        "metadata": {},
+                        "files_count": 0,
+                        "upload_time": ""
+                    })
+            
+            return models
+            
+        except Exception as e:
+            logger.error(f"Failed to list models: {e}")
+            return []
+    
+    def _upload_file(self, local_path: str, remote_path: str) -> None:
+        """Upload file to storage."""
+        if self.config.provider == "s3":
+            self._client.upload_file(local_path, self.config.bucket_name, remote_path)
+        elif self.config.provider == "gcs":
+            bucket = self._client.bucket(self.config.bucket_name)
+            blob = bucket.blob(remote_path)
+            blob.upload_from_filename(local_path)
+        else:
+            raise NotImplementedError(f"Upload not implemented for {self.config.provider}")
+    
+    def _upload_content(self, content: str, remote_path: str) -> None:
+        """Upload string content to storage."""
+        if self.config.provider == "s3":
+            self._client.put_object(
+                Bucket=self.config.bucket_name,
+                Key=remote_path,
+                Body=content.encode('utf-8')
+            )
+        elif self.config.provider == "gcs":
+            bucket = self._client.bucket(self.config.bucket_name)
+            blob = bucket.blob(remote_path)
+            blob.upload_from_string(content)
+        else:
+            raise NotImplementedError(f"Upload not implemented for {self.config.provider}")
+    
+    def _download_file(self, remote_path: str, local_path: str) -> None:
+        """Download file from storage."""
+        if self.config.provider == "s3":
+            self._client.download_file(self.config.bucket_name, remote_path, local_path)
+        elif self.config.provider == "gcs":
+            bucket = self._client.bucket(self.config.bucket_name)
+            blob = bucket.blob(remote_path)
+            blob.download_to_filename(local_path)
+        else:
+            raise NotImplementedError(f"Download not implemented for {self.config.provider}")
+    
+    def _download_content(self, remote_path: str) -> str:
+        """Download content as string."""
+        if self.config.provider == "s3":
+            response = self._client.get_object(Bucket=self.config.bucket_name, Key=remote_path)
+            return response['Body'].read().decode('utf-8')
+        elif self.config.provider == "gcs":
+            bucket = self._client.bucket(self.config.bucket_name)
+            blob = bucket.blob(remote_path)
+            return blob.download_as_text()
+        else:
+            raise NotImplementedError(f"Download not implemented for {self.config.provider}")
+    
+    def _list_objects(self, prefix: str) -> List[str]:
+        """List objects with given prefix."""
+        if self.config.provider == "s3":
+            response = self._client.list_objects_v2(
+                Bucket=self.config.bucket_name,
+                Prefix=prefix
+            )
+            return [obj['Key'] for obj in response.get('Contents', [])]
+        elif self.config.provider == "gcs":
+            bucket = self._client.bucket(self.config.bucket_name)
+            blobs = bucket.list_blobs(prefix=prefix)
+            return [blob.name for blob in blobs]
+        else:
+            raise NotImplementedError(f"List objects not implemented for {self.config.provider}")
+    
+    def _list_prefixes(self, prefix: str) -> List[str]:
+        """List prefixes (directories) under given prefix."""
+        if self.config.provider == "s3":
+            response = self._client.list_objects_v2(
+                Bucket=self.config.bucket_name,
+                Prefix=prefix,
+                Delimiter='/'
+            )
+            return [cp['Prefix'] for cp in response.get('CommonPrefixes', [])]
+        elif self.config.provider == "gcs":
+            # GCS doesn't have true directories, so we simulate by grouping by prefix
+            bucket = self._client.bucket(self.config.bucket_name)
+            blobs = bucket.list_blobs(prefix=prefix)
+            prefixes = set()
+            for blob in blobs:
+                parts = blob.name[len(prefix):].split('/')
+                if len(parts) > 1:
+                    prefixes.add(f"{prefix}{parts[0]}/")
+            return list(prefixes)
+        else:
+            raise NotImplementedError(f"List prefixes not implemented for {self.config.provider}")
+    
+    def _get_public_url(self, remote_path: str) -> str:
+        """Get public URL for object."""
+        if self.config.provider == "s3":
+            return f"https://{self.config.bucket_name}.s3.{self.config.region}.amazonaws.com/{remote_path}"
+        elif self.config.provider == "gcs":
+            return f"https://storage.googleapis.com/{self.config.bucket_name}/{remote_path}"
+        else:
+            return f"{self.config.provider}://{self.config.bucket_name}/{remote_path}"
+    
+    def _get_object_size(self, remote_path: str) -> int:
+        """Get object size in bytes."""
+        try:
+            if self.config.provider == "s3":
+                response = self._client.head_object(Bucket=self.config.bucket_name, Key=remote_path)
+                return response['ContentLength']
+            elif self.config.provider == "gcs":
+                bucket = self._client.bucket(self.config.bucket_name)
+                blob = bucket.blob(remote_path)
+                blob.reload()
+                return blob.size
+        except:
+            return 0
+    
+    def _get_object_modified_time(self, remote_path: str) -> str:
+        """Get object last modified time."""
+        try:
+            if self.config.provider == "s3":
+                response = self._client.head_object(Bucket=self.config.bucket_name, Key=remote_path)
+                return response['LastModified'].isoformat()
+            elif self.config.provider == "gcs":
+                bucket = self._client.bucket(self.config.bucket_name)
+                blob = bucket.blob(remote_path)
+                blob.reload()
+                return blob.time_created.isoformat()
+        except:
+            return ""
+    
+    def _get_timestamp(self) -> str:
+        """Get current timestamp."""
+        from datetime import datetime
+        return datetime.utcnow().isoformat() 

isa_model/training/core/__init__.py

@@ -0,0 +1,23 @@
+"""
+Core Training Components for ISA Model SDK
+
+This module provides the core training functionality:
+- Base training classes and interfaces
+- Configuration management
+- Training utilities
+"""
+
+from .trainer import BaseTrainer, SFTTrainer
+from .config import TrainingConfig, LoRAConfig, DatasetConfig
+from .dataset import DatasetManager
+from .utils import TrainingUtils
+
+__all__ = [
+    'BaseTrainer',
+    'SFTTrainer', 
+    'TrainingConfig',
+    'LoRAConfig',
+    'DatasetConfig',
+    'DatasetManager',
+    'TrainingUtils'
+] 

isa_model/training/core/config.py

@@ -0,0 +1,181 @@
+"""
+Training Configuration Classes
+
+Defines configuration classes for different training scenarios.
+"""
+
+from dataclasses import dataclass, field
+from typing import Optional, Dict, Any, List
+from pathlib import Path
+
+
+@dataclass
+class LoRAConfig:
+    """LoRA (Low-Rank Adaptation) configuration."""
+    
+    use_lora: bool = True
+    lora_rank: int = 8
+    lora_alpha: int = 16
+    lora_dropout: float = 0.05
+    lora_target_modules: Optional[List[str]] = None
+    
+    def __post_init__(self):
+        if self.lora_target_modules is None:
+            # Default target modules for most transformer models
+            self.lora_target_modules = ["q_proj", "v_proj", "k_proj", "o_proj"]
+
+
+@dataclass
+class DatasetConfig:
+    """Dataset configuration."""
+    
+    dataset_path: str
+    dataset_format: str = "alpaca"  # alpaca, sharegpt, custom
+    max_length: int = 1024
+    validation_split: float = 0.1
+    preprocessing_num_workers: int = 4
+    
+    def __post_init__(self):
+        if not Path(self.dataset_path).exists() and not self.dataset_path.startswith("http"):
+            # Assume it's a HuggingFace dataset name
+            pass
+
+
+@dataclass
+class TrainingConfig:
+    """Main training configuration."""
+    
+    # Model configuration
+    model_name: str
+    output_dir: str
+    
+    # Training hyperparameters
+    num_epochs: int = 3
+    batch_size: int = 4
+    learning_rate: float = 2e-5
+    warmup_steps: int = 100
+    weight_decay: float = 0.01
+    gradient_accumulation_steps: int = 1
+    max_grad_norm: float = 1.0
+    
+    # Training strategy
+    training_type: str = "sft"  # sft, dpo, rlhf
+    fp16: bool = True
+    bf16: bool = False
+    gradient_checkpointing: bool = True
+    
+    # Saving and logging
+    save_steps: int = 500
+    logging_steps: int = 10
+    eval_steps: int = 500
+    save_total_limit: int = 3
+    
+    # LoRA configuration
+    lora_config: Optional[LoRAConfig] = field(default_factory=LoRAConfig)
+    
+    # Dataset configuration
+    dataset_config: Optional[DatasetConfig] = None
+    
+    # Additional parameters
+    extra_params: Dict[str, Any] = field(default_factory=dict)
+    
+    def __post_init__(self):
+        # Create output directory if it doesn't exist
+        Path(self.output_dir).mkdir(parents=True, exist_ok=True)
+        
+        # Set BF16 for newer GPUs, FP16 for older ones
+        if self.bf16:
+            self.fp16 = False
+    
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert config to dictionary."""
+        config_dict = {}
+        
+        for key, value in self.__dict__.items():
+            if key in ['lora_config', 'dataset_config']:
+                if value is not None:
+                    config_dict[key] = value.__dict__
+                else:
+                    config_dict[key] = None
+            else:
+                config_dict[key] = value
+        
+        return config_dict
+    
+    @classmethod
+    def from_dict(cls, config_dict: Dict[str, Any]) -> 'TrainingConfig':
+        """Create config from dictionary."""
+        # Handle nested configs
+        if 'lora_config' in config_dict and config_dict['lora_config'] is not None:
+            config_dict['lora_config'] = LoRAConfig(**config_dict['lora_config'])
+        
+        if 'dataset_config' in config_dict and config_dict['dataset_config'] is not None:
+            config_dict['dataset_config'] = DatasetConfig(**config_dict['dataset_config'])
+        
+        return cls(**config_dict)
+
+
+@dataclass
+class RunPodConfig:
+    """RunPod cloud training configuration."""
+    
+    api_key: str
+    template_id: str
+    gpu_type: str = "NVIDIA RTX A6000"
+    gpu_count: int = 1
+    container_disk_in_gb: int = 50
+    volume_in_gb: int = 100
+    max_runtime_hours: int = 24
+    
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dictionary."""
+        return self.__dict__.copy()
+
+
+@dataclass
+class StorageConfig:
+    """Cloud storage configuration."""
+    
+    provider: str  # s3, gcs, local
+    bucket_name: Optional[str] = None
+    region: Optional[str] = None
+    access_key: Optional[str] = None
+    secret_key: Optional[str] = None
+    service_account_path: Optional[str] = None
+    
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dictionary."""
+        return self.__dict__.copy()
+
+
+@dataclass 
+class JobConfig:
+    """Training job configuration for cloud training."""
+    
+    model_name: str
+    dataset_source: str
+    job_name: Optional[str] = None
+    description: Optional[str] = None
+    
+    # Training parameters
+    training_type: str = "sft"
+    num_epochs: int = 3
+    batch_size: int = 4
+    learning_rate: float = 2e-5
+    max_length: int = 1024
+    
+    # LoRA parameters
+    use_lora: bool = True
+    lora_rank: int = 8
+    lora_alpha: int = 16
+    lora_dropout: float = 0.05
+    
+    # Storage parameters
+    save_model_to_storage: bool = True
+    model_name_in_storage: Optional[str] = None
+    upload_to_hf: bool = False
+    hf_model_name: Optional[str] = None
+    
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dictionary."""
+        return self.__dict__.copy()