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

snowflake/ml/monitoring/_client/model_monitor_sql_client.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, List, Mapping, Optional
+from typing import Any, Mapping, Optional
 
 from snowflake import snowpark
 from snowflake.ml._internal.utils import (
@@ -15,7 +15,7 @@ MODEL_JSON_MODEL_NAME_FIELD = "model_name"
 MODEL_JSON_VERSION_NAME_FIELD = "version_name"
 
 
-def _build_sql_list_from_columns(columns: List[sql_identifier.SqlIdentifier]) -> str:
+def _build_sql_list_from_columns(columns: list[sql_identifier.SqlIdentifier]) -> str:
     sql_list = ", ".join([f"'{column}'" for column in columns])
     return f"({sql_list})"
 
@@ -60,17 +60,17 @@ class ModelMonitorSQLClient:
         function_name: str,
         warehouse_name: sql_identifier.SqlIdentifier,
         timestamp_column: sql_identifier.SqlIdentifier,
-        id_columns: List[sql_identifier.SqlIdentifier],
-        prediction_score_columns: List[sql_identifier.SqlIdentifier],
-        prediction_class_columns: List[sql_identifier.SqlIdentifier],
-        actual_score_columns: List[sql_identifier.SqlIdentifier],
-        actual_class_columns: List[sql_identifier.SqlIdentifier],
+        id_columns: list[sql_identifier.SqlIdentifier],
+        prediction_score_columns: list[sql_identifier.SqlIdentifier],
+        prediction_class_columns: list[sql_identifier.SqlIdentifier],
+        actual_score_columns: list[sql_identifier.SqlIdentifier],
+        actual_class_columns: list[sql_identifier.SqlIdentifier],
         refresh_interval: str,
         aggregation_window: str,
         baseline_database: Optional[sql_identifier.SqlIdentifier] = None,
         baseline_schema: Optional[sql_identifier.SqlIdentifier] = None,
         baseline: Optional[sql_identifier.SqlIdentifier] = None,
-        statement_params: Optional[Dict[str, Any]] = None,
+        statement_params: Optional[dict[str, Any]] = None,
     ) -> None:
         baseline_sql = ""
         if baseline:
@@ -103,7 +103,7 @@ class ModelMonitorSQLClient:
         database_name: Optional[sql_identifier.SqlIdentifier] = None,
         schema_name: Optional[sql_identifier.SqlIdentifier] = None,
         monitor_name: sql_identifier.SqlIdentifier,
-        statement_params: Optional[Dict[str, Any]] = None,
+        statement_params: Optional[dict[str, Any]] = None,
     ) -> None:
         search_database_name = database_name or self._database_name
         search_schema_name = schema_name or self._schema_name
@@ -116,8 +116,8 @@ class ModelMonitorSQLClient:
     def show_model_monitors(
         self,
         *,
-        statement_params: Optional[Dict[str, Any]] = None,
-    ) -> List[snowpark.Row]:
+        statement_params: Optional[dict[str, Any]] = None,
+    ) -> list[snowpark.Row]:
         fully_qualified_schema_name = ".".join([self._database_name.identifier(), self._schema_name.identifier()])
         return (
             query_result_checker.SqlResultValidator(
@@ -135,7 +135,7 @@ class ModelMonitorSQLClient:
         database_name: Optional[sql_identifier.SqlIdentifier] = None,
         schema_name: Optional[sql_identifier.SqlIdentifier] = None,
         monitor_name: sql_identifier.SqlIdentifier,
-        statement_params: Optional[Dict[str, Any]] = None,
+        statement_params: Optional[dict[str, Any]] = None,
     ) -> bool:
         search_database_name = database_name or self._database_name
         search_schema_name = schema_name or self._schema_name
@@ -153,7 +153,7 @@ class ModelMonitorSQLClient:
     def validate_monitor_warehouse(
         self,
         warehouse_name: sql_identifier.SqlIdentifier,
-        statement_params: Optional[Dict[str, Any]] = None,
+        statement_params: Optional[dict[str, Any]] = None,
     ) -> None:
         """Validate warehouse provided for monitoring exists.
 
@@ -177,11 +177,11 @@ class ModelMonitorSQLClient:
         *,
         source_column_schema: Mapping[str, types.DataType],
         timestamp_column: sql_identifier.SqlIdentifier,
-        prediction_score_columns: List[sql_identifier.SqlIdentifier],
-        prediction_class_columns: List[sql_identifier.SqlIdentifier],
-        actual_score_columns: List[sql_identifier.SqlIdentifier],
-        actual_class_columns: List[sql_identifier.SqlIdentifier],
-        id_columns: List[sql_identifier.SqlIdentifier],
+        prediction_score_columns: list[sql_identifier.SqlIdentifier],
+        prediction_class_columns: list[sql_identifier.SqlIdentifier],
+        actual_score_columns: list[sql_identifier.SqlIdentifier],
+        actual_class_columns: list[sql_identifier.SqlIdentifier],
+        id_columns: list[sql_identifier.SqlIdentifier],
     ) -> None:
         """Ensures all columns exist in the source table.
 
@@ -221,11 +221,11 @@ class ModelMonitorSQLClient:
         source_schema: Optional[sql_identifier.SqlIdentifier],
         source: sql_identifier.SqlIdentifier,
         timestamp_column: sql_identifier.SqlIdentifier,
-        prediction_score_columns: List[sql_identifier.SqlIdentifier],
-        prediction_class_columns: List[sql_identifier.SqlIdentifier],
-        actual_score_columns: List[sql_identifier.SqlIdentifier],
-        actual_class_columns: List[sql_identifier.SqlIdentifier],
-        id_columns: List[sql_identifier.SqlIdentifier],
+        prediction_score_columns: list[sql_identifier.SqlIdentifier],
+        prediction_class_columns: list[sql_identifier.SqlIdentifier],
+        actual_score_columns: list[sql_identifier.SqlIdentifier],
+        actual_class_columns: list[sql_identifier.SqlIdentifier],
+        id_columns: list[sql_identifier.SqlIdentifier],
     ) -> None:
         source_database = source_database or self._database_name
         source_schema = source_schema or self._schema_name
@@ -250,7 +250,7 @@ class ModelMonitorSQLClient:
         self,
         operation: str,
         monitor_name: sql_identifier.SqlIdentifier,
-        statement_params: Optional[Dict[str, Any]] = None,
+        statement_params: Optional[dict[str, Any]] = None,
     ) -> None:
         if operation not in {"SUSPEND", "RESUME"}:
             raise ValueError(f"Operation {operation} not supported for altering Dynamic Tables")
@@ -263,7 +263,7 @@ class ModelMonitorSQLClient:
     def suspend_monitor(
         self,
         monitor_name: sql_identifier.SqlIdentifier,
-        statement_params: Optional[Dict[str, Any]] = None,
+        statement_params: Optional[dict[str, Any]] = None,
     ) -> None:
         self._alter_monitor(
             operation="SUSPEND",
@@ -274,7 +274,7 @@ class ModelMonitorSQLClient:
     def resume_monitor(
         self,
         monitor_name: sql_identifier.SqlIdentifier,
-        statement_params: Optional[Dict[str, Any]] = None,
+        statement_params: Optional[dict[str, Any]] = None,
     ) -> None:
         self._alter_monitor(
             operation="RESUME",

snowflake/ml/monitoring/_manager/model_monitor_manager.py

@@ -1,5 +1,5 @@
 import json
-from typing import Any, Dict, List, Optional
+from typing import Any, Optional
 
 from snowflake import snowpark
 from snowflake.ml._internal.utils import sql_identifier
@@ -20,7 +20,7 @@ class ModelMonitorManager:
         database_name: sql_identifier.SqlIdentifier,
         schema_name: sql_identifier.SqlIdentifier,
         *,
-        statement_params: Optional[Dict[str, Any]] = None,
+        statement_params: Optional[dict[str, Any]] = None,
     ) -> None:
         """
         Opens a ModelMonitorManager for a given database and schema.
@@ -64,7 +64,7 @@ class ModelMonitorManager:
             f"Found: {existing_target_methods}."
         )
 
-    def _build_column_list_from_input(self, columns: Optional[List[str]]) -> List[sql_identifier.SqlIdentifier]:
+    def _build_column_list_from_input(self, columns: Optional[list[str]]) -> list[sql_identifier.SqlIdentifier]:
         return [sql_identifier.SqlIdentifier(column_name) for column_name in columns] if columns else []
 
     def add_monitor(
@@ -172,7 +172,7 @@ class ModelMonitorManager:
         """
         rows = self._model_monitor_client.show_model_monitors(statement_params=self.statement_params)
 
-        def model_match_fn(model_details: Dict[str, str]) -> bool:
+        def model_match_fn(model_details: dict[str, str]) -> bool:
             return (
                 model_details[model_monitor_sql_client.MODEL_JSON_MODEL_NAME_FIELD] == model_version.model_name
                 and model_details[model_monitor_sql_client.MODEL_JSON_VERSION_NAME_FIELD] == model_version.version_name
@@ -215,7 +215,7 @@ class ModelMonitorManager:
             name=monitor_name_id,
         )
 
-    def show_model_monitors(self) -> List[snowpark.Row]:
+    def show_model_monitors(self) -> list[snowpark.Row]:
         """Show all model monitors in the registry.
 
         Returns:

snowflake/ml/monitoring/entities/model_monitor_config.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass
-from typing import List, Optional
+from typing import Optional
 
 from snowflake.ml.model._client.model import model_version_impl
 
@@ -14,20 +14,20 @@ class ModelMonitorSourceConfig:
     timestamp_column: str
     """Name of column in the source containing timestamp."""
 
-    id_columns: List[str]
+    id_columns: list[str]
     """List of columns in the source containing unique identifiers."""
 
-    prediction_score_columns: Optional[List[str]] = None
+    prediction_score_columns: Optional[list[str]] = None
     """List of columns in the source containing prediction scores.
     Can be regression scores for regression models and probability scores for classification models."""
 
-    prediction_class_columns: Optional[List[str]] = None
+    prediction_class_columns: Optional[list[str]] = None
     """List of columns in the source containing prediction classes for classification models."""
 
-    actual_score_columns: Optional[List[str]] = None
+    actual_score_columns: Optional[list[str]] = None
     """List of columns in the source containing actual scores."""
 
-    actual_class_columns: Optional[List[str]] = None
+    actual_class_columns: Optional[list[str]] = None
     """List of columns in the source containing actual classes for classification models."""
 
     baseline: Optional[str] = None

snowflake/ml/registry/_manager/model_manager.py

@@ -1,11 +1,10 @@
-import os
 from types import ModuleType
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any, Optional, Union
 
 import pandas as pd
 from absl.logging import logging
 
-from snowflake.ml._internal import platform_capabilities, telemetry
+from snowflake.ml._internal import env, platform_capabilities, telemetry
 from snowflake.ml._internal.exceptions import error_codes, exceptions
 from snowflake.ml._internal.human_readable_id import hrid_generator
 from snowflake.ml._internal.utils import sql_identifier
@@ -14,7 +13,6 @@ from snowflake.ml.model._client.model import model_impl, model_version_impl
 from snowflake.ml.model._client.ops import metadata_ops, model_ops, service_ops
 from snowflake.ml.model._model_composer import model_composer
 from snowflake.ml.model._packager.model_meta import model_meta
-from snowflake.ml.modeling._internal import constants
 from snowflake.snowpark import exceptions as snowpark_exceptions, session
 
 logger = logging.getLogger(__name__)
@@ -45,20 +43,21 @@ class ModelManager:
         model_name: str,
         version_name: Optional[str] = None,
         comment: Optional[str] = None,
-        metrics: Optional[Dict[str, Any]] = 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.SupportedTargetPlatformType]] = None,
+        metrics: Optional[dict[str, Any]] = 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.SupportedTargetPlatformType]] = None,
         python_version: Optional[str] = None,
-        signatures: Optional[Dict[str, model_signature.ModelSignature]] = None,
+        signatures: Optional[dict[str, model_signature.ModelSignature]] = None,
         sample_input_data: Optional[model_types.SupportedDataType] = None,
-        user_files: Optional[Dict[str, List[str]]] = None,
-        code_paths: Optional[List[str]] = None,
-        ext_modules: Optional[List[ModuleType]] = None,
+        user_files: Optional[dict[str, list[str]]] = None,
+        code_paths: Optional[list[str]] = None,
+        ext_modules: Optional[list[ModuleType]] = None,
         task: model_types.Task = model_types.Task.UNKNOWN,
         options: Optional[model_types.ModelSaveOption] = None,
-        statement_params: Optional[Dict[str, Any]] = None,
+        statement_params: Optional[dict[str, Any]] = None,
     ) -> model_version_impl.ModelVersion:
 
         database_name_id, schema_name_id, model_name_id = self._parse_fully_qualified_name(model_name)
@@ -131,6 +130,7 @@ class ModelManager:
             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,
             signatures=signatures,
@@ -150,20 +150,21 @@ class ModelManager:
         model_name: str,
         version_name: str,
         comment: Optional[str] = None,
-        metrics: Optional[Dict[str, Any]] = 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.SupportedTargetPlatformType]] = None,
+        metrics: Optional[dict[str, Any]] = 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.SupportedTargetPlatformType]] = None,
         python_version: Optional[str] = None,
-        signatures: Optional[Dict[str, model_signature.ModelSignature]] = None,
+        signatures: Optional[dict[str, model_signature.ModelSignature]] = None,
         sample_input_data: Optional[model_types.SupportedDataType] = None,
-        user_files: Optional[Dict[str, List[str]]] = None,
-        code_paths: Optional[List[str]] = None,
-        ext_modules: Optional[List[ModuleType]] = None,
+        user_files: Optional[dict[str, list[str]]] = None,
+        code_paths: Optional[list[str]] = None,
+        ext_modules: Optional[list[ModuleType]] = None,
         task: model_types.Task = model_types.Task.UNKNOWN,
         options: Optional[model_types.ModelSaveOption] = None,
-        statement_params: Optional[Dict[str, Any]] = None,
+        statement_params: Optional[dict[str, Any]] = None,
     ) -> model_version_impl.ModelVersion:
         database_name_id, schema_name_id, model_name_id = sql_identifier.parse_fully_qualified_name(model_name)
         version_name_id = sql_identifier.SqlIdentifier(version_name)
@@ -212,7 +213,7 @@ class ModelManager:
             platforms = [model_types.TargetPlatform(platform) for platform in target_platforms]
         else:
             # Default the target platform to SPCS if not specified when running in ML runtime
-            if os.getenv(constants.IN_ML_RUNTIME_ENV_VAR):
+            if env.IN_ML_RUNTIME:
                 logger.info(
                     "Logging the model on Container Runtime for ML without specifying `target_platforms`. "
                     'Default to `target_platforms=["SNOWPARK_CONTAINER_SERVICES"]`.'
@@ -253,6 +254,7 @@ class ModelManager:
             conda_dependencies=conda_dependencies,
             pip_requirements=pip_requirements,
             artifact_repository_map=artifact_repository_map,
+            resource_constraint=resource_constraint,
             target_platforms=platforms,
             python_version=python_version,
             user_files=user_files,
@@ -314,7 +316,7 @@ class ModelManager:
         self,
         model_name: str,
         *,
-        statement_params: Optional[Dict[str, Any]] = None,
+        statement_params: Optional[dict[str, Any]] = None,
     ) -> model_impl.Model:
         database_name_id, schema_name_id, model_name_id = self._parse_fully_qualified_name(model_name)
         if self._model_ops.validate_existence(
@@ -342,8 +344,8 @@ class ModelManager:
     def models(
         self,
         *,
-        statement_params: Optional[Dict[str, Any]] = None,
-    ) -> List[model_impl.Model]:
+        statement_params: Optional[dict[str, Any]] = None,
+    ) -> list[model_impl.Model]:
         model_names = self._model_ops.list_models_or_versions(
             database_name=None,
             schema_name=None,
@@ -361,7 +363,7 @@ class ModelManager:
     def show_models(
         self,
         *,
-        statement_params: Optional[Dict[str, Any]] = None,
+        statement_params: Optional[dict[str, Any]] = None,
     ) -> pd.DataFrame:
         rows = self._model_ops.show_models_or_versions(
             database_name=None,
@@ -374,7 +376,7 @@ class ModelManager:
         self,
         model_name: str,
         *,
-        statement_params: Optional[Dict[str, Any]] = None,
+        statement_params: Optional[dict[str, Any]] = None,
     ) -> None:
         database_name_id, schema_name_id, model_name_id = self._parse_fully_qualified_name(model_name)
 
@@ -387,7 +389,7 @@ class ModelManager:
 
     def _parse_fully_qualified_name(
         self, model_name: str
-    ) -> Tuple[
+    ) -> tuple[
         Optional[sql_identifier.SqlIdentifier], Optional[sql_identifier.SqlIdentifier], sql_identifier.SqlIdentifier
     ]:
         try:

snowflake/ml/registry/registry.py

@@ -1,6 +1,6 @@
 import warnings
 from types import ModuleType
-from typing import Any, Dict, List, Optional, Union, overload
+from typing import Any, Optional, Union, overload
 
 import pandas as pd
 
@@ -36,7 +36,7 @@ class Registry:
         *,
         database_name: Optional[str] = None,
         schema_name: Optional[str] = None,
-        options: Optional[Dict[str, Any]] = None,
+        options: Optional[dict[str, Any]] = None,
     ) -> None:
         """Opens a registry within a pre-created Snowflake schema.
 
@@ -107,17 +107,18 @@ class Registry:
         model_name: str,
         version_name: Optional[str] = None,
         comment: Optional[str] = None,
-        metrics: Optional[Dict[str, Any]] = 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.SupportedTargetPlatformType]] = None,
+        metrics: Optional[dict[str, Any]] = 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.SupportedTargetPlatformType]] = None,
         python_version: Optional[str] = None,
-        signatures: Optional[Dict[str, model_signature.ModelSignature]] = None,
+        signatures: Optional[dict[str, model_signature.ModelSignature]] = None,
         sample_input_data: Optional[model_types.SupportedDataType] = None,
-        user_files: Optional[Dict[str, List[str]]] = None,
-        code_paths: Optional[List[str]] = None,
-        ext_modules: Optional[List[ModuleType]] = None,
+        user_files: Optional[dict[str, list[str]]] = None,
+        code_paths: Optional[list[str]] = None,
+        ext_modules: Optional[list[ModuleType]] = None,
         task: model_types.Task = model_types.Task.UNKNOWN,
         options: Optional[model_types.ModelSaveOption] = None,
     ) -> ModelVersion:
@@ -152,6 +153,7 @@ class Registry:
                 Format: {channel_name: artifact_repository_name}, where:
                    - channel_name: The name of the Conda package channel (e.g., 'condaforge') or 'pip' for pip packages.
                    - artifact_repository_name: The name or URL of the repository to fetch packages from.
+            resource_constraint: Mapping of resource constraint keys and values, e.g. {"architecture": "x86"}.
             target_platforms: List of target platforms to run the model. The only acceptable inputs are a combination of
                 {"WAREHOUSE", "SNOWPARK_CONTAINER_SERVICES"}. Defaults to None.
             python_version: Python version in which the model is run. Defaults to None.
@@ -232,6 +234,7 @@ class Registry:
             "conda_dependencies",
             "pip_requirements",
             "artifact_repository_map",
+            "resource_constraint",
             "target_platforms",
             "python_version",
             "signatures",
@@ -244,17 +247,18 @@ class Registry:
         model_name: str,
         version_name: Optional[str] = None,
         comment: Optional[str] = None,
-        metrics: Optional[Dict[str, Any]] = 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.SupportedTargetPlatformType]] = None,
+        metrics: Optional[dict[str, Any]] = 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.SupportedTargetPlatformType]] = None,
         python_version: Optional[str] = None,
-        signatures: Optional[Dict[str, model_signature.ModelSignature]] = None,
+        signatures: Optional[dict[str, model_signature.ModelSignature]] = None,
         sample_input_data: Optional[model_types.SupportedDataType] = None,
-        user_files: Optional[Dict[str, List[str]]] = None,
-        code_paths: Optional[List[str]] = None,
-        ext_modules: Optional[List[ModuleType]] = None,
+        user_files: Optional[dict[str, list[str]]] = None,
+        code_paths: Optional[list[str]] = None,
+        ext_modules: Optional[list[ModuleType]] = None,
         task: model_types.Task = model_types.Task.UNKNOWN,
         options: Optional[model_types.ModelSaveOption] = None,
     ) -> ModelVersion:
@@ -289,6 +293,7 @@ class Registry:
                 Format: {channel_name: artifact_repository_name}, where:
                    - channel_name: The name of the Conda package channel (e.g., 'condaforge') or 'pip' for pip packages.
                    - artifact_repository_name: The name or URL of the repository to fetch packages from.
+            resource_constraint: Mapping of resource constraint keys and values, e.g. {"architecture": "x86"}.
             target_platforms: List of target platforms to run the model. The only acceptable inputs are a combination of
                 {"WAREHOUSE", "SNOWPARK_CONTAINER_SERVICES"}. Defaults to None.
             python_version: Python version in which the model is run. Defaults to None.
@@ -373,6 +378,7 @@ class Registry:
                 conda_dependencies,
                 pip_requirements,
                 artifact_repository_map,
+                resource_constraint,
                 target_platforms,
                 python_version,
                 signatures,
@@ -407,6 +413,7 @@ class Registry:
             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,
             signatures=signatures,
@@ -442,7 +449,7 @@ class Registry:
         project=_TELEMETRY_PROJECT,
         subproject=_MODEL_TELEMETRY_SUBPROJECT,
     )
-    def models(self) -> List[Model]:
+    def models(self) -> list[Model]:
         """Get all models in the schema where the registry is opened.
 
         Returns:
@@ -568,7 +575,7 @@ class Registry:
         subproject=telemetry.TelemetrySubProject.MONITORING.value,
     )
     @snowpark._internal.utils.private_preview(version=model_monitor_version.SNOWFLAKE_ML_MONITORING_MIN_VERSION)
-    def show_model_monitors(self) -> List[snowpark.Row]:
+    def show_model_monitors(self) -> list[snowpark.Row]:
         """Show all model monitors in the registry.
 
         Returns:

snowflake/ml/utils/authentication.py

@@ -1,7 +1,7 @@
 import http
 import logging
 from datetime import timedelta
-from typing import Dict, Optional
+from typing import Optional
 
 import requests
 from cryptography.hazmat.primitives.asymmetric import types
@@ -10,7 +10,7 @@ from requests import auth
 from snowflake.ml._internal.utils import jwt_generator
 
 logger = logging.getLogger(__name__)
-_JWT_TOKEN_CACHE: Dict[str, Dict[int, str]] = {}
+_JWT_TOKEN_CACHE: dict[str, dict[int, str]] = {}
 
 
 def get_jwt_token_generator(

snowflake/ml/utils/connection_params.py

@@ -1,6 +1,6 @@
 import configparser
 import os
-from typing import Dict, Optional, Union
+from typing import Optional, Union
 
 from absl import logging
 from cryptography.hazmat import backends
@@ -76,7 +76,7 @@ def _load_pem_to_der(private_key_path: str) -> bytes:
     )
 
 
-def _connection_properties_from_env() -> Dict[str, str]:
+def _connection_properties_from_env() -> dict[str, str]:
     """Returns a dict with all possible login related env variables."""
     sf_conn_prop = {
         # Mandatory fields
@@ -104,7 +104,7 @@ def _connection_properties_from_env() -> Dict[str, str]:
     return sf_conn_prop
 
 
-def _load_from_snowsql_config_file(connection_name: str, login_file: str = "") -> Dict[str, str]:
+def _load_from_snowsql_config_file(connection_name: str, login_file: str = "") -> dict[str, str]:
     """Loads the dictionary from snowsql config file."""
     snowsql_config_file = login_file if login_file else os.path.expanduser(_DEFAULT_CONNECTION_FILE)
     if not os.path.exists(snowsql_config_file):
@@ -133,7 +133,7 @@ def _load_from_snowsql_config_file(connection_name: str, login_file: str = "") -
 
 
 @snowpark._internal.utils.private_preview(version="0.2.0")
-def SnowflakeLoginOptions(connection_name: str = "", login_file: Optional[str] = None) -> Dict[str, Union[str, bytes]]:
+def SnowflakeLoginOptions(connection_name: str = "", login_file: Optional[str] = None) -> dict[str, Union[str, bytes]]:
     """Returns a dict that can be used directly into snowflake python connector or Snowpark session config.
 
     NOTE: Token/Auth information is sideloaded in all cases above, if provided in following order:
@@ -164,7 +164,7 @@ def SnowflakeLoginOptions(connection_name: str = "", login_file: Optional[str] =
     Raises:
         Exception: if none of config file and environment variable are present.
     """
-    conn_prop: Dict[str, Union[str, bytes]] = {}
+    conn_prop: dict[str, Union[str, bytes]] = {}
     login_file = login_file or os.path.expanduser(_DEFAULT_CONNECTION_FILE)
     # If login file exists, use this exclusively.
     if os.path.exists(login_file):

snowflake/ml/utils/sparse.py

@@ -1,6 +1,6 @@
 import collections
 import json
-from typing import List, Optional
+from typing import Optional
 
 import pandas as pd
 from pandas import arrays as pandas_arrays
@@ -9,7 +9,7 @@ from pandas.core.arrays import sparse as pandas_sparse
 from snowflake.snowpark import DataFrame
 
 
-def _pandas_to_sparse_pandas(pandas_df: pd.DataFrame, sparse_cols: List[str]) -> Optional[pd.DataFrame]:
+def _pandas_to_sparse_pandas(pandas_df: pd.DataFrame, sparse_cols: list[str]) -> Optional[pd.DataFrame]:
     """Convert the pandas df into pandas df with multiple SparseArray columns."""
     num_rows = pandas_df.shape[0]
     if num_rows == 0:
@@ -52,8 +52,9 @@ def _pandas_to_sparse_pandas(pandas_df: pd.DataFrame, sparse_cols: List[str]) ->
     return pandas_df
 
 
-def to_pandas_with_sparse(df: DataFrame, sparse_cols: List[str]) -> pd.DataFrame:
-    """Load a Snowpark df with sparse columns represented in JSON strings into pandas df with multiple SparseArray columns.
+def to_pandas_with_sparse(df: DataFrame, sparse_cols: list[str]) -> pd.DataFrame:
+    """Load a Snowpark df with sparse columns represented in JSON strings into pandas df with multiple SparseArray
+        columns.
 
        For example, for below input:
        ----------------------------------------------

snowflake/ml/utils/sql_client.py

@@ -1,5 +1,4 @@
 from enum import Enum
-from typing import Dict
 
 
 class CreationOption(Enum):
@@ -13,7 +12,7 @@ class CreationMode:
         self.if_not_exists = if_not_exists
         self.or_replace = or_replace
 
-    def get_ddl_phrases(self) -> Dict[CreationOption, str]:
+    def get_ddl_phrases(self) -> dict[CreationOption, str]:
         if_not_exists_sql = " IF NOT EXISTS" if self.if_not_exists else ""
         or_replace_sql = " OR REPLACE" if self.or_replace else ""
         return {

snowflake/ml/version.py

@@ -1 +1,2 @@
-VERSION="1.8.2"
+# This is parsed by regex in conda recipe meta file. Make sure not to break it.
+VERSION = "1.8.3"