ds-agent-cli 0.1.0

package/src/storage/helpers.py

@@ -0,0 +1,116 @@
+"""
+Helper utilities for artifact storage integration
+"""
+
+import os
+import tempfile
+import joblib
+from typing import Any, Dict, Optional
+from pathlib import Path
+
+
+def save_model_with_store(
+    model_data: Any,
+    filename: str,
+    metadata: Optional[Dict[str, Any]] = None,
+    fallback_dir: str = "./outputs/models"
+) -> str:
+    """
+    Save model using artifact store if available, otherwise use fallback path.
+    
+    Args:
+        model_data: Model object or dict to save
+        filename: Name of the model file (e.g., "model.pkl")
+        metadata: Optional metadata to attach
+        fallback_dir: Directory to use if artifact store unavailable
+        
+    Returns:
+        Path where model was saved
+    """
+    try:
+        from storage import get_artifact_store
+        store = get_artifact_store()
+        
+        # Save to temp file first
+        with tempfile.NamedTemporaryFile(mode='wb', suffix='.pkl', delete=False) as tmp:
+            joblib.dump(model_data, tmp.name)
+            model_path = store.save_model(tmp.name, metadata=metadata)
+        os.unlink(tmp.name)
+        
+        return model_path
+        
+    except ImportError:
+        # Fallback to local path
+        model_path = os.path.join(fallback_dir, filename)
+        Path(model_path).parent.mkdir(parents=True, exist_ok=True)
+        joblib.dump(model_data, model_path)
+        return model_path
+
+
+def save_plot_with_store(
+    plot_path: str,
+    metadata: Optional[Dict[str, Any]] = None
+) -> str:
+    """
+    Save plot using artifact store if available.
+    
+    Args:
+        plot_path: Path to existing plot file
+        metadata: Optional metadata to attach
+        
+    Returns:
+        Path where plot was saved
+    """
+    try:
+        from storage import get_artifact_store
+        store = get_artifact_store()
+        return store.save_plot(plot_path, metadata=metadata)
+    except ImportError:
+        # Already saved locally
+        return plot_path
+
+
+def save_report_with_store(
+    report_path: str,
+    metadata: Optional[Dict[str, Any]] = None
+) -> str:
+    """
+    Save report using artifact store if available.
+    
+    Args:
+        report_path: Path to existing report file
+        metadata: Optional metadata to attach
+        
+    Returns:
+        Path where report was saved
+    """
+    try:
+        from storage import get_artifact_store
+        store = get_artifact_store()
+        return store.save_report(report_path, metadata=metadata)
+    except ImportError:
+        # Already saved locally
+        return report_path
+
+
+def save_data_with_store(
+    data_path: str,
+    metadata: Optional[Dict[str, Any]] = None
+) -> str:
+    """
+    Save data file using artifact store if available.
+    
+    Args:
+        data_path: Path to existing data file
+        metadata: Optional metadata to attach
+        
+    Returns:
+        Path where data was saved
+    """
+    try:
+        from storage import get_artifact_store
+        store = get_artifact_store()
+        return store.save_data(data_path, metadata=metadata)
+    except ImportError:
+        # Already saved locally
+        return data_path