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

snowflake/ml/jobs/job.py

@@ -1,5 +1,5 @@
 import time
-from typing import Any, Dict, Generic, List, Optional, TypeVar, cast
+from typing import Any, Generic, Literal, Optional, TypeVar, Union, cast, overload
 
 import yaml
 
@@ -18,11 +18,11 @@ class MLJob(Generic[T]):
     def __init__(
         self,
         id: str,
-        service_spec: Optional[Dict[str, Any]] = None,
+        service_spec: Optional[dict[str, Any]] = None,
         session: Optional[snowpark.Session] = None,
     ) -> None:
         self._id = id
-        self._service_spec_cached: Optional[Dict[str, Any]] = service_spec
+        self._service_spec_cached: Optional[dict[str, Any]] = service_spec
         self._session = session or sp_context.get_active_session()
 
         self._status: types.JOB_STATUS = "PENDING"
@@ -42,18 +42,18 @@ class MLJob(Generic[T]):
         return self._status
 
     @property
-    def _service_spec(self) -> Dict[str, Any]:
+    def _service_spec(self) -> dict[str, Any]:
         """Get the job's service spec."""
         if not self._service_spec_cached:
             self._service_spec_cached = _get_service_spec(self._session, self.id)
         return self._service_spec_cached
 
     @property
-    def _container_spec(self) -> Dict[str, Any]:
+    def _container_spec(self) -> dict[str, Any]:
         """Get the job's main container spec."""
         containers = self._service_spec["spec"]["containers"]
         container_spec = next(c for c in containers if c["name"] == constants.DEFAULT_CONTAINER_NAME)
-        return cast(Dict[str, Any], container_spec)
+        return cast(dict[str, Any], container_spec)
 
     @property
     def _stage_path(self) -> str:
@@ -70,8 +70,17 @@ class MLJob(Generic[T]):
             raise RuntimeError(f"Job {self.id} doesn't have a result path configured")
         return f"{self._stage_path}/{result_path}"
 
-    @snowpark._internal.utils.private_preview(version="1.7.4")
-    def get_logs(self, limit: int = -1, instance_id: Optional[int] = None) -> str:
+    @overload
+    def get_logs(self, limit: int = -1, instance_id: Optional[int] = None, *, as_list: Literal[True]) -> list[str]:
+        ...
+
+    @overload
+    def get_logs(self, limit: int = -1, instance_id: Optional[int] = None, *, as_list: Literal[False] = False) -> str:
+        ...
+
+    def get_logs(
+        self, limit: int = -1, instance_id: Optional[int] = None, *, as_list: bool = False
+    ) -> Union[str, list[str]]:
         """
         Return the job's execution logs.
 
@@ -79,15 +88,17 @@ class MLJob(Generic[T]):
             limit: The maximum number of lines to return. Negative values are treated as no limit.
             instance_id: Optional instance ID to get logs from a specific instance.
                          If not provided, returns logs from the head node.
+            as_list: If True, returns logs as a list of lines. Otherwise, returns logs as a single string.
 
         Returns:
             The job's execution logs.
         """
         logs = _get_logs(self._session, self.id, limit, instance_id)
         assert isinstance(logs, str)  # mypy
+        if as_list:
+            return logs.splitlines()
         return logs
 
-    @snowpark._internal.utils.private_preview(version="1.7.4")
     def show_logs(self, limit: int = -1, instance_id: Optional[int] = None) -> None:
         """
         Display the job's execution logs.
@@ -97,9 +108,8 @@ class MLJob(Generic[T]):
             instance_id: Optional instance ID to get logs from a specific instance.
                          If not provided, displays logs from the head node.
         """
-        print(self.get_logs(limit, instance_id))  # noqa: T201: we need to print here.
+        print(self.get_logs(limit, instance_id, as_list=False))  # noqa: T201: we need to print here.
 
-    @snowpark._internal.utils.private_preview(version="1.7.4")
     @telemetry.send_api_usage_telemetry(project=_PROJECT, func_params_to_log=["timeout"])
     def wait(self, timeout: float = -1) -> types.JOB_STATUS:
         """
@@ -167,10 +177,10 @@ def _get_status(session: snowpark.Session, job_id: str, instance_id: Optional[in
 
 
 @telemetry.send_api_usage_telemetry(project=_PROJECT, func_params_to_log=["job_id"])
-def _get_service_spec(session: snowpark.Session, job_id: str) -> Dict[str, Any]:
+def _get_service_spec(session: snowpark.Session, job_id: str) -> dict[str, Any]:
     """Retrieve job execution service spec."""
     (row,) = session.sql("DESCRIBE SERVICE IDENTIFIER(?)", params=[job_id]).collect()
-    return cast(Dict[str, Any], yaml.safe_load(row["spec"]))
+    return cast(dict[str, Any], yaml.safe_load(row["spec"]))
 
 
 @telemetry.send_api_usage_telemetry(project=_PROJECT, func_params_to_log=["job_id", "limit", "instance_id"])
@@ -192,7 +202,7 @@ def _get_logs(session: snowpark.Session, job_id: str, limit: int = -1, instance_
         instance_id = _get_head_instance_id(session, job_id)
 
     # Assemble params: [job_id, instance_id, container_name, (optional) limit]
-    params: List[Any] = [
+    params: list[Any] = [
         job_id,
         0 if instance_id is None else instance_id,
         constants.DEFAULT_CONTAINER_NAME,

snowflake/ml/jobs/manager.py

@@ -1,16 +1,7 @@
+import logging
 import pathlib
 import textwrap
-from typing import (
-    Any,
-    Callable,
-    Dict,
-    List,
-    Literal,
-    Optional,
-    TypeVar,
-    Union,
-    overload,
-)
+from typing import Any, Callable, Literal, Optional, TypeVar, Union, overload
 from uuid import uuid4
 
 import yaml
@@ -23,13 +14,14 @@ from snowflake.ml.jobs._utils import payload_utils, spec_utils
 from snowflake.snowpark.context import get_active_session
 from snowflake.snowpark.exceptions import SnowparkSQLException
 
+logger = logging.getLogger(__name__)
+
 _PROJECT = "MLJob"
 JOB_ID_PREFIX = "MLJOB_"
 
 T = TypeVar("T")
 
 
-@snowpark._internal.utils.private_preview(version="1.7.4")
 @telemetry.send_api_usage_telemetry(project=_PROJECT, func_params_to_log=["limit", "scope"])
 def list_jobs(
     limit: int = 10,
@@ -69,7 +61,6 @@ def list_jobs(
     return df
 
 
-@snowpark._internal.utils.private_preview(version="1.7.4")
 @telemetry.send_api_usage_telemetry(project=_PROJECT)
 def get_job(job_id: str, session: Optional[snowpark.Session] = None) -> jb.MLJob[Any]:
     """Retrieve a job service from the backend."""
@@ -93,7 +84,6 @@ def get_job(job_id: str, session: Optional[snowpark.Session] = None) -> jb.MLJob
         raise
 
 
-@snowpark._internal.utils.private_preview(version="1.7.4")
 @telemetry.send_api_usage_telemetry(project=_PROJECT)
 def delete_job(job: Union[str, jb.MLJob[Any]], session: Optional[snowpark.Session] = None) -> None:
     """Delete a job service from the backend. Status and logs will be lost."""
@@ -106,19 +96,18 @@ def delete_job(job: Union[str, jb.MLJob[Any]], session: Optional[snowpark.Sessio
     session.sql("DROP SERVICE IDENTIFIER(?)", params=(job_id,)).collect()
 
 
-@snowpark._internal.utils.private_preview(version="1.7.4")
 @telemetry.send_api_usage_telemetry(project=_PROJECT)
 def submit_file(
     file_path: str,
     compute_pool: str,
     *,
     stage_name: str,
-    args: Optional[List[str]] = None,
-    env_vars: Optional[Dict[str, str]] = None,
-    pip_requirements: Optional[List[str]] = None,
-    external_access_integrations: Optional[List[str]] = None,
+    args: Optional[list[str]] = None,
+    env_vars: Optional[dict[str, str]] = None,
+    pip_requirements: Optional[list[str]] = None,
+    external_access_integrations: Optional[list[str]] = None,
     query_warehouse: Optional[str] = None,
-    spec_overrides: Optional[Dict[str, Any]] = None,
+    spec_overrides: Optional[dict[str, Any]] = None,
     num_instances: Optional[int] = None,
     enable_metrics: bool = False,
     session: Optional[snowpark.Session] = None,
@@ -159,7 +148,6 @@ def submit_file(
     )
 
 
-@snowpark._internal.utils.private_preview(version="1.7.4")
 @telemetry.send_api_usage_telemetry(project=_PROJECT)
 def submit_directory(
     dir_path: str,
@@ -167,12 +155,12 @@ def submit_directory(
     *,
     entrypoint: str,
     stage_name: str,
-    args: Optional[List[str]] = None,
-    env_vars: Optional[Dict[str, str]] = None,
-    pip_requirements: Optional[List[str]] = None,
-    external_access_integrations: Optional[List[str]] = None,
+    args: Optional[list[str]] = None,
+    env_vars: Optional[dict[str, str]] = None,
+    pip_requirements: Optional[list[str]] = None,
+    external_access_integrations: Optional[list[str]] = None,
     query_warehouse: Optional[str] = None,
-    spec_overrides: Optional[Dict[str, Any]] = None,
+    spec_overrides: Optional[dict[str, Any]] = None,
     num_instances: Optional[int] = None,
     enable_metrics: bool = False,
     session: Optional[snowpark.Session] = None,
@@ -222,12 +210,12 @@ def _submit_job(
     *,
     stage_name: str,
     entrypoint: Optional[str] = None,
-    args: Optional[List[str]] = None,
-    env_vars: Optional[Dict[str, str]] = None,
-    pip_requirements: Optional[List[str]] = None,
-    external_access_integrations: Optional[List[str]] = None,
+    args: Optional[list[str]] = None,
+    env_vars: Optional[dict[str, str]] = None,
+    pip_requirements: Optional[list[str]] = None,
+    external_access_integrations: Optional[list[str]] = None,
     query_warehouse: Optional[str] = None,
-    spec_overrides: Optional[Dict[str, Any]] = None,
+    spec_overrides: Optional[dict[str, Any]] = None,
     num_instances: Optional[int] = None,
     enable_metrics: bool = False,
     session: Optional[snowpark.Session] = None,
@@ -242,12 +230,12 @@ def _submit_job(
     *,
     stage_name: str,
     entrypoint: Optional[str] = None,
-    args: Optional[List[str]] = None,
-    env_vars: Optional[Dict[str, str]] = None,
-    pip_requirements: Optional[List[str]] = None,
-    external_access_integrations: Optional[List[str]] = None,
+    args: Optional[list[str]] = None,
+    env_vars: Optional[dict[str, str]] = None,
+    pip_requirements: Optional[list[str]] = None,
+    external_access_integrations: Optional[list[str]] = None,
     query_warehouse: Optional[str] = None,
-    spec_overrides: Optional[Dict[str, Any]] = None,
+    spec_overrides: Optional[dict[str, Any]] = None,
     num_instances: Optional[int] = None,
     enable_metrics: bool = False,
     session: Optional[snowpark.Session] = None,
@@ -263,6 +251,8 @@ def _submit_job(
         # TODO: Log lengths of args, env_vars, and spec_overrides values
         "pip_requirements",
         "external_access_integrations",
+        "num_instances",
+        "enable_metrics",
     ],
 )
 def _submit_job(
@@ -271,12 +261,12 @@ def _submit_job(
     *,
     stage_name: str,
     entrypoint: Optional[str] = None,
-    args: Optional[List[str]] = None,
-    env_vars: Optional[Dict[str, str]] = None,
-    pip_requirements: Optional[List[str]] = None,
-    external_access_integrations: Optional[List[str]] = None,
+    args: Optional[list[str]] = None,
+    env_vars: Optional[dict[str, str]] = None,
+    pip_requirements: Optional[list[str]] = None,
+    external_access_integrations: Optional[list[str]] = None,
     query_warehouse: Optional[str] = None,
-    spec_overrides: Optional[Dict[str, Any]] = None,
+    spec_overrides: Optional[dict[str, Any]] = None,
     num_instances: Optional[int] = None,
     enable_metrics: bool = False,
     session: Optional[snowpark.Session] = None,
@@ -305,6 +295,12 @@ def _submit_job(
     Raises:
         RuntimeError: If required Snowflake features are not enabled.
     """
+    # Display warning about PrPr parameters
+    if num_instances is not None:
+        logger.warning(
+            "_submit_job() parameter 'num_instances' is in private preview since 1.8.2. Do not use it in production.",
+        )
+
     session = session or get_active_session()
     job_id = f"{JOB_ID_PREFIX}{str(uuid4()).replace('-', '_').upper()}"
     stage_name = "@" + stage_name.lstrip("@").rstrip("/")

snowflake/ml/lineage/lineage_node.py

@@ -1,6 +1,6 @@
 import json
 from datetime import datetime
-from typing import TYPE_CHECKING, Dict, List, Literal, Optional, Set, Type, Union
+from typing import TYPE_CHECKING, Literal, Optional, Union
 
 from snowflake import snowpark
 from snowflake.ml._internal import telemetry
@@ -12,7 +12,7 @@ if TYPE_CHECKING:
     from snowflake.ml.model._client.model import model_version_impl
 
 _PROJECT = "LINEAGE"
-DOMAIN_LINEAGE_REGISTRY: Dict[str, Type["LineageNode"]] = {}
+DOMAIN_LINEAGE_REGISTRY: dict[str, type["LineageNode"]] = {}
 
 
 class LineageNode:
@@ -87,8 +87,8 @@ class LineageNode:
     def lineage(
         self,
         direction: Literal["upstream", "downstream"] = "downstream",
-        domain_filter: Optional[Set[Literal["feature_view", "dataset", "model", "table", "view"]]] = None,
-    ) -> List[Union["feature_view.FeatureView", "dataset.Dataset", "model_version_impl.ModelVersion", "LineageNode"]]:
+        domain_filter: Optional[set[Literal["feature_view", "dataset", "model", "table", "view"]]] = None,
+    ) -> list[Union["feature_view.FeatureView", "dataset.Dataset", "model_version_impl.ModelVersion", "LineageNode"]]:
         """
         Retrieves the lineage nodes connected to this node.
 
@@ -109,7 +109,7 @@ class LineageNode:
         if domain_filter is not None:
             domain_filter = {d.lower() for d in domain_filter}  # type: ignore[misc]
 
-        lineage_nodes: List["LineageNode"] = []
+        lineage_nodes: list["LineageNode"] = []
         for row in df.collect():
             lineage_object = (
                 json.loads(row["TARGET_OBJECT"])

snowflake/ml/model/_client/model/model_impl.py

@@ -1,4 +1,4 @@
-from typing import Dict, List, Optional, Union
+from typing import Optional, Union
 
 import pandas as pd
 
@@ -224,7 +224,7 @@ class Model:
         project=_TELEMETRY_PROJECT,
         subproject=_TELEMETRY_SUBPROJECT,
     )
-    def versions(self) -> List[model_version_impl.ModelVersion]:
+    def versions(self) -> list[model_version_impl.ModelVersion]:
         """Get all versions in the model.
 
         Returns:
@@ -298,7 +298,7 @@ class Model:
         project=_TELEMETRY_PROJECT,
         subproject=_TELEMETRY_SUBPROJECT,
     )
-    def show_tags(self) -> Dict[str, str]:
+    def show_tags(self) -> dict[str, str]:
         """Get a dictionary showing the tag and its value attached to the model.
 
         Returns:

snowflake/ml/model/_client/model/model_version_impl.py

@@ -2,10 +2,11 @@ import enum
 import pathlib
 import tempfile
 import warnings
-from typing import Any, Callable, Dict, List, Optional, Union, overload
+from typing import Any, Callable, Optional, Union, overload
 
 import pandas as pd
 
+from snowflake import snowpark
 from snowflake.ml._internal import telemetry
 from snowflake.ml._internal.utils import sql_identifier
 from snowflake.ml.lineage import lineage_node
@@ -32,7 +33,7 @@ class ModelVersion(lineage_node.LineageNode):
     _service_ops: service_ops.ServiceOperator
     _model_name: sql_identifier.SqlIdentifier
     _version_name: sql_identifier.SqlIdentifier
-    _functions: List[model_manifest_schema.ModelFunctionInfo]
+    _functions: list[model_manifest_schema.ModelFunctionInfo]
 
     def __init__(self) -> None:
         raise RuntimeError("ModelVersion's initializer is not meant to be used. Use `version` from model instead.")
@@ -152,7 +153,7 @@ class ModelVersion(lineage_node.LineageNode):
         project=_TELEMETRY_PROJECT,
         subproject=_TELEMETRY_SUBPROJECT,
     )
-    def show_metrics(self) -> Dict[str, Any]:
+    def show_metrics(self) -> dict[str, Any]:
         """Show all metrics logged with the model version.
 
         Returns:
@@ -293,7 +294,7 @@ class ModelVersion(lineage_node.LineageNode):
             statement_params=statement_params,
         )
 
-    def _get_functions(self) -> List[model_manifest_schema.ModelFunctionInfo]:
+    def _get_functions(self) -> list[model_manifest_schema.ModelFunctionInfo]:
         statement_params = telemetry.get_statement_params(
             project=_TELEMETRY_PROJECT,
             subproject=_TELEMETRY_SUBPROJECT,
@@ -327,7 +328,7 @@ class ModelVersion(lineage_node.LineageNode):
         project=_TELEMETRY_PROJECT,
         subproject=_TELEMETRY_SUBPROJECT,
     )
-    def show_functions(self) -> List[model_manifest_schema.ModelFunctionInfo]:
+    def show_functions(self) -> list[model_manifest_schema.ModelFunctionInfo]:
         """Show all functions information in a model version that is callable.
 
         Returns:
@@ -405,11 +406,6 @@ class ModelVersion(lineage_node.LineageNode):
             strict_input_validation: Enable stricter validation for the input data. This will result value range based
                 type validation to make sure your input data won't overflow when providing to the model.
 
-        Raises:
-            ValueError: When no method with the corresponding name is available.
-            ValueError: When there are more than 1 target methods available in the model but no function name specified.
-            ValueError: When the partition column is not a valid Snowflake identifier.
-
         Returns:
             The prediction data. It would be the same type dataframe as your input.
         """
@@ -422,29 +418,7 @@ class ModelVersion(lineage_node.LineageNode):
             # Partition column must be a valid identifier
             partition_column = sql_identifier.SqlIdentifier(partition_column)
 
-        functions: List[model_manifest_schema.ModelFunctionInfo] = self._functions
-
-        if function_name:
-            req_method_name = sql_identifier.SqlIdentifier(function_name).identifier()
-            find_method: Callable[[model_manifest_schema.ModelFunctionInfo], bool] = (
-                lambda method: method["name"] == req_method_name
-            )
-            target_function_info = next(
-                filter(find_method, functions),
-                None,
-            )
-            if target_function_info is None:
-                raise ValueError(
-                    f"There is no method with name {function_name} available in the model"
-                    f" {self.fully_qualified_model_name} version {self.version_name}"
-                )
-        elif len(functions) != 1:
-            raise ValueError(
-                f"There are more than 1 target methods available in the model {self.fully_qualified_model_name}"
-                f" version {self.version_name}. Please specify a `function_name` when calling the `run` method."
-            )
-        else:
-            target_function_info = functions[0]
+        target_function_info = self._get_function_info(function_name=function_name)
 
         if service_name:
             database_name_id, schema_name_id, service_name_id = sql_identifier.parse_fully_qualified_name(service_name)
@@ -475,6 +449,33 @@ class ModelVersion(lineage_node.LineageNode):
                 is_partitioned=target_function_info["is_partitioned"],
             )
 
+    def _get_function_info(self, function_name: Optional[str]) -> model_manifest_schema.ModelFunctionInfo:
+        functions: list[model_manifest_schema.ModelFunctionInfo] = self._functions
+
+        if function_name:
+            req_method_name = sql_identifier.SqlIdentifier(function_name).identifier()
+            find_method: Callable[[model_manifest_schema.ModelFunctionInfo], bool] = (
+                lambda method: method["name"] == req_method_name
+            )
+            target_function_info = next(
+                filter(find_method, functions),
+                None,
+            )
+            if target_function_info is None:
+                raise ValueError(
+                    f"There is no method with name {function_name} available in the model"
+                    f" {self.fully_qualified_model_name} version {self.version_name}"
+                )
+        elif len(functions) != 1:
+            raise ValueError(
+                f"There are more than 1 target methods available in the model {self.fully_qualified_model_name}"
+                f" version {self.version_name}. Please specify a `function_name` when calling the `run` method."
+            )
+        else:
+            target_function_info = functions[0]
+
+        return target_function_info
+
     @telemetry.send_api_usage_telemetry(
         project=_TELEMETRY_PROJECT, subproject=_TELEMETRY_SUBPROJECT, func_params_to_log=["export_mode"]
     )
@@ -684,7 +685,7 @@ class ModelVersion(lineage_node.LineageNode):
         num_workers: Optional[int] = None,
         max_batch_rows: Optional[int] = None,
         force_rebuild: bool = False,
-        build_external_access_integrations: Optional[List[str]] = None,
+        build_external_access_integrations: Optional[list[str]] = None,
         block: bool = True,
     ) -> Union[str, async_job.AsyncJob]:
         """Create an inference service with the given spec.
@@ -751,7 +752,7 @@ class ModelVersion(lineage_node.LineageNode):
         max_batch_rows: Optional[int] = None,
         force_rebuild: bool = False,
         build_external_access_integration: Optional[str] = None,
-        build_external_access_integrations: Optional[List[str]] = None,
+        build_external_access_integrations: Optional[list[str]] = None,
         block: bool = True,
     ) -> Union[str, async_job.AsyncJob]:
         """Create an inference service with the given spec.
@@ -914,5 +915,72 @@ class ModelVersion(lineage_node.LineageNode):
             statement_params=statement_params,
         )
 
+    @snowpark._internal.utils.private_preview(version="1.8.3")
+    @telemetry.send_api_usage_telemetry(
+        project=_TELEMETRY_PROJECT,
+        subproject=_TELEMETRY_SUBPROJECT,
+    )
+    def run_job(
+        self,
+        X: Union[pd.DataFrame, "dataframe.DataFrame"],
+        *,
+        job_name: str,
+        compute_pool: str,
+        image_repo: str,
+        output_table_name: str,
+        function_name: Optional[str] = None,
+        cpu_requests: Optional[str] = None,
+        memory_requests: Optional[str] = None,
+        gpu_requests: Optional[Union[str, int]] = None,
+        num_workers: Optional[int] = None,
+        max_batch_rows: Optional[int] = None,
+        force_rebuild: bool = False,
+        build_external_access_integrations: Optional[list[str]] = None,
+    ) -> Union[pd.DataFrame, dataframe.DataFrame]:
+        statement_params = telemetry.get_statement_params(
+            project=_TELEMETRY_PROJECT,
+            subproject=_TELEMETRY_SUBPROJECT,
+        )
+        target_function_info = self._get_function_info(function_name=function_name)
+        job_db_id, job_schema_id, job_id = sql_identifier.parse_fully_qualified_name(job_name)
+        image_repo_db_id, image_repo_schema_id, image_repo_id = sql_identifier.parse_fully_qualified_name(image_repo)
+        output_table_db_id, output_table_schema_id, output_table_id = sql_identifier.parse_fully_qualified_name(
+            output_table_name
+        )
+        warehouse = self._service_ops._session.get_current_warehouse()
+        assert warehouse, "No active warehouse selected in the current session."
+        return self._service_ops.invoke_job_method(
+            target_method=target_function_info["target_method"],
+            signature=target_function_info["signature"],
+            X=X,
+            database_name=None,
+            schema_name=None,
+            model_name=self._model_name,
+            version_name=self._version_name,
+            job_database_name=job_db_id,
+            job_schema_name=job_schema_id,
+            job_name=job_id,
+            compute_pool_name=sql_identifier.SqlIdentifier(compute_pool),
+            warehouse_name=sql_identifier.SqlIdentifier(warehouse),
+            image_repo_database_name=image_repo_db_id,
+            image_repo_schema_name=image_repo_schema_id,
+            image_repo_name=image_repo_id,
+            output_table_database_name=output_table_db_id,
+            output_table_schema_name=output_table_schema_id,
+            output_table_name=output_table_id,
+            cpu_requests=cpu_requests,
+            memory_requests=memory_requests,
+            gpu_requests=gpu_requests,
+            num_workers=num_workers,
+            max_batch_rows=max_batch_rows,
+            force_rebuild=force_rebuild,
+            build_external_access_integrations=(
+                None
+                if build_external_access_integrations is None
+                else [sql_identifier.SqlIdentifier(eai) for eai in build_external_access_integrations]
+            ),
+            statement_params=statement_params,
+        )
+
 
 lineage_node.DOMAIN_LINEAGE_REGISTRY["model"] = ModelVersion

snowflake/ml/model/_client/ops/metadata_ops.py

@@ -1,5 +1,5 @@
 import json
-from typing import Any, Dict, Optional, TypedDict
+from typing import Any, Optional, TypedDict
 
 from typing_extensions import NotRequired
 
@@ -14,7 +14,7 @@ MODEL_VERSION_METADATA_SCHEMA_VERSION = "2024-01-01"
 
 
 class ModelVersionMetadataSchema(TypedDict):
-    metrics: NotRequired[Dict[str, Any]]
+    metrics: NotRequired[dict[str, Any]]
 
 
 class MetadataOperator:
@@ -44,7 +44,7 @@ class MetadataOperator:
         )
 
     @staticmethod
-    def _parse(metadata_dict: Dict[str, Any]) -> ModelVersionMetadataSchema:
+    def _parse(metadata_dict: dict[str, Any]) -> ModelVersionMetadataSchema:
         loaded_metadata_schema_version = metadata_dict.get("snowpark_ml_schema_version", None)
         if loaded_metadata_schema_version is None:
             return ModelVersionMetadataSchema(metrics={})
@@ -65,8 +65,8 @@ class MetadataOperator:
         schema_name: Optional[sql_identifier.SqlIdentifier],
         model_name: sql_identifier.SqlIdentifier,
         version_name: sql_identifier.SqlIdentifier,
-        statement_params: Optional[Dict[str, Any]] = None,
-    ) -> Dict[str, Any]:
+        statement_params: Optional[dict[str, Any]] = None,
+    ) -> dict[str, Any]:
         version_info_list = self._model_client.show_versions(
             database_name=database_name,
             schema_name=schema_name,
@@ -89,7 +89,7 @@ class MetadataOperator:
         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,
     ) -> ModelVersionMetadataSchema:
         metadata_dict = self._get_current_metadata_dict(
             database_name=database_name,
@@ -108,7 +108,7 @@ class MetadataOperator:
         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:
         metadata_dict = self._get_current_metadata_dict(
             database_name=database_name,