Merge pull request #127 from slaclab/mlflow-floats

pluflou · web-flow · commit c77dfed7414b · 2025-07-16T13:43:27.000-07:00
Adjust logging to support non-torch custom models
diff --git a/lume_model/base.py b/lume_model/base.py
@@ -9,7 +9,6 @@
 import yaml
 import numpy as np
 from pydantic import BaseModel, ConfigDict, field_validator
-import torch  # TODO: for torch.Tensor type hinting, but may need to make more general in mlflow class
 
 from lume_model.variables import ScalarVariable, get_variable, ConfigEnum
 from lume_model.utils import (
@@ -456,7 +455,6 @@ def from_yaml(cls, yaml_obj: [str, TextIOWrapper]):
 
     def register_to_mlflow(
         self,
-        input_dict: dict[str, Union[float, torch.Tensor]],
         artifact_path: str,
         registered_model_name: str | None = None,
         tags: dict[str, Any] | None = None,
@@ -477,7 +475,6 @@ def register_to_mlflow(
         https://mlflow.org/docs/latest/getting-started/intro-quickstart/ for more info.
 
         Args:
-            input_dict: Input dictionary to infer the model signature.
             artifact_path: Path to store the model in MLflow.
             registered_model_name: Name of the registered model in MLflow. Optional.
             tags: Tags to add to the MLflow model. Optional.
@@ -493,7 +490,6 @@ def register_to_mlflow(
         """
         return register_model(
             self,
-            input_dict,
             artifact_path,
             registered_model_name,
             tags,
diff --git a/lume_model/mlflow_utils.py b/lume_model/mlflow_utils.py
@@ -1,9 +1,9 @@
 import os
 import warnings
-from typing import Any, Union
+from typing import Any
 from contextlib import nullcontext
 
-from torch import Tensor, nn
+from torch import nn
 
 try:
     import mlflow
@@ -20,7 +20,6 @@
 
 def register_model(
     lume_model,
-    input: dict[str, Union[float, Tensor]] | Tensor,
     artifact_path: str,
     registered_model_name: str | None = None,
     tags: dict[str, Any] | None = None,
@@ -40,9 +39,10 @@ def register_model(
     a tracking server, set the environment variable MLFLOW_TRACKING_URI, e.g. a local port/path. See
     https://mlflow.org/docs/latest/getting-started/intro-quickstart/ for more info.
 
+    Note that at the moment, this does not log artifacts for custom models other than the YAML dump file.
+
     Args:
         lume_model: LumeModel to register.
-        input: Input dictionary to infer the model signature.
         artifact_path: Path to store the model in MLflow.
         registered_model_name: Name of the registered model in MLflow.
         tags: Tags to add to the MLflow model.
@@ -70,30 +70,19 @@ def register_model(
         else nullcontext()
     )
     with ctx:
-        # Define the signature of the model
         if isinstance(lume_model, nn.Module):
-            signature = mlflow.models.infer_signature(
-                input.numpy(), lume_model(Tensor(input)).detach().numpy()
-            )
             model_info = mlflow.pytorch.log_model(
                 pytorch_model=lume_model,
                 artifact_path=artifact_path,
-                signature=signature,
                 registered_model_name=registered_model_name,
                 **kwargs,
             )
         else:
             # Create pyfunc model for MLflow to be able to log/load the model
             pf_model = create_mlflow_model(lume_model)
-            # Adjust the input to match the expected input format
-            # Must be one of `numpy.ndarray`, `List[numpy.ndarray]`, `Dict[str, numpy.ndarray]` or `pandas.DataFrame`
-            input = {key: value.numpy() for key, value in input.items()}
-            signature = mlflow.models.infer_signature(input, pf_model.predict(input))
             model_info = mlflow.pyfunc.log_model(
                 python_model=pf_model,
                 artifact_path=artifact_path,
-                signature=signature,
-                input_example=input,
                 registered_model_name=registered_model_name,
                 **kwargs,
             )
@@ -106,23 +95,35 @@ def register_model(
 
             lume_model.dump(f"{name}.yml", save_jit=save_jit)
             mlflow.log_artifact(f"{name}.yml", artifact_path)
-            mlflow.log_artifact(f"{name}_model.pt", artifact_path)
             os.remove(f"{name}.yml")
-            os.remove(f"{name}_model.pt")
-            if save_jit:
-                mlflow.log_artifact(f"{name}_model.jit", artifact_path)
-                os.remove(f"{name}_model.jit")
-
-            # Get and log the input and output transformers
-            lume_model = (
-                lume_model._model if isinstance(lume_model, nn.Module) else lume_model
-            )
-            for i in range(len(lume_model.input_transformers)):
-                mlflow.log_artifact(f"{name}_input_transformers_{i}.pt", artifact_path)
-                os.remove(f"{name}_input_transformers_{i}.pt")
-            for i in range(len(lume_model.output_transformers)):
-                mlflow.log_artifact(f"{name}_output_transformers_{i}.pt", artifact_path)
-                os.remove(f"{name}_output_transformers_{i}.pt")
+
+            from lume_model.models import registered_models
+
+            if type(lume_model) in registered_models:
+                # all registered models are torch models at the moment
+                # may change in the future
+                mlflow.log_artifact(f"{name}_model.pt", artifact_path)
+                os.remove(f"{name}_model.pt")
+                if save_jit:
+                    mlflow.log_artifact(f"{name}_model.jit", artifact_path)
+                    os.remove(f"{name}_model.jit")
+
+                # Get and log the input and output transformers
+                lume_model = (
+                    lume_model._model
+                    if isinstance(lume_model, nn.Module)
+                    else lume_model
+                )
+                for i in range(len(lume_model.input_transformers)):
+                    mlflow.log_artifact(
+                        f"{name}_input_transformers_{i}.pt", artifact_path
+                    )
+                    os.remove(f"{name}_input_transformers_{i}.pt")
+                for i in range(len(lume_model.output_transformers)):
+                    mlflow.log_artifact(
+                        f"{name}_output_transformers_{i}.pt", artifact_path
+                    )
+                    os.remove(f"{name}_output_transformers_{i}.pt")
 
     if (tags or alias or version_tags) and registered_model_name:
         from mlflow import MlflowClient
@@ -167,18 +168,13 @@ class PyFuncModel(mlflow.pyfunc.PythonModel):
 
     # Disable type hint validation for the predict method to avoid annoying warnings
     # since we have type validation in the lume-model itself.
-    # If we need to implement this, this may be helpful:
-    # g
     _skip_type_hint_validation = True
 
     def __init__(self, model):
         self.model = model
 
     def predict(self, model_input):
         """Evaluate the model with the given input."""
-        # Convert input to the format expected by the model
-        # TODO: this isn't very general but type validation in torch modules requires this. May need to adjust.
-        model_input = {key: Tensor(value) for key, value in model_input.items()}
         return self.model.evaluate(model_input)
 
     def save_model(self):
diff --git a/lume_model/models/torch_module.py b/lume_model/models/torch_module.py
@@ -198,7 +198,6 @@ def _validate_input(x: torch.Tensor) -> torch.Tensor:
 
     def register_to_mlflow(
         self,
-        input: torch.Tensor,
         artifact_path: str,
         registered_model_name: str | None = None,
         tags: dict[str, Any] | None = None,
@@ -219,7 +218,6 @@ def register_to_mlflow(
         https://mlflow.org/docs/latest/getting-started/intro-quickstart/ for more info.
 
         Args:
-            input: Input tensor to infer the model signature.
             artifact_path: Path to store the model in MLflow.
             registered_model_name: Name of the registered model in MLflow. Optional.
             tags: Tags to add to the MLflow model. Optional.
@@ -235,7 +233,6 @@ def register_to_mlflow(
         """
         return register_model(
             self,
-            input,
             artifact_path,
             registered_model_name,
             tags,
diff --git a/pyproject.toml b/pyproject.toml
@@ -24,19 +24,16 @@ dependencies = [
     "pydantic",
     "numpy",
     "pyyaml",
+    "torch",
+    "botorch",
     "mlflow"
 ]
 dynamic = ["version"]
 [tool.setuptools_scm]
 version_file = "lume_model/_version.py"
 
 [project.optional-dependencies]
-torch = [
-    "botorch",
-    "torch"
-]
 dev = [
-    "lume-model[torch]",
     "pre-commit",
     "pytest"
 ]
