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

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

@@ -1,5 +1,5 @@
 import pathlib
-from typing import List, Optional, Union
+from typing import Any, Optional, Union, overload
 
 import yaml
 
@@ -19,6 +19,7 @@ class ModelDeploymentSpec:
     def __init__(self, workspace_path: Optional[pathlib.Path] = None) -> None:
         self.workspace_path = workspace_path
 
+    @overload
     def save(
         self,
         *,
@@ -26,91 +27,214 @@ class ModelDeploymentSpec:
         schema_name: sql_identifier.SqlIdentifier,
         model_name: sql_identifier.SqlIdentifier,
         version_name: sql_identifier.SqlIdentifier,
-        service_database_name: Optional[sql_identifier.SqlIdentifier],
-        service_schema_name: Optional[sql_identifier.SqlIdentifier],
+        service_database_name: Optional[sql_identifier.SqlIdentifier] = None,
+        service_schema_name: Optional[sql_identifier.SqlIdentifier] = None,
         service_name: sql_identifier.SqlIdentifier,
+        inference_compute_pool_name: sql_identifier.SqlIdentifier,
         image_build_compute_pool_name: sql_identifier.SqlIdentifier,
-        service_compute_pool_name: sql_identifier.SqlIdentifier,
         image_repo_database_name: Optional[sql_identifier.SqlIdentifier],
         image_repo_schema_name: Optional[sql_identifier.SqlIdentifier],
         image_repo_name: sql_identifier.SqlIdentifier,
+        cpu: Optional[str],
+        memory: Optional[str],
+        gpu: Optional[Union[str, int]],
+        num_workers: Optional[int],
+        max_batch_rows: Optional[int],
+        force_rebuild: bool,
+        external_access_integrations: Optional[list[sql_identifier.SqlIdentifier]],
+        # service spec
         ingress_enabled: bool,
         max_instances: int,
+    ) -> str:
+        ...
+
+    @overload
+    def save(
+        self,
+        *,
+        database_name: sql_identifier.SqlIdentifier,
+        schema_name: sql_identifier.SqlIdentifier,
+        model_name: sql_identifier.SqlIdentifier,
+        version_name: sql_identifier.SqlIdentifier,
+        job_database_name: Optional[sql_identifier.SqlIdentifier] = None,
+        job_schema_name: Optional[sql_identifier.SqlIdentifier] = None,
+        job_name: sql_identifier.SqlIdentifier,
+        inference_compute_pool_name: sql_identifier.SqlIdentifier,
+        image_build_compute_pool_name: sql_identifier.SqlIdentifier,
+        image_repo_database_name: Optional[sql_identifier.SqlIdentifier],
+        image_repo_schema_name: Optional[sql_identifier.SqlIdentifier],
+        image_repo_name: sql_identifier.SqlIdentifier,
+        cpu: Optional[str],
+        memory: Optional[str],
+        gpu: Optional[Union[str, int]],
+        num_workers: Optional[int],
+        max_batch_rows: Optional[int],
+        force_rebuild: bool,
+        external_access_integrations: Optional[list[sql_identifier.SqlIdentifier]],
+        # job spec
+        warehouse: sql_identifier.SqlIdentifier,
+        target_method: str,
+        input_table_database_name: Optional[sql_identifier.SqlIdentifier] = None,
+        input_table_schema_name: Optional[sql_identifier.SqlIdentifier] = None,
+        input_table_name: sql_identifier.SqlIdentifier,
+        output_table_database_name: Optional[sql_identifier.SqlIdentifier] = None,
+        output_table_schema_name: Optional[sql_identifier.SqlIdentifier] = None,
+        output_table_name: sql_identifier.SqlIdentifier,
+    ) -> str:
+        ...
+
+    def save(
+        self,
+        *,
+        database_name: sql_identifier.SqlIdentifier,
+        schema_name: sql_identifier.SqlIdentifier,
+        model_name: sql_identifier.SqlIdentifier,
+        version_name: sql_identifier.SqlIdentifier,
+        service_database_name: Optional[sql_identifier.SqlIdentifier] = None,
+        service_schema_name: Optional[sql_identifier.SqlIdentifier] = None,
+        service_name: Optional[sql_identifier.SqlIdentifier] = None,
+        job_database_name: Optional[sql_identifier.SqlIdentifier] = None,
+        job_schema_name: Optional[sql_identifier.SqlIdentifier] = None,
+        job_name: Optional[sql_identifier.SqlIdentifier] = None,
+        inference_compute_pool_name: sql_identifier.SqlIdentifier,
+        image_build_compute_pool_name: sql_identifier.SqlIdentifier,
+        image_repo_database_name: Optional[sql_identifier.SqlIdentifier],
+        image_repo_schema_name: Optional[sql_identifier.SqlIdentifier],
+        image_repo_name: sql_identifier.SqlIdentifier,
         cpu: Optional[str],
         memory: Optional[str],
         gpu: Optional[Union[str, int]],
         num_workers: Optional[int],
         max_batch_rows: Optional[int],
         force_rebuild: bool,
-        external_access_integrations: Optional[List[sql_identifier.SqlIdentifier]],
+        external_access_integrations: Optional[list[sql_identifier.SqlIdentifier]],
+        # service spec
+        ingress_enabled: Optional[bool] = None,
+        max_instances: Optional[int] = None,
+        # job spec
+        warehouse: Optional[sql_identifier.SqlIdentifier] = None,
+        target_method: Optional[str] = None,
+        input_table_database_name: Optional[sql_identifier.SqlIdentifier] = None,
+        input_table_schema_name: Optional[sql_identifier.SqlIdentifier] = None,
+        input_table_name: Optional[sql_identifier.SqlIdentifier] = None,
+        output_table_database_name: Optional[sql_identifier.SqlIdentifier] = None,
+        output_table_schema_name: Optional[sql_identifier.SqlIdentifier] = None,
+        output_table_name: Optional[sql_identifier.SqlIdentifier] = None,
     ) -> str:
         # create the deployment spec
         # models spec
         fq_model_name = identifier.get_schema_level_object_identifier(
             database_name.identifier(), schema_name.identifier(), model_name.identifier()
         )
-        model_dict = model_deployment_spec_schema.ModelDict(name=fq_model_name, version=version_name.identifier())
+        model = model_deployment_spec_schema.Model(name=fq_model_name, version=version_name.identifier())
 
         # image_build spec
         saved_image_repo_database = image_repo_database_name or database_name
         saved_image_repo_schema = image_repo_schema_name or schema_name
         fq_image_repo_name = identifier.get_schema_level_object_identifier(
-            saved_image_repo_database.identifier(), saved_image_repo_schema.identifier(), image_repo_name.identifier()
+            db=saved_image_repo_database.identifier(),
+            schema=saved_image_repo_schema.identifier(),
+            object_name=image_repo_name.identifier(),
         )
-        image_build_dict: model_deployment_spec_schema.ImageBuildDict = {
-            "compute_pool": image_build_compute_pool_name.identifier(),
-            "image_repo": fq_image_repo_name,
-            "force_rebuild": force_rebuild,
-        }
-        if external_access_integrations is not None:
-            image_build_dict["external_access_integrations"] = [
-                eai.identifier() for eai in external_access_integrations
-            ]
 
-        # service spec
-        saved_service_database = service_database_name or database_name
-        saved_service_schema = service_schema_name or schema_name
-        fq_service_name = identifier.get_schema_level_object_identifier(
-            saved_service_database.identifier(), saved_service_schema.identifier(), service_name.identifier()
+        image_build = model_deployment_spec_schema.ImageBuild(
+            compute_pool=image_build_compute_pool_name.identifier(),
+            image_repo=fq_image_repo_name,
+            force_rebuild=force_rebuild,
+            external_access_integrations=(
+                [eai.identifier() for eai in external_access_integrations] if external_access_integrations else None
+            ),
         )
-        service_dict = model_deployment_spec_schema.ServiceDict(
-            name=fq_service_name,
-            compute_pool=service_compute_pool_name.identifier(),
-            ingress_enabled=ingress_enabled,
-            max_instances=max_instances,
-        )
-        if cpu:
-            service_dict["cpu"] = cpu
 
+        # universal base inference spec in service and job
+        base_inference_spec: dict[str, Any] = {}
+        if cpu:
+            base_inference_spec["cpu"] = cpu
         if memory:
-            service_dict["memory"] = memory
-
+            base_inference_spec["memory"] = memory
         if gpu:
             if isinstance(gpu, int):
                 gpu_str = str(gpu)
             else:
                 gpu_str = gpu
-            service_dict["gpu"] = gpu_str
-
+            base_inference_spec["gpu"] = gpu_str
         if num_workers:
-            service_dict["num_workers"] = num_workers
-
+            base_inference_spec["num_workers"] = num_workers
         if max_batch_rows:
-            service_dict["max_batch_rows"] = max_batch_rows
+            base_inference_spec["max_batch_rows"] = max_batch_rows
+
+        if service_name:  # service spec
+            assert ingress_enabled, "ingress_enabled is required for service spec"
+            assert max_instances, "max_instances is required for service spec"
+            saved_service_database = service_database_name or database_name
+            saved_service_schema = service_schema_name or schema_name
+            fq_service_name = identifier.get_schema_level_object_identifier(
+                saved_service_database.identifier(), saved_service_schema.identifier(), service_name.identifier()
+            )
+            service = model_deployment_spec_schema.Service(
+                name=fq_service_name,
+                compute_pool=inference_compute_pool_name.identifier(),
+                ingress_enabled=ingress_enabled,
+                max_instances=max_instances,
+                **base_inference_spec,
+            )
+
+            # model deployment spec
+            model_deployment_spec: Union[
+                model_deployment_spec_schema.ModelServiceDeploymentSpec,
+                model_deployment_spec_schema.ModelJobDeploymentSpec,
+            ] = model_deployment_spec_schema.ModelServiceDeploymentSpec(
+                models=[model],
+                image_build=image_build,
+                service=service,
+            )
+        else:  # job spec
+            assert job_name, "job_name is required for job spec"
+            assert warehouse, "warehouse is required for job spec"
+            assert target_method, "target_method is required for job spec"
+            assert input_table_name, "input_table_name is required for job spec"
+            assert output_table_name, "output_table_name is required for job spec"
+            saved_job_database = job_database_name or database_name
+            saved_job_schema = job_schema_name or schema_name
+            input_table_database_name = input_table_database_name or database_name
+            input_table_schema_name = input_table_schema_name or schema_name
+            output_table_database_name = output_table_database_name or database_name
+            output_table_schema_name = output_table_schema_name or schema_name
+            fq_job_name = identifier.get_schema_level_object_identifier(
+                saved_job_database.identifier(), saved_job_schema.identifier(), job_name.identifier()
+            )
+            fq_input_table_name = identifier.get_schema_level_object_identifier(
+                input_table_database_name.identifier(),
+                input_table_schema_name.identifier(),
+                input_table_name.identifier(),
+            )
+            fq_output_table_name = identifier.get_schema_level_object_identifier(
+                output_table_database_name.identifier(),
+                output_table_schema_name.identifier(),
+                output_table_name.identifier(),
+            )
+            job = model_deployment_spec_schema.Job(
+                name=fq_job_name,
+                compute_pool=inference_compute_pool_name.identifier(),
+                warehouse=warehouse.identifier(),
+                target_method=target_method,
+                input_table_name=fq_input_table_name,
+                output_table_name=fq_output_table_name,
+                **base_inference_spec,
+            )
+
+            # model deployment spec
+            model_deployment_spec = model_deployment_spec_schema.ModelJobDeploymentSpec(
+                models=[model],
+                image_build=image_build,
+                job=job,
+            )
 
-        # model deployment spec
-        model_deployment_spec_dict = model_deployment_spec_schema.ModelDeploymentSpecDict(
-            models=[model_dict],
-            image_build=image_build_dict,
-            service=service_dict,
-        )
-
-        # Anchors are not supported in the server, avoid that.
-        yaml.SafeDumper.ignore_aliases = lambda *args: True  # type: ignore[method-assign]
         if self.workspace_path is None:
-            return yaml.safe_dump(model_deployment_spec_dict)
+            return yaml.safe_dump(model_deployment_spec.model_dump(exclude_none=True))
+
         # save the yaml
         file_path = self.workspace_path / self.DEPLOY_SPEC_FILE_REL_PATH
         with file_path.open("w", encoding="utf-8") as f:
-            yaml.safe_dump(model_deployment_spec_dict, f)
+            yaml.safe_dump(model_deployment_spec.model_dump(exclude_none=True), f)
         return str(file_path.resolve())

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

@@ -1,33 +1,53 @@
-from typing import List, TypedDict
+from typing import Optional
 
-from typing_extensions import NotRequired, Required
+from pydantic import BaseModel
 
 
-class ModelDict(TypedDict):
-    name: Required[str]
-    version: Required[str]
+class Model(BaseModel):
+    name: str
+    version: str
 
 
-class ImageBuildDict(TypedDict):
-    compute_pool: Required[str]
-    image_repo: Required[str]
-    force_rebuild: Required[bool]
-    external_access_integrations: NotRequired[List[str]]
+class ImageBuild(BaseModel):
+    compute_pool: str
+    image_repo: str
+    force_rebuild: bool
+    external_access_integrations: Optional[list[str]] = None
 
 
-class ServiceDict(TypedDict):
-    name: Required[str]
-    compute_pool: Required[str]
-    ingress_enabled: Required[bool]
-    max_instances: Required[int]
-    cpu: NotRequired[str]
-    memory: NotRequired[str]
-    gpu: NotRequired[str]
-    num_workers: NotRequired[int]
-    max_batch_rows: NotRequired[int]
+class Service(BaseModel):
+    name: str
+    compute_pool: str
+    ingress_enabled: bool
+    max_instances: int
+    cpu: Optional[str] = None
+    memory: Optional[str] = None
+    gpu: Optional[str] = None
+    num_workers: Optional[int] = None
+    max_batch_rows: Optional[int] = None
 
 
-class ModelDeploymentSpecDict(TypedDict):
-    models: Required[List[ModelDict]]
-    image_build: Required[ImageBuildDict]
-    service: Required[ServiceDict]
+class Job(BaseModel):
+    name: str
+    compute_pool: str
+    cpu: Optional[str] = None
+    memory: Optional[str] = None
+    gpu: Optional[str] = None
+    num_workers: Optional[int] = None
+    max_batch_rows: Optional[int] = None
+    warehouse: str
+    target_method: str
+    input_table_name: str
+    output_table_name: str
+
+
+class ModelServiceDeploymentSpec(BaseModel):
+    models: list[Model]
+    image_build: ImageBuild
+    service: Service
+
+
+class ModelJobDeploymentSpec(BaseModel):
+    models: list[Model]
+    image_build: ImageBuild
+    job: Job

snowflake/ml/model/_client/sql/model.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
@@ -24,8 +24,8 @@ class ModelSQLClient(_base._BaseSQLClient):
         schema_name: Optional[sql_identifier.SqlIdentifier],
         model_name: Optional[sql_identifier.SqlIdentifier] = None,
         validate_result: bool = True,
-        statement_params: Optional[Dict[str, Any]] = None,
-    ) -> List[row.Row]:
+        statement_params: Optional[dict[str, Any]] = None,
+    ) -> list[row.Row]:
         actual_database_name = database_name or self._database_name
         actual_schema_name = schema_name or self._schema_name
         fully_qualified_schema_name = ".".join([actual_database_name.identifier(), actual_schema_name.identifier()])
@@ -57,8 +57,8 @@ class ModelSQLClient(_base._BaseSQLClient):
         version_name: Optional[sql_identifier.SqlIdentifier] = None,
         validate_result: bool = True,
         check_model_details: bool = False,
-        statement_params: Optional[Dict[str, Any]] = None,
-    ) -> List[row.Row]:
+        statement_params: Optional[dict[str, Any]] = None,
+    ) -> list[row.Row]:
         like_sql = ""
         if version_name:
             like_sql = f" LIKE '{version_name.resolved()}'"
@@ -90,7 +90,7 @@ class ModelSQLClient(_base._BaseSQLClient):
         schema_name: Optional[sql_identifier.SqlIdentifier],
         model_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,
@@ -107,7 +107,7 @@ class ModelSQLClient(_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,
+        statement_params: Optional[dict[str, Any]] = None,
     ) -> None:
         query_result_checker.SqlResultValidator(
             self._session,
@@ -124,7 +124,7 @@ class ModelSQLClient(_base._BaseSQLClient):
         new_model_db: Optional[sql_identifier.SqlIdentifier],
         new_model_schema: Optional[sql_identifier.SqlIdentifier],
         new_model_name: sql_identifier.SqlIdentifier,
-        statement_params: Optional[Dict[str, Any]] = None,
+        statement_params: Optional[dict[str, Any]] = None,
     ) -> None:
         # Use registry's database and schema if a non fully qualified new model name is provided.
         new_fully_qualified_name = self.fully_qualified_object_name(new_model_db, new_model_schema, new_model_name)

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,7 +1,7 @@
 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
@@ -47,7 +47,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 +76,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 +95,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
@@ -181,7 +181,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 +206,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 +231,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 +245,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(