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

snowflake/ml/fileset/sfcfs.py

@@ -1,7 +1,7 @@
 import collections
 import logging
 from functools import partial
-from typing import Any, Callable, Dict, List, Optional, Tuple, Type, Union, cast
+from typing import Any, Callable, Optional, Union, cast
 
 import fsspec
 
@@ -100,7 +100,7 @@ class SFFileSystem(fsspec.AbstractFileSystem):
             raise ValueError("Either sf_connection or snowpark_session has to be non-empty!")
         self._conn = self._session._conn._conn  # Telemetry wrappers expect connection under `conn_attr_name="_conn"``
         self._kwargs = kwargs
-        self._stage_fs_set: Dict[Tuple[str, str, str], stage_fs.SFStageFileSystem] = {}
+        self._stage_fs_set: dict[tuple[str, str, str], stage_fs.SFStageFileSystem] = {}
 
         super().__init__(**kwargs)
 
@@ -133,7 +133,7 @@ class SFFileSystem(fsspec.AbstractFileSystem):
         assert isinstance(session, snowpark.Session)
         return session
 
-    def __reduce__(self) -> Tuple[Callable[[], Type["SFFileSystem"]], Tuple[()], Dict[str, Any]]:
+    def __reduce__(self) -> tuple[Callable[[], type["SFFileSystem"]], tuple[()], dict[str, Any]]:
         """Returns a state dictionary for use in serialization.
 
         Returns:
@@ -145,7 +145,7 @@ class SFFileSystem(fsspec.AbstractFileSystem):
 
         return partial(self.__class__, **{_RECREATE_FROM_SERIALIZED: True}), (), state_dictionary
 
-    def __setstate__(self, state_dict: Dict[str, Any]) -> None:
+    def __setstate__(self, state_dict: dict[str, Any]) -> None:
         """Sets the dictionary state at deserialization time, and rebuilds a snowflake connection.
 
         Args:
@@ -191,7 +191,7 @@ class SFFileSystem(fsspec.AbstractFileSystem):
         func_params_to_log=["detail"],
         conn_attr_name="_conn",
     )
-    def ls(self, path: str, detail: bool = False, **kwargs: Any) -> Union[List[str], List[Dict[str, Any]]]:
+    def ls(self, path: str, detail: bool = False, **kwargs: Any) -> Union[list[str], list[dict[str, Any]]]:
         """Override fsspec `ls` method. List single "directory" with or without details.
 
         Args:
@@ -214,14 +214,14 @@ class SFFileSystem(fsspec.AbstractFileSystem):
         file_path = self._parse_file_path(path)
         stage_fs = self._get_stage_fs(file_path)
         stage_path_list = stage_fs.ls(file_path.filepath, detail=True, **kwargs)
-        stage_path_list = cast(List[Dict[str, Any]], stage_path_list)
+        stage_path_list = cast(list[dict[str, Any]], stage_path_list)
         return self._decorate_ls_res(stage_fs, stage_path_list, detail)
 
     @telemetry.send_api_usage_telemetry(
         project=_PROJECT,
         conn_attr_name="_conn",
     )
-    def optimize_read(self, files: Optional[List[str]] = None) -> None:
+    def optimize_read(self, files: Optional[list[str]] = None) -> None:
         """Prefetch and cache the presigned urls for all the given files to speed up the file opening.
 
         All the files introduced here will have their urls cached. Further open() on any of cached urls will lead to a
@@ -232,8 +232,8 @@ class SFFileSystem(fsspec.AbstractFileSystem):
         """
         if not files:
             return
-        stage_fs_dict: Dict[str, stage_fs.SFStageFileSystem] = {}
-        stage_file_paths: Dict[str, List[str]] = collections.defaultdict(list)
+        stage_fs_dict: dict[str, stage_fs.SFStageFileSystem] = {}
+        stage_file_paths: dict[str, list[str]] = collections.defaultdict(list)
         for file in files:
             path_info = self._parse_file_path(file)
             fs = self._get_stage_fs(path_info)
@@ -271,11 +271,11 @@ class SFFileSystem(fsspec.AbstractFileSystem):
         project=_PROJECT,
         conn_attr_name="_conn",
     )
-    def info(self, path: str, **kwargs: Any) -> Dict[str, Any]:
+    def info(self, path: str, **kwargs: Any) -> dict[str, Any]:
         """Override fsspec `info` method. Give details of entry at path."""
         file_path = self._parse_file_path(path)
         stage_fs = self._get_stage_fs(file_path)
-        res: Dict[str, Any] = stage_fs.info(file_path.filepath, **kwargs)
+        res: dict[str, Any] = stage_fs.info(file_path.filepath, **kwargs)
         if res:
             res["name"] = self._stage_path_to_absolute_path(stage_fs, res["name"])
         return res
@@ -283,9 +283,9 @@ class SFFileSystem(fsspec.AbstractFileSystem):
     def _decorate_ls_res(
         self,
         stage_fs: stage_fs.SFStageFileSystem,
-        stage_path_list: List[Dict[str, Any]],
+        stage_path_list: list[dict[str, Any]],
         detail: bool,
-    ) -> Union[List[str], List[Dict[str, Any]]]:
+    ) -> Union[list[str], list[dict[str, Any]]]:
         """Add the stage location as the prefix of file names returned by ls() of stagefs"""
         for path in stage_path_list:
             path["name"] = self._stage_path_to_absolute_path(stage_fs, path["name"])

snowflake/ml/fileset/stage_fs.py

@@ -2,7 +2,7 @@ import inspect
 import logging
 import time
 from dataclasses import dataclass
-from typing import Any, Dict, List, Optional, Tuple, Union, cast
+from typing import Any, Optional, Union, cast
 
 import fsspec
 from fsspec.implementations import http as httpfs
@@ -44,7 +44,7 @@ class _PresignedUrl:
         return not self.expire_at or time.time() > self.expire_at - headroom_sec
 
 
-def _get_httpfs_kwargs(**kwargs: Any) -> Dict[str, Any]:
+def _get_httpfs_kwargs(**kwargs: Any) -> dict[str, Any]:
     """Extract kwargs that are meaningful to HTTPFileSystem."""
     httpfs_related_keys = [
         "block_size",
@@ -124,7 +124,7 @@ class SFStageFileSystem(fsspec.AbstractFileSystem):
         self._db = db
         self._schema = schema
         self._stage = stage
-        self._url_cache: Dict[str, _PresignedUrl] = {}
+        self._url_cache: dict[str, _PresignedUrl] = {}
 
         httpfs_kwargs = _get_httpfs_kwargs(**kwargs)
         self._fs = httpfs.HTTPFileSystem(**httpfs_kwargs)
@@ -145,7 +145,7 @@ class SFStageFileSystem(fsspec.AbstractFileSystem):
         project=_PROJECT,
         func_params_to_log=["detail"],
     )
-    def ls(self, path: str, detail: bool = False) -> Union[List[str], List[Dict[str, Any]]]:
+    def ls(self, path: str, detail: bool = False) -> Union[list[str], list[dict[str, Any]]]:
         """Override fsspec `ls` method. List single "directory" with or without details.
 
         Args:
@@ -169,7 +169,7 @@ class SFStageFileSystem(fsspec.AbstractFileSystem):
             loc = self.stage_name
             path = path.lstrip("/")
             async_job: snowpark.AsyncJob = self._session.sql(f"LIST '{loc}/{path}'").collect(block=False)
-            objects: List[snowpark.Row] = _resolve_async_job(async_job)
+            objects: list[snowpark.Row] = _resolve_async_job(async_job)
         except snowpark_exceptions.SnowparkSQLException as e:
             if e.sql_error_code == fileset_errors.ERRNO_DOMAIN_NOT_EXIST:
                 raise snowml_exceptions.SnowflakeMLException(
@@ -192,7 +192,7 @@ class SFStageFileSystem(fsspec.AbstractFileSystem):
     @telemetry.send_api_usage_telemetry(
         project=_PROJECT,
     )
-    def optimize_read(self, files: Optional[List[str]] = None) -> None:
+    def optimize_read(self, files: Optional[list[str]] = None) -> None:
         """Prefetch and cache the presigned urls for all the given files to speed up the read performance.
 
         All the files introduced here will have their urls cached. Further open() on any of cached urls will lead to a
@@ -271,7 +271,7 @@ class SFStageFileSystem(fsspec.AbstractFileSystem):
                 original_exception=fileset_errors.StageFileNotFoundError(f"Stage file {path} doesn't exist."),
             )
 
-    def _open_with_snowpark(self, path: str, **kwargs: Dict[str, Any]) -> fsspec.spec.AbstractBufferedFile:
+    def _open_with_snowpark(self, path: str, **kwargs: dict[str, Any]) -> fsspec.spec.AbstractBufferedFile:
         """Open the a file for reading using snowflake.snowpark.file_operation
 
         Args:
@@ -299,7 +299,7 @@ class SFStageFileSystem(fsspec.AbstractFileSystem):
                     original_exception=e,
                 )
 
-    def _parse_list_result(self, list_result: List[snowpark.Row], search_path: str) -> List[Dict[str, Any]]:
+    def _parse_list_result(self, list_result: list[snowpark.Row], search_path: str) -> list[dict[str, Any]]:
         """Convert the result from LIST query to the expected format of fsspec ls() method.
 
         Note that Snowflake LIST query has different behavior with ls(). LIST query will return all the stage files
@@ -318,7 +318,7 @@ class SFStageFileSystem(fsspec.AbstractFileSystem):
         Returns:
             A list of dict, where each dict contains key-value pairs as the properties of a file.
         """
-        files: Dict[str, Dict[str, Any]] = {}
+        files: dict[str, dict[str, Any]] = {}
         search_path = search_path.strip("/")
         for row in list_result:
             name, size, md5, last_modified = row["name"], row["size"], row["md5"], row["last_modified"]
@@ -360,7 +360,7 @@ class SFStageFileSystem(fsspec.AbstractFileSystem):
 
     def _add_file_info_helper(
         self,
-        files: Dict[str, Dict[str, Any]],
+        files: dict[str, dict[str, Any]],
         object_path: str,
         file_size: int,
         file_type: str,
@@ -379,12 +379,12 @@ class SFStageFileSystem(fsspec.AbstractFileSystem):
         )
 
     def _fetch_presigned_urls(
-        self, files: List[str], url_lifetime: float = _PRESIGNED_URL_LIFETIME_SEC
-    ) -> List[Tuple[str, str]]:
+        self, files: list[str], url_lifetime: float = _PRESIGNED_URL_LIFETIME_SEC
+    ) -> list[tuple[str, str]]:
         """Fetch presigned urls for the given files."""
         file_df = self._session.create_dataframe(files).to_df("name")
         try:
-            presigned_urls: List[Tuple[str, str]] = file_df.select_expr(
+            presigned_urls: list[tuple[str, str]] = file_df.select_expr(
                 f"name, get_presigned_url('{self.stage_name}', name, {url_lifetime}) as url"
             ).collect(
                 statement_params=telemetry.get_function_usage_statement_params(
@@ -418,10 +418,10 @@ def _match_error_code(ex: snowpark_exceptions.SnowparkSQLException, error_code:
 
 
 @snowflake_plan.SnowflakePlan.Decorator.wrap_exception  # type: ignore[misc]
-def _resolve_async_job(async_job: snowpark.AsyncJob) -> List[snowpark.Row]:
+def _resolve_async_job(async_job: snowpark.AsyncJob) -> list[snowpark.Row]:
     # Make sure Snowpark exceptions are properly caught and converted by wrap_exception wrapper
     try:
-        query_result = cast(List[snowpark.Row], async_job.result("row"))
+        query_result = cast(list[snowpark.Row], async_job.result("row"))
         return query_result
     except snowpark_errors.DatabaseError as e:
         # HACK: Snowpark surfaces a generic exception if query doesn't complete immediately

snowflake/ml/jobs/_utils/interop_utils.py

@@ -10,7 +10,7 @@ import traceback
 from collections import namedtuple
 from dataclasses import dataclass
 from types import TracebackType
-from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Type, Union, cast
+from typing import Any, Callable, Optional, Union, cast
 
 from snowflake import snowpark
 from snowflake.snowpark import exceptions as sp_exceptions
@@ -33,7 +33,7 @@ class ExecutionResult:
     def success(self) -> bool:
         return self.exception is None
 
-    def to_dict(self) -> Dict[str, Any]:
+    def to_dict(self) -> dict[str, Any]:
         """Return the serializable dictionary."""
         if isinstance(self.exception, BaseException):
             exc_type = type(self.exception)
@@ -50,7 +50,7 @@ class ExecutionResult:
         }
 
     @classmethod
-    def from_dict(cls, result_dict: Dict[str, Any]) -> "ExecutionResult":
+    def from_dict(cls, result_dict: dict[str, Any]) -> "ExecutionResult":
         if not isinstance(result_dict.get("success"), bool):
             raise ValueError("Invalid result dictionary")
 
@@ -242,11 +242,11 @@ def _install_sys_excepthook() -> None:
         original_excepthook = sys.excepthook
 
         def custom_excepthook(
-            exc_type: Type[BaseException],
+            exc_type: type[BaseException],
             exc_value: BaseException,
             exc_tb: Optional[TracebackType],
             *,
-            seen_exc_ids: Optional[Set[int]] = None,
+            seen_exc_ids: Optional[set[int]] = None,
         ) -> None:
             if seen_exc_ids is None:
                 seen_exc_ids = set()
@@ -331,7 +331,7 @@ def _install_ipython_hook() -> bool:
     except ImportError:
         return False
 
-    def parse_traceback_str(traceback_str: str) -> List[Tuple[str, int, str, str]]:
+    def parse_traceback_str(traceback_str: str) -> list[tuple[str, int, str, str]]:
         return [
             (m.group("filename"), int(m.group("lineno")), m.group("name"), m.group("line"))
             for m in re.finditer(_TRACEBACK_ENTRY_PATTERN, traceback_str)
@@ -342,13 +342,13 @@ def _install_ipython_hook() -> bool:
 
         def custom_format_exception_as_a_whole(
             self: VerboseTB,
-            etype: Type[BaseException],
+            etype: type[BaseException],
             evalue: Optional[BaseException],
             etb: Optional[TracebackType],
             number_of_lines_of_context: int,
             tb_offset: Optional[int],
             **kwargs: Any,
-        ) -> List[List[str]]:
+        ) -> list[list[str]]:
             if (remote_err := _retrieve_remote_error_info(evalue)) and isinstance(remote_err, RemoteError):
                 # Implementation forked from IPython.core.ultratb.VerboseTB.format_exception_as_a_whole
                 head = self.prepare_header(remote_err.exc_type, long_version=False).replace(
@@ -388,7 +388,7 @@ def _install_ipython_hook() -> bool:
             etb: Optional[TracebackType],
             tb_offset: Optional[int] = None,
             **kwargs: Any,
-        ) -> List[str]:
+        ) -> list[str]:
             if (remote_err := _retrieve_remote_error_info(evalue)) and isinstance(remote_err, RemoteError):
                 tb_list = [
                     (m.group("filename"), m.group("lineno"), m.group("name"), m.group("line"))
@@ -400,7 +400,7 @@ def _install_ipython_hook() -> bool:
                         "(most recent call last)",
                         "(from remote execution)",
                     )
-                return cast(List[str], out_list)
+                return cast(list[str], out_list)
             return original_structured_traceback(  # type: ignore[no-any-return]
                 self, etype, evalue, etb, tb_offset, **kwargs
             )

snowflake/ml/jobs/_utils/payload_utils.py

@@ -6,17 +6,7 @@ import pickle
 import sys
 import textwrap
 from pathlib import Path, PurePath
-from typing import (
-    Any,
-    Callable,
-    List,
-    Optional,
-    Type,
-    Union,
-    cast,
-    get_args,
-    get_origin,
-)
+from typing import Any, Callable, Optional, Union, cast, get_args, get_origin
 
 import cloudpickle as cp
 
@@ -277,7 +267,7 @@ class JobPayload:
         source: Union[str, Path, Callable[..., Any]],
         entrypoint: Optional[Union[str, Path]] = None,
         *,
-        pip_requirements: Optional[List[str]] = None,
+        pip_requirements: Optional[list[str]] = None,
     ) -> None:
         self.source = Path(source) if isinstance(source, str) else source
         self.entrypoint = Path(entrypoint) if isinstance(entrypoint, str) else entrypoint
@@ -364,7 +354,7 @@ class JobPayload:
                     auto_compress=False,
                 )
 
-        python_entrypoint: List[Union[str, PurePath]] = [
+        python_entrypoint: list[Union[str, PurePath]] = [
             PurePath("mljob_launcher.py"),
             entrypoint.file_path.relative_to(source),
         ]
@@ -381,7 +371,7 @@ class JobPayload:
         )
 
 
-def _get_parameter_type(param: inspect.Parameter) -> Optional[Type[object]]:
+def _get_parameter_type(param: inspect.Parameter) -> Optional[type[object]]:
     # Unwrap Optional type annotations
     param_type = param.annotation
     if get_origin(param_type) is Union and len(get_args(param_type)) == 2 and type(None) in get_args(param_type):
@@ -390,10 +380,10 @@ def _get_parameter_type(param: inspect.Parameter) -> Optional[Type[object]]:
     # Return None for empty type annotations
     if param_type == inspect.Parameter.empty:
         return None
-    return cast(Type[object], param_type)
+    return cast(type[object], param_type)
 
 
-def _validate_parameter_type(param_type: Type[object], param_name: str) -> None:
+def _validate_parameter_type(param_type: type[object], param_name: str) -> None:
     # Validate param_type is a supported type
     if param_type not in _SUPPORTED_ARG_TYPES:
         raise ValueError(

snowflake/ml/jobs/_utils/scripts/mljob_launcher.py

@@ -1,4 +1,5 @@
 import argparse
+import copy
 import importlib.util
 import json
 import os
@@ -7,7 +8,7 @@ import sys
 import traceback
 import warnings
 from pathlib import Path
-from typing import Any, Dict, Optional
+from typing import Any, Optional
 
 import cloudpickle
 
@@ -27,7 +28,7 @@ except ImportError:
     from dataclasses import dataclass
 
     @dataclass(frozen=True)
-    class ExecutionResult:
+    class ExecutionResult:  # type: ignore[no-redef]
         result: Optional[Any] = None
         exception: Optional[BaseException] = None
 
@@ -35,7 +36,7 @@ except ImportError:
         def success(self) -> bool:
             return self.exception is None
 
-        def to_dict(self) -> Dict[str, Any]:
+        def to_dict(self) -> dict[str, Any]:
             """Return the serializable dictionary."""
             if isinstance(self.exception, BaseException):
                 exc_type = type(self.exception)
@@ -136,7 +137,9 @@ def main(script_path: str, *script_args: Any, script_main_func: Optional[str] =
         while tb and tb.tb_frame.f_code.co_filename in skip_files:
             # Skip any frames preceding user script execution
             tb = tb.tb_next
-        result_obj = ExecutionResult(exception=e.with_traceback(tb))
+        cleaned_ex = copy.copy(e)  # Need to create a mutable copy of exception to set __traceback__
+        cleaned_ex = cleaned_ex.with_traceback(tb)
+        result_obj = ExecutionResult(exception=cleaned_ex)
         raise
     finally:
         result_dict = result_obj.to_dict()

snowflake/ml/jobs/_utils/scripts/signal_workers.py

@@ -9,7 +9,7 @@ import logging
 import socket
 import sys
 import time
-from typing import Any, Dict, List, Set
+from typing import Any
 
 import ray
 from constants import (
@@ -33,34 +33,34 @@ class ShutdownSignal:
         self.acknowledged_workers = set()
         logging.info(f"ShutdownSignal actor created on {self.hostname}")
 
-    def request_shutdown(self) -> Dict[str, Any]:
+    def request_shutdown(self) -> dict[str, Any]:
         """Signal workers to shut down"""
         self.shutdown_requested = True
         self.timestamp = time.time()
         logging.info(f"Shutdown requested by head node at {self.timestamp}")
         return {"status": "shutdown_requested", "timestamp": self.timestamp, "host": self.hostname}
 
-    def should_shutdown(self) -> Dict[str, Any]:
+    def should_shutdown(self) -> dict[str, Any]:
         """Check if shutdown has been requested"""
         return {"shutdown": self.shutdown_requested, "timestamp": self.timestamp, "host": self.hostname}
 
-    def ping(self) -> Dict[str, Any]:
+    def ping(self) -> dict[str, Any]:
         """Simple method to test connectivity"""
         return {"status": "alive", "host": self.hostname}
 
-    def acknowledge_shutdown(self, worker_id: str) -> Dict[str, Any]:
+    def acknowledge_shutdown(self, worker_id: str) -> dict[str, Any]:
         """Worker acknowledges it has received the shutdown signal and is terminating"""
         self.acknowledged_workers.add(worker_id)
         logging.info(f"Worker {worker_id} acknowledged shutdown. Total acknowledged: {len(self.acknowledged_workers)}")
 
         return {"status": "acknowledged", "worker_id": worker_id, "acknowledged_count": len(self.acknowledged_workers)}
 
-    def get_acknowledgment_workers(self) -> Set[str]:
+    def get_acknowledgment_workers(self) -> set[str]:
         """Get the set of workers who have acknowledged shutdown"""
         return self.acknowledged_workers
 
 
-def get_worker_node_ids() -> List[str]:
+def get_worker_node_ids() -> list[str]:
     """Get the IDs of all active worker nodes.
 
     Returns:
@@ -127,7 +127,7 @@ def verify_shutdown(shutdown_signal: ActorHandle) -> None:
     logging.debug(f"Shutdown status check: {check}")
 
 
-def wait_for_acknowledgments(shutdown_signal: ActorHandle, worker_node_ids: List[str], wait_time: int) -> None:
+def wait_for_acknowledgments(shutdown_signal: ActorHandle, worker_node_ids: list[str], wait_time: int) -> None:
     """Wait for workers to acknowledge shutdown.
 
     Args:

snowflake/ml/jobs/_utils/spec_utils.py

@@ -1,7 +1,7 @@
 import logging
 from math import ceil
 from pathlib import PurePath
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Optional, Union
 
 from snowflake import snowpark
 from snowflake.ml._internal.utils import snowflake_env
@@ -15,10 +15,7 @@ def _get_node_resources(session: snowpark.Session, compute_pool: str) -> types.C
     if not rows:
         raise ValueError(f"Compute pool '{compute_pool}' not found")
     instance_family: str = rows[0]["instance_family"]
-
-    # Get the cloud we're using (AWS, Azure, etc)
-    region = snowflake_env.get_regions(session)[snowflake_env.get_current_region_id(session)]
-    cloud = region["cloud"]
+    cloud = snowflake_env.get_current_cloud(session, default=snowflake_env.SnowflakeCloudType.AWS)
 
     return (
         constants.COMMON_INSTANCE_FAMILIES.get(instance_family)
@@ -26,22 +23,14 @@ def _get_node_resources(session: snowpark.Session, compute_pool: str) -> types.C
     )
 
 
-def _get_image_spec(session: snowpark.Session, compute_pool: str, image_tag: Optional[str] = None) -> types.ImageSpec:
+def _get_image_spec(session: snowpark.Session, compute_pool: str) -> types.ImageSpec:
     # Retrieve compute pool node resources
     resources = _get_node_resources(session, compute_pool=compute_pool)
 
     # Use MLRuntime image
     image_repo = constants.DEFAULT_IMAGE_REPO
     image_name = constants.DEFAULT_IMAGE_GPU if resources.gpu > 0 else constants.DEFAULT_IMAGE_CPU
-
-    # Try to pull latest image tag from server side if possible
-    if not image_tag:
-        query_result = session.sql("SHOW PARAMETERS LIKE 'constants.RUNTIME_BASE_IMAGE_TAG' IN ACCOUNT").collect()
-        if query_result:
-            image_tag = query_result[0]["value"]
-
-    if image_tag is None:
-        image_tag = constants.DEFAULT_IMAGE_TAG
+    image_tag = constants.DEFAULT_IMAGE_TAG
 
     # TODO: Should each instance consume the entire pod?
     return types.ImageSpec(
@@ -54,9 +43,9 @@ def _get_image_spec(session: snowpark.Session, compute_pool: str, image_tag: Opt
 
 
 def generate_spec_overrides(
-    environment_vars: Optional[Dict[str, str]] = None,
-    custom_overrides: Optional[Dict[str, Any]] = None,
-) -> Dict[str, Any]:
+    environment_vars: Optional[dict[str, str]] = None,
+    custom_overrides: Optional[dict[str, Any]] = None,
+) -> dict[str, Any]:
     """
     Generate a dictionary of service specification overrides.
 
@@ -68,7 +57,7 @@ def generate_spec_overrides(
         Resulting service specifiation patch dict. Empty if no overrides were supplied.
     """
     # Generate container level overrides
-    container_spec: Dict[str, Any] = {
+    container_spec: dict[str, Any] = {
         "name": constants.DEFAULT_CONTAINER_NAME,
     }
     if environment_vars:
@@ -95,10 +84,10 @@ def generate_service_spec(
     session: snowpark.Session,
     compute_pool: str,
     payload: types.UploadedPayload,
-    args: Optional[List[str]] = None,
+    args: Optional[list[str]] = None,
     num_instances: Optional[int] = None,
     enable_metrics: bool = False,
-) -> Dict[str, Any]:
+) -> dict[str, Any]:
     """
     Generate a service specification for a job.
 
@@ -117,11 +106,11 @@ def generate_service_spec(
     image_spec = _get_image_spec(session, compute_pool)
 
     # Set resource requests/limits, including nvidia.com/gpu quantity if applicable
-    resource_requests: Dict[str, Union[str, int]] = {
+    resource_requests: dict[str, Union[str, int]] = {
         "cpu": f"{int(image_spec.resource_requests.cpu * 1000)}m",
         "memory": f"{image_spec.resource_limits.memory}Gi",
     }
-    resource_limits: Dict[str, Union[str, int]] = {
+    resource_limits: dict[str, Union[str, int]] = {
         "cpu": f"{int(image_spec.resource_requests.cpu * 1000)}m",
         "memory": f"{image_spec.resource_limits.memory}Gi",
     }
@@ -130,8 +119,8 @@ def generate_service_spec(
         resource_limits["nvidia.com/gpu"] = image_spec.resource_limits.gpu
 
     # Add local volumes for ephemeral logs and artifacts
-    volumes: List[Dict[str, str]] = []
-    volume_mounts: List[Dict[str, str]] = []
+    volumes: list[dict[str, str]] = []
+    volume_mounts: list[dict[str, str]] = []
     for volume_name, mount_path in [
         ("system-logs", "/var/log/managedservices/system/mlrs"),
         ("user-logs", "/var/log/managedservices/user/mlrs"),
@@ -302,11 +291,11 @@ def merge_patch(base: Any, patch: Any, display_name: str = "") -> Any:
 
 
 def _merge_lists_of_dicts(
-    base: List[Dict[str, Any]],
-    patch: List[Dict[str, Any]],
+    base: list[dict[str, Any]],
+    patch: list[dict[str, Any]],
     merge_key: str = "name",
     display_name: str = "",
-) -> List[Dict[str, Any]]:
+) -> list[dict[str, Any]]:
     """
     Attempts to merge lists of dicts by matching on a merge key (default "name").
     - If the merge key is missing, the behavior falls back to overwriting the list.

snowflake/ml/jobs/_utils/types.py

@@ -1,6 +1,6 @@
 from dataclasses import dataclass
 from pathlib import PurePath
-from typing import List, Literal, Optional, Union
+from typing import Literal, Optional, Union
 
 JOB_STATUS = Literal[
     "PENDING",
@@ -21,7 +21,7 @@ class PayloadEntrypoint:
 class UploadedPayload:
     # TODO: Include manifest of payload files for validation
     stage_path: PurePath
-    entrypoint: List[Union[str, PurePath]]
+    entrypoint: list[Union[str, PurePath]]
 
 
 @dataclass(frozen=True)

snowflake/ml/jobs/decorators.py

@@ -1,6 +1,6 @@
 import copy
 import functools
-from typing import Callable, Dict, List, Optional, TypeVar
+from typing import Callable, Optional, TypeVar
 
 from typing_extensions import ParamSpec
 
@@ -15,16 +15,15 @@ _Args = ParamSpec("_Args")
 _ReturnValue = TypeVar("_ReturnValue")
 
 
-@snowpark._internal.utils.private_preview(version="1.7.4")
 @telemetry.send_api_usage_telemetry(project=_PROJECT)
 def remote(
     compute_pool: str,
     *,
     stage_name: str,
-    pip_requirements: Optional[List[str]] = None,
-    external_access_integrations: Optional[List[str]] = None,
+    pip_requirements: Optional[list[str]] = None,
+    external_access_integrations: Optional[list[str]] = None,
     query_warehouse: Optional[str] = None,
-    env_vars: Optional[Dict[str, str]] = None,
+    env_vars: Optional[dict[str, str]] = None,
     num_instances: Optional[int] = None,
     enable_metrics: bool = False,
     session: Optional[snowpark.Session] = None,