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

snowflake/ml/_internal/file_utils.py

@@ -11,18 +11,7 @@ import sys
 import tarfile
 import tempfile
 import zipfile
-from typing import (
-    Any,
-    Callable,
-    Dict,
-    Generator,
-    List,
-    Literal,
-    Optional,
-    Set,
-    Tuple,
-    Union,
-)
+from typing import Any, Callable, Generator, Literal, Optional, Union
 from urllib import parse
 
 import cloudpickle
@@ -37,7 +26,7 @@ GENERATED_PY_FILE_EXT = (".pyc", ".pyo", ".pyd", ".pyi")
 def copytree(
     src: "Union[str, os.PathLike[str]]",
     dst: "Union[str, os.PathLike[str]]",
-    ignore: Optional[Callable[..., Set[str]]] = None,
+    ignore: Optional[Callable[..., set[str]]] = None,
     dirs_exist_ok: bool = False,
 ) -> "Union[str, os.PathLike[str]]":
     """This is a forked version of shutil.copytree that remove all copystat, to make sure it works in Sproc.
@@ -170,7 +159,7 @@ def zip_python_package(zipfile_path: str, package_name: str, ignore_generated_py
 
 
 def hash_directory(
-    directory: Union[str, pathlib.Path], *, ignore_hidden: bool = False, excluded_files: Optional[List[str]] = None
+    directory: Union[str, pathlib.Path], *, ignore_hidden: bool = False, excluded_files: Optional[list[str]] = None
 ) -> str:
     """Hash the **content** of a folder recursively using SHA-1.
 
@@ -186,7 +175,7 @@ def hash_directory(
         excluded_files = []
 
     def _update_hash_from_dir(
-        directory: Union[str, pathlib.Path], hash: "hashlib._Hash", *, ignore_hidden: bool, excluded_files: List[str]
+        directory: Union[str, pathlib.Path], hash: "hashlib._Hash", *, ignore_hidden: bool, excluded_files: list[str]
     ) -> "hashlib._Hash":
         assert pathlib.Path(directory).is_dir(), "Provided path is not a directory."
         for path in sorted(pathlib.Path(directory).iterdir(), key=lambda p: str(p).lower()):
@@ -208,7 +197,7 @@ def hash_directory(
     ).hexdigest()
 
 
-def get_all_modules(dirname: str, prefix: str = "") -> List[str]:
+def get_all_modules(dirname: str, prefix: str = "") -> list[str]:
     modules = [mod.name for mod in pkgutil.iter_modules([dirname], prefix=prefix)]
     subdirs = [f.path for f in os.scandir(dirname) if f.is_dir()]
     for sub_dirname in subdirs:
@@ -248,7 +237,7 @@ def _create_tar_gz_stream(source_dir: str, arcname: Optional[str] = None) -> Gen
         yield output_stream
 
 
-def get_package_path(package_name: str, strategy: Literal["first", "last"] = "first") -> Tuple[str, str]:
+def get_package_path(package_name: str, strategy: Literal["first", "last"] = "first") -> tuple[str, str]:
     """[Obsolete]Return the path to where a package is defined and its start location.
     Example 1: snowflake.ml -> path/to/site-packages/snowflake/ml, path/to/site-packages
     Example 2: zip_imported_module -> path/to/some/zipfile.zip/zip_imported_module, path/to/some/zipfile.zip
@@ -267,7 +256,7 @@ def get_package_path(package_name: str, strategy: Literal["first", "last"] = "fi
     return pkg_path, pkg_start_path
 
 
-def stage_object(session: snowpark.Session, object: object, stage_location: str) -> List[snowpark.PutResult]:
+def stage_object(session: snowpark.Session, object: object, stage_location: str) -> list[snowpark.PutResult]:
     temp_file = tempfile.NamedTemporaryFile(delete=False)
     temp_file_path = temp_file.name
     temp_file.close()
@@ -279,7 +268,7 @@ def stage_object(session: snowpark.Session, object: object, stage_location: str)
 
 
 def stage_file_exists(
-    session: snowpark.Session, stage_location: str, file_name: str, statement_params: Dict[str, Any]
+    session: snowpark.Session, stage_location: str, file_name: str, statement_params: dict[str, Any]
 ) -> bool:
     try:
         res = session.sql(f"list {stage_location}/{file_name}").collect(statement_params=statement_params)
@@ -297,7 +286,7 @@ def upload_directory_to_stage(
     local_path: pathlib.Path,
     stage_path: Union[pathlib.PurePosixPath, parse.ParseResult],
     *,
-    statement_params: Optional[Dict[str, Any]] = None,
+    statement_params: Optional[dict[str, Any]] = None,
 ) -> None:
     """Upload a local folder recursively to a stage and keep the structure.
 
@@ -350,7 +339,7 @@ def download_directory_from_stage(
     stage_path: pathlib.PurePosixPath,
     local_path: pathlib.Path,
     *,
-    statement_params: Optional[Dict[str, Any]] = None,
+    statement_params: Optional[dict[str, Any]] = None,
 ) -> None:
     """Upload a folder in stage recursively to a folder in local and keep the structure.
 

snowflake/ml/_internal/human_readable_id/hrid_generator_base.py

@@ -15,7 +15,6 @@ In this module you will find:
 
 import math
 from abc import ABC, abstractmethod
-from typing import Dict, List, Tuple
 
 
 class HRIDBase(ABC):
@@ -28,12 +27,11 @@ class HRIDBase(ABC):
     @abstractmethod
     def __id_generator__(self) -> int:
         """The generator to use to generate new IDs. The implementer needs to provide this."""
-        pass
 
-    __hrid_structure__: Tuple[str, ...]
+    __hrid_structure__: tuple[str, ...]
     """The HRID structure to be generated. The implementer needs to provide this."""
 
-    __hrid_words__: Dict[str, Tuple[str, ...]]
+    __hrid_words__: dict[str, tuple[str, ...]]
     """The mapping between the HRID parts and the words to use. The implementer needs to provide this."""
 
     __separator__ = "_"
@@ -82,7 +80,7 @@ class HRIDBase(ABC):
             hrid.append(str(values[idxs[i]]))
         return self.__separator__.join(hrid)
 
-    def generate(self) -> Tuple[int, str]:
+    def generate(self) -> tuple[int, str]:
         """Generate an ID and the corresponding HRID.
 
         Returns:
@@ -92,7 +90,7 @@ class HRIDBase(ABC):
         hrid = self.id_to_hrid(id)
         return (id, hrid)
 
-    def _id_to_idxs(self, id: int) -> List[int]:
+    def _id_to_idxs(self, id: int) -> list[int]:
         """Take the ID and convert it to indices into the HRID words.
 
         Args:
@@ -109,7 +107,7 @@ class HRIDBase(ABC):
             idxs.append((id & mask) >> shift)
         return idxs
 
-    def _hrid_to_idxs(self, hrid: str) -> List[int]:
+    def _hrid_to_idxs(self, hrid: str) -> list[int]:
         """Take the HRID and convert it to indices into the HRID words.
 
         Args:

snowflake/ml/_internal/init_utils.py

@@ -2,10 +2,9 @@ import importlib
 import inspect
 import pkgutil
 from types import FunctionType
-from typing import Dict
 
 
-def fetch_classes_from_modules_in_pkg_dir(pkg_dir: str, pkg_name: str) -> Dict[str, type]:
+def fetch_classes_from_modules_in_pkg_dir(pkg_dir: str, pkg_name: str) -> dict[str, type]:
     """Finds classes defined all the python modules in the given package directory.
 
     Args:
@@ -36,7 +35,7 @@ def fetch_classes_from_modules_in_pkg_dir(pkg_dir: str, pkg_name: str) -> Dict[s
     return exportable_classes
 
 
-def fetch_functions_from_modules_in_pkg_dir(pkg_dir: str, pkg_name: str) -> Dict[str, FunctionType]:
+def fetch_functions_from_modules_in_pkg_dir(pkg_dir: str, pkg_name: str) -> dict[str, FunctionType]:
     """Finds functions defined all the python modules in the given package directory.
 
     Args:

snowflake/ml/_internal/lineage/lineage_utils.py

@@ -1,6 +1,6 @@
 import copy
 import functools
-from typing import Any, Callable, List, Optional, get_args
+from typing import Any, Callable, Optional, get_args
 
 from snowflake import snowpark
 from snowflake.ml.data import data_source
@@ -9,7 +9,7 @@ _DATA_SOURCES_ATTR = "_data_sources"
 
 
 def _wrap_func(
-    fn: Callable[..., snowpark.DataFrame], data_sources: List[data_source.DataSource]
+    fn: Callable[..., snowpark.DataFrame], data_sources: list[data_source.DataSource]
 ) -> Callable[..., snowpark.DataFrame]:
     """Wrap a DataFrame transform function to propagate data_sources to derived DataFrames."""
 
@@ -34,9 +34,9 @@ def _wrap_class_func(fn: Callable[..., snowpark.DataFrame]) -> Callable[..., sno
     return wrapped
 
 
-def get_data_sources(*args: Any) -> Optional[List[data_source.DataSource]]:
+def get_data_sources(*args: Any) -> Optional[list[data_source.DataSource]]:
     """Helper method for extracting data sources attribute from DataFrames in an argument list"""
-    result: Optional[List[data_source.DataSource]] = None
+    result: Optional[list[data_source.DataSource]] = None
     for arg in args:
         srcs = getattr(arg, _DATA_SOURCES_ATTR, None)
         if isinstance(srcs, list) and all(isinstance(s, get_args(data_source.DataSource)) for s in srcs):
@@ -46,7 +46,7 @@ def get_data_sources(*args: Any) -> Optional[List[data_source.DataSource]]:
     return result
 
 
-def set_data_sources(obj: Any, data_sources: Optional[List[data_source.DataSource]]) -> None:
+def set_data_sources(obj: Any, data_sources: Optional[list[data_source.DataSource]]) -> None:
     """Helper method for attaching data sources to an object"""
     if data_sources:
         assert all(isinstance(ds, get_args(data_source.DataSource)) for ds in data_sources)
@@ -54,7 +54,7 @@ def set_data_sources(obj: Any, data_sources: Optional[List[data_source.DataSourc
 
 
 def patch_dataframe(
-    df: snowpark.DataFrame, data_sources: List[data_source.DataSource], inplace: bool = False
+    df: snowpark.DataFrame, data_sources: list[data_source.DataSource], inplace: bool = False
 ) -> snowpark.DataFrame:
     """
     Monkey patch a DataFrame to add attach the provided data_sources as an attribute of the DataFrame.

snowflake/ml/_internal/platform_capabilities.py

@@ -1,6 +1,6 @@
 import json
 from contextlib import contextmanager
-from typing import Any, Dict, Optional
+from typing import Any, Optional
 
 from absl import logging
 
@@ -11,6 +11,9 @@ from snowflake.snowpark import (
     session as snowpark_session,
 )
 
+LIVE_COMMIT_PARAMETER = "ENABLE_LIVE_VERSION_IN_SDK"
+INLINE_DEPLOYMENT_SPEC_PARAMETER = "ENABLE_INLINE_DEPLOYMENT_SPEC"
+
 
 class PlatformCapabilities:
     """Class that retrieves platform feature values for the currently running server.
@@ -18,18 +21,18 @@ class PlatformCapabilities:
     Example usage:
     ```
     pc = PlatformCapabilities.get_instance(session)
-    if pc.is_nested_function_enabled():
-        # Nested functions are enabled.
-        print("Nested functions are enabled.")
+    if pc.is_inlined_deployment_spec_enabled():
+        # Inline deployment spec is enabled.
+        print("Inline deployment spec is enabled.")
     else:
-        # Nested functions are disabled.
-        print("Nested functions are disabled or not supported.")
+        # Inline deployment spec is disabled.
+        print("Inline deployment spec is disabled or not supported.")
     ```
     """
 
     _instance: Optional["PlatformCapabilities"] = None
     # Used for unittesting only. This is to avoid the need to mock the session object or reaching out to Snowflake
-    _mock_features: Optional[Dict[str, Any]] = None
+    _mock_features: Optional[dict[str, Any]] = None
 
     @classmethod
     def get_instance(cls, session: Optional[snowpark_session.Session] = None) -> "PlatformCapabilities":
@@ -41,7 +44,7 @@ class PlatformCapabilities:
         return cls._instance
 
     @classmethod
-    def set_mock_features(cls, features: Optional[Dict[str, Any]] = None) -> None:
+    def set_mock_features(cls, features: Optional[dict[str, Any]] = None) -> None:
         cls._mock_features = features
 
     @classmethod
@@ -50,9 +53,11 @@ class PlatformCapabilities:
 
     # For contextmanager, we need to have return type Iterator[Never]. However, Never type is introduced only in
     # Python 3.11. So, we are ignoring the type for this method.
+    _dummy_features: dict[str, Any] = {"dummy": "dummy"}
+
     @classmethod  # type: ignore[arg-type]
     @contextmanager
-    def mock_features(cls, features: Dict[str, Any]) -> None:  # type: ignore[misc]
+    def mock_features(cls, features: dict[str, Any] = _dummy_features) -> None:  # type: ignore[misc]
         logging.debug(f"Setting mock features: {features}")
         cls.set_mock_features(features)
         try:
@@ -61,17 +66,14 @@ class PlatformCapabilities:
             logging.debug(f"Clearing mock features: {features}")
             cls.clear_mock_features()
 
-    def is_nested_function_enabled(self) -> bool:
-        return self._get_bool_feature("SPCS_MODEL_ENABLE_EMBEDDED_SERVICE_FUNCTIONS", False)
-
     def is_inlined_deployment_spec_enabled(self) -> bool:
-        return self._get_bool_feature("ENABLE_INLINE_DEPLOYMENT_SPEC", False)
+        return self._get_bool_feature(INLINE_DEPLOYMENT_SPEC_PARAMETER, False)
 
     def is_live_commit_enabled(self) -> bool:
-        return self._get_bool_feature("ENABLE_BUNDLE_MODULE_CHECKOUT", False)
+        return self._get_bool_feature(LIVE_COMMIT_PARAMETER, False)
 
     @staticmethod
-    def _get_features(session: snowpark_session.Session) -> Dict[str, Any]:
+    def _get_features(session: snowpark_session.Session) -> dict[str, Any]:
         try:
             result = (
                 query_result_checker.SqlResultValidator(
@@ -99,7 +101,7 @@ class PlatformCapabilities:
         return {}
 
     def __init__(
-        self, *, session: Optional[snowpark_session.Session] = None, features: Optional[Dict[str, Any]] = None
+        self, *, session: Optional[snowpark_session.Session] = None, features: Optional[dict[str, Any]] = None
     ) -> None:
         # This is for testing purposes only.
         if features:

snowflake/ml/_internal/telemetry.py

@@ -8,25 +8,13 @@ import sys
 import time
 import traceback
 import types
-from typing import (
-    Any,
-    Callable,
-    Dict,
-    Iterable,
-    List,
-    Mapping,
-    Optional,
-    Set,
-    Tuple,
-    TypeVar,
-    Union,
-    cast,
-)
+from typing import Any, Callable, Iterable, Mapping, Optional, TypeVar, Union, cast
 
 from typing_extensions import ParamSpec
 
 from snowflake import connector
 from snowflake.connector import telemetry as connector_telemetry, time_util
+from snowflake.ml import version as snowml_version
 from snowflake.ml._internal import env
 from snowflake.ml._internal.exceptions import (
     error_codes,
@@ -99,13 +87,13 @@ class _TelemetrySourceType(enum.Enum):
     AUGMENT_TELEMETRY = "SNOWML_AUGMENT_TELEMETRY"
 
 
-_statement_params_context_var: contextvars.ContextVar[Dict[str, str]] = contextvars.ContextVar("statement_params")
+_statement_params_context_var: contextvars.ContextVar[dict[str, str]] = contextvars.ContextVar("statement_params")
 
 
 class _StatementParamsPatchManager:
     def __init__(self) -> None:
-        self._patch_cache: Set[server_connection.ServerConnection] = set()
-        self._context_var: contextvars.ContextVar[Dict[str, str]] = _statement_params_context_var
+        self._patch_cache: set[server_connection.ServerConnection] = set()
+        self._context_var: contextvars.ContextVar[dict[str, str]] = _statement_params_context_var
 
     def apply_patches(self) -> None:
         try:
@@ -117,7 +105,7 @@ class _StatementParamsPatchManager:
         except snowpark_exceptions.SnowparkSessionException:
             pass
 
-    def set_statement_params(self, statement_params: Dict[str, str]) -> None:
+    def set_statement_params(self, statement_params: dict[str, str]) -> None:
         # Only set value if not already set in context
         if not self._context_var.get({}):
             self._context_var.set(statement_params)
@@ -152,7 +140,6 @@ class _StatementParamsPatchManager:
                 if throw_on_patch_fail:  # primarily used for testing
                     raise
                 # TODO: Log a warning, this probably means there was a breaking change in Snowpark/SnowflakeConnection
-                pass
 
     def _patch_with_statement_params(
         self, target: object, function_name: str, param_name: str = "statement_params"
@@ -197,10 +184,10 @@ class _StatementParamsPatchManager:
 
         setattr(target, function_name, wrapper)
 
-    def __getstate__(self) -> Dict[str, Any]:
+    def __getstate__(self) -> dict[str, Any]:
         return {}
 
-    def __setstate__(self, state: Dict[str, Any]) -> None:
+    def __setstate__(self, state: dict[str, Any]) -> None:
         # unpickling does not call __init__ by default, do it manually here
         self.__init__()  # type: ignore[misc]
 
@@ -210,7 +197,7 @@ _patch_manager = _StatementParamsPatchManager()
 
 def get_statement_params(
     project: str, subproject: Optional[str] = None, class_name: Optional[str] = None
-) -> Dict[str, Any]:
+) -> dict[str, Any]:
     """
     Get telemetry statement parameters.
 
@@ -231,8 +218,8 @@ def get_statement_params(
 
 
 def add_statement_params_custom_tags(
-    statement_params: Optional[Dict[str, Any]], custom_tags: Mapping[str, Any]
-) -> Dict[str, Any]:
+    statement_params: Optional[dict[str, Any]], custom_tags: Mapping[str, Any]
+) -> dict[str, Any]:
     """
     Add custom_tags to existing statement_params.  Overwrite keys in custom_tags dict that already exist.
     If existing statement_params are not provided, do nothing as the information cannot be effectively tracked.
@@ -246,7 +233,7 @@ def add_statement_params_custom_tags(
     """
     if not statement_params:
         return {}
-    existing_custom_tags: Dict[str, Any] = statement_params.pop(TelemetryField.KEY_CUSTOM_TAGS.value, {})
+    existing_custom_tags: dict[str, Any] = statement_params.pop(TelemetryField.KEY_CUSTOM_TAGS.value, {})
     existing_custom_tags.update(custom_tags)
     # NOTE: This can be done with | operator after upgrade from py3.8
     return {
@@ -289,17 +276,17 @@ def get_function_usage_statement_params(
     *,
     function_category: str = TelemetryField.FUNC_CAT_USAGE.value,
     function_name: Optional[str] = None,
-    function_parameters: Optional[Dict[str, Any]] = None,
+    function_parameters: Optional[dict[str, Any]] = None,
     api_calls: Optional[
-        List[
+        list[
             Union[
-                Dict[str, Union[Callable[..., Any], str]],
+                dict[str, Union[Callable[..., Any], str]],
                 Union[Callable[..., Any], str],
             ]
         ]
     ] = None,
-    custom_tags: Optional[Dict[str, Union[bool, int, str, float]]] = None,
-) -> Dict[str, Any]:
+    custom_tags: Optional[dict[str, Union[bool, int, str, float]]] = None,
+) -> dict[str, Any]:
     """
     Get function usage statement parameters.
 
@@ -321,12 +308,12 @@ def get_function_usage_statement_params(
         >>> df.collect(statement_params=statement_params)
     """
     telemetry_type = f"{env.SOURCE.lower()}_{TelemetryField.TYPE_FUNCTION_USAGE.value}"
-    statement_params: Dict[str, Any] = {
+    statement_params: dict[str, Any] = {
         connector_telemetry.TelemetryField.KEY_SOURCE.value: env.SOURCE,
         TelemetryField.KEY_PROJECT.value: project,
         TelemetryField.KEY_SUBPROJECT.value: subproject,
         TelemetryField.KEY_OS.value: env.OS,
-        TelemetryField.KEY_VERSION.value: env.VERSION,
+        TelemetryField.KEY_VERSION.value: snowml_version.VERSION,
         TelemetryField.KEY_PYTHON_VERSION.value: env.PYTHON_VERSION,
         connector_telemetry.TelemetryField.KEY_TYPE.value: telemetry_type,
         TelemetryField.KEY_CATEGORY.value: function_category,
@@ -339,7 +326,7 @@ def get_function_usage_statement_params(
     if api_calls:
         statement_params[TelemetryField.KEY_API_CALLS.value] = []
         for api_call in api_calls:
-            if isinstance(api_call, Dict):
+            if isinstance(api_call, dict):
                 telemetry_api_call = api_call.copy()
                 # convert Callable to str
                 for field, api in api_call.items():
@@ -388,7 +375,7 @@ def send_custom_usage(
     *,
     telemetry_type: str,
     subproject: Optional[str] = None,
-    data: Optional[Dict[str, Any]] = None,
+    data: Optional[dict[str, Any]] = None,
     **kwargs: Any,
 ) -> None:
     active_session = next(iter(session._get_active_sessions()))
@@ -409,17 +396,17 @@ def send_api_usage_telemetry(
     api_calls_extractor: Optional[
         Callable[
             ...,
-            List[
+            list[
                 Union[
-                    Dict[str, Union[Callable[..., Any], str]],
+                    dict[str, Union[Callable[..., Any], str]],
                     Union[Callable[..., Any], str],
                 ]
             ],
         ]
     ] = None,
-    sfqids_extractor: Optional[Callable[..., List[str]]] = None,
+    sfqids_extractor: Optional[Callable[..., list[str]]] = None,
     subproject_extractor: Optional[Callable[[Any], str]] = None,
-    custom_tags: Optional[Dict[str, Union[bool, int, str, float]]] = None,
+    custom_tags: Optional[dict[str, Union[bool, int, str, float]]] = None,
 ) -> Callable[[Callable[_Args, _ReturnValue]], Callable[_Args, _ReturnValue]]:
     """
     Decorator that sends API usage telemetry and adds function usage statement parameters to the dataframe returned by
@@ -454,7 +441,7 @@ def send_api_usage_telemetry(
         def wrap(*args: Any, **kwargs: Any) -> _ReturnValue:
             params = _get_func_params(func, func_params_to_log, args, kwargs) if func_params_to_log else None
 
-            api_calls: List[Union[Dict[str, Union[Callable[..., Any], str]], Callable[..., Any], str]] = []
+            api_calls: list[Union[dict[str, Union[Callable[..., Any], str]], Callable[..., Any], str]] = []
             if api_calls_extractor:
                 extracted_api_calls = api_calls_extractor(args[0])
                 for api_call in extracted_api_calls:
@@ -484,7 +471,7 @@ def send_api_usage_telemetry(
                 custom_tags=custom_tags,
             )
 
-            def update_stmt_params_if_snowpark_df(obj: _ReturnValue, statement_params: Dict[str, Any]) -> _ReturnValue:
+            def update_stmt_params_if_snowpark_df(obj: _ReturnValue, statement_params: dict[str, Any]) -> _ReturnValue:
                 """
                 Update SnowML function usage statement parameters to the object if it is a Snowpark DataFrame.
                 Used to track APIs returning a Snowpark DataFrame.
@@ -614,7 +601,7 @@ def _get_full_func_name(func: Callable[..., Any]) -> str:
 
 def _get_func_params(
     func: Callable[..., Any], func_params_to_log: Optional[Iterable[str]], args: Any, kwargs: Any
-) -> Dict[str, Any]:
+) -> dict[str, Any]:
     """
     Get function parameters.
 
@@ -639,7 +626,7 @@ def _get_func_params(
     return params
 
 
-def _extract_arg_value(field: str, func_spec: inspect.FullArgSpec, args: Any, kwargs: Any) -> Tuple[bool, Any]:
+def _extract_arg_value(field: str, func_spec: inspect.FullArgSpec, args: Any, kwargs: Any) -> tuple[bool, Any]:
     """
     Function to extract a specified argument value.
 
@@ -702,11 +689,11 @@ class _SourceTelemetryClient:
         self.source: str = env.SOURCE
         self.project: Optional[str] = project
         self.subproject: Optional[str] = subproject
-        self.version = env.VERSION
+        self.version = snowml_version.VERSION
         self.python_version: str = env.PYTHON_VERSION
         self.os: str = env.OS
 
-    def _send(self, msg: Dict[str, Any], timestamp: Optional[int] = None) -> None:
+    def _send(self, msg: dict[str, Any], timestamp: Optional[int] = None) -> None:
         """
         Add telemetry data to a batch in connector client.
 
@@ -720,7 +707,7 @@ class _SourceTelemetryClient:
             telemetry_data = connector_telemetry.TelemetryData(message=msg, timestamp=timestamp)
             self._telemetry.try_add_log_to_batch(telemetry_data)
 
-    def _create_basic_telemetry_data(self, telemetry_type: str) -> Dict[str, Any]:
+    def _create_basic_telemetry_data(self, telemetry_type: str) -> dict[str, Any]:
         message = {
             connector_telemetry.TelemetryField.KEY_SOURCE.value: self.source,
             TelemetryField.KEY_PROJECT.value: self.project,
@@ -738,10 +725,10 @@ class _SourceTelemetryClient:
         func_name: str,
         function_category: str,
         duration: float,
-        func_params: Optional[Dict[str, Any]] = None,
-        api_calls: Optional[List[Dict[str, Any]]] = None,
-        sfqids: Optional[List[Any]] = None,
-        custom_tags: Optional[Dict[str, Union[bool, int, str, float]]] = None,
+        func_params: Optional[dict[str, Any]] = None,
+        api_calls: Optional[list[dict[str, Any]]] = None,
+        sfqids: Optional[list[Any]] = None,
+        custom_tags: Optional[dict[str, Union[bool, int, str, float]]] = None,
         error: Optional[str] = None,
         error_code: Optional[str] = None,
         stack_trace: Optional[str] = None,
@@ -761,7 +748,7 @@ class _SourceTelemetryClient:
             error_code: Error code.
             stack_trace: Error stack trace.
         """
-        data: Dict[str, Any] = {
+        data: dict[str, Any] = {
             TelemetryField.KEY_FUNC_NAME.value: func_name,
             TelemetryField.KEY_CATEGORY.value: function_category,
         }
@@ -775,7 +762,7 @@ class _SourceTelemetryClient:
             data[TelemetryField.KEY_CUSTOM_TAGS.value] = custom_tags
 
         telemetry_type = f"{self.source.lower()}_{TelemetryField.TYPE_FUNCTION_USAGE.value}"
-        message: Dict[str, Any] = {
+        message: dict[str, Any] = {
             **self._create_basic_telemetry_data(telemetry_type),
             TelemetryField.KEY_DATA.value: data,
             TelemetryField.KEY_DURATION.value: duration,
@@ -795,7 +782,7 @@ class _SourceTelemetryClient:
             self._telemetry.send_batch()
 
 
-def get_sproc_statement_params_kwargs(sproc: Callable[..., Any], statement_params: Dict[str, Any]) -> Dict[str, Any]:
+def get_sproc_statement_params_kwargs(sproc: Callable[..., Any], statement_params: dict[str, Any]) -> dict[str, Any]:
     """
     Get statement_params keyword argument for sproc call.
 

snowflake/ml/_internal/type_utils.py

@@ -11,7 +11,7 @@ T = TypeVar("T")
 class LazyType(Generic[T]):
     """Utility type to help defer need of importing."""
 
-    def __init__(self, klass: Union[str, Type[T]]) -> None:
+    def __init__(self, klass: Union[str, type[T]]) -> None:
         self.qualname = ""
         if isinstance(klass, str):
             parts = klass.rsplit(".", 1)
@@ -30,7 +30,7 @@ class LazyType(Generic[T]):
         return self.isinstance(obj)
 
     @classmethod
-    def from_type(cls, typ_: Union["LazyType[T]", Type[T]]) -> "LazyType[T]":
+    def from_type(cls, typ_: Union["LazyType[T]", type[T]]) -> "LazyType[T]":
         if isinstance(typ_, LazyType):
             return typ_
         return cls(typ_)
@@ -48,7 +48,7 @@ class LazyType(Generic[T]):
     def __repr__(self) -> str:
         return f'LazyType("{self.module}", "{self.qualname}")'
 
-    def get_class(self) -> Type[T]:
+    def get_class(self) -> type[T]:
         if self._runtime_class is None:
             try:
                 m = importlib.import_module(self.module)

snowflake/ml/_internal/utils/db_utils.py

@@ -1,5 +1,5 @@
 from enum import Enum
-from typing import Any, Dict, Optional
+from typing import Any, Optional
 
 from snowflake.ml._internal.utils import query_result_checker, sql_identifier
 from snowflake.snowpark import session
@@ -19,7 +19,7 @@ def db_object_exists(
     *,
     database_name: Optional[sql_identifier.SqlIdentifier] = None,
     schema_name: Optional[sql_identifier.SqlIdentifier] = None,
-    statement_params: Optional[Dict[str, Any]] = None,
+    statement_params: Optional[dict[str, Any]] = None,
 ) -> bool:
     """Check if object exists in database.