snowflake-ml-python 1.8.2__py3-none-any.whl → 1.8.4__py3-none-any.whl

snowflake/ml/model/_client/sql/model_version.py

@@ -1,7 +1,7 @@
 import json
 import pathlib
 import textwrap
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any, Optional
 from urllib.parse import ParseResult
 
 from snowflake.ml._internal.utils import (
@@ -34,7 +34,7 @@ class ModelVersionSQLClient(_base._BaseSQLClient):
         model_name: sql_identifier.SqlIdentifier,
         version_name: sql_identifier.SqlIdentifier,
         stage_path: str,
-        statement_params: Optional[Dict[str, Any]] = None,
+        statement_params: Optional[dict[str, Any]] = None,
     ) -> None:
         query_result_checker.SqlResultValidator(
             self._session,
@@ -56,7 +56,7 @@ class ModelVersionSQLClient(_base._BaseSQLClient):
         schema_name: Optional[sql_identifier.SqlIdentifier],
         model_name: sql_identifier.SqlIdentifier,
         version_name: sql_identifier.SqlIdentifier,
-        statement_params: Optional[Dict[str, Any]] = None,
+        statement_params: Optional[dict[str, Any]] = None,
     ) -> None:
         fq_source_model_name = self.fully_qualified_object_name(
             source_database_name, source_schema_name, source_model_name
@@ -78,7 +78,7 @@ class ModelVersionSQLClient(_base._BaseSQLClient):
         schema_name: Optional[sql_identifier.SqlIdentifier],
         model_name: sql_identifier.SqlIdentifier,
         version_name: sql_identifier.SqlIdentifier,
-        statement_params: Optional[Dict[str, Any]] = None,
+        statement_params: Optional[dict[str, Any]] = None,
     ) -> None:
         sql = (
             f"CREATE MODEL {self.fully_qualified_object_name(database_name, schema_name, model_name)}"
@@ -97,7 +97,7 @@ class ModelVersionSQLClient(_base._BaseSQLClient):
         schema_name: Optional[sql_identifier.SqlIdentifier],
         model_name: sql_identifier.SqlIdentifier,
         version_name: sql_identifier.SqlIdentifier,
-        statement_params: Optional[Dict[str, Any]] = None,
+        statement_params: Optional[dict[str, Any]] = None,
     ) -> None:
         sql = (
             f"ALTER MODEL {self.fully_qualified_object_name(database_name, schema_name, model_name)}"
@@ -116,7 +116,7 @@ class ModelVersionSQLClient(_base._BaseSQLClient):
         schema_name: Optional[sql_identifier.SqlIdentifier],
         model_name: sql_identifier.SqlIdentifier,
         version_name: sql_identifier.SqlIdentifier,
-        statement_params: Optional[Dict[str, Any]] = None,
+        statement_params: Optional[dict[str, Any]] = None,
     ) -> None:
         sql = (
             f"ALTER MODEL {self.fully_qualified_object_name(database_name, schema_name, model_name)}"
@@ -138,7 +138,7 @@ class ModelVersionSQLClient(_base._BaseSQLClient):
         model_name: sql_identifier.SqlIdentifier,
         version_name: sql_identifier.SqlIdentifier,
         stage_path: str,
-        statement_params: Optional[Dict[str, Any]] = None,
+        statement_params: Optional[dict[str, Any]] = None,
     ) -> None:
         query_result_checker.SqlResultValidator(
             self._session,
@@ -160,7 +160,7 @@ class ModelVersionSQLClient(_base._BaseSQLClient):
         schema_name: Optional[sql_identifier.SqlIdentifier],
         model_name: sql_identifier.SqlIdentifier,
         version_name: sql_identifier.SqlIdentifier,
-        statement_params: Optional[Dict[str, Any]] = None,
+        statement_params: Optional[dict[str, Any]] = None,
     ) -> None:
         fq_source_model_name = self.fully_qualified_object_name(
             source_database_name, source_schema_name, source_model_name
@@ -182,7 +182,7 @@ class ModelVersionSQLClient(_base._BaseSQLClient):
         schema_name: Optional[sql_identifier.SqlIdentifier],
         model_name: sql_identifier.SqlIdentifier,
         version_name: sql_identifier.SqlIdentifier,
-        statement_params: Optional[Dict[str, Any]] = None,
+        statement_params: Optional[dict[str, Any]] = None,
     ) -> None:
         query_result_checker.SqlResultValidator(
             self._session,
@@ -201,7 +201,7 @@ class ModelVersionSQLClient(_base._BaseSQLClient):
         model_name: sql_identifier.SqlIdentifier,
         version_name: sql_identifier.SqlIdentifier,
         alias_name: sql_identifier.SqlIdentifier,
-        statement_params: Optional[Dict[str, Any]] = None,
+        statement_params: Optional[dict[str, Any]] = None,
     ) -> None:
         query_result_checker.SqlResultValidator(
             self._session,
@@ -219,7 +219,7 @@ class ModelVersionSQLClient(_base._BaseSQLClient):
         schema_name: Optional[sql_identifier.SqlIdentifier],
         model_name: sql_identifier.SqlIdentifier,
         version_or_alias_name: sql_identifier.SqlIdentifier,
-        statement_params: Optional[Dict[str, Any]] = None,
+        statement_params: Optional[dict[str, Any]] = None,
     ) -> None:
         query_result_checker.SqlResultValidator(
             self._session,
@@ -239,8 +239,8 @@ class ModelVersionSQLClient(_base._BaseSQLClient):
         version_name: sql_identifier.SqlIdentifier,
         file_path: pathlib.PurePosixPath,
         is_dir: bool = False,
-        statement_params: Optional[Dict[str, Any]] = None,
-    ) -> List[row.Row]:
+        statement_params: Optional[dict[str, Any]] = None,
+    ) -> list[row.Row]:
         # Workaround for snowURL bug.
         trailing_slash = "/" if is_dir else ""
 
@@ -276,7 +276,7 @@ class ModelVersionSQLClient(_base._BaseSQLClient):
         version_name: sql_identifier.SqlIdentifier,
         file_path: pathlib.PurePosixPath,
         target_path: pathlib.Path,
-        statement_params: Optional[Dict[str, Any]] = None,
+        statement_params: Optional[dict[str, Any]] = None,
     ) -> pathlib.Path:
         stage_location = pathlib.PurePosixPath(
             self.fully_qualified_object_name(database_name, schema_name, model_name),
@@ -310,8 +310,8 @@ class ModelVersionSQLClient(_base._BaseSQLClient):
         schema_name: Optional[sql_identifier.SqlIdentifier],
         model_name: sql_identifier.SqlIdentifier,
         version_name: sql_identifier.SqlIdentifier,
-        statement_params: Optional[Dict[str, Any]] = None,
-    ) -> List[row.Row]:
+        statement_params: Optional[dict[str, Any]] = None,
+    ) -> list[row.Row]:
         res = query_result_checker.SqlResultValidator(
             self._session,
             (
@@ -331,7 +331,7 @@ class ModelVersionSQLClient(_base._BaseSQLClient):
         model_name: sql_identifier.SqlIdentifier,
         version_name: sql_identifier.SqlIdentifier,
         comment: str,
-        statement_params: Optional[Dict[str, Any]] = None,
+        statement_params: Optional[dict[str, Any]] = None,
     ) -> None:
         query_result_checker.SqlResultValidator(
             self._session,
@@ -351,9 +351,9 @@ class ModelVersionSQLClient(_base._BaseSQLClient):
         version_name: sql_identifier.SqlIdentifier,
         method_name: sql_identifier.SqlIdentifier,
         input_df: dataframe.DataFrame,
-        input_args: List[sql_identifier.SqlIdentifier],
-        returns: List[Tuple[str, spt.DataType, sql_identifier.SqlIdentifier]],
-        statement_params: Optional[Dict[str, Any]] = None,
+        input_args: list[sql_identifier.SqlIdentifier],
+        returns: list[tuple[str, spt.DataType, sql_identifier.SqlIdentifier]],
+        statement_params: Optional[dict[str, Any]] = None,
     ) -> dataframe.DataFrame:
         with_statements = []
         if len(input_df.queries["queries"]) == 1 and len(input_df.queries["post_actions"]) == 0:
@@ -433,10 +433,10 @@ class ModelVersionSQLClient(_base._BaseSQLClient):
         version_name: sql_identifier.SqlIdentifier,
         method_name: sql_identifier.SqlIdentifier,
         input_df: dataframe.DataFrame,
-        input_args: List[sql_identifier.SqlIdentifier],
-        returns: List[Tuple[str, spt.DataType, sql_identifier.SqlIdentifier]],
+        input_args: list[sql_identifier.SqlIdentifier],
+        returns: list[tuple[str, spt.DataType, sql_identifier.SqlIdentifier]],
         partition_column: Optional[sql_identifier.SqlIdentifier],
-        statement_params: Optional[Dict[str, Any]] = None,
+        statement_params: Optional[dict[str, Any]] = None,
         is_partitioned: bool = True,
     ) -> dataframe.DataFrame:
         with_statements = []
@@ -529,13 +529,13 @@ class ModelVersionSQLClient(_base._BaseSQLClient):
 
     def set_metadata(
         self,
-        metadata_dict: Dict[str, Any],
+        metadata_dict: dict[str, Any],
         *,
         database_name: Optional[sql_identifier.SqlIdentifier],
         schema_name: Optional[sql_identifier.SqlIdentifier],
         model_name: sql_identifier.SqlIdentifier,
         version_name: sql_identifier.SqlIdentifier,
-        statement_params: Optional[Dict[str, Any]] = None,
+        statement_params: Optional[dict[str, Any]] = None,
     ) -> None:
         json_metadata = json.dumps(metadata_dict)
         query_result_checker.SqlResultValidator(
@@ -554,7 +554,7 @@ class ModelVersionSQLClient(_base._BaseSQLClient):
         schema_name: Optional[sql_identifier.SqlIdentifier],
         model_name: sql_identifier.SqlIdentifier,
         version_name: sql_identifier.SqlIdentifier,
-        statement_params: Optional[Dict[str, Any]] = None,
+        statement_params: Optional[dict[str, Any]] = None,
     ) -> None:
         query_result_checker.SqlResultValidator(
             self._session,

snowflake/ml/model/_client/sql/service.py

@@ -1,10 +1,9 @@
 import enum
 import json
 import textwrap
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any, Optional, Union
 
 from snowflake import snowpark
-from snowflake.ml._internal import platform_capabilities
 from snowflake.ml._internal.utils import (
     identifier,
     query_result_checker,
@@ -47,7 +46,7 @@ class ServiceSQLClient(_base._BaseSQLClient):
         gpu: Optional[Union[str, int]],
         force_rebuild: bool,
         external_access_integration: sql_identifier.SqlIdentifier,
-        statement_params: Optional[Dict[str, Any]] = None,
+        statement_params: Optional[dict[str, Any]] = None,
     ) -> None:
         actual_image_repo_database = image_repo_database_name or self._database_name
         actual_image_repo_schema = image_repo_schema_name or self._schema_name
@@ -76,8 +75,8 @@ class ServiceSQLClient(_base._BaseSQLClient):
         stage_path: Optional[str] = None,
         model_deployment_spec_yaml_str: Optional[str] = None,
         model_deployment_spec_file_rel_path: Optional[str] = None,
-        statement_params: Optional[Dict[str, Any]] = None,
-    ) -> Tuple[str, snowpark.AsyncJob]:
+        statement_params: Optional[dict[str, Any]] = None,
+    ) -> tuple[str, snowpark.AsyncJob]:
         assert model_deployment_spec_yaml_str or model_deployment_spec_file_rel_path
         if model_deployment_spec_yaml_str:
             sql_str = f"CALL SYSTEM$DEPLOY_MODEL('{model_deployment_spec_yaml_str}')"
@@ -95,9 +94,9 @@ class ServiceSQLClient(_base._BaseSQLClient):
         service_name: sql_identifier.SqlIdentifier,
         method_name: sql_identifier.SqlIdentifier,
         input_df: dataframe.DataFrame,
-        input_args: List[sql_identifier.SqlIdentifier],
-        returns: List[Tuple[str, spt.DataType, sql_identifier.SqlIdentifier]],
-        statement_params: Optional[Dict[str, Any]] = None,
+        input_args: list[sql_identifier.SqlIdentifier],
+        returns: list[tuple[str, spt.DataType, sql_identifier.SqlIdentifier]],
+        statement_params: Optional[dict[str, Any]] = None,
     ) -> dataframe.DataFrame:
         with_statements = []
         actual_database_name = database_name or self._database_name
@@ -133,18 +132,10 @@ class ServiceSQLClient(_base._BaseSQLClient):
             input_args_sql = ", ".join(f"'{arg}', {arg.identifier()}" for arg in input_args)
             args_sql = f"object_construct_keep_null({input_args_sql})"
 
-        if platform_capabilities.PlatformCapabilities.get_instance().is_nested_function_enabled():
-            fully_qualified_service_name = self.fully_qualified_object_name(
-                actual_database_name, actual_schema_name, service_name
-            )
-            fully_qualified_function_name = f"{fully_qualified_service_name}!{method_name.identifier()}"
-        else:
-            function_name = identifier.concat_names([service_name.identifier(), "_", method_name.identifier()])
-            fully_qualified_function_name = identifier.get_schema_level_object_identifier(
-                actual_database_name.identifier(),
-                actual_schema_name.identifier(),
-                function_name,
-            )
+        fully_qualified_service_name = self.fully_qualified_object_name(
+            actual_database_name, actual_schema_name, service_name
+        )
+        fully_qualified_function_name = f"{fully_qualified_service_name}!{method_name.identifier()}"
 
         sql = textwrap.dedent(
             f"""{with_sql}
@@ -181,7 +172,7 @@ class ServiceSQLClient(_base._BaseSQLClient):
         service_name: sql_identifier.SqlIdentifier,
         instance_id: str = "0",
         container_name: str,
-        statement_params: Optional[Dict[str, Any]] = None,
+        statement_params: Optional[dict[str, Any]] = None,
     ) -> str:
         system_func = "SYSTEM$GET_SERVICE_LOGS"
         rows = (
@@ -206,8 +197,8 @@ class ServiceSQLClient(_base._BaseSQLClient):
         schema_name: Optional[sql_identifier.SqlIdentifier],
         service_name: sql_identifier.SqlIdentifier,
         include_message: bool = False,
-        statement_params: Optional[Dict[str, Any]] = None,
-    ) -> Tuple[ServiceStatus, Optional[str]]:
+        statement_params: Optional[dict[str, Any]] = None,
+    ) -> tuple[ServiceStatus, Optional[str]]:
         system_func = "SYSTEM$GET_SERVICE_STATUS"
         rows = (
             query_result_checker.SqlResultValidator(
@@ -231,7 +222,7 @@ class ServiceSQLClient(_base._BaseSQLClient):
         database_name: Optional[sql_identifier.SqlIdentifier],
         schema_name: Optional[sql_identifier.SqlIdentifier],
         service_name: sql_identifier.SqlIdentifier,
-        statement_params: Optional[Dict[str, Any]] = None,
+        statement_params: Optional[dict[str, Any]] = None,
     ) -> None:
         query_result_checker.SqlResultValidator(
             self._session,
@@ -245,8 +236,8 @@ class ServiceSQLClient(_base._BaseSQLClient):
         database_name: Optional[sql_identifier.SqlIdentifier],
         schema_name: Optional[sql_identifier.SqlIdentifier],
         service_name: sql_identifier.SqlIdentifier,
-        statement_params: Optional[Dict[str, Any]] = None,
-    ) -> List[row.Row]:
+        statement_params: Optional[dict[str, Any]] = None,
+    ) -> list[row.Row]:
         fully_qualified_service_name = self.fully_qualified_object_name(database_name, schema_name, service_name)
         res = (
             query_result_checker.SqlResultValidator(

snowflake/ml/model/_client/sql/stage.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, Optional
+from typing import Any, Optional
 
 from snowflake.ml._internal.utils import query_result_checker, sql_identifier
 from snowflake.ml.model._client.sql import _base
@@ -11,7 +11,7 @@ class StageSQLClient(_base._BaseSQLClient):
         database_name: Optional[sql_identifier.SqlIdentifier],
         schema_name: Optional[sql_identifier.SqlIdentifier],
         stage_name: sql_identifier.SqlIdentifier,
-        statement_params: Optional[Dict[str, Any]] = None,
+        statement_params: Optional[dict[str, Any]] = None,
     ) -> None:
         query_result_checker.SqlResultValidator(
             self._session,

snowflake/ml/model/_client/sql/tag.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, List, Optional
+from typing import Any, Optional
 
 from snowflake.ml._internal.utils import query_result_checker, sql_identifier
 from snowflake.ml.model._client.sql import _base
@@ -16,7 +16,7 @@ class ModuleTagSQLClient(_base._BaseSQLClient):
         tag_schema_name: Optional[sql_identifier.SqlIdentifier],
         tag_name: sql_identifier.SqlIdentifier,
         tag_value: str,
-        statement_params: Optional[Dict[str, Any]] = None,
+        statement_params: Optional[dict[str, Any]] = None,
     ) -> None:
         fq_model_name = self.fully_qualified_object_name(database_name, schema_name, model_name)
         fq_tag_name = self.fully_qualified_object_name(tag_database_name, tag_schema_name, tag_name)
@@ -35,7 +35,7 @@ class ModuleTagSQLClient(_base._BaseSQLClient):
         tag_database_name: Optional[sql_identifier.SqlIdentifier],
         tag_schema_name: Optional[sql_identifier.SqlIdentifier],
         tag_name: sql_identifier.SqlIdentifier,
-        statement_params: Optional[Dict[str, Any]] = None,
+        statement_params: Optional[dict[str, Any]] = None,
     ) -> None:
         fq_model_name = self.fully_qualified_object_name(database_name, schema_name, model_name)
         fq_tag_name = self.fully_qualified_object_name(tag_database_name, tag_schema_name, tag_name)
@@ -54,7 +54,7 @@ class ModuleTagSQLClient(_base._BaseSQLClient):
         tag_database_name: Optional[sql_identifier.SqlIdentifier],
         tag_schema_name: Optional[sql_identifier.SqlIdentifier],
         tag_name: sql_identifier.SqlIdentifier,
-        statement_params: Optional[Dict[str, Any]] = None,
+        statement_params: Optional[dict[str, Any]] = None,
     ) -> row.Row:
         fq_model_name = self.fully_qualified_object_name(database_name, schema_name, model_name)
         fq_tag_name = self.fully_qualified_object_name(tag_database_name, tag_schema_name, tag_name)
@@ -75,8 +75,8 @@ class ModuleTagSQLClient(_base._BaseSQLClient):
         database_name: Optional[sql_identifier.SqlIdentifier],
         schema_name: Optional[sql_identifier.SqlIdentifier],
         model_name: sql_identifier.SqlIdentifier,
-        statement_params: Optional[Dict[str, Any]] = None,
-    ) -> List[row.Row]:
+        statement_params: Optional[dict[str, Any]] = None,
+    ) -> list[row.Row]:
         fq_model_name = self.fully_qualified_object_name(database_name, schema_name, model_name)
         actual_database_name = database_name or self._database_name
         return (

snowflake/ml/model/_model_composer/model_composer.py

@@ -3,13 +3,14 @@ import tempfile
 import uuid
 import warnings
 from types import ModuleType
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Optional, Union
 from urllib import parse
 
 from absl import logging
 from packaging import requirements
 
 from snowflake import snowpark
+from snowflake.ml import version as snowml_version
 from snowflake.ml._internal import env as snowml_env, env_utils, file_utils
 from snowflake.ml._internal.lineage import lineage_utils
 from snowflake.ml.data import data_source
@@ -43,7 +44,7 @@ class ModelComposer:
         session: Session,
         stage_path: str,
         *,
-        statement_params: Optional[Dict[str, Any]] = None,
+        statement_params: Optional[dict[str, Any]] = None,
         save_location: Optional[str] = None,
     ) -> None:
         self.session = session
@@ -122,17 +123,18 @@ class ModelComposer:
         *,
         name: str,
         model: model_types.SupportedModelType,
-        signatures: Optional[Dict[str, model_signature.ModelSignature]] = None,
+        signatures: Optional[dict[str, model_signature.ModelSignature]] = None,
         sample_input_data: Optional[model_types.SupportedDataType] = None,
-        metadata: 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,
-        target_platforms: Optional[List[model_types.TargetPlatform]] = None,
+        metadata: 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,
+        target_platforms: Optional[list[model_types.TargetPlatform]] = None,
         python_version: Optional[str] = None,
-        user_files: Optional[Dict[str, List[str]]] = None,
-        ext_modules: Optional[List[ModuleType]] = None,
-        code_paths: Optional[List[str]] = None,
+        user_files: Optional[dict[str, list[str]]] = None,
+        ext_modules: Optional[list[ModuleType]] = None,
+        code_paths: Optional[list[str]] = None,
         task: model_types.Task = model_types.Task.UNKNOWN,
         options: Optional[model_types.ModelSaveOption] = None,
     ) -> model_meta.ModelMetadata:
@@ -140,40 +142,63 @@ class ModelComposer:
         conda_dep_dict = env_utils.validate_conda_dependency_string_list(
             conda_dependencies if conda_dependencies else []
         )
-        is_warehouse_runnable = (
-            not conda_dep_dict
-            or all(
-                chan == env_utils.DEFAULT_CHANNEL_NAME or chan == env_utils.SNOWFLAKE_CONDA_CHANNEL_URL
-                for chan in conda_dep_dict
-            )
-        ) and (not pip_requirements)
-        disable_explainability = (
-            target_platforms and model_types.TargetPlatform.SNOWPARK_CONTAINER_SERVICES in target_platforms
-        ) or (not is_warehouse_runnable)
-
-        if disable_explainability and options and options.get("enable_explainability", False):
-            warnings.warn(
-                ("The model can be deployed to Snowpark Container Services only if `enable_explainability=False`."),
-                category=UserWarning,
-                stacklevel=2,
+
+        enable_explainability = None
+
+        if options:
+            enable_explainability = options.get("enable_explainability", None)
+
+        # skip everything if user said False explicitly
+        if enable_explainability is None or enable_explainability is True:
+            is_warehouse_runnable = (
+                not conda_dep_dict
+                or all(
+                    chan == env_utils.DEFAULT_CHANNEL_NAME or chan == env_utils.SNOWFLAKE_CONDA_CHANNEL_URL
+                    for chan in conda_dep_dict
+                )
+            ) and (not pip_requirements)
+
+            only_spcs = (
+                target_platforms
+                and len(target_platforms) == 1
+                and model_types.TargetPlatform.SNOWPARK_CONTAINER_SERVICES in target_platforms
             )
+            if only_spcs or (not is_warehouse_runnable):
+                # if only SPCS and user asked for explainability we fail
+                if enable_explainability is True:
+                    raise ValueError(
+                        "`enable_explainability` cannot be set to True when the model is not runnable in WH "
+                        "or the target platforms include SPCS."
+                    )
+                elif not options:  # explicitly set flag to false in these cases if not specified
+                    options = model_types.BaseModelSaveOption()
+                    options["enable_explainability"] = False
+            elif (
+                target_platforms
+                and len(target_platforms) > 1
+                and model_types.TargetPlatform.SNOWPARK_CONTAINER_SERVICES in target_platforms
+            ):  # if both then only available for WH
+                if enable_explainability is True:
+                    warnings.warn(
+                        ("Explain function will only be available for model deployed to warehouse."),
+                        category=UserWarning,
+                        stacklevel=2,
+                    )
 
         if not options:
             options = model_types.BaseModelSaveOption()
-            if disable_explainability:
-                options["enable_explainability"] = False
 
         if not snowpark_utils.is_in_stored_procedure():  # type: ignore[no-untyped-call]
             snowml_matched_versions = env_utils.get_matched_package_versions_in_information_schema(
                 self.session,
-                reqs=[requirements.Requirement(f"{env_utils.SNOWPARK_ML_PKG_NAME}=={snowml_env.VERSION}")],
+                reqs=[requirements.Requirement(f"{env_utils.SNOWPARK_ML_PKG_NAME}=={snowml_version.VERSION}")],
                 python_version=python_version or snowml_env.PYTHON_VERSION,
                 statement_params=self._statement_params,
             ).get(env_utils.SNOWPARK_ML_PKG_NAME, [])
 
             if len(snowml_matched_versions) < 1 and options.get("embed_local_ml_library", False) is False:
                 logging.info(
-                    f"Local snowflake-ml-python library has version {snowml_env.VERSION},"
+                    f"Local snowflake-ml-python library has version {snowml_version.VERSION},"
                     " which is not available in the Snowflake server, embedding local ML library automatically."
                 )
                 options["embed_local_ml_library"] = True
@@ -187,6 +212,7 @@ class ModelComposer:
             conda_dependencies=conda_dependencies,
             pip_requirements=pip_requirements,
             artifact_repository_map=artifact_repository_map,
+            resource_constraint=resource_constraint,
             target_platforms=target_platforms,
             python_version=python_version,
             ext_modules=ext_modules,
@@ -226,7 +252,7 @@ class ModelComposer:
 
     def _get_data_sources(
         self, model: model_types.SupportedModelType, sample_input_data: Optional[model_types.SupportedDataType] = None
-    ) -> Optional[List[data_source.DataSource]]:
+    ) -> Optional[list[data_source.DataSource]]:
         data_sources = lineage_utils.get_data_sources(model)
         if not data_sources and sample_input_data is not None:
             data_sources = lineage_utils.get_data_sources(sample_input_data)

snowflake/ml/model/_model_composer/model_manifest/model_manifest.py

@@ -2,7 +2,7 @@ import collections
 import logging
 import pathlib
 import warnings
-from typing import Dict, List, Optional, cast
+from typing import Optional, cast
 
 import yaml
 
@@ -45,10 +45,10 @@ class ModelManifest:
         self,
         model_meta: model_meta_api.ModelMetadata,
         model_rel_path: pathlib.PurePosixPath,
-        user_files: Optional[Dict[str, List[str]]] = None,
+        user_files: Optional[dict[str, list[str]]] = None,
         options: Optional[type_hints.ModelSaveOption] = None,
-        data_sources: Optional[List[data_source.DataSource]] = None,
-        target_platforms: Optional[List[type_hints.TargetPlatform]] = None,
+        data_sources: Optional[list[data_source.DataSource]] = None,
+        target_platforms: Optional[list[type_hints.TargetPlatform]] = None,
     ) -> None:
         if options is None:
             options = {}
@@ -78,12 +78,13 @@ class ModelManifest:
             logger.info(f"Conda dependencies: {runtime_to_use.runtime_env.conda_dependencies}")
             logger.info(f"Pip requirements: {runtime_to_use.runtime_env.pip_requirements}")
             logger.info(f"artifact_repository_map: {runtime_to_use.runtime_env.artifact_repository_map}")
+            logger.info(f"resource_constraint: {runtime_to_use.runtime_env.resource_constraint}")
         runtime_dict = runtime_to_use.save(
             self.workspace_path, default_channel_override=env_utils.SNOWFLAKE_CONDA_CHANNEL_URL
         )
 
         self.function_generator = function_generator.FunctionGenerator(model_dir_rel_path=model_rel_path)
-        self.methods: List[model_method.ModelMethod] = []
+        self.methods: list[model_method.ModelMethod] = []
 
         for target_method in model_meta.signatures.keys():
             method = model_method.ModelMethod(
@@ -100,7 +101,7 @@ class ModelManifest:
 
             self.methods.append(method)
 
-        self.user_files: List[model_user_file.ModelUserFile] = []
+        self.user_files: list[model_user_file.ModelUserFile] = []
 
         if user_files is not None:
             for subdirectory, paths in user_files.items():
@@ -127,16 +128,19 @@ class ModelManifest:
         if model_meta.env.artifact_repository_map:
             dependencies["artifact_repository_map"] = runtime_dict["dependencies"]["artifact_repository_map"]
 
+        runtime = model_manifest_schema.ModelRuntimeDict(
+            language="PYTHON",
+            version=runtime_to_use.runtime_env.python_version,
+            imports=runtime_dict["imports"],
+            dependencies=dependencies,
+        )
+
+        if runtime_dict["resource_constraint"]:
+            runtime["resource_constraint"] = runtime_dict["resource_constraint"]
+
         manifest_dict = model_manifest_schema.ModelManifestDict(
             manifest_version=model_manifest_schema.MODEL_MANIFEST_VERSION,
-            runtimes={
-                self._DEFAULT_RUNTIME_NAME: model_manifest_schema.ModelRuntimeDict(
-                    language="PYTHON",
-                    version=runtime_to_use.runtime_env.python_version,
-                    imports=runtime_dict["imports"],
-                    dependencies=dependencies,
-                )
-            },
+            runtimes={self._DEFAULT_RUNTIME_NAME: runtime},
             methods=[
                 method.save(
                     self.workspace_path,
@@ -178,8 +182,8 @@ class ModelManifest:
         return res
 
     def _extract_lineage_info(
-        self, data_sources: Optional[List[data_source.DataSource]]
-    ) -> List[model_manifest_schema.LineageSourceDict]:
+        self, data_sources: Optional[list[data_source.DataSource]]
+    ) -> list[model_manifest_schema.LineageSourceDict]:
         result = []
         if data_sources:
             for source in data_sources:

snowflake/ml/model/_model_composer/model_manifest/model_manifest_schema.py

@@ -1,6 +1,6 @@
 # This files contains schema definition of what will be written into MANIFEST.yml
 import enum
-from typing import Any, Dict, List, Literal, Optional, TypedDict, Union
+from typing import Any, Literal, Optional, TypedDict, Union
 
 from typing_extensions import NotRequired, Required
 
@@ -20,14 +20,15 @@ class ModelMethodFunctionTypes(enum.Enum):
 class ModelRuntimeDependenciesDict(TypedDict):
     conda: NotRequired[str]
     pip: NotRequired[str]
-    artifact_repository_map: NotRequired[Optional[Dict[str, str]]]
+    artifact_repository_map: NotRequired[Optional[dict[str, str]]]
 
 
 class ModelRuntimeDict(TypedDict):
     language: Required[Literal["PYTHON"]]
     version: Required[str]
-    imports: Required[List[str]]
+    imports: Required[list[str]]
     dependencies: Required[ModelRuntimeDependenciesDict]
+    resource_constraint: NotRequired[Optional[dict[str, str]]]
 
 
 class ModelMethodSignatureField(TypedDict):
@@ -43,8 +44,8 @@ class ModelFunctionMethodDict(TypedDict):
     runtime: Required[str]
     type: Required[str]
     handler: Required[str]
-    inputs: Required[List[ModelMethodSignatureFieldWithName]]
-    outputs: Required[Union[List[ModelMethodSignatureField], List[ModelMethodSignatureFieldWithName]]]
+    inputs: Required[list[ModelMethodSignatureFieldWithName]]
+    outputs: Required[Union[list[ModelMethodSignatureField], list[ModelMethodSignatureFieldWithName]]]
 
 
 ModelMethodDict = ModelFunctionMethodDict
@@ -71,12 +72,12 @@ class ModelFunctionInfo(TypedDict):
 class ModelFunctionInfoDict(TypedDict):
     name: Required[str]
     target_method: Required[str]
-    signature: Required[Dict[str, Any]]
+    signature: Required[dict[str, Any]]
 
 
 class SnowparkMLDataDict(TypedDict):
     schema_version: Required[str]
-    functions: Required[List[ModelFunctionInfoDict]]
+    functions: Required[list[ModelFunctionInfoDict]]
 
 
 class LineageSourceTypes(enum.Enum):
@@ -92,9 +93,9 @@ class LineageSourceDict(TypedDict):
 
 class ModelManifestDict(TypedDict):
     manifest_version: Required[str]
-    runtimes: Required[Dict[str, ModelRuntimeDict]]
-    methods: Required[List[ModelMethodDict]]
-    user_data: NotRequired[Dict[str, Any]]
-    user_files: NotRequired[List[str]]
-    lineage_sources: NotRequired[List[LineageSourceDict]]
-    target_platforms: NotRequired[List[str]]
+    runtimes: Required[dict[str, ModelRuntimeDict]]
+    methods: Required[list[ModelMethodDict]]
+    user_data: NotRequired[dict[str, Any]]
+    user_files: NotRequired[list[str]]
+    lineage_sources: NotRequired[list[LineageSourceDict]]
+    target_platforms: NotRequired[list[str]]