apache-hamilton 1.90.0.dev0__py3-none-any.whl

hamilton/plugins/kedro_extensions.py

@@ -0,0 +1,117 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import dataclasses
+from collections.abc import Collection
+from typing import Any
+
+from kedro.io import DataCatalog
+
+from hamilton import registry
+from hamilton.io.data_adapters import DataLoader, DataSaver
+
+
+@dataclasses.dataclass
+class KedroSaver(DataSaver):
+    """Use Kedro DataCatalog and Dataset to save results
+    ref: https://docs.kedro.org/en/stable/data/advanced_data_catalog_usage.html
+
+    .. code-block:: python
+
+        from kedro.framework.session import KedroSession
+
+        with KedroSession.create() as session:
+            context = session.load_context()
+            catalog = context.catalog
+
+        dr.materialize(
+            to.kedro(
+                id="my_dataset__kedro",
+                dependencies=["my_dataset"],
+                dataset_name="my_dataset",
+                catalog=catalog,
+            )
+        )
+    """
+
+    dataset_name: str
+    catalog: DataCatalog
+
+    @classmethod
+    def applicable_types(cls) -> Collection[type]:
+        return [Any]
+
+    def save_data(self, data: Any) -> dict[str, Any]:
+        self.catalog.save(self.dataset_name, data)
+        return dict(success=True)
+
+    @classmethod
+    def name(cls) -> str:
+        return "kedro"
+
+
+@dataclasses.dataclass
+class KedroLoader(DataLoader):
+    """Use Kedro DataCatalog and Dataset to load data
+    ref: https://docs.kedro.org/en/stable/data/advanced_data_catalog_usage.html
+
+    .. code-block:: python
+
+        from kedro.framework.session import KedroSession
+
+        with KedroSession.create() as session:
+            context = session.load_context()
+            catalog = context.catalog
+
+        dr.materialize(
+            from_.kedro(
+                target="input_table",
+                dataset_name="input_table",
+                catalog=catalog
+            )
+        )
+    """
+
+    dataset_name: str
+    catalog: DataCatalog
+    version: str | None = None
+
+    @classmethod
+    def applicable_types(cls) -> Collection[type]:
+        return [Any]
+
+    def load_data(self, type_: type) -> tuple[Any, dict[str, Any]]:
+        data = self.catalog.load(self.dataset_name, self.version)
+        metadata = dict(dataset_name=self.dataset_name, version=self.version)
+        return data, metadata
+
+    @classmethod
+    def name(cls) -> str:
+        return "kedro"
+
+
+def register_data_loaders():
+    for loader in [
+        KedroSaver,
+        KedroLoader,
+    ]:
+        registry.register_adapter(loader)
+
+
+register_data_loaders()
+
+COLUMN_FRIENDLY_DF_TYPE = False

hamilton/plugins/lightgbm_extensions.py

@@ -0,0 +1,99 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import dataclasses
+from collections.abc import Collection
+from pathlib import Path
+from typing import Any, Literal, Union
+
+try:
+    import lightgbm
+except ImportError as e:
+    raise NotImplementedError("LightGBM is not installed.") from e
+
+
+from hamilton import registry
+from hamilton.io import utils
+from hamilton.io.data_adapters import DataLoader, DataSaver
+
+LIGHTGBM_MODEL_TYPES = [lightgbm.LGBMModel, lightgbm.Booster, lightgbm.CVBooster]
+LIGHTGBM_MODEL_TYPES_ANNOTATION = Union[lightgbm.LGBMModel, lightgbm.Booster, lightgbm.CVBooster]
+
+
+@dataclasses.dataclass
+class LightGBMFileWriter(DataSaver):
+    """Write LighGBM models and boosters to a file"""
+
+    path: str | Path
+    num_iteration: int | None = None
+    start_iteration: int = 0
+    importance_type: Literal["split", "gain"] = "split"
+
+    @classmethod
+    def applicable_types(cls) -> Collection[type]:
+        return LIGHTGBM_MODEL_TYPES
+
+    def save_data(self, data: LIGHTGBM_MODEL_TYPES_ANNOTATION) -> dict[str, Any]:
+        if isinstance(data, lightgbm.LGBMModel):
+            data = data.booster_
+
+        data.save_model(
+            filename=self.path,
+            num_iteration=self.num_iteration,
+            start_iteration=self.start_iteration,
+            importance_type=self.importance_type,
+        )
+        return utils.get_file_metadata(self.path)
+
+    @classmethod
+    def name(cls) -> str:
+        return "file"
+
+
+@dataclasses.dataclass
+class LightGBMFileReader(DataLoader):
+    """Load LighGBM models and boosters from a file"""
+
+    path: str | Path
+
+    @classmethod
+    def applicable_types(cls) -> Collection[type]:
+        return LIGHTGBM_MODEL_TYPES
+
+    def load_data(
+        self, type_: type
+    ) -> tuple[lightgbm.Booster | lightgbm.CVBooster, dict[str, Any]]:
+        model = type_(model_file=self.path)
+        metadata = utils.get_file_metadata(self.path)
+        return model, metadata
+
+    @classmethod
+    def name(cls) -> str:
+        return "file"
+
+
+def register_data_loaders():
+    for loader in [
+        LightGBMFileReader,
+        LightGBMFileWriter,
+    ]:
+        registry.register_adapter(loader)
+
+
+register_data_loaders()
+
+COLUMN_FRIENDLY_DF_TYPE = False

hamilton/plugins/matplotlib_extensions.py

@@ -0,0 +1,108 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import dataclasses
+from collections.abc import Collection
+from os import PathLike
+from typing import IO, Any
+
+try:
+    from matplotlib.artist import Artist
+    from matplotlib.figure import Figure
+    from matplotlib.transforms import Bbox
+except ImportError as e:
+    raise NotImplementedError("Matplotlib is not installed.") from e
+
+from hamilton import registry
+from hamilton.io import utils
+from hamilton.io.data_adapters import DataSaver
+
+
+@dataclasses.dataclass
+class MatplotlibWriter(DataSaver):
+    """Write Matplotlib figure as static image format
+    ref: https://matplotlib.org/stable/api/figure_api.html#matplotlib.figure.Figure
+    """
+
+    path: str | PathLike | IO
+    dpi: float | str | None = None
+    format: str | None = None
+    metadata: dict | None = None
+    bbox_inches: str | Bbox | None = None
+    pad_inches: float | str | None = None
+    facecolor: str | float | tuple | None = None
+    edgecolor: str | float | tuple | None = None
+    backend: str | None = None
+    orientation: str | None = None
+    papertype: str | None = None
+    transparent: bool | None = None
+    bbox_extra_artists: list[Artist] | None = None
+    pil_kwargs: dict | None = None
+
+    def _get_saving_kwargs(self) -> dict:
+        kwargs = {}
+        if self.format is not None:
+            kwargs["format"] = self.format
+        if self.metadata is not None:
+            kwargs["metadata"] = self.metadata
+        if self.bbox_inches is not None:
+            kwargs["bbox_inches"] = self.bbox_inches
+        if self.pad_inches is not None:
+            kwargs["pad_inches"] = self.pad_inches
+        if self.facecolor is not None:
+            kwargs["facecolor"] = self.facecolor
+        if self.edgecolor is not None:
+            kwargs["edgecolor"] = self.edgecolor
+        if self.backend is not None:
+            kwargs["backend"] = self.backend
+        if self.orientation is not None:
+            kwargs["orientation"] = self.orientation
+        if self.papertype is not None:
+            kwargs["papertype"] = self.papertype
+        if self.transparent is not None:
+            kwargs["transparent"] = self.transparent
+        if self.bbox_extra_artists is not None:
+            kwargs["bbox_extra_artists"] = self.bbox_extra_artists
+        if self.pil_kwargs is not None:
+            kwargs["pil_kwargs"] = self.pil_kwargs
+
+        return kwargs
+
+    def save_data(self, data: Figure) -> dict[str, Any]:
+        data.savefig(fname=self.path, **self._get_saving_kwargs())
+        # TODO make utils.get_file_metadata() safer for when self.path is IO type
+        return utils.get_file_metadata(self.path)
+
+    @classmethod
+    def applicable_types(cls) -> Collection[type]:
+        return [Figure]
+
+    @classmethod
+    def name(cls) -> str:
+        return "plt"
+
+
+def register_data_savers():
+    for saver in [
+        MatplotlibWriter,
+    ]:
+        registry.register_adapter(saver)
+
+
+register_data_savers()
+
+COLUMN_FRIENDLY_DF_TYPE = False

hamilton/plugins/mlflow_extensions.py

@@ -0,0 +1,216 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import dataclasses
+import pathlib
+from collections.abc import Collection
+from types import ModuleType
+from typing import Any, Literal
+
+try:
+    import mlflow
+except ImportError as e:
+    raise NotImplementedError("MLFlow is not installed.") from e
+
+from hamilton import registry
+from hamilton.io.data_adapters import DataLoader, DataSaver
+
+
+@dataclasses.dataclass
+class MLFlowModelSaver(DataSaver):
+    """Save model to the MLFlow tracking server using `.log_model()`
+
+    :param path: Run relative path to store model. Will constitute the model URI.
+    :param register_as: If not None, register the model under the specified name.
+    :param flavor: Library format to save the model (sklearn, xgboost, etc.). Automatically inferred if None.
+    :param run_id: Log model to a specific run. Leave to `None` if using the `MLFlowTracker`
+    :param mlflow_kwargs: Arguments for `.log_model()`. Can be flavor-specific.
+    """
+
+    path: str | pathlib.Path = "model"
+    register_as: str | None = None
+    flavor: str | ModuleType | None = None
+    run_id: str | None = None
+    mlflow_kwargs: dict[str, Any] = None
+
+    def __post_init__(self):
+        self.mlflow_kwargs = self.mlflow_kwargs or {}
+
+    @classmethod
+    def name(cls) -> str:
+        return "mlflow"
+
+    @classmethod
+    def applicable_types(cls) -> Collection[type]:
+        return [Any]
+
+    def save_data(self, data) -> dict[str, Any]:
+        if self.flavor:
+            flavor = self.flavor
+        else:
+            # infer the flavor from the base module of the data class
+            # for example, extract `sklearn` from `sklearn.linear_model._base`
+            flavor, _, _ = data.__module__.partition(".")
+
+        # pass flavor as module, this suports standard flavors `like mlflow.sklearn`
+        # but also supports custom flavors like `my_flavor` that implements save(), log(), load()
+        if isinstance(flavor, ModuleType):
+            flavor_module = flavor
+        else:
+            # retrieve the `mlflow.FLAVOR` submodule to `.log_model()`
+            try:
+                flavor_module = getattr(mlflow, flavor)
+            except ImportError as e:
+                raise ImportError(f"Flavor {flavor} is unsupported by MLFlow") from e
+
+        # handle `run_id` and active run conflicts
+        if mlflow.active_run() and self.run_id:
+            if mlflow.active_run().info.run_id != self.run_id:
+                raise RuntimeError(
+                    "The MLFlowModelSaver `run_id` doesn't match the active `run_id`\n",
+                    "Set `run_id=None` to save to the active MLFlow run.",
+                )
+
+        # save to active run
+        if mlflow.active_run():
+            model_info = flavor_module.log_model(data, self.path, **self.mlflow_kwargs)
+        # create a run with `run_id` and save to it
+        else:
+            with mlflow.start_run(run_id=self.run_id):
+                model_info = flavor_module.log_model(data, self.path, **self.mlflow_kwargs)
+
+        # create metadata from ModelInfo object
+        metadata = {k.strip("_"): v for k, v in model_info.__dict__.items()}
+
+        if self.register_as:
+            model_version = mlflow.register_model(
+                model_uri=metadata["model_uri"], name=self.register_as
+            )
+            # update metadata with the registered ModelVersion
+            # there's a contract between this key and the MLFlowTracker's
+            # post_node_execute() reads this metadata
+            metadata["registered_model"] = {
+                k.strip("_"): v for k, v in model_version.__dict__.items()
+            }
+
+        return metadata
+
+
+@dataclasses.dataclass
+class MLFlowModelLoader(DataLoader):
+    """Load model from the MLFlow tracking server or model registry using .load_model()
+    You can pass a model URI or the necessary metadata to retrieve the model
+
+    :param model_uri: Model location starting as `runs:/` for tracking or `models:/` for registry
+    :param mode: `tracking` or registry`. tracking needs `run_id` and `path`. registry needs `model_name` and `version` or `version_alias`.
+    :param run_id: Run id of the model on the tracking server
+    :param path: Run relative path where the model is stored
+    :param model_name: Name of the registered model (equivalent to `register_as` in model saver)
+    :param version: Version of the registered model. Can pass as string `v1` or integer `1`
+    :param version_alias: Version alias of the registered model. Specify either this or `version`
+    :param flavor: Library format to load the model (sklearn, xgboost, etc.). Automatically inferred if None.
+    :param mlflow_kwargs: Arguments for `.load_model()`. Can be flavor-specific.
+    """
+
+    model_uri: str | None = None
+    mode: Literal["tracking", "registry"] = "tracking"
+    run_id: str | None = None
+    path: str | pathlib.Path = "model"
+    model_name: str | None = None
+    version: str | int | None = None
+    version_alias: str | None = None
+    flavor: ModuleType | str | None = None
+    mlflow_kwargs: dict[str, Any] = None
+
+    # __post_init__ is required to set kwargs as empty dict because
+    # can't set: kwargs: Dict[str, Any] = dataclasses.field(default_factory=dict)
+    # otherwise raises `InvalidDecoratorException` because materializer factory check
+    # for all params being set and `kwargs` would be unset until instantiation.
+    def __post_init__(self):
+        self.mlflow_kwargs = self.mlflow_kwargs or {}
+
+        if self.model_uri:
+            return
+
+        if self.mode == "tracking":
+            if (not self.run_id) or (not self.path):
+                raise ValueError("Using `mode='tracking'` requires passing `run_id` and `path`")
+
+            self.model_uri = f"runs:/{self.run_id}/{self.path}"
+
+        elif self.mode == "registry":
+            if not self.model_name:
+                raise ValueError("Using `mode='registry` requires passing `model_name`")
+
+            if not (bool(self.version) ^ bool(self.version_alias)):
+                raise ValueError(
+                    "If using `mode='registry'` requires passing `version` OR `version_alias"
+                )
+
+            if self.version:
+                self.model_uri = f"models:/{self.model_name}/{self.version}"
+            elif self.version_alias:
+                self.model_uri = f"models:/{self.model_name}@{self.version_alias}"
+
+    @classmethod
+    def name(cls) -> str:
+        return "mlflow"
+
+    @classmethod
+    def applicable_types(cls) -> Collection[type]:
+        return [Any]
+
+    def load_data(self, type_: type) -> tuple[Any, dict[str, Any]]:
+        model_info = mlflow.models.model.get_model_info(self.model_uri)
+        metadata = {k.strip("_"): v for k, v in model_info.__dict__.items()}
+
+        flavor = self.flavor
+        # if flavor not explicitly passed, retrieve flavor from the ModelInfo
+        if flavor is None:
+            # prioritize the library specific flavor. Default to `pyfunc` if none available.
+            try:
+                flavor = next(f for f in metadata["flavors"].keys() if f != "python_function")
+            except StopIteration:
+                flavor = "pyfunc"
+
+        # pass flavor as module, this suports standard flavors `like mlflow.sklearn`
+        # but also supports custom flavors like `my_flavor` that implements save(), log(), load()
+        if isinstance(flavor, ModuleType):
+            flavor_module = flavor
+        else:
+            # retrieve the `mlflow.FLAVOR` submodule to `.log_model()`
+            try:
+                flavor_module = getattr(mlflow, flavor)
+            except ImportError as e:
+                raise ImportError(f"Flavor {flavor} is unsupported by MLFlow") from e
+
+        model = flavor_module.load_model(model_uri=self.model_uri, **self.mlflow_kwargs)
+        return model, metadata
+
+
+def register_data_loaders():
+    """Function to register the data loaders for this extension."""
+    for loader in [
+        MLFlowModelSaver,
+        MLFlowModelLoader,
+    ]:
+        registry.register_adapter(loader)
+
+
+register_data_loaders()
+
+COLUMN_FRIENDLY_DF_TYPE = False

hamilton/plugins/numpy_extensions.py

@@ -0,0 +1,105 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import dataclasses
+import pathlib
+from collections.abc import Collection
+from typing import IO, Any
+
+try:
+    import numpy as np
+except ImportError as e:
+    raise NotImplementedError("Numpy is not installed.") from e
+
+from typing import Literal
+
+from hamilton import registry
+from hamilton.io import utils
+from hamilton.io.data_adapters import DataLoader, DataSaver
+
+
+@dataclasses.dataclass
+class NumpyNpyWriter(DataSaver):
+    """Write Numpy multidimensional arrays to custom .npy format
+    ref: https://numpy.org/doc/stable/reference/routines.io.html
+    """
+
+    path: str | pathlib.Path | IO
+    allow_pickle: bool | None = None
+    fix_imports: bool | None = None
+
+    def save_data(self, data: np.ndarray) -> dict[str, Any]:
+        kwargs = dict(file=self.path, arr=data, allow_pickle=self.allow_pickle)
+        if np.__version__ < "2.4" and self.fix_imports is not None:
+            kwargs["fix_imports"] = self.fix_imports
+        np.save(**kwargs)
+        return utils.get_file_metadata(self.path)
+
+    @classmethod
+    def applicable_types(cls) -> Collection[type]:
+        return [np.ndarray]
+
+    @classmethod
+    def name(cls) -> str:
+        return "npy"
+
+
+@dataclasses.dataclass
+class NumpyNpyReader(DataLoader):
+    """Read Numpy multidimensional arrays from custom .npy format
+    ref: https://numpy.org/doc/stable/reference/routines.io.html
+    """
+
+    path: str | pathlib.Path | IO
+    mmap_mode: str | None = None
+    allow_pickle: bool | None = None
+    fix_imports: bool | None = None
+    encoding: Literal["ASCII", "latin1", "bytes"] = "ASCII"
+
+    @classmethod
+    def applicable_types(cls) -> Collection[type]:
+        return [np.ndarray]
+
+    def load_data(self, type_: type) -> tuple[np.ndarray, dict[str, Any]]:
+        kwargs = dict(
+            file=self.path,
+            mmap_mode=self.mmap_mode,
+            allow_pickle=self.allow_pickle,
+            encoding=self.encoding,
+        )
+        if np.__version__ < "2.4" and self.fix_imports is not None:
+            kwargs["fix_imports"] = self.fix_imports
+        array = np.load(**kwargs)
+        metadata = utils.get_file_metadata(self.path)
+        return array, metadata
+
+    @classmethod
+    def name(cls) -> str:
+        return "npy"
+
+
+def register_data_loaders():
+    for loader in [
+        NumpyNpyWriter,
+        NumpyNpyReader,
+    ]:
+        registry.register_adapter(loader)
+
+
+register_data_loaders()
+
+COLUMN_FRIENDLY_DF_TYPE = False