File size: 3,114 Bytes
226ac39
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
"""
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