chalkpy 2.89.22__py3-none-any.whl → 2.95.3__py3-none-any.whl

chalk/ml/model_hooks.py

@@ -1,68 +1,385 @@
-from typing import Any, Callable, Dict, Optional, Tuple
+from typing import TYPE_CHECKING, Any, Dict, Optional, Protocol, Tuple
 
-import pyarrow as pa
+from chalk.ml.utils import ModelClass, ModelEncoding, ModelType
 
-from chalk.ml.utils import ModelEncoding, ModelType
+if TYPE_CHECKING:
+    from chalk.features.resolver import ResourceHint
 
 
-def load_xgb_classifier(f: str):
-    import xgboost  # pyright: ignore[reportMissingImports]
+class ModelInference(Protocol):
+    """Abstract base class for model loading and inference."""
 
-    model = xgboost.XGBClassifier()
-    model.load_model(f)
-    return model
+    def load_model(self, path: str, resource_hint: Optional["ResourceHint"] = None) -> Any:
+        """Load a model from the given path."""
+        pass
 
+    def predict(self, model: Any, X: Any) -> Any:
+        """Run inference on the model with input X."""
+        pass
 
-def load_xgb_regressor(f: str):
-    import xgboost  # pyright: ignore[reportMissingImports]
+    def prepare_input(self, feature_table: Any) -> Any:
+        """Convert PyArrow table to model input format.
 
-    model = xgboost.XGBRegressor()
-    model.load_model(f)
-    return model
+        Default implementation converts to numpy array via __array__().
+        Override for model-specific input formats (e.g., ONNX struct arrays).
+        """
+        return feature_table.__array__()
 
+    def extract_output(self, result: Any, output_feature_name: str) -> Any:
+        """Extract single output from model result.
 
-def load_pytorch_model(f: str):
-    import torch  # pyright: ignore[reportMissingImports]
+        Default implementation returns result as-is (for single outputs).
+        Override for models with structured outputs (e.g., ONNX struct arrays).
+        """
+        return result
 
-    torch.set_grad_enabled(False)
-    model = torch.jit.load(f)
-    model.input_to_tensor = lambda X: torch.from_numpy(X.__array__()).float()
-    return model
 
+class XGBoostClassifierInference(ModelInference):
+    """Model inference for XGBoost classifiers."""
 
-MODEL_HOOKS: Dict[Tuple[ModelType, ModelEncoding, Optional[str]], Callable[[str], Any]] = {
-    (ModelType.PYTORCH, ModelEncoding.PICKLE, None): load_pytorch_model,
-    (ModelType.SKLEARN, ModelEncoding.PICKLE, None): lambda f: __import__("joblib").load(f),
-    (ModelType.TENSORFLOW, ModelEncoding.HDF5, None): lambda f: __import__("tensorflow").keras.models.load_model(f),
-    (ModelType.TENSORFLOW, ModelEncoding.SAFETENSOR, None): lambda f: __import__("tensorflow").keras.models.load_model(
-        f
-    ),
-    (ModelType.XGBOOST, ModelEncoding.JSON, None): load_xgb_regressor,
-    (ModelType.XGBOOST, ModelEncoding.JSON, "classifier"): load_xgb_classifier,
-    (ModelType.XGBOOST, ModelEncoding.JSON, "regressor"): load_xgb_regressor,
-    (ModelType.LIGHTGBM, ModelEncoding.TEXT, None): lambda f: __import__("lightgbm").Booster(model_file=f),
-    (ModelType.CATBOOST, ModelEncoding.CBM, None): lambda f: __import__("catboost").CatBoost().load_model(f),
-    (ModelType.ONNX, ModelEncoding.PROTOBUF, None): lambda f: __import__("onnxruntime").InferenceSession(f),
-}
+    def load_model(self, path: str, resource_hint: Optional["ResourceHint"] = None) -> Any:
+        import xgboost  # pyright: ignore[reportMissingImports]
 
+        model = xgboost.XGBClassifier()
+        model.load_model(path)
+        return model
 
-def pytorch_predict(model: Any, X: Any):
-    outputs = model(model.input_to_tensor(X))
-    result = outputs.detach().numpy().astype("float64")
-    result = result.squeeze()
-    # Convert 0-dimensional array to scalar, or ensure we have a proper 1D array
-    if result.ndim == 0:
-        return result.item()
-    return result
+    def predict(self, model: Any, X: Any) -> Any:
+        return model.predict(X)
 
 
-PREDICT_HOOKS: Dict[Tuple[ModelType, ModelEncoding, Optional[str]], Callable[[Any, pa.Table], Any]] = {
-    (ModelType.PYTORCH, ModelEncoding.PICKLE, None): pytorch_predict,
-    (ModelType.SKLEARN, ModelEncoding.PICKLE, None): lambda model, X: model.predict(X),
-    (ModelType.TENSORFLOW, ModelEncoding.HDF5, None): lambda model, X: model.predict(X),
-    (ModelType.TENSORFLOW, ModelEncoding.SAFETENSOR, None): lambda model, X: model.predict(X),
-    (ModelType.XGBOOST, ModelEncoding.JSON, None): lambda model, X: model.predict(X),
-    (ModelType.LIGHTGBM, ModelEncoding.TEXT, None): lambda model, X: model.predict(X),
-    (ModelType.CATBOOST, ModelEncoding.CBM, None): lambda model, X: model.predict(X),
-    (ModelType.ONNX, ModelEncoding.PROTOBUF, None): lambda model, X: model.run(None, {"input": X.astype("float32")})[0],
-}
+class XGBoostRegressorInference(ModelInference):
+    """Model inference for XGBoost regressors."""
+
+    def load_model(self, path: str, resource_hint: Optional["ResourceHint"] = None) -> Any:
+        import xgboost  # pyright: ignore[reportMissingImports]
+
+        model = xgboost.XGBRegressor()
+        model.load_model(path)
+        return model
+
+    def predict(self, model: Any, X: Any) -> Any:
+        return model.predict(X)
+
+
+class PyTorchInference(ModelInference):
+    """Model inference for PyTorch models."""
+
+    def load_model(self, path: str, resource_hint: Optional["ResourceHint"] = None) -> Any:
+        import torch  # pyright: ignore[reportMissingImports]
+
+        torch.set_grad_enabled(False)
+
+        # Load the model
+        model = torch.jit.load(path)
+
+        # If resource_hint is "gpu", move model to GPU
+        if resource_hint == "gpu" and torch.cuda.is_available():
+            device = torch.device("cuda")
+            model = model.to(device)
+            model.input_to_tensor = lambda X: torch.from_numpy(X).float().to(device)
+        else:
+            model.input_to_tensor = lambda X: torch.from_numpy(X).float()
+
+        return model
+
+    def predict(self, model: Any, X: Any) -> Any:
+        outputs = model(model.input_to_tensor(X))
+        result = outputs.detach().cpu().numpy().astype("float64")
+        result = result.squeeze()
+
+        # Convert 0-dimensional array to scalar, or ensure we have a proper 1D array
+        if result.ndim == 0:
+            return result.item()
+
+        return result
+
+
+class SklearnInference(ModelInference):
+    """Model inference for scikit-learn models."""
+
+    def load_model(self, path: str, resource_hint: Optional["ResourceHint"] = None) -> Any:
+        import joblib  # pyright: ignore[reportMissingImports]
+
+        return joblib.load(path)
+
+    def predict(self, model: Any, X: Any) -> Any:
+        return model.predict(X)
+
+
+class TensorFlowInference(ModelInference):
+    """Model inference for TensorFlow models."""
+
+    def load_model(self, path: str, resource_hint: Optional["ResourceHint"] = None) -> Any:
+        import tensorflow  # pyright: ignore[reportMissingImports]
+
+        return tensorflow.keras.models.load_model(path)
+
+    def predict(self, model: Any, X: Any) -> Any:
+        return model.predict(X)
+
+
+class LightGBMInference(ModelInference):
+    """Model inference for LightGBM models."""
+
+    def load_model(self, path: str, resource_hint: Optional["ResourceHint"] = None) -> Any:
+        import lightgbm  # pyright: ignore[reportMissingImports]
+
+        return lightgbm.Booster(model_file=path)
+
+    def predict(self, model: Any, X: Any) -> Any:
+        return model.predict(X)
+
+
+class CatBoostInference(ModelInference):
+    """Model inference for CatBoost models."""
+
+    def load_model(self, path: str, resource_hint: Optional["ResourceHint"] = None) -> Any:
+        import catboost  # pyright: ignore[reportMissingImports]
+
+        return catboost.CatBoost().load_model(path)
+
+    def predict(self, model: Any, X: Any) -> Any:
+        return model.predict(X)
+
+
+class ONNXInference(ModelInference):
+    """Model inference for ONNX models with struct input/output support."""
+
+    def load_model(self, path: str, resource_hint: Optional["ResourceHint"] = None) -> Any:
+        import onnxruntime  # pyright: ignore[reportMissingImports]
+
+        # Conditionally add CUDAExecutionProvider based on resource_hint
+        providers = (
+            ["CUDAExecutionProvider", "CPUExecutionProvider"] if resource_hint == "gpu" else ["CPUExecutionProvider"]
+        )
+        return onnxruntime.InferenceSession(path, providers=providers)
+
+    def prepare_input(self, feature_table: Any) -> Any:
+        """Convert PyArrow table to struct array for ONNX models."""
+        import pyarrow as pa
+
+        # Get arrays for each column, combining chunks if necessary
+        arrays = []
+        for i in range(feature_table.num_columns):
+            col = feature_table.column(i)
+            if isinstance(col, pa.ChunkedArray):
+                arrays.append(col.combine_chunks())
+            else:
+                arrays.append(col)
+
+        # Create fields from schema, preserving original field names
+        # Field names should match ONNX input names exactly
+        fields = []
+        for field in feature_table.schema:
+            fields.append(pa.field(field.name, field.type))
+
+        # Create struct array where each row is a struct with named fields
+        return pa.StructArray.from_arrays(arrays, fields=fields)
+
+    def extract_output(self, result: Any, output_feature_name: str) -> Any:
+        """Extract single field from ONNX struct output."""
+        import pyarrow as pa
+
+        if not isinstance(result, (pa.StructArray, pa.ChunkedArray)):
+            return result
+
+        struct_type = result.type if isinstance(result, pa.StructArray) else result.chunk(0).type
+
+        # Find matching field by name, or use first field
+        field_index = None
+        for i, field in enumerate(struct_type):
+            if field.name == output_feature_name:
+                field_index = i
+                break
+
+        return result.field(field_index if field_index is not None else 0)
+
+    def predict(self, model: Any, X: Any) -> Any:
+        """Run ONNX inference with struct input/output."""
+        # Get ONNX model input/output names
+        input_names = [inp.name for inp in model.get_inputs()]
+        output_names = [out.name for out in model.get_outputs()]
+
+        # Convert struct input to ONNX input dict
+        input_dict = self._struct_to_inputs(X, input_names)
+
+        # Run ONNX inference
+        outputs = model.run(output_names, input_dict)
+
+        # Always return outputs as struct array
+        return self._outputs_to_struct(output_names, outputs)
+
+    def _struct_to_inputs(self, struct_array: Any, input_names: list) -> dict:
+        """Extract ONNX inputs from struct array by matching field names.
+
+        Struct field names must match ONNX input names (supports list/Tensor types).
+        If ONNX expects a single input but struct has multiple scalar fields,
+        stack them into a 2D array.
+        """
+        import numpy as np
+        import pyarrow as pa
+
+        if isinstance(struct_array, pa.ChunkedArray):
+            struct_array = struct_array.combine_chunks()
+
+        input_dict = {}
+        struct_fields = {field.name: i for i, field in enumerate(struct_array.type)}
+
+        # Check if struct field names match ONNX input names
+        fields_match = all(input_name in struct_fields for input_name in input_names)
+
+        if not fields_match:
+            # Special case 1: ONNX expects single input and struct has single field
+            # Use that field regardless of name mismatch
+            if len(input_names) == 1 and len(struct_fields) == 1:
+                field_data = struct_array.field(0)
+                input_dict[input_names[0]] = self._arrow_to_numpy(field_data)
+                return input_dict
+
+            # Special case 2: ONNX expects single input, but struct has multiple scalar fields
+            # Stack them into a 2D array [batch_size, num_fields]
+            if len(input_names) == 1 and len(struct_fields) > 1:
+                # Check if all fields are scalar (not nested lists)
+                all_scalar = all(
+                    not pa.types.is_list(struct_array.type[i].type)
+                    and not pa.types.is_large_list(struct_array.type[i].type)
+                    for i in range(len(struct_array.type))
+                )
+
+                if all_scalar:
+                    # Stack all fields into a single 2D array
+                    columns = []
+                    for i in range(len(struct_array.type)):
+                        field_data = struct_array.field(i)
+                        col_array = self._arrow_to_numpy(field_data)
+                        columns.append(col_array)
+
+                    # Stack columns horizontally to create [batch_size, num_features]
+                    stacked = np.column_stack(columns)
+                    input_dict[input_names[0]] = stacked
+                    return input_dict
+
+            raise ValueError(
+                f"ONNX inputs {input_names} not found in struct fields {list(struct_fields.keys())}. "
+                + "Struct field names must match ONNX input names."
+            )
+
+        # Direct mapping: struct fields match ONNX inputs (for Tensor/list types or named inputs)
+        for input_name in input_names:
+            field_data = struct_array.field(struct_fields[input_name])
+            input_dict[input_name] = self._arrow_to_numpy(field_data)
+
+        return input_dict
+
+    def _arrow_to_numpy(self, arrow_array: Any) -> Any:
+        """Convert Arrow array (including nested lists) to dense numpy array."""
+        import numpy as np
+        import pyarrow as pa
+
+        if isinstance(arrow_array, pa.ChunkedArray):
+            arrow_array = arrow_array.combine_chunks()
+
+        # Convert to Python list, then numpy - handles all cases (nested lists, flat arrays, etc.)
+        return np.array(arrow_array.to_pylist(), dtype=np.float32)
+
+    def _outputs_to_struct(self, output_names: list, outputs: list) -> Any:
+        """Convert ONNX outputs to PyArrow struct array."""
+        import pyarrow as pa
+
+        if not outputs:
+            raise ValueError("ONNX model returned no outputs")
+
+        # Convert each output to Arrow array with proper type
+        fields = []
+        arrays = []
+
+        for name, output_array in zip(output_names, outputs):
+            arrow_array = self._numpy_to_arrow_array(output_array)
+            fields.append(pa.field(name, arrow_array.type))
+            arrays.append(arrow_array)
+
+        return pa.StructArray.from_arrays(arrays, fields=fields)
+
+    def _numpy_to_arrow_array(self, arr: Any) -> Any:
+        """Convert numpy array to PyArrow array (possibly nested list)."""
+        import pyarrow as pa
+
+        # PyArrow can infer the correct nested list type from Python lists
+        # Shape (batch, dim1, dim2, ...) -> list[list[...]]
+        return pa.array(arr.tolist())
+
+
+class ModelInferenceRegistry:
+    """Registry for model inference implementations."""
+
+    def __init__(self):
+        super().__init__()
+        self._registry: Dict[Tuple[ModelType, ModelEncoding, Optional[ModelClass]], ModelInference] = {}
+
+    def register(
+        self,
+        model_type: ModelType,
+        encoding: ModelEncoding,
+        model_class: Optional[ModelClass],
+        inference: ModelInference,
+    ) -> None:
+        """Register a model inference implementation."""
+        self._registry[(model_type, encoding, model_class)] = inference
+
+    def register_for_all_classes(
+        self,
+        model_type: ModelType,
+        encoding: ModelEncoding,
+        inference: ModelInference,
+    ) -> None:
+        """Register inference for None, CLASSIFICATION, and REGRESSION variants."""
+        self.register(model_type, encoding, None, inference)
+        self.register(model_type, encoding, ModelClass.CLASSIFICATION, inference)
+        self.register(model_type, encoding, ModelClass.REGRESSION, inference)
+
+    def get(
+        self,
+        model_type: ModelType,
+        encoding: ModelEncoding,
+        model_class: Optional[ModelClass] = None,
+    ) -> Optional[ModelInference]:
+        """Get a model inference implementation from the registry."""
+        return self._registry.get((model_type, encoding, model_class), None)
+
+    def get_loader(
+        self,
+        model_type: ModelType,
+        encoding: ModelEncoding,
+        model_class: Optional[ModelClass] = None,
+    ):
+        """Get the load_model function for a given configuration."""
+        inference = self.get(model_type, encoding, model_class)
+        return inference.load_model if inference else None
+
+    def get_predictor(
+        self,
+        model_type: ModelType,
+        encoding: ModelEncoding,
+        model_class: Optional[ModelClass] = None,
+    ):
+        """Get the predict function for a given configuration."""
+        inference = self.get(model_type, encoding, model_class)
+        return inference.predict if inference else None
+
+
+# Global registry instance
+MODEL_REGISTRY = ModelInferenceRegistry()
+
+# Register all model types
+MODEL_REGISTRY.register_for_all_classes(ModelType.PYTORCH, ModelEncoding.PICKLE, PyTorchInference())
+MODEL_REGISTRY.register_for_all_classes(ModelType.SKLEARN, ModelEncoding.PICKLE, SklearnInference())
+MODEL_REGISTRY.register_for_all_classes(ModelType.TENSORFLOW, ModelEncoding.HDF5, TensorFlowInference())
+MODEL_REGISTRY.register_for_all_classes(ModelType.LIGHTGBM, ModelEncoding.TEXT, LightGBMInference())
+MODEL_REGISTRY.register_for_all_classes(ModelType.CATBOOST, ModelEncoding.CBM, CatBoostInference())
+MODEL_REGISTRY.register_for_all_classes(ModelType.ONNX, ModelEncoding.PROTOBUF, ONNXInference())
+
+# XGBoost requires different implementations for classification vs regression
+MODEL_REGISTRY.register(ModelType.XGBOOST, ModelEncoding.JSON, None, XGBoostRegressorInference())
+MODEL_REGISTRY.register(ModelType.XGBOOST, ModelEncoding.JSON, ModelClass.CLASSIFICATION, XGBoostClassifierInference())
+MODEL_REGISTRY.register(ModelType.XGBOOST, ModelEncoding.JSON, ModelClass.REGRESSION, XGBoostRegressorInference())

chalk/ml/model_reference.py

@@ -3,12 +3,22 @@ from __future__ import annotations
 import inspect
 import os
 from datetime import datetime
+from typing import TYPE_CHECKING
 
 from chalk.ml.model_version import ModelVersion
-from chalk.ml.utils import REGISTRY_METADATA_FILE, get_model_spec, model_encoding_from_proto, model_type_from_proto
+from chalk.ml.utils import (
+    ModelClass,
+    get_model_spec,
+    get_registry_metadata_file,
+    model_encoding_from_proto,
+    model_type_from_proto,
+)
 from chalk.utils.object_inspect import get_source_object_starting
 from chalk.utils.source_parsing import should_skip_source_code_parsing
 
+if TYPE_CHECKING:
+    from chalk.features.resolver import ResourceHint
+
 
 class ModelReference:
     def __init__(
@@ -18,6 +28,8 @@ class ModelReference:
         version: int | None = None,
         alias: str | None = None,
         as_of_date: datetime | None = None,
+        resource_hint: "ResourceHint | None" = None,
+        resource_group: str | None = None,
     ):
         """Specifies the model version that should be loaded into the deployment.
 
@@ -68,6 +80,8 @@ class ModelReference:
         self.as_of_date = as_of_date
         self.alias = alias
         self.identifier = identifier
+        self.resource_hint = resource_hint
+        self.resource_group = resource_group
 
         self.filename = filename
         self.source_line_start = source_line_start
@@ -89,7 +103,8 @@ class ModelReference:
         MODEL_REFERENCE_REGISTRY[(name, identifier)] = self
 
         # Only load model if the metadata file exists, which only happens in deployed environments
-        if REGISTRY_METADATA_FILE is not None and os.path.exists(REGISTRY_METADATA_FILE):
+        registry_metadata_file = get_registry_metadata_file()
+        if registry_metadata_file is not None and os.path.exists(registry_metadata_file):
             model_artifact_metadata = get_model_spec(model_name=name, identifier=identifier)
 
             mv = ModelVersion(
@@ -100,6 +115,11 @@ class ModelReference:
                 identifier=identifier,
                 model_type=model_type_from_proto(model_artifact_metadata.spec.model_type),
                 model_encoding=model_encoding_from_proto(model_artifact_metadata.spec.model_encoding),
+                model_class=ModelClass(model_artifact_metadata.spec.model_class)
+                if model_artifact_metadata.spec.model_class
+                else None,
+                resource_hint=resource_hint,
+                resource_group=resource_group,
             )
 
             from chalk.features.hooks import before_all
@@ -107,14 +127,22 @@ class ModelReference:
             def hook():
                 mv.load_model()
 
-            before_all(hook)
+            before_all(hook, resource_hint=resource_hint, resource_group=resource_group)
 
             self.model_version = mv
         else:
-            self.model_version = ModelVersion(name=name, identifier=identifier)
+            self.model_version = ModelVersion(
+                name=name, identifier=identifier, resource_hint=resource_hint, resource_group=resource_group
+            )
 
     @classmethod
-    def as_of(cls, name: str, when: datetime) -> ModelVersion:
+    def as_of(
+        cls,
+        name: str,
+        when: datetime,
+        resource_hint: "ResourceHint | None" = None,
+        resource_group: str | None = None,
+    ) -> ModelVersion:
         """Creates a ModelReference for a specific point in time.
 
         Parameters
@@ -123,6 +151,11 @@ class ModelReference:
             The name of the model.
         when
             The datetime to use for creating the model version identifier.
+        resource_hint
+            Whether this model loading is bound by CPU, I/O, or GPU.
+        resource_group
+            The resource group for the model: this is used to isolate execution
+            onto a separate pod (or set of nodes), such as on a GPU-enabled node.
 
         Returns
         -------
@@ -134,13 +167,20 @@ class ModelReference:
         >>> import datetime
         >>> timestamp = datetime.datetime(2023, 10, 15, 14, 30, 0)
         >>> model = ModelReference.as_of("fraud_model", timestamp)
+        >>> model = ModelReference.as_of("fraud_model", timestamp, resource_hint="gpu", resource_group="gpu-group")
         """
 
-        mr = ModelReference(name=name, as_of_date=when)
+        mr = ModelReference(name=name, as_of_date=when, resource_hint=resource_hint, resource_group=resource_group)
         return mr.model_version
 
     @classmethod
-    def from_version(cls, name: str, version: int) -> ModelVersion:
+    def from_version(
+        cls,
+        name: str,
+        version: int,
+        resource_hint: "ResourceHint | None" = None,
+        resource_group: str | None = None,
+    ) -> ModelVersion:
         """Creates a ModelReference using a numeric version identifier.
 
         Parameters
@@ -149,6 +189,11 @@ class ModelReference:
             The name of the model.
         version
             The version number. Must be a non-negative integer.
+        resource_hint
+            Whether this model loading is bound by CPU, I/O, or GPU.
+        resource_group
+            The resource group for the model: this is used to isolate execution
+            onto a separate pod (or set of nodes), such as on a GPU-enabled node.
 
         Returns
         -------
@@ -163,15 +208,22 @@ class ModelReference:
         Examples
         --------
         >>> model = ModelReference.from_version("fraud_model", 1)
+        >>> model = ModelReference.from_version("fraud_model", 1, resource_hint="gpu", resource_group="gpu-group")
         """
         if version < 0:
             raise ValueError("Version number must be a non-negative integer.")
 
-        mr = ModelReference(name=name, version=version)
+        mr = ModelReference(name=name, version=version, resource_hint=resource_hint, resource_group=resource_group)
         return mr.model_version
 
     @classmethod
-    def from_alias(cls, name: str, alias: str) -> ModelVersion:
+    def from_alias(
+        cls,
+        name: str,
+        alias: str,
+        resource_hint: "ResourceHint | None" = None,
+        resource_group: str | None = None,
+    ) -> ModelVersion:
         """Creates a ModelReference using an alias identifier.
 
         Parameters
@@ -180,6 +232,11 @@ class ModelReference:
             The name of the model.
         alias
             The alias string. Must be non-empty.
+        resource_hint
+            Whether this model loading is bound by CPU, I/O, or GPU.
+        resource_group
+            The resource group for the model: this is used to isolate execution
+            onto a separate pod (or set of nodes), such as on a GPU-enabled node.
 
         Returns
         -------
@@ -194,11 +251,12 @@ class ModelReference:
         Examples
         --------
         >>> model = ModelReference.from_alias("fraud_model", "latest")
+        >>> model = ModelReference.from_alias("fraud_model", "latest", resource_hint="gpu", resource_group="gpu-group")
         """
         if not alias:
             raise ValueError("Alias must be a non-empty string.")
 
-        mr = ModelReference(name=name, alias=alias)
+        mr = ModelReference(name=name, alias=alias, resource_hint=resource_hint, resource_group=resource_group)
         return mr.model_version