Spaces:
Sleeping
Sleeping
File size: 9,292 Bytes
0dbcacb a9f0e5a | 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 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 | """Utilities for connecting to and interacting with Hopsworks Feature Store."""
import os
import logging
from typing import Optional
import pandas as pd
import hopsworks
from dotenv import load_dotenv
logger = logging.getLogger(__name__)
# Load environment variables
load_dotenv()
def connect_to_hopsworks(api_key: Optional[str] = None, project_name: Optional[str] = None):
"""
Connect to Hopsworks.
Args:
api_key: Hopsworks API key. If None, reads from HOPSWORKS_API_KEY env variable.
project_name: Hopsworks project name. If None, reads from HOPSWORKS_PROJECT_NAME env variable.
Returns:
Hopsworks project object
Raises:
ValueError: If API key or project name is not provided
"""
# Get API key from parameter or environment
api_key = api_key or os.getenv("HOPSWORKS_API_KEY")
if not api_key:
raise ValueError(
"Hopsworks API key not provided. "
"Set HOPSWORKS_API_KEY environment variable or pass api_key parameter."
)
# Get project name from parameter or environment
project_name = project_name or os.getenv("HOPSWORKS_PROJECT")
logger.info(f"Connecting to Hopsworks project: {project_name or 'default'}")
# Login to Hopsworks
try:
if project_name:
project = hopsworks.login(
api_key_value=api_key,
project=project_name,
engine="python" # Use Python engine (serverless, no cert download)
)
else:
project = hopsworks.login(
api_key_value=api_key,
engine="python"
)
except Exception as e:
logger.error(f"Failed to connect to Hopsworks: {e}")
logger.info("Trying to connect without specifying project...")
project = hopsworks.login(
api_key_value=api_key,
engine="python"
)
logger.info(f"Successfully connected to Hopsworks project: {project.name}")
return project
def get_or_create_feature_group(
project,
name: str,
version: int = 1,
):
"""
Get existing feature group or create new one if it doesn't exist.
Args:
project: Hopsworks project object
name: Feature group name
version: Feature group version
description: Description of the feature group
primary_key: List of column names to use as primary key
event_time: Column name to use as event time
online_enabled: Whether to enable online feature serving
Returns:
Feature group object
"""
fs = project.get_feature_store()
try:
# Try to get existing feature group
fg = fs.get_feature_group(name=name, version=version)
logger.info(f"Retrieved existing feature group: {name} (version {version})")
return fg
except Exception:
# Feature group doesn't exist, will need to create it
logger.info(f"Feature group {name} (version {version}) not found, will create on first insert")
return None
def upload_dataframe_to_feature_group(
project,
df: pd.DataFrame,
feature_group_name: str,
version: int = 1,
description: str = "",
primary_key: list = None,
event_time: Optional[str] = None,
online_enabled: bool = False,
write_options: dict = None
):
"""
Upload a DataFrame to a Hopsworks feature group.
Args:
project: Hopsworks project object
df: Pandas DataFrame to upload
feature_group_name: Name of the feature group
version: Feature group version
description: Description of the feature group
primary_key: List of column names to use as primary key
event_time: Column name to use as event time
online_enabled: Whether to enable online feature serving
write_options: Additional write options (e.g., {"wait_for_job": False})
Returns:
Feature group object
"""
fs = project.get_feature_store()
logger.info(f"Uploading DataFrame to feature group: {feature_group_name} (version {version})")
logger.info(f"DataFrame shape: {df.shape}")
# Create or get feature group
fg = fs.get_or_create_feature_group(
name=feature_group_name,
version=version,
description=description,
primary_key=primary_key or [],
event_time=event_time,
online_enabled=online_enabled
)
# Insert data
write_options = write_options or {"wait_for_job": True}
fg.insert(df, write_options=write_options)
logger.info(f"Successfully uploaded {len(df)} rows to {feature_group_name}")
return fg
def read_feature_group(
project,
feature_group_name: str,
version: int = 1,
online: bool = False
) -> pd.DataFrame:
"""
Read data from a Hopsworks feature group.
Args:
project: Hopsworks project object
feature_group_name: Name of the feature group
version: Feature group version
online: Whether to read from online feature store
Returns:
Pandas DataFrame with feature group data
"""
fs = project.get_feature_store()
logger.info(f"Reading feature group: {feature_group_name} (version {version})")
fg = fs.get_feature_group(name=feature_group_name, version=version)
if online:
df = fg.read(online=True)
else:
df = fg.read()
logger.info(f"Read {len(df)} rows from {feature_group_name}")
return df
def create_feature_view(
project,
name: str,
version: int = 1,
description: str = "",
query=None,
labels: list = None
):
"""
Create a feature view for training datasets.
Args:
project: Hopsworks project object
name: Feature view name
version: Feature view version
description: Description of the feature view
query: Query object to define feature selection
labels: List of label column names
Returns:
Feature view object
"""
fs = project.get_feature_store()
logger.info(f"Creating feature view: {name} (version {version})")
fv = fs.create_feature_view(
name=name,
version=version,
description=description,
query=query,
labels=labels or []
)
logger.info(f"Successfully created feature view: {name}")
return fv
def save_model_to_registry(
project,
model,
model_name: str,
metrics: dict = None,
description: str = "",
model_schema: dict = None,
scaler=None,
feature_names: list = None
):
"""
Save a trained model to Hopsworks Model Registry with all artifacts.
Args:
project: Hopsworks project object
model: Trained model object
model_name: Name for the model in registry
metrics: Dictionary of model metrics
description: Model description
model_schema: Optional model schema
scaler: Optional scaler object to save with model
feature_names: Optional list of feature names
Returns:
Model registry object
"""
import joblib
import os
import tempfile
logger.info(f"Saving model to Hopsworks Model Registry: {model_name}")
# Get model registry
mr = project.get_model_registry()
# Create temporary directory for model artifacts
with tempfile.TemporaryDirectory() as tmpdir:
# Save model using joblib
model_path = os.path.join(tmpdir, "model.pkl")
joblib.dump(model, model_path)
logger.info(f"Model saved to temporary path: {model_path}")
# Save scaler if provided
if scaler is not None:
scaler_path = os.path.join(tmpdir, "scaler.pkl")
joblib.dump(scaler, scaler_path)
logger.info(f"Scaler saved to temporary path: {scaler_path}")
# Save feature names if provided
if feature_names is not None:
feature_names_path = os.path.join(tmpdir, "feature_names.txt")
with open(feature_names_path, 'w') as f:
f.write('\n'.join(feature_names))
logger.info(f"Feature names saved to temporary path: {feature_names_path}")
# Save hyperparameters from description if they exist
if description and "Best Parameters:" in description:
import json
# Try to extract and save params as JSON for easier loading
params_path = os.path.join(tmpdir, "hyperparameters.txt")
with open(params_path, 'w') as f:
f.write(description)
logger.info(f"Hyperparameters saved to temporary path: {params_path}")
# Create model in registry
try:
model_registry = mr.python.create_model(
name=model_name,
metrics=metrics or {},
description=description,
input_example=None,
model_schema=model_schema
)
# Save all model artifacts (model, scaler, feature_names)
model_registry.save(tmpdir)
logger.info(f"Successfully saved model '{model_name}' with all artifacts to Model Registry")
return model_registry
except Exception as e:
logger.error(f"Failed to save model to registry: {e}")
raise
|