isa-model 0.0.1__py3-none-any.whl

isa_model/scripts/inference_tracker.py

@@ -0,0 +1,283 @@
+"""
+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

@@ -0,0 +1,379 @@
+"""
+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