snowflake-ml-python 1.8.1__py3-none-any.whl → 1.8.3__py3-none-any.whl

snowflake/ml/model/_packager/model_handlers/lightgbm.py

@@ -1,16 +1,6 @@
 import os
 import warnings
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    Callable,
-    Dict,
-    Optional,
-    Type,
-    Union,
-    cast,
-    final,
-)
+from typing import TYPE_CHECKING, Any, Callable, Optional, Union, cast, final
 
 import cloudpickle
 import numpy as np
@@ -41,7 +31,7 @@ class LGBMModelHandler(_base.BaseModelHandler[Union["lightgbm.Booster", "lightgb
     HANDLER_TYPE = "lightgbm"
     HANDLER_VERSION = "2024-03-19"
     _MIN_SNOWPARK_ML_VERSION = "1.3.1"
-    _HANDLER_MIGRATOR_PLANS: Dict[str, Type[base_migrator.BaseModelHandlerMigrator]] = {}
+    _HANDLER_MIGRATOR_PLANS: dict[str, type[base_migrator.BaseModelHandlerMigrator]] = {}
 
     MODEL_BLOB_FILE_OR_DIR = "model.pkl"
     DEFAULT_TARGET_METHODS = ["predict", "predict_proba"]
@@ -215,7 +205,7 @@ class LGBMModelHandler(_base.BaseModelHandler[Union["lightgbm.Booster", "lightgb
         def _create_custom_model(
             raw_model: Union["lightgbm.Booster", "lightgbm.LGBMModel"],
             model_meta: model_meta_api.ModelMetadata,
-        ) -> Type[custom_model.CustomModel]:
+        ) -> type[custom_model.CustomModel]:
             def fn_factory(
                 raw_model: Union["lightgbm.Booster", "lightgbm.LGBMModel"],
                 signature: model_signature.ModelSignature,
@@ -250,7 +240,7 @@ class LGBMModelHandler(_base.BaseModelHandler[Union["lightgbm.Booster", "lightgb
 
                 return fn
 
-            type_method_dict: Dict[str, Any] = {"_raw_model": raw_model}
+            type_method_dict: dict[str, Any] = {"_raw_model": raw_model}
             for target_method_name, sig in model_meta.signatures.items():
                 type_method_dict[target_method_name] = fn_factory(raw_model, sig, target_method_name)
 

snowflake/ml/model/_packager/model_handlers/mlflow.py

@@ -1,7 +1,7 @@
 import os
 import pathlib
 import tempfile
-from typing import TYPE_CHECKING, Callable, Dict, Optional, Type, cast, final
+from typing import TYPE_CHECKING, Callable, Optional, cast, final
 
 import pandas as pd
 from typing_extensions import TypeGuard, Unpack
@@ -61,7 +61,7 @@ class MLFlowHandler(_base.BaseModelHandler["mlflow.pyfunc.PyFuncModel"]):
     HANDLER_TYPE = "mlflow"
     HANDLER_VERSION = "2023-12-01"
     _MIN_SNOWPARK_ML_VERSION = "1.0.12"
-    _HANDLER_MIGRATOR_PLANS: Dict[str, Type[base_migrator.BaseModelHandlerMigrator]] = {}
+    _HANDLER_MIGRATOR_PLANS: dict[str, type[base_migrator.BaseModelHandlerMigrator]] = {}
 
     MODEL_BLOB_FILE_OR_DIR = "model"
     _DEFAULT_TARGET_METHOD = "predict"
@@ -204,7 +204,7 @@ class MLFlowHandler(_base.BaseModelHandler["mlflow.pyfunc.PyFuncModel"]):
         def _create_custom_model(
             raw_model: "mlflow.pyfunc.PyFuncModel",
             model_meta: model_meta_api.ModelMetadata,
-        ) -> Type[custom_model.CustomModel]:
+        ) -> type[custom_model.CustomModel]:
             def fn_factory(
                 raw_model: "mlflow.pyfunc.PyFuncModel",
                 signature: model_signature.ModelSignature,

snowflake/ml/model/_packager/model_handlers/pytorch.py

@@ -1,6 +1,6 @@
 import os
 import sys
-from typing import TYPE_CHECKING, Callable, Dict, Optional, Type, cast, final
+from typing import TYPE_CHECKING, Callable, Optional, cast, final
 
 import cloudpickle
 import pandas as pd
@@ -38,7 +38,7 @@ class PyTorchHandler(_base.BaseModelHandler["torch.nn.Module"]):
     HANDLER_TYPE = "pytorch"
     HANDLER_VERSION = "2025-03-01"
     _MIN_SNOWPARK_ML_VERSION = "1.8.0"
-    _HANDLER_MIGRATOR_PLANS: Dict[str, Type[base_migrator.BaseModelHandlerMigrator]] = {
+    _HANDLER_MIGRATOR_PLANS: dict[str, type[base_migrator.BaseModelHandlerMigrator]] = {
         "2023-12-01": pytorch_migrator_2023_12_01.PyTorchHandlerMigrator20231201
     }
 
@@ -151,7 +151,7 @@ class PyTorchHandler(_base.BaseModelHandler["torch.nn.Module"]):
         model_meta.env.include_if_absent(
             [model_env.ModelDependency(requirement="pytorch", pip_name="torch")], check_local_version=True
         )
-        model_meta.env.cuda_version = kwargs.get("cuda_version", model_env.DEFAULT_CUDA_VERSION)
+        model_meta.env.cuda_version = kwargs.get("cuda_version", handlers_utils.get_default_cuda_version())
 
     @classmethod
     def load_model(
@@ -188,7 +188,7 @@ class PyTorchHandler(_base.BaseModelHandler["torch.nn.Module"]):
         def _create_custom_model(
             raw_model: "torch.nn.Module",
             model_meta: model_meta_api.ModelMetadata,
-        ) -> Type[custom_model.CustomModel]:
+        ) -> type[custom_model.CustomModel]:
             multiple_inputs = cast(
                 model_meta_schema.PyTorchModelBlobOptions, model_meta.models[model_meta.name].options
             )["multiple_inputs"]

snowflake/ml/model/_packager/model_handlers/sentence_transformers.py

@@ -1,7 +1,7 @@
 import inspect
 import logging
 import os
-from typing import TYPE_CHECKING, Callable, Dict, Optional, Type, cast, final
+from typing import TYPE_CHECKING, Callable, Optional, cast, final
 
 import pandas as pd
 from typing_extensions import TypeGuard, Unpack
@@ -24,7 +24,7 @@ if TYPE_CHECKING:
 logger = logging.getLogger(__name__)
 
 
-def _validate_sentence_transformers_signatures(sigs: Dict[str, model_signature.ModelSignature]) -> None:
+def _validate_sentence_transformers_signatures(sigs: dict[str, model_signature.ModelSignature]) -> None:
     if list(sigs.keys()) != ["encode"]:
         raise ValueError("target_methods can only be ['encode']")
 
@@ -48,7 +48,7 @@ class SentenceTransformerHandler(_base.BaseModelHandler["sentence_transformers.S
     HANDLER_TYPE = "sentence_transformers"
     HANDLER_VERSION = "2024-03-15"
     _MIN_SNOWPARK_ML_VERSION = "1.3.1"
-    _HANDLER_MIGRATOR_PLANS: Dict[str, Type[base_migrator.BaseModelHandlerMigrator]] = {}
+    _HANDLER_MIGRATOR_PLANS: dict[str, type[base_migrator.BaseModelHandlerMigrator]] = {}
 
     MODEL_BLOB_FILE_OR_DIR = "model"
     DEFAULT_TARGET_METHODS = ["encode"]
@@ -166,7 +166,7 @@ class SentenceTransformerHandler(_base.BaseModelHandler["sentence_transformers.S
             ],
             check_local_version=True,
         )
-        model_meta.env.cuda_version = kwargs.get("cuda_version", model_env.DEFAULT_CUDA_VERSION)
+        model_meta.env.cuda_version = kwargs.get("cuda_version", handlers_utils.get_default_cuda_version())
 
     @staticmethod
     def _get_device_config(**kwargs: Unpack[model_types.SentenceTransformersLoadOptions]) -> Optional[str]:
@@ -224,7 +224,7 @@ class SentenceTransformerHandler(_base.BaseModelHandler["sentence_transformers.S
         def _create_custom_model(
             raw_model: "sentence_transformers.SentenceTransformer",
             model_meta: model_meta_api.ModelMetadata,
-        ) -> Type[custom_model.CustomModel]:
+        ) -> type[custom_model.CustomModel]:
             batch_size = cast(
                 model_meta_schema.SentenceTransformersModelBlobOptions, model_meta.models[model_meta.name].options
             ).get("batch_size", None)

snowflake/ml/model/_packager/model_handlers/sklearn.py

@@ -1,13 +1,13 @@
 import os
 import warnings
-from typing import TYPE_CHECKING, Callable, Dict, Optional, Type, Union, cast, final
+from typing import TYPE_CHECKING, Callable, Optional, Union, cast, final
 
 import cloudpickle
 import numpy as np
 import pandas as pd
 from typing_extensions import TypeGuard, Unpack
 
-from snowflake.ml._internal import type_utils
+from snowflake.ml._internal import env, type_utils
 from snowflake.ml.model import custom_model, model_signature, type_hints as model_types
 from snowflake.ml.model._packager.model_env import model_env
 from snowflake.ml.model._packager.model_handlers import _base, _utils as handlers_utils
@@ -19,7 +19,6 @@ from snowflake.ml.model._packager.model_meta import (
 )
 from snowflake.ml.model._packager.model_task import model_task_utils
 from snowflake.ml.model._signatures import numpy_handler, utils as model_signature_utils
-from snowflake.ml.modeling._internal.constants import IN_ML_RUNTIME_ENV_VAR
 
 if TYPE_CHECKING:
     import sklearn.base
@@ -49,7 +48,7 @@ class SKLModelHandler(_base.BaseModelHandler[Union["sklearn.base.BaseEstimator",
     HANDLER_TYPE = "sklearn"
     HANDLER_VERSION = "2023-12-01"
     _MIN_SNOWPARK_ML_VERSION = "1.0.12"
-    _HANDLER_MIGRATOR_PLANS: Dict[str, Type[base_migrator.BaseModelHandlerMigrator]] = {}
+    _HANDLER_MIGRATOR_PLANS: dict[str, type[base_migrator.BaseModelHandlerMigrator]] = {}
 
     DEFAULT_TARGET_METHODS = [
         "predict",
@@ -113,7 +112,7 @@ class SKLModelHandler(_base.BaseModelHandler[Union["sklearn.base.BaseEstimator",
                 raise ValueError("Sample input data is required to enable explainability.")
 
         # If this is a pipeline and we are in the container runtime, check for distributed estimator.
-        if os.getenv(IN_ML_RUNTIME_ENV_VAR) and isinstance(model, sklearn.pipeline.Pipeline):
+        if env.IN_ML_RUNTIME and isinstance(model, sklearn.pipeline.Pipeline):
             model = _unpack_container_runtime_pipeline(model)
 
         if not is_sub_model:
@@ -265,7 +264,7 @@ class SKLModelHandler(_base.BaseModelHandler[Union["sklearn.base.BaseEstimator",
         def _create_custom_model(
             raw_model: Union["sklearn.base.BaseEstimator", "sklearn.pipeline.Pipeline"],
             model_meta: model_meta_api.ModelMetadata,
-        ) -> Type[custom_model.CustomModel]:
+        ) -> type[custom_model.CustomModel]:
             def fn_factory(
                 raw_model: Union["sklearn.base.BaseEstimator", "sklearn.pipeline.Pipeline"],
                 signature: model_signature.ModelSignature,

snowflake/ml/model/_packager/model_handlers/snowmlmodel.py

@@ -1,6 +1,6 @@
 import os
 import warnings
-from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Type, cast, final
+from typing import TYPE_CHECKING, Any, Callable, Optional, cast, final
 
 import cloudpickle
 import numpy as np
@@ -36,7 +36,7 @@ class SnowMLModelHandler(_base.BaseModelHandler["BaseEstimator"]):
     HANDLER_TYPE = "snowml"
     HANDLER_VERSION = "2023-12-01"
     _MIN_SNOWPARK_ML_VERSION = "1.0.12"
-    _HANDLER_MIGRATOR_PLANS: Dict[str, Type[base_migrator.BaseModelHandlerMigrator]] = {}
+    _HANDLER_MIGRATOR_PLANS: dict[str, type[base_migrator.BaseModelHandlerMigrator]] = {}
 
     DEFAULT_TARGET_METHODS = ["predict", "transform", "predict_proba", "predict_log_proba", "decision_function"]
     EXPLAIN_TARGET_METHODS = ["predict", "predict_proba", "predict_log_proba"]
@@ -264,7 +264,7 @@ class SnowMLModelHandler(_base.BaseModelHandler["BaseEstimator"]):
         def _create_custom_model(
             raw_model: "BaseEstimator",
             model_meta: model_meta_api.ModelMetadata,
-        ) -> Type[custom_model.CustomModel]:
+        ) -> type[custom_model.CustomModel]:
             def fn_factory(
                 raw_model: "BaseEstimator",
                 signature: model_signature.ModelSignature,

snowflake/ml/model/_packager/model_handlers/tensorflow.py

@@ -1,5 +1,5 @@
 import os
-from typing import TYPE_CHECKING, Callable, Dict, Optional, Type, cast, final
+from typing import TYPE_CHECKING, Callable, Optional, cast, final
 
 import pandas as pd
 from packaging import version
@@ -38,7 +38,7 @@ class TensorFlowHandler(_base.BaseModelHandler["tensorflow.Module"]):
     HANDLER_TYPE = "tensorflow"
     HANDLER_VERSION = "2025-03-01"
     _MIN_SNOWPARK_ML_VERSION = "1.8.0"
-    _HANDLER_MIGRATOR_PLANS: Dict[str, Type[base_migrator.BaseModelHandlerMigrator]] = {
+    _HANDLER_MIGRATOR_PLANS: dict[str, type[base_migrator.BaseModelHandlerMigrator]] = {
         "2023-12-01": tensorflow_migrator_2023_12_01.TensorflowHandlerMigrator20231201,
         "2025-01-01": tensorflow_migrator_2025_01_01.TensorflowHandlerMigrator20250101,
     }
@@ -188,7 +188,7 @@ class TensorFlowHandler(_base.BaseModelHandler["tensorflow.Module"]):
             dependencies,
             check_local_version=True,
         )
-        model_meta.env.cuda_version = kwargs.get("cuda_version", model_env.DEFAULT_CUDA_VERSION)
+        model_meta.env.cuda_version = kwargs.get("cuda_version", handlers_utils.get_default_cuda_version())
 
     @classmethod
     def load_model(
@@ -230,7 +230,7 @@ class TensorFlowHandler(_base.BaseModelHandler["tensorflow.Module"]):
         def _create_custom_model(
             raw_model: "tensorflow.Module",
             model_meta: model_meta_api.ModelMetadata,
-        ) -> Type[custom_model.CustomModel]:
+        ) -> type[custom_model.CustomModel]:
             multiple_inputs = cast(
                 model_meta_schema.TensorflowModelBlobOptions, model_meta.models[model_meta.name].options
             )["multiple_inputs"]

snowflake/ml/model/_packager/model_handlers/torchscript.py

@@ -1,5 +1,5 @@
 import os
-from typing import TYPE_CHECKING, Callable, Dict, Optional, Type, cast, final
+from typing import TYPE_CHECKING, Callable, Optional, cast, final
 
 import pandas as pd
 from typing_extensions import TypeGuard, Unpack
@@ -36,7 +36,7 @@ class TorchScriptHandler(_base.BaseModelHandler["torch.jit.ScriptModule"]):
     HANDLER_TYPE = "torchscript"
     HANDLER_VERSION = "2025-03-01"
     _MIN_SNOWPARK_ML_VERSION = "1.8.0"
-    _HANDLER_MIGRATOR_PLANS: Dict[str, Type[base_migrator.BaseModelHandlerMigrator]] = {
+    _HANDLER_MIGRATOR_PLANS: dict[str, type[base_migrator.BaseModelHandlerMigrator]] = {
         "2023-12-01": torchscript_migrator_2023_12_01.TorchScriptHandlerMigrator20231201
     }
 
@@ -141,7 +141,7 @@ class TorchScriptHandler(_base.BaseModelHandler["torch.jit.ScriptModule"]):
         model_meta.env.include_if_absent(
             [model_env.ModelDependency(requirement="pytorch", pip_name="torch")], check_local_version=True
         )
-        model_meta.env.cuda_version = kwargs.get("cuda_version", model_env.DEFAULT_CUDA_VERSION)
+        model_meta.env.cuda_version = kwargs.get("cuda_version", handlers_utils.get_default_cuda_version())
 
     @classmethod
     def load_model(
@@ -181,7 +181,7 @@ class TorchScriptHandler(_base.BaseModelHandler["torch.jit.ScriptModule"]):
         def _create_custom_model(
             raw_model: "torch.jit.ScriptModule",
             model_meta: model_meta_api.ModelMetadata,
-        ) -> Type[custom_model.CustomModel]:
+        ) -> type[custom_model.CustomModel]:
             def fn_factory(
                 raw_model: "torch.jit.ScriptModule",
                 signature: model_signature.ModelSignature,

snowflake/ml/model/_packager/model_handlers/xgboost.py

@@ -1,17 +1,7 @@
 # mypy: disable-error-code="import"
 import os
 import warnings
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    Callable,
-    Dict,
-    Optional,
-    Type,
-    Union,
-    cast,
-    final,
-)
+from typing import TYPE_CHECKING, Any, Callable, Optional, Union, cast, final
 
 import numpy as np
 import pandas as pd
@@ -44,7 +34,7 @@ class XGBModelHandler(_base.BaseModelHandler[Union["xgboost.Booster", "xgboost.X
     HANDLER_TYPE = "xgboost"
     HANDLER_VERSION = "2023-12-01"
     _MIN_SNOWPARK_ML_VERSION = "1.0.12"
-    _HANDLER_MIGRATOR_PLANS: Dict[str, Type[base_migrator.BaseModelHandlerMigrator]] = {}
+    _HANDLER_MIGRATOR_PLANS: dict[str, type[base_migrator.BaseModelHandlerMigrator]] = {}
 
     MODEL_BLOB_FILE_OR_DIR = "model.ubj"
     DEFAULT_TARGET_METHODS = ["predict", "predict_proba"]
@@ -175,7 +165,7 @@ class XGBModelHandler(_base.BaseModelHandler[Union["xgboost.Booster", "xgboost.X
         if enable_explainability:
             model_meta.env.include_if_absent([model_env.ModelDependency(requirement="shap>=0.46.0", pip_name="shap")])
             model_meta.explain_algorithm = model_meta_schema.ModelExplainAlgorithm.SHAP
-        model_meta.env.cuda_version = kwargs.get("cuda_version", model_env.DEFAULT_CUDA_VERSION)
+        model_meta.env.cuda_version = kwargs.get("cuda_version", handlers_utils.get_default_cuda_version())
 
     @classmethod
     def load_model(
@@ -227,7 +217,7 @@ class XGBModelHandler(_base.BaseModelHandler[Union["xgboost.Booster", "xgboost.X
         def _create_custom_model(
             raw_model: Union["xgboost.Booster", "xgboost.XGBModel"],
             model_meta: model_meta_api.ModelMetadata,
-        ) -> Type[custom_model.CustomModel]:
+        ) -> type[custom_model.CustomModel]:
             def fn_factory(
                 raw_model: Union["xgboost.Booster", "xgboost.XGBModel"],
                 signature: model_signature.ModelSignature,
@@ -261,7 +251,7 @@ class XGBModelHandler(_base.BaseModelHandler[Union["xgboost.Booster", "xgboost.X
                     return explain_fn
                 return fn
 
-            type_method_dict: Dict[str, Any] = {"_raw_model": raw_model}
+            type_method_dict: dict[str, Any] = {"_raw_model": raw_model}
             for target_method_name, sig in model_meta.signatures.items():
                 type_method_dict[target_method_name] = fn_factory(raw_model, sig, target_method_name)
 

snowflake/ml/model/_packager/model_meta/model_blob_meta.py

@@ -1,4 +1,4 @@
-from typing import Dict, cast
+from typing import cast
 
 from typing_extensions import Unpack
 
@@ -25,7 +25,7 @@ class ModelBlobMeta:
         self.handler_version = kwargs["handler_version"]
         self.function_properties = kwargs.get("function_properties", {})
 
-        self.artifacts: Dict[str, str] = {}
+        self.artifacts: dict[str, str] = {}
         artifacts = kwargs.get("artifacts", None)
         if artifacts:
             self.artifacts = artifacts

snowflake/ml/model/_packager/model_meta/model_meta.py

@@ -6,31 +6,26 @@ import zipfile
 from contextlib import contextmanager
 from datetime import datetime
 from types import ModuleType
-from typing import Any, Dict, Generator, List, Optional, TypedDict
+from typing import Any, Generator, Optional, TypedDict
 
 import cloudpickle
 import yaml
 from packaging import requirements, version
 from typing_extensions import Required
 
-from snowflake.ml._internal import env as snowml_env, env_utils, file_utils
+from snowflake.ml import version as snowml_version
+from snowflake.ml._internal import env_utils, file_utils
 from snowflake.ml.model import model_signature, type_hints as model_types
 from snowflake.ml.model._packager.model_env import model_env
-from snowflake.ml.model._packager.model_meta import (
-    _packaging_requirements,
-    model_blob_meta,
-    model_meta_schema,
-)
+from snowflake.ml.model._packager.model_meta import model_blob_meta, model_meta_schema
 from snowflake.ml.model._packager.model_meta_migrator import migrator_plans
 from snowflake.ml.model._packager.model_runtime import model_runtime
 
 MODEL_METADATA_FILE = "model.yaml"
 MODEL_CODE_DIR = "code"
 
-_PACKAGING_REQUIREMENTS = [
-    str(env_utils.get_package_spec_with_supported_ops_only(requirements.Requirement(r)))
-    for r in _packaging_requirements.REQUIREMENTS
-]
+_PACKAGING_REQUIREMENTS = ["cloudpickle"]
+
 _SNOWFLAKE_PKG_NAME = "snowflake"
 _SNOWFLAKE_ML_PKG_NAME = f"{_SNOWFLAKE_PKG_NAME}.ml"
 
@@ -41,14 +36,16 @@ def create_model_metadata(
     model_dir_path: str,
     name: str,
     model_type: model_types.SupportedModelHandlerType,
-    signatures: Optional[Dict[str, model_signature.ModelSignature]] = None,
-    function_properties: Optional[Dict[str, Dict[str, Any]]] = None,
-    metadata: Optional[Dict[str, str]] = None,
-    code_paths: Optional[List[str]] = None,
-    ext_modules: Optional[List[ModuleType]] = None,
-    conda_dependencies: Optional[List[str]] = None,
-    pip_requirements: Optional[List[str]] = None,
-    artifact_repository_map: Optional[Dict[str, str]] = None,
+    signatures: Optional[dict[str, model_signature.ModelSignature]] = None,
+    function_properties: Optional[dict[str, dict[str, Any]]] = None,
+    metadata: Optional[dict[str, str]] = None,
+    code_paths: Optional[list[str]] = None,
+    ext_modules: Optional[list[ModuleType]] = None,
+    conda_dependencies: Optional[list[str]] = None,
+    pip_requirements: Optional[list[str]] = None,
+    artifact_repository_map: Optional[dict[str, str]] = None,
+    resource_constraint: Optional[dict[str, str]] = None,
+    target_platforms: Optional[list[model_types.TargetPlatform]] = None,
     python_version: Optional[str] = None,
     task: model_types.Task = model_types.Task.UNKNOWN,
     **kwargs: Any,
@@ -69,6 +66,8 @@ def create_model_metadata(
         conda_dependencies: List of conda requirements for running the model. Defaults to None.
         pip_requirements: List of pip Python packages requirements for running the model. Defaults to None.
         artifact_repository_map: A dict mapping from package channel to artifact repository name.
+        resource_constraint: Mapping of resource constraint keys and values, e.g. {"architecture": "x86"}.
+        target_platforms: List of target platforms to run the model.
         python_version: A string of python version where model is run. Used for user override. If specified as None,
             current version would be captured. Defaults to None.
         task: The task of the Model Version. It is an enum class Task with values TABULAR_REGRESSION,
@@ -101,16 +100,19 @@ def create_model_metadata(
         else:
             raise ValueError("`snowflake.ml` is imported via a way that embedding local ML library is not supported.")
 
+    prefer_pip = target_platforms == [model_types.TargetPlatform.SNOWPARK_CONTAINER_SERVICES]
     env = _create_env_for_model_metadata(
         conda_dependencies=conda_dependencies,
         pip_requirements=pip_requirements,
         artifact_repository_map=artifact_repository_map,
+        resource_constraint=resource_constraint,
         python_version=python_version,
         embed_local_ml_library=embed_local_ml_library,
+        prefer_pip=prefer_pip,
     )
 
     if embed_local_ml_library:
-        env.snowpark_ml_version = f"{snowml_env.VERSION}+{file_utils.hash_directory(path_to_copy)}"
+        env.snowpark_ml_version = f"{snowml_version.VERSION}+{file_utils.hash_directory(path_to_copy)}"
 
     model_meta = ModelMetadata(
         name=name,
@@ -152,20 +154,23 @@ def create_model_metadata(
 
 def _create_env_for_model_metadata(
     *,
-    conda_dependencies: Optional[List[str]] = None,
-    pip_requirements: Optional[List[str]] = None,
-    artifact_repository_map: Optional[Dict[str, str]] = None,
+    conda_dependencies: Optional[list[str]] = None,
+    pip_requirements: Optional[list[str]] = None,
+    artifact_repository_map: Optional[dict[str, str]] = None,
+    resource_constraint: Optional[dict[str, str]] = None,
     python_version: Optional[str] = None,
     embed_local_ml_library: bool = False,
+    prefer_pip: bool = False,
 ) -> model_env.ModelEnv:
-    env = model_env.ModelEnv()
+    env = model_env.ModelEnv(prefer_pip=prefer_pip)
 
     # Mypy doesn't like getter and setter have different types. See python/mypy #3004
     env.conda_dependencies = conda_dependencies  # type: ignore[assignment]
     env.pip_requirements = pip_requirements  # type: ignore[assignment]
     env.artifact_repository_map = artifact_repository_map
+    env.resource_constraint = resource_constraint
     env.python_version = python_version  # type: ignore[assignment]
-    env.snowpark_ml_version = snowml_env.VERSION
+    env.snowpark_ml_version = snowml_version.VERSION
 
     requirements_to_add = _PACKAGING_REQUIREMENTS
 
@@ -237,20 +242,20 @@ class ModelMetadata:
         name: str,
         env: model_env.ModelEnv,
         model_type: model_types.SupportedModelHandlerType,
-        runtimes: Optional[Dict[str, model_runtime.ModelRuntime]] = None,
-        signatures: Optional[Dict[str, model_signature.ModelSignature]] = None,
-        function_properties: Optional[Dict[str, Dict[str, Any]]] = None,
-        user_files: Optional[Dict[str, List[str]]] = None,
-        metadata: Optional[Dict[str, str]] = None,
+        runtimes: Optional[dict[str, model_runtime.ModelRuntime]] = None,
+        signatures: Optional[dict[str, model_signature.ModelSignature]] = None,
+        function_properties: Optional[dict[str, dict[str, Any]]] = None,
+        user_files: Optional[dict[str, list[str]]] = None,
+        metadata: Optional[dict[str, str]] = None,
         creation_timestamp: Optional[str] = None,
         min_snowpark_ml_version: Optional[str] = None,
-        models: Optional[Dict[str, model_blob_meta.ModelBlobMeta]] = None,
+        models: Optional[dict[str, model_blob_meta.ModelBlobMeta]] = None,
         original_metadata_version: Optional[str] = model_meta_schema.MODEL_METADATA_VERSION,
         task: model_types.Task = model_types.Task.UNKNOWN,
         explain_algorithm: Optional[model_meta_schema.ModelExplainAlgorithm] = None,
     ) -> None:
         self.name = name
-        self.signatures: Dict[str, model_signature.ModelSignature] = dict()
+        self.signatures: dict[str, model_signature.ModelSignature] = dict()
         if signatures:
             self.signatures = signatures
         self.function_properties = function_properties or {}
@@ -265,7 +270,7 @@ class ModelMetadata:
             else model_meta_schema.MODEL_METADATA_MIN_SNOWPARK_ML_VERSION
         )
 
-        self.models: Dict[str, model_blob_meta.ModelBlobMeta] = dict()
+        self.models: dict[str, model_blob_meta.ModelBlobMeta] = dict()
         if models:
             self.models = models
 
@@ -286,7 +291,7 @@ class ModelMetadata:
         self._min_snowpark_ml_version = max(self._min_snowpark_ml_version, parsed_min_snowpark_ml_version)
 
     @property
-    def runtimes(self) -> Dict[str, model_runtime.ModelRuntime]:
+    def runtimes(self) -> dict[str, model_runtime.ModelRuntime]:
         if self._runtimes and "cpu" in self._runtimes:
             return self._runtimes
         runtimes = {
@@ -353,11 +358,11 @@ class ModelMetadata:
 
         loaded_meta_min_snowpark_ml_version = loaded_meta.get("min_snowpark_ml_version", None)
         if not loaded_meta_min_snowpark_ml_version or (
-            version.parse(loaded_meta_min_snowpark_ml_version) > version.parse(snowml_env.VERSION)
+            version.parse(loaded_meta_min_snowpark_ml_version) > version.parse(snowml_version.VERSION)
         ):
             raise RuntimeError(
                 f"The minimal version required to load the model is {loaded_meta_min_snowpark_ml_version}, "
-                f"while current version of Snowpark ML library is {snowml_env.VERSION}."
+                f"while current version of Snowpark ML library is {snowml_version.VERSION}."
             )
         return model_meta_schema.ModelMetadataDict(
             creation_timestamp=loaded_meta["creation_timestamp"],
@@ -400,7 +405,7 @@ class ModelMetadata:
         env = model_env.ModelEnv()
         env.load_from_dict(pathlib.Path(model_dir_path), model_dict["env"])
 
-        runtimes: Optional[Dict[str, model_runtime.ModelRuntime]]
+        runtimes: Optional[dict[str, model_runtime.ModelRuntime]]
         if model_dict.get("runtimes", None):
             runtimes = {
                 name: model_runtime.ModelRuntime.load(pathlib.Path(model_dir_path), name, env, runtime_dict)

snowflake/ml/model/_packager/model_meta/model_meta_schema.py

@@ -1,7 +1,7 @@
 # This files contains schema definition of what will be written into model.yml
 # Changing this file should lead to a change of the schema version.
 from enum import Enum
-from typing import Any, Dict, List, Optional, TypedDict, Union
+from typing import Any, Optional, TypedDict, Union
 
 from typing_extensions import NotRequired, Required
 
@@ -18,18 +18,20 @@ class FunctionProperties(Enum):
 class ModelRuntimeDependenciesDict(TypedDict):
     conda: Required[str]
     pip: Required[str]
-    artifact_repository_map: NotRequired[Optional[Dict[str, str]]]
+    artifact_repository_map: NotRequired[Optional[dict[str, str]]]
 
 
 class ModelRuntimeDict(TypedDict):
-    imports: Required[List[str]]
+    imports: Required[list[str]]
     dependencies: Required[ModelRuntimeDependenciesDict]
+    resource_constraint: NotRequired[Optional[dict[str, str]]]
 
 
 class ModelEnvDict(TypedDict):
     conda: Required[str]
     pip: Required[str]
-    artifact_repository_map: NotRequired[Optional[Dict[str, str]]]
+    artifact_repository_map: NotRequired[Optional[dict[str, str]]]
+    resource_constraint: NotRequired[Optional[dict[str, str]]]
     python_version: Required[str]
     cuda_version: NotRequired[Optional[str]]
     snowpark_ml_version: Required[str]
@@ -102,25 +104,25 @@ class ModelBlobMetadataDict(TypedDict):
     model_type: Required[type_hints.SupportedModelHandlerType]
     path: Required[str]
     handler_version: Required[str]
-    function_properties: NotRequired[Dict[str, Dict[str, Any]]]
-    artifacts: NotRequired[Dict[str, str]]
+    function_properties: NotRequired[dict[str, dict[str, Any]]]
+    artifacts: NotRequired[dict[str, str]]
     options: NotRequired[ModelBlobOptions]
 
 
 class ModelMetadataDict(TypedDict):
     creation_timestamp: Required[str]
     env: Required[ModelEnvDict]
-    runtimes: NotRequired[Dict[str, ModelRuntimeDict]]
-    metadata: NotRequired[Optional[Dict[str, str]]]
+    runtimes: NotRequired[dict[str, ModelRuntimeDict]]
+    metadata: NotRequired[Optional[dict[str, str]]]
     model_type: Required[type_hints.SupportedModelHandlerType]
-    models: Required[Dict[str, ModelBlobMetadataDict]]
+    models: Required[dict[str, ModelBlobMetadataDict]]
     name: Required[str]
-    signatures: Required[Dict[str, Dict[str, Any]]]
+    signatures: Required[dict[str, dict[str, Any]]]
     version: Required[str]
     min_snowpark_ml_version: Required[str]
     task: Required[str]
     explainability: NotRequired[Optional[ExplainabilityMetadataDict]]
-    function_properties: NotRequired[Dict[str, Dict[str, Any]]]
+    function_properties: NotRequired[dict[str, dict[str, Any]]]
 
 
 class ModelExplainAlgorithm(Enum):

snowflake/ml/model/_packager/model_meta_migrator/base_migrator.py

@@ -1,6 +1,6 @@
 import copy
 from abc import abstractmethod
-from typing import Any, Dict, Protocol, final
+from typing import Any, Protocol, final
 
 from snowflake.ml._internal import migrator_utils
 
@@ -11,13 +11,13 @@ class _BaseModelMetaMigratorProtocol(Protocol):
 
     @staticmethod
     @abstractmethod
-    def upgrade(original_meta_dict: Dict[str, Any]) -> Dict[str, Any]:
+    def upgrade(original_meta_dict: dict[str, Any]) -> dict[str, Any]:
         raise NotImplementedError
 
 
 class BaseModelMetaMigrator(_BaseModelMetaMigratorProtocol):
     @final
-    def try_upgrade(self, original_meta_dict: Dict[str, Any]) -> Dict[str, Any]:
+    def try_upgrade(self, original_meta_dict: dict[str, Any]) -> dict[str, Any]:
         loaded_meta_version = original_meta_dict.get("version", None)
         if not loaded_meta_version or str(loaded_meta_version) != self.source_version:
             raise NotImplementedError(