openstef-models 4.2.2__tar.gz → 4.2.3__tar.gz

{openstef_models-4.2.2 → openstef_models-4.2.3}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.4
+Metadata-Version: 2.5
 Name: openstef-models
-Version: 4.2.2
+Version: 4.2.3
 Summary: Core models for OpenSTEF
 Project-URL: Documentation, https://openstef.github.io/openstef/index.html
 Project-URL: Homepage, https://lfenergy.org/projects/openstef/

{openstef_models-4.2.2 → openstef_models-4.2.3}/pyproject.toml

@@ -7,7 +7,7 @@ requires = [ "hatchling" ]
 
 [project]
 name = "openstef-models"
-version = "4.2.2"
+version = "4.2.3"
 description = "Core models for OpenSTEF"
 readme = "README.md"
 keywords = [ "energy", "forecasting", "machinelearning" ]

{openstef_models-4.2.2 → openstef_models-4.2.3}/src/openstef_models/integrations/mlflow/mlflow_storage.py

@@ -22,7 +22,7 @@ from urllib.parse import urlparse
 from mlflow import MlflowClient
 from mlflow.entities import Metric, Param, Run
 from mlflow.exceptions import MlflowException
-from pydantic import Field, PrivateAttr
+from pydantic import ConfigDict, Field, PrivateAttr
 
 from openstef_core.base_model import BaseConfig
 from openstef_core.exceptions import ModelNotFoundError
@@ -31,7 +31,7 @@ from openstef_models.integrations.joblib import JoblibModelSerializer
 from openstef_models.mixins import ModelIdentifier, ModelSerializer
 
 
-def normalize_tracking_uri(uri: str) -> str:
+def normalize_tracking_uri(uri: str, exceptions: Sequence[str] = []) -> str:
     r"""Normalize a tracking URI to a file:/// URI when it refers to a local path.
 
     MLflow's model registry rejects bare Windows paths (e.g. ``D:\mlflow``) because
@@ -44,10 +44,13 @@ def normalize_tracking_uri(uri: str) -> str:
 
     Args:
         uri: Raw tracking URI string, may be a path or a proper URI.
+        exceptions: List of URI to treat as urls even if they have no scheme
 
     Returns:
         A ``file:///`` URI for local paths, or the original URI for remote schemes.
     """
+    if uri in exceptions:
+        return uri
     scheme = urlparse(uri).scheme
     # Empty scheme → relative/absolute POSIX path.
     # Single-letter scheme → Windows drive letter (e.g. "D" from "D:\\...").
@@ -63,9 +66,33 @@ class MLFlowStorage(BaseConfig):
     Handles creation, storage, and retrieval of MLflow runs including models,
     training data, metrics, and hyperparameters. Organizes artifacts locally
     before uploading to MLflow tracking server.
+
+    ``tracking_uri`` configures the Tracking backend that stores experiments and
+    runs. ``registry_uri`` optionally configures a separate Model Registry
+    backend; if it is omitted, MLflow uses the registry associated with the
+    tracking backend. ``artifact_location`` configures the artifact root for
+    experiments created by this instance.
+
+    Typical configurations:
+
+    tracking_uri="sqlite:///mlflow.db", # Local SQLite tracking backend with separate local artifacts
+    tracking_uri="https://mlflow.example.com", # Remote MLflow Tracking Server with S3 artifacts
+    tracking_uri="databricks", # Databricks MLflow Tracking Server with DBFS artifacts
+
+    A caller may provide ``client`` to use an already configured
+    :class:`mlflow.MlflowClient`. When supplied, it takes precedence over
+    ``tracking_uri`` and ``registry_uri`` and is excluded from serialized model
+    configuration.
     """
 
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
     tracking_uri: str = Field(default="./mlflow", description="MLflow tracking server URI.")
+    registry_uri: str | None = Field(default=None, description="MLflow tracking server URI.")
+
+    # Excluded from model_dump(), JSON serialization, and config persistence.
+    client: MlflowClient | None = Field(default=None, exclude=True, repr=False)
+
     local_artifacts_path: Path = Field(
         default=Path("./mlflow_artifacts_local"), description="Local path for storing MLflow artifacts before upload."
     )
@@ -86,6 +113,11 @@ class MLFlowStorage(BaseConfig):
 
     model_serializer: ModelSerializer = Field(default_factory=JoblibModelSerializer)
 
+    uri_exceptions: Sequence[str] = Field(
+        default=["databricks"],
+        description="List of URI treat as urls even if they have no scheme",
+    )
+
     _client: MlflowClient = PrivateAttr()
     _logger: logging.Logger = PrivateAttr(default=logging.getLogger(__name__))
 
@@ -95,8 +127,17 @@ class MLFlowStorage(BaseConfig):
             # Suppress MLflow's stdout messages (emoji URLs)
             os.environ.setdefault("MLFLOW_SUPPRESS_PRINTING_URL_TO_STDOUT", "true")
             os.environ.setdefault("MLFLOW_ENABLE_ARTIFACTS_PROGRESS_BAR", "false")
+
+        if self.client is not None:
+            self._client = self.client
+            return
+
         self.tracking_uri = normalize_tracking_uri(self.tracking_uri)
-        self._client = MlflowClient(tracking_uri=self.tracking_uri)
+        self.registry_uri = normalize_tracking_uri(self.registry_uri) if self.registry_uri else None
+        self._client = MlflowClient(
+            tracking_uri=self.tracking_uri,
+            registry_uri=self.registry_uri,
+        )
 
     def create_run(
         self,

{openstef_models-4.2.2 → openstef_models-4.2.3}/src/openstef_models/models/forecasting/median_forecaster.py

@@ -200,7 +200,10 @@ class MedianForecaster(Forecaster, ExplainableForecaster, ContributionsMixin):
         lag_df = input_data.reindex(new_index, fill_value=np.nan)[self._feature_names]
 
         # Convert the lag DataFrame to NumPy arrays for faster processing.
-        lag_array = lag_df.to_numpy()
+        # Copy explicitly: pandas Copy-on-Write may hand back a read-only view of the
+        # DataFrame's underlying block, but the autoregressive loop below fills the
+        # array in place (via _fill_diagonal_with_median), so it must own a writable copy.
+        lag_array = lag_df.to_numpy(copy=True)
         # Initialize the prediction array with NaNs.
         prediction = np.full(lag_array.shape[0], np.nan)
 

{openstef_models-4.2.2 → openstef_models-4.2.3}/src/openstef_models/utils/data_split.py

@@ -201,7 +201,7 @@ def _get_extreme_days(
         raise TypeError("target_series must have a DatetimeIndex.")
 
     # Compute daily min and max once
-    daily_agg: pd.DataFrame = target_series.resample("1D").agg(["min", "max"])  # type: ignore
+    daily_agg: pd.DataFrame = target_series.resample("1D").agg(["min", "max"])
     n_days = len(daily_agg)
     n_extremes = max(int(fraction * n_days), 2)
 

{openstef_models-4.2.2 → openstef_models-4.2.3}/tests/unit/integrations/mlflow/test_mlflow_storage.py

@@ -7,6 +7,7 @@ from __future__ import annotations
 from typing import TYPE_CHECKING, cast
 
 import pytest
+from mlflow import MlflowClient
 
 from openstef_core.mixins import HyperParams, Stateful
 from openstef_models.integrations.mlflow import MLFlowStorage
@@ -49,6 +50,21 @@ def model_id() -> str:
     return "test_model_123"
 
 
+def test_construct_from_mlflow(tmp_path: Path, model_id: str):
+    """Test that experiment_name_prefix is prepended to experiment names."""
+    # Arrange
+    mlflow_client = MlflowClient(tracking_uri=f"sqlite:///{tmp_path / 'mlflow.db'}")
+
+    # Act
+    storage = MLFlowStorage(
+        tracking_uri=f"sqlite:///{tmp_path / 'mlflow.db'}",
+        client=mlflow_client,
+    )
+
+    # Assert
+    assert storage._client is mlflow_client
+
+
 def test_create_run(storage: MLFlowStorage, model_id: str):
     """Test that create_run converts HyperParams to MLflow parameters."""
     # Arrange
@@ -194,12 +210,15 @@ def test_search_run__returns_matching_run(storage: MLFlowStorage, model_id: str)
         pytest.param("https://mlflow.example.com:5000", "https://mlflow.example.com:5000", id="https-port"),
         pytest.param("sqlite:///mlflow.db", "sqlite:///mlflow.db", id="sqlite"),
         pytest.param("postgresql://user:pass@host/db", "postgresql://user:pass@host/db", id="postgresql"),
+        pytest.param("databricks", "databricks", id="databricks"),
     ],
 )
 def test_normalize_tracking_uri(uri: str, expected: str):
     """Local paths become file:/// URIs; remote/database URIs pass through unchanged."""
     # Act
-    result = normalize_tracking_uri(uri)
+    result = normalize_tracking_uri(
+        uri, exceptions=["databricks"]
+    )  # databricks is a special case that should not be normalized
 
     # Assert
     if expected == "file:///":

{openstef_models-4.2.2 → openstef_models-4.2.3}/tests/unit/models/forecasting/conftest.py

@@ -30,7 +30,7 @@ def sample_forecast_input_dataset() -> ForecastInputDataset:
                 "feature2": feature_2,
                 "feature3": feature_3,
             },
-            index=pd.date_range(start=start_date, periods=num_samples, freq="1d"),
+            index=pd.date_range(start=start_date, periods=num_samples, freq="1D"),
         ),
         sample_interval=timedelta(days=1),
         target_column="load",

{openstef_models-4.2.2 → openstef_models-4.2.3}/tests/unit/models/forecasting/test_median_forecaster.py

@@ -158,7 +158,12 @@ def test_median_handles_all_missing_data():
             {
                 "quantile_P50": [],
             },
-            index=pd.DatetimeIndex([], freq="h"),
+            # Slice `index` rather than building a bare `DatetimeIndex([], freq="h")` so the
+            # (empty) expected index shares its dtype resolution with the one produced by the
+            # real prediction pipeline (pandas infers datetime64 resolution from context, and
+            # an index with no data to infer from does not necessarily match one derived from
+            # actual timestamps).
+            index=index[:0],
         ),
         sample_interval=training_input_data.sample_interval,
         forecast_start=training_input_data.forecast_start,

{openstef_models-4.2.2 → openstef_models-4.2.3}/tests/unit/transforms/general/test_dimensionality_reducer.py

@@ -33,7 +33,7 @@ def sample_forecast_input_dataset() -> TimeSeriesDataset:
                 "feature2": feature_2,
                 "feature3": feature_3,
             },
-            index=pd.date_range(start=start_date, periods=num_samples, freq="1d"),
+            index=pd.date_range(start=start_date, periods=num_samples, freq="1D"),
         ),
         sample_interval=timedelta(days=1),
     )