isa-model 0.3.9__py3-none-any.whl → 0.4.0__py3-none-any.whl

isa_model/scripts/inference_tracker.py

@@ -1,283 +0,0 @@
-"""
-MLflow tracker for inference workflows.
-"""
-
-import os
-import json
-import time
-import logging
-from typing import Dict, List, Optional, Any, Union
-from contextlib import contextmanager
-
-from .mlflow_manager import MLflowManager, ExperimentType
-from .model_registry import ModelRegistry, ModelStage, ModelVersion
-
-
-logger = logging.getLogger(__name__)
-
-
-class InferenceTracker:
-    """
-    Tracker for model inference workflows.
-    
-    This class provides utilities to track model inference using MLflow,
-    including performance metrics and input/output logging.
-    
-    Example:
-        ```python
-        # Initialize tracker
-        tracker = InferenceTracker(
-            tracking_uri="http://localhost:5000"
-        )
-        
-        # Get model from registry
-        model_version = tracker.get_production_model("llama-7b")
-        
-        # Track inference
-        with tracker.track_inference(
-            model_name="llama-7b",
-            model_version=model_version.version
-        ):
-            # Start timer
-            start_time = time.time()
-            
-            # Generate text
-            output = model.generate(prompt)
-            
-            # Log inference
-            tracker.log_inference(
-                input=prompt,
-                output=output,
-                latency_ms=(time.time() - start_time) * 1000
-            )
-        ```
-    """
-    
-    def __init__(
-        self,
-        tracking_uri: Optional[str] = None,
-        artifact_uri: Optional[str] = None,
-        registry_uri: Optional[str] = None
-    ):
-        """
-        Initialize the inference tracker.
-        
-        Args:
-            tracking_uri: URI for MLflow tracking server
-            artifact_uri: URI for MLflow artifacts
-            registry_uri: URI for MLflow model registry
-        """
-        self.mlflow_manager = MLflowManager(
-            tracking_uri=tracking_uri,
-            artifact_uri=artifact_uri,
-            registry_uri=registry_uri
-        )
-        self.model_registry = ModelRegistry(
-            tracking_uri=tracking_uri,
-            registry_uri=registry_uri
-        )
-        self.current_run_info = {}
-        self.inference_samples = []
-        
-    def get_production_model(self, model_name: str) -> Optional[ModelVersion]:
-        """
-        Get the production version of a model.
-        
-        Args:
-            model_name: Name of the model
-            
-        Returns:
-            Production ModelVersion or None if not found
-        """
-        return self.model_registry.get_latest_model_version(
-            name=model_name,
-            stage=ModelStage.PRODUCTION
-        )
-        
-    def get_staging_model(self, model_name: str) -> Optional[ModelVersion]:
-        """
-        Get the staging version of a model.
-        
-        Args:
-            model_name: Name of the model
-            
-        Returns:
-            Staging ModelVersion or None if not found
-        """
-        return self.model_registry.get_latest_model_version(
-            name=model_name,
-            stage=ModelStage.STAGING
-        )
-        
-    @contextmanager
-    def track_inference(
-        self,
-        model_name: str,
-        model_version: Optional[str] = None,
-        batch_size: Optional[int] = None,
-        tags: Optional[Dict[str, str]] = None
-    ):
-        """
-        Track model inference with MLflow.
-        
-        Args:
-            model_name: Name of the model
-            model_version: Version of the model
-            batch_size: Batch size for inference
-            tags: Tags for the run
-            
-        Yields:
-            Dictionary with run information
-        """
-        run_info = {
-            "model_name": model_name,
-            "model_version": model_version,
-            "batch_size": batch_size,
-            "start_time": time.time(),
-            "metrics": {}
-        }
-        
-        # Prepare tags
-        if tags is None:
-            tags = {}
-            
-        tags["model_name"] = model_name
-        if model_version:
-            tags["model_version"] = model_version
-            
-        if batch_size:
-            tags["batch_size"] = str(batch_size)
-            
-        # Start the MLflow run
-        with self.mlflow_manager.start_run(
-            experiment_type=ExperimentType.INFERENCE,
-            model_name=model_name,
-            tags=tags
-        ) as run:
-            run_info["run_id"] = run.info.run_id
-            run_info["experiment_id"] = run.info.experiment_id
-            
-            # Reset inference samples
-            self.inference_samples = []
-            
-            self.current_run_info = run_info
-            try:
-                yield run_info
-                
-                # Calculate and log summary metrics
-                self._log_summary_metrics()
-                
-                # Save inference samples
-                if self.inference_samples:
-                    self._save_inference_samples()
-                    
-            finally:
-                run_info["end_time"] = time.time()
-                run_info["duration"] = run_info["end_time"] - run_info["start_time"]
-                
-                # Log duration
-                self.mlflow_manager.log_metrics({
-                    "duration_seconds": run_info["duration"]
-                })
-                
-                self.current_run_info = {}
-                
-    def log_inference(
-        self,
-        input: str,
-        output: str,
-        latency_ms: Optional[float] = None,
-        token_count: Optional[int] = None,
-        tokens_per_second: Optional[float] = None,
-        metadata: Optional[Dict[str, Any]] = None
-    ) -> None:
-        """
-        Log an inference sample.
-        
-        Args:
-            input: Input prompt
-            output: Generated output
-            latency_ms: Latency in milliseconds
-            token_count: Number of tokens generated
-            tokens_per_second: Tokens per second
-            metadata: Additional metadata
-        """
-        if not self.current_run_info:
-            logger.warning("No active run. Inference will not be logged.")
-            return
-            
-        sample = {
-            "input": input,
-            "output": output,
-            "timestamp": time.time()
-        }
-        
-        if latency_ms is not None:
-            sample["latency_ms"] = latency_ms
-            
-        if token_count is not None:
-            sample["token_count"] = token_count
-            
-        if tokens_per_second is not None:
-            sample["tokens_per_second"] = tokens_per_second
-            
-        if metadata:
-            sample["metadata"] = metadata
-            
-        self.inference_samples.append(sample)
-        
-        # Log individual metrics
-        metrics = {}
-        if latency_ms is not None:
-            metrics["latency_ms"] = latency_ms
-            
-        if token_count is not None:
-            metrics["token_count"] = token_count
-            
-        if tokens_per_second is not None:
-            metrics["tokens_per_second"] = tokens_per_second
-            
-        if metrics:
-            self.mlflow_manager.log_metrics(metrics)
-    
-    def _log_summary_metrics(self) -> None:
-        """Log summary metrics based on all inference samples."""
-        if not self.inference_samples:
-            return
-            
-        latencies = [s.get("latency_ms") for s in self.inference_samples if "latency_ms" in s]
-        token_counts = [s.get("token_count") for s in self.inference_samples if "token_count" in s]
-        tokens_per_second = [s.get("tokens_per_second") for s in self.inference_samples if "tokens_per_second" in s]
-        
-        metrics = {
-            "inference_count": len(self.inference_samples)
-        }
-        
-        if latencies:
-            metrics["avg_latency_ms"] = sum(latencies) / len(latencies)
-            metrics["min_latency_ms"] = min(latencies)
-            metrics["max_latency_ms"] = max(latencies)
-            
-        if token_counts:
-            metrics["avg_token_count"] = sum(token_counts) / len(token_counts)
-            metrics["total_tokens"] = sum(token_counts)
-            
-        if tokens_per_second:
-            metrics["avg_tokens_per_second"] = sum(tokens_per_second) / len(tokens_per_second)
-            
-        self.mlflow_manager.log_metrics(metrics)
-        
-    def _save_inference_samples(self) -> None:
-        """Save inference samples as an artifact."""
-        import tempfile
-        
-        with tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.json') as f:
-            json.dump(self.inference_samples, f, indent=2)
-            temp_path = f.name
-            
-        self.mlflow_manager.log_artifact(temp_path, "inference_samples.json")
-        
-        try:
-            os.remove(temp_path)
-        except:
-            pass 

isa_model/scripts/mlflow_manager.py

@@ -1,379 +0,0 @@
-"""
-MLflow manager for experiment tracking and model management.
-"""
-
-import os
-import logging
-from enum import Enum
-from typing import Dict, List, Optional, Any, Union
-import mlflow
-from mlflow.tracking import MlflowClient
-
-logger = logging.getLogger(__name__)
-
-
-class ExperimentType(str, Enum):
-    """Types of experiments that can be tracked."""
-    
-    TRAINING = "training"
-    FINETUNING = "finetuning"
-    REINFORCEMENT_LEARNING = "rl"
-    INFERENCE = "inference"
-    EVALUATION = "evaluation"
-
-
-class MLflowManager:
-    """
-    Manager class for MLflow operations.
-    
-    This class provides methods to set up MLflow, track experiments,
-    log metrics, and manage models.
-    
-    Example:
-        ```python
-        # Initialize MLflow manager
-        mlflow_manager = MLflowManager(
-            tracking_uri="http://localhost:5000",
-            artifact_uri="s3://bucket/artifacts"
-        )
-        
-        # Set up experiment and start run
-        with mlflow_manager.start_run(
-            experiment_type=ExperimentType.FINETUNING,
-            model_name="llama-7b"
-        ) as run:
-            # Log parameters
-            mlflow_manager.log_params({
-                "learning_rate": 2e-5,
-                "batch_size": 8
-            })
-            
-            # Train model...
-            
-            # Log metrics
-            mlflow_manager.log_metrics({
-                "accuracy": 0.95,
-                "loss": 0.02
-            })
-            
-            # Log model
-            mlflow_manager.log_model(
-                model_path="/path/to/model",
-                name="finetuned-llama-7b"
-            )
-        ```
-    """
-    
-    def __init__(
-        self,
-        tracking_uri: Optional[str] = None,
-        artifact_uri: Optional[str] = None,
-        registry_uri: Optional[str] = None
-    ):
-        """
-        Initialize the MLflow manager.
-        
-        Args:
-            tracking_uri: URI for MLflow tracking server
-            artifact_uri: URI for MLflow artifacts
-            registry_uri: URI for MLflow model registry
-        """
-        self.tracking_uri = tracking_uri or os.environ.get("MLFLOW_TRACKING_URI", "")
-        self.artifact_uri = artifact_uri or os.environ.get("MLFLOW_ARTIFACT_URI", "")
-        self.registry_uri = registry_uri or os.environ.get("MLFLOW_REGISTRY_URI", "")
-        
-        self._setup_mlflow()
-        self.client = MlflowClient(tracking_uri=self.tracking_uri, registry_uri=self.registry_uri)
-        self.active_run = None
-        
-    def _setup_mlflow(self) -> None:
-        """Set up MLflow configuration."""
-        if self.tracking_uri:
-            mlflow.set_tracking_uri(self.tracking_uri)
-            logger.info(f"Set MLflow tracking URI to {self.tracking_uri}")
-            
-        if self.registry_uri:
-            mlflow.set_registry_uri(self.registry_uri)
-            logger.info(f"Set MLflow registry URI to {self.registry_uri}")
-    
-    def create_experiment(
-        self,
-        experiment_type: ExperimentType,
-        model_name: str,
-        tags: Optional[Dict[str, str]] = None
-    ) -> str:
-        """
-        Create a new experiment if it doesn't exist.
-        
-        Args:
-            experiment_type: Type of experiment
-            model_name: Name of the model
-            tags: Tags for the experiment
-            
-        Returns:
-            ID of the experiment
-        """
-        experiment_name = f"{model_name}_{experiment_type.value}"
-        
-        # Get experiment if exists, create if not
-        experiment = mlflow.get_experiment_by_name(experiment_name)
-        if experiment is None:
-            experiment_id = mlflow.create_experiment(
-                name=experiment_name,
-                artifact_location=self.artifact_uri if self.artifact_uri else None,
-                tags=tags
-            )
-            logger.info(f"Created new experiment: {experiment_name} (ID: {experiment_id})")
-        else:
-            experiment_id = experiment.experiment_id
-            logger.info(f"Using existing experiment: {experiment_name} (ID: {experiment_id})")
-            
-        return experiment_id
-    
-    def start_run(
-        self,
-        experiment_type: ExperimentType,
-        model_name: str,
-        run_name: Optional[str] = None,
-        tags: Optional[Dict[str, str]] = None,
-        nested: bool = False
-    ) -> mlflow.ActiveRun:
-        """
-        Start a new MLflow run.
-        
-        Args:
-            experiment_type: Type of experiment
-            model_name: Name of the model
-            run_name: Name for the run
-            tags: Tags for the run
-            nested: Whether this is a nested run
-            
-        Returns:
-            MLflow active run context
-        """
-        experiment_id = self.create_experiment(experiment_type, model_name)
-        
-        if not run_name:
-            import datetime
-            timestamp = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
-            run_name = f"{model_name}_{experiment_type.value}_{timestamp}"
-            
-        self.active_run = mlflow.start_run(
-            experiment_id=experiment_id,
-            run_name=run_name,
-            tags=tags,
-            nested=nested
-        )
-        
-        logger.info(f"Started MLflow run: {run_name} (ID: {self.active_run.info.run_id})")
-        return self.active_run
-    
-    def end_run(self) -> None:
-        """End the current MLflow run."""
-        if mlflow.active_run():
-            run_id = mlflow.active_run().info.run_id
-            mlflow.end_run()
-            logger.info(f"Ended MLflow run: {run_id}")
-            self.active_run = None
-    
-    def log_params(self, params: Dict[str, Any]) -> None:
-        """
-        Log parameters to the current run.
-        
-        Args:
-            params: Dictionary of parameters to log
-        """
-        if not mlflow.active_run():
-            logger.warning("No active run. Parameters will not be logged.")
-            return
-            
-        mlflow.log_params(params)
-        logger.debug(f"Logged parameters: {params}")
-    
-    def log_metrics(self, metrics: Dict[str, float], step: Optional[int] = None) -> None:
-        """
-        Log metrics to the current run.
-        
-        Args:
-            metrics: Dictionary of metrics to log
-            step: Step value for the metrics
-        """
-        if not mlflow.active_run():
-            logger.warning("No active run. Metrics will not be logged.")
-            return
-            
-        mlflow.log_metrics(metrics, step=step)
-        logger.debug(f"Logged metrics: {metrics}")
-    
-    def log_model(
-        self,
-        model_path: str,
-        name: str,
-        flavor: str = "pyfunc",
-        **kwargs
-    ) -> str:
-        """
-        Log a model to MLflow.
-        
-        Args:
-            model_path: Path to the model
-            name: Name for the logged model
-            flavor: MLflow model flavor
-            **kwargs: Additional arguments for model logging
-            
-        Returns:
-            Path where the model is logged
-        """
-        if not mlflow.active_run():
-            logger.warning("No active run. Model will not be logged.")
-            return ""
-            
-        log_func = getattr(mlflow, f"log_{flavor}")
-        if not log_func:
-            logger.warning(f"Unsupported model flavor: {flavor}. Using pyfunc instead.")
-            log_func = mlflow.pyfunc.log_model
-            
-        artifact_path = f"models/{name}"
-        logged_model = log_func(
-            artifact_path=artifact_path,
-            path=model_path,
-            **kwargs
-        )
-        
-        logger.info(f"Logged model: {name} at {artifact_path}")
-        return artifact_path
-    
-    def log_artifact(self, local_path: str, artifact_path: Optional[str] = None) -> None:
-        """
-        Log an artifact to MLflow.
-        
-        Args:
-            local_path: Local path to the artifact
-            artifact_path: Path for the artifact in MLflow
-        """
-        if not mlflow.active_run():
-            logger.warning("No active run. Artifact will not be logged.")
-            return
-            
-        mlflow.log_artifact(local_path, artifact_path)
-        logger.debug(f"Logged artifact: {local_path} to {artifact_path or 'root'}")
-    
-    def log_artifacts(self, local_dir: str, artifact_path: Optional[str] = None) -> None:
-        """
-        Log multiple artifacts to MLflow.
-        
-        Args:
-            local_dir: Local directory containing artifacts
-            artifact_path: Path for the artifacts in MLflow
-        """
-        if not mlflow.active_run():
-            logger.warning("No active run. Artifacts will not be logged.")
-            return
-            
-        mlflow.log_artifacts(local_dir, artifact_path)
-        logger.debug(f"Logged artifacts from directory: {local_dir} to {artifact_path or 'root'}")
-    
-    def get_run(self, run_id: str) -> Optional[mlflow.entities.Run]:
-        """
-        Get a run by ID.
-        
-        Args:
-            run_id: ID of the run
-            
-        Returns:
-            MLflow run entity or None if not found
-        """
-        try:
-            return self.client.get_run(run_id)
-        except mlflow.exceptions.MlflowException as e:
-            logger.error(f"Failed to get run {run_id}: {e}")
-            return None
-    
-    def search_runs(
-        self,
-        experiment_ids: List[str],
-        filter_string: Optional[str] = None,
-        max_results: int = 100
-    ) -> List[mlflow.entities.Run]:
-        """
-        Search for runs in the given experiments.
-        
-        Args:
-            experiment_ids: List of experiment IDs
-            filter_string: Filter string for the search
-            max_results: Maximum number of results to return
-            
-        Returns:
-            List of MLflow run entities
-        """
-        try:
-            return self.client.search_runs(
-                experiment_ids=experiment_ids,
-                filter_string=filter_string,
-                max_results=max_results
-            )
-        except mlflow.exceptions.MlflowException as e:
-            logger.error(f"Failed to search runs: {e}")
-            return []
-            
-    def get_experiment_id_by_name(self, experiment_name: str) -> Optional[str]:
-        """
-        Get experiment ID by name.
-        
-        Args:
-            experiment_name: Name of the experiment
-            
-        Returns:
-            Experiment ID or None if not found
-        """
-        experiment = mlflow.get_experiment_by_name(experiment_name)
-        if experiment:
-            return experiment.experiment_id
-        return None
-        
-    def set_tracking_tag(self, key: str, value: str) -> None:
-        """
-        Set a tag for the current run.
-        
-        Args:
-            key: Tag key
-            value: Tag value
-        """
-        if not mlflow.active_run():
-            logger.warning("No active run. Tag will not be set.")
-            return
-            
-        mlflow.set_tag(key, value)
-        logger.debug(f"Set tag: {key}={value}")
-        
-    def create_model_version(
-        self,
-        name: str,
-        source: str,
-        description: Optional[str] = None,
-        tags: Optional[Dict[str, str]] = None
-    ) -> Optional[str]:
-        """
-        Create a new model version in the registry.
-        
-        Args:
-            name: Name of the registered model
-            source: Source path of the model
-            description: Description for the model version
-            tags: Tags for the model version
-            
-        Returns:
-            Version of the created model or None if creation failed
-        """
-        try:
-            version = self.client.create_model_version(
-                name=name,
-                source=source,
-                description=description,
-                tags=tags
-            )
-            logger.info(f"Created model version: {name} v{version.version}")
-            return version.version
-        except mlflow.exceptions.MlflowException as e:
-            logger.error(f"Failed to create model version: {e}")
-            return None