snowflake-ml-python 1.7.1__py3-none-any.whl → 1.7.3__py3-none-any.whl

snowflake/ml/monitoring/entities/model_monitor_config.py

@@ -6,23 +6,49 @@ from snowflake.ml.model._client.model import model_version_impl
 
 @dataclass
 class ModelMonitorSourceConfig:
+    """Configuration for the source of data to be monitored."""
+
     source: str
+    """Name of table or view containing monitoring data."""
+
     timestamp_column: str
+    """Name of column in the source containing timestamp."""
+
     id_columns: List[str]
+    """List of columns in the source containing unique identifiers."""
+
     prediction_score_columns: Optional[List[str]] = None
+    """List of columns in the source containing prediction scores.
+    Can be regression scores for regression models and probability scores for classification models."""
+
     prediction_class_columns: Optional[List[str]] = None
+    """List of columns in the source containing prediction classes for classification models."""
+
     actual_score_columns: Optional[List[str]] = None
+    """List of columns in the source containing actual scores."""
+
     actual_class_columns: Optional[List[str]] = None
+    """List of columns in the source containing actual classes for classification models."""
+
     baseline: Optional[str] = None
+    """Name of table containing the baseline data."""
 
 
 @dataclass
 class ModelMonitorConfig:
+    """Configuration for the Model Monitor."""
+
     model_version: model_version_impl.ModelVersion
+    """Model version to monitor."""
 
-    # Python model function name
     model_function_name: str
+    """Function name in the model to monitor."""
+
     background_compute_warehouse_name: str
-    # TODO: Add support for pythonic notion of time.
+    """Name of the warehouse to use for background compute."""
+
     refresh_interval: str = "1 hour"
+    """Interval at which to refresh the monitoring data."""
+
     aggregation_window: str = "1 day"
+    """Window for aggregating monitoring data."""

snowflake/ml/monitoring/model_monitor.py

@@ -1,5 +1,7 @@
+from snowflake import snowpark
 from snowflake.ml._internal import telemetry
 from snowflake.ml._internal.utils import sql_identifier
+from snowflake.ml.monitoring import model_monitor_version
 from snowflake.ml.monitoring._client import model_monitor_sql_client
 
 
@@ -9,13 +11,8 @@ class ModelMonitor:
     name: sql_identifier.SqlIdentifier
     _model_monitor_client: model_monitor_sql_client.ModelMonitorSQLClient
 
-    statement_params = telemetry.get_statement_params(
-        telemetry.TelemetryProject.MLOPS.value,
-        telemetry.TelemetrySubProject.MONITORING.value,
-    )
-
     def __init__(self) -> None:
-        raise RuntimeError("ModelMonitor's initializer is not meant to be used.")
+        raise RuntimeError("Model Monitor's initializer is not meant to be used.")
 
     @classmethod
     def _ref(
@@ -28,10 +25,28 @@ class ModelMonitor:
         self._model_monitor_client = model_monitor_client
         return self
 
+    @telemetry.send_api_usage_telemetry(
+        project=telemetry.TelemetryProject.MLOPS.value,
+        subproject=telemetry.TelemetrySubProject.MONITORING.value,
+    )
+    @snowpark._internal.utils.private_preview(version=model_monitor_version.SNOWFLAKE_ML_MONITORING_MIN_VERSION)
     def suspend(self) -> None:
-        """Suspend pipeline for ModelMonitor"""
-        self._model_monitor_client.suspend_monitor(self.name, statement_params=self.statement_params)
-
+        """Suspend the Model Monitor"""
+        statement_params = telemetry.get_statement_params(
+            telemetry.TelemetryProject.MLOPS.value,
+            telemetry.TelemetrySubProject.MONITORING.value,
+        )
+        self._model_monitor_client.suspend_monitor(self.name, statement_params=statement_params)
+
+    @telemetry.send_api_usage_telemetry(
+        project=telemetry.TelemetryProject.MLOPS.value,
+        subproject=telemetry.TelemetrySubProject.MONITORING.value,
+    )
+    @snowpark._internal.utils.private_preview(version=model_monitor_version.SNOWFLAKE_ML_MONITORING_MIN_VERSION)
     def resume(self) -> None:
-        """Resume pipeline for ModelMonitor"""
-        self._model_monitor_client.resume_monitor(self.name, statement_params=self.statement_params)
+        """Resume the Model Monitor"""
+        statement_params = telemetry.get_statement_params(
+            telemetry.TelemetryProject.MLOPS.value,
+            telemetry.TelemetrySubProject.MONITORING.value,
+        )
+        self._model_monitor_client.resume_monitor(self.name, statement_params=statement_params)

snowflake/ml/registry/_manager/model_manager.py

@@ -1,13 +1,13 @@
 from types import ModuleType
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 import pandas as pd
 from absl.logging import logging
-from packaging import version
 
 from snowflake.ml._internal import telemetry
+from snowflake.ml._internal.exceptions import error_codes, exceptions
 from snowflake.ml._internal.human_readable_id import hrid_generator
-from snowflake.ml._internal.utils import snowflake_env, sql_identifier
+from snowflake.ml._internal.utils import sql_identifier
 from snowflake.ml.model import model_signature, type_hints as model_types
 from snowflake.ml.model._client.model import model_impl, model_version_impl
 from snowflake.ml.model._client.ops import metadata_ops, model_ops, service_ops
@@ -50,14 +50,40 @@ class ModelManager:
         python_version: Optional[str] = None,
         signatures: Optional[Dict[str, model_signature.ModelSignature]] = None,
         sample_input_data: Optional[model_types.SupportedDataType] = None,
+        user_files: Optional[Dict[str, List[str]]] = None,
         code_paths: Optional[List[str]] = None,
         ext_modules: Optional[List[ModuleType]] = None,
         task: model_types.Task = model_types.Task.UNKNOWN,
         options: Optional[model_types.ModelSaveOption] = None,
         statement_params: Optional[Dict[str, Any]] = None,
     ) -> model_version_impl.ModelVersion:
-        if not version_name:
-            version_name = self._hrid_generator.generate()[1]
+
+        database_name_id, schema_name_id, model_name_id = self._parse_fully_qualified_name(model_name)
+
+        model_exists = self._model_ops.validate_existence(
+            database_name=database_name_id,
+            schema_name=schema_name_id,
+            model_name=model_name_id,
+            statement_params=statement_params,
+        )
+
+        if version_name is None:
+            if model_exists:
+                versions = self._model_ops.list_models_or_versions(
+                    database_name=database_name_id,
+                    schema_name=schema_name_id,
+                    model_name=model_name_id,
+                    statement_params=statement_params,
+                )
+                for _ in range(1000):
+                    hrid = self._hrid_generator.generate()[1]
+                    if sql_identifier.SqlIdentifier(hrid) not in versions:
+                        version_name = hrid
+                        break
+                if version_name is None:
+                    raise RuntimeError("Random version name generation failed.")
+            else:
+                version_name = self._hrid_generator.generate()[1]
 
         if isinstance(model, model_version_impl.ModelVersion):
             (
@@ -75,10 +101,24 @@ class ModelManager:
                 schema_name=None,
                 model_name=sql_identifier.SqlIdentifier(model_name),
                 version_name=sql_identifier.SqlIdentifier(version_name),
+                model_exists=model_exists,
                 statement_params=statement_params,
             )
             return self.get_model(model_name=model_name, statement_params=statement_params).version(version_name)
 
+        version_name_id = sql_identifier.SqlIdentifier(version_name)
+        if model_exists and self._model_ops.validate_existence(
+            database_name=database_name_id,
+            schema_name=schema_name_id,
+            model_name=model_name_id,
+            version_name=version_name_id,
+            statement_params=statement_params,
+        ):
+            raise ValueError(
+                f"Model {model_name} version {version_name} already existed. "
+                + "To auto-generate `version_name`, skip that argument."
+            )
+
         return self._log_model(
             model=model,
             model_name=model_name,
@@ -91,6 +131,7 @@ class ModelManager:
             python_version=python_version,
             signatures=signatures,
             sample_input_data=sample_input_data,
+            user_files=user_files,
             code_paths=code_paths,
             ext_modules=ext_modules,
             task=task,
@@ -103,7 +144,7 @@ class ModelManager:
         model: model_types.SupportedModelType,
         *,
         model_name: str,
-        version_name: Optional[str] = None,
+        version_name: str,
         comment: Optional[str] = None,
         metrics: Optional[Dict[str, Any]] = None,
         conda_dependencies: Optional[List[str]] = None,
@@ -112,6 +153,7 @@ class ModelManager:
         python_version: Optional[str] = None,
         signatures: Optional[Dict[str, model_signature.ModelSignature]] = None,
         sample_input_data: Optional[model_types.SupportedDataType] = None,
+        user_files: Optional[Dict[str, List[str]]] = None,
         code_paths: Optional[List[str]] = None,
         ext_modules: Optional[List[ModuleType]] = None,
         task: model_types.Task = model_types.Task.UNKNOWN,
@@ -119,28 +161,8 @@ class ModelManager:
         statement_params: Optional[Dict[str, Any]] = None,
     ) -> model_version_impl.ModelVersion:
         database_name_id, schema_name_id, model_name_id = sql_identifier.parse_fully_qualified_name(model_name)
-
-        if not version_name:
-            version_name = self._hrid_generator.generate()[1]
         version_name_id = sql_identifier.SqlIdentifier(version_name)
 
-        if self._model_ops.validate_existence(
-            database_name=database_name_id,
-            schema_name=schema_name_id,
-            model_name=model_name_id,
-            statement_params=statement_params,
-        ) and self._model_ops.validate_existence(
-            database_name=database_name_id,
-            schema_name=schema_name_id,
-            model_name=model_name_id,
-            version_name=version_name_id,
-            statement_params=statement_params,
-        ):
-            raise ValueError(
-                f"Model {model_name} version {version_name} already existed. "
-                + "To auto-generate `version_name`, skip that argument."
-            )
-
         stage_path = self._model_ops.prepare_model_stage_path(
             database_name=database_name_id,
             schema_name=schema_name_id,
@@ -148,13 +170,10 @@ class ModelManager:
         )
 
         platforms = None
-        # TODO(jbahk): Remove the version check after Snowflake 8.40.0 release
         # User specified target platforms are defaulted to None and will not show up in the generated manifest.
-        # In the backend, we attempt to create a model for all platforms (WH, SPCS) regardless by default.
-        if snowflake_env.get_current_snowflake_version(self._model_ops._session) >= version.parse("8.40.0"):
+        if target_platforms:
             # Convert any string target platforms to TargetPlatform objects
-            if target_platforms:
-                platforms = [model_types.TargetPlatform(platform) for platform in target_platforms]
+            platforms = [model_types.TargetPlatform(platform) for platform in target_platforms]
 
         logger.info("Start packaging and uploading your model. It might take some time based on the size of the model.")
 
@@ -170,6 +189,7 @@ class ModelManager:
             pip_requirements=pip_requirements,
             target_platforms=platforms,
             python_version=python_version,
+            user_files=user_files,
             code_paths=code_paths,
             ext_modules=ext_modules,
             options=options,
@@ -229,7 +249,7 @@ class ModelManager:
         *,
         statement_params: Optional[Dict[str, Any]] = None,
     ) -> model_impl.Model:
-        database_name_id, schema_name_id, model_name_id = sql_identifier.parse_fully_qualified_name(model_name)
+        database_name_id, schema_name_id, model_name_id = self._parse_fully_qualified_name(model_name)
         if self._model_ops.validate_existence(
             database_name=database_name_id,
             schema_name=schema_name_id,
@@ -289,7 +309,7 @@ class ModelManager:
         *,
         statement_params: Optional[Dict[str, Any]] = None,
     ) -> None:
-        database_name_id, schema_name_id, model_name_id = sql_identifier.parse_fully_qualified_name(model_name)
+        database_name_id, schema_name_id, model_name_id = self._parse_fully_qualified_name(model_name)
 
         self._model_ops.delete_model_or_version(
             database_name=database_name_id,
@@ -297,3 +317,20 @@ class ModelManager:
             model_name=model_name_id,
             statement_params=statement_params,
         )
+
+    def _parse_fully_qualified_name(
+        self, model_name: str
+    ) -> Tuple[
+        Optional[sql_identifier.SqlIdentifier], Optional[sql_identifier.SqlIdentifier], sql_identifier.SqlIdentifier
+    ]:
+        try:
+            return sql_identifier.parse_fully_qualified_name(model_name)
+        except ValueError:
+            raise exceptions.SnowflakeMLException(
+                error_code=error_codes.INVALID_ARGUMENT,
+                original_exception=ValueError(
+                    f"The model_name `{model_name}` cannot be parsed as a SQL identifier. Alphanumeric characters and "
+                    "underscores are permitted. See https://docs.snowflake.com/en/sql-reference/identifiers-syntax for "
+                    "more information."
+                ),
+            )

snowflake/ml/registry/registry.py

@@ -117,41 +117,49 @@ class Registry:
         options: Optional[model_types.ModelSaveOption] = None,
     ) -> ModelVersion:
         """
-        Log a model with various parameters and metadata.
+        Log a model with various parameters and metadata, or a ModelVersion object.
 
         Args:
-            model: Model object of supported types such as Scikit-learn, XGBoost, LightGBM, Snowpark ML,
-                PyTorch, TorchScript, Tensorflow, Tensorflow Keras, MLFlow, HuggingFace Pipeline,
-                Sentence Transformers, or Custom Model.
-            model_name: Name to identify the model.
+            model: Supported model or ModelVersion object.
+                - Supported model: Model object of supported types such as Scikit-learn, XGBoost, LightGBM, Snowpark ML,
+                PyTorch, TorchScript, Tensorflow, Tensorflow Keras, MLFlow, HuggingFace Pipeline, Sentence Transformers,
+                or Custom Model.
+                - ModelVersion: Source ModelVersion object used to create the new ModelVersion object.
+            model_name: Name to identify the model. This must be a valid Snowflake SQL Identifier. Alphanumeric
+                characters and underscores are permitted.
+                See https://docs.snowflake.com/en/sql-reference/identifiers-syntax for more.
             version_name: Version identifier for the model. Combination of model_name and version_name must be unique.
                 If not specified, a random name will be generated.
             comment: Comment associated with the model version. Defaults to None.
             metrics: A JSON serializable dictionary containing metrics linked to the model version. Defaults to None.
-            signatures: Model data signatures for inputs and outputs for various target methods. If it is None,
-                sample_input_data would be used to infer the signatures for those models that cannot automatically
-                infer the signature. Defaults to None.
-            sample_input_data: Sample input data to infer model signatures from.
-                It would also be used as background data in explanation and to capture data lineage. Defaults to None.
             conda_dependencies: List of Conda package specifications. Use "[channel::]package [operator version]" syntax
                 to specify a dependency. It is a recommended way to specify your dependencies using conda. When channel
                 is not specified, Snowflake Anaconda Channel will be used. Defaults to None.
             pip_requirements: List of Pip package specifications. Defaults to None.
-                Currently it is not supported since Model can only executed in Snowflake Warehouse where all
-                dependencies are required to be retrieved from Snowflake Anaconda Channel.
+                Models with pip requirements are currently only runnable in Snowpark Container Services.
+                See https://docs.snowflake.com/en/developer-guide/snowflake-ml/model-registry/container for more.
+                Models with pip requirements specified will not be executable in Snowflake Warehouse where all
+                dependencies must be retrieved from Snowflake Anaconda Channel.
             target_platforms: List of target platforms to run the model. The only acceptable inputs are a combination of
                 {"WAREHOUSE", "SNOWPARK_CONTAINER_SERVICES"}. Defaults to None.
             python_version: Python version in which the model is run. Defaults to None.
+            signatures: Model data signatures for inputs and outputs for various target methods. If it is None,
+                sample_input_data would be used to infer the signatures for those models that cannot automatically
+                infer the signature. If not None, sample_input_data should not be specified. Defaults to None.
+            sample_input_data: Sample input data to infer model signatures from.
+                It would also be used as background data in explanation and to capture data lineage. Defaults to None.
             code_paths: List of directories containing code to import. Defaults to None.
             ext_modules: List of external modules to pickle with the model object.
                 Only supported when logging the following types of model:
                 Scikit-learn, Snowpark ML, PyTorch, TorchScript and Custom Model. Defaults to None.
             options (Dict[str, Any], optional): Additional model saving options.
+
                 Model Saving Options include:
+
                 - embed_local_ml_library: Embed local Snowpark ML into the code directory or folder.
                     Override to True if the local Snowpark ML version is not available in the Snowflake Anaconda
                     Channel. Otherwise, defaults to False
-                - relax_version: Whether or not relax the version constraints of the dependencies when running in the
+                - relax_version: Whether to relax the version constraints of the dependencies when running in the
                     Warehouse. It detects any ==x.y.z in specifiers and replaced with >=x.y, <(x+1). Defaults to True.
                 - function_type: Set the method function type globally. To set method function types individually see
                   function_type in model_options.
@@ -163,7 +171,10 @@ class Registry:
                     - max_batch_size: Maximum batch size that the method could accept in the Snowflake Warehouse.
                         Defaults to None, determined automatically by Snowflake.
                     - function_type: One of supported model method function types (FUNCTION or TABLE_FUNCTION).
+        Returns:
+            ModelVersion: ModelVersion object corresponding to the model just logged.
         """
+
         ...
 
     @overload
@@ -214,6 +225,7 @@ class Registry:
         python_version: Optional[str] = None,
         signatures: Optional[Dict[str, model_signature.ModelSignature]] = None,
         sample_input_data: Optional[model_types.SupportedDataType] = None,
+        user_files: Optional[Dict[str, List[str]]] = None,
         code_paths: Optional[List[str]] = None,
         ext_modules: Optional[List[ModuleType]] = None,
         task: model_types.Task = model_types.Task.UNKNOWN,
@@ -228,25 +240,31 @@ class Registry:
                 PyTorch, TorchScript, Tensorflow, Tensorflow Keras, MLFlow, HuggingFace Pipeline, Sentence Transformers,
                 or Custom Model.
                 - ModelVersion: Source ModelVersion object used to create the new ModelVersion object.
-            model_name: Name to identify the model.
+            model_name: Name to identify the model. This must be a valid Snowflake SQL Identifier. Alphanumeric
+                characters and underscores are permitted.
+                See https://docs.snowflake.com/en/sql-reference/identifiers-syntax for more.
             version_name: Version identifier for the model. Combination of model_name and version_name must be unique.
                 If not specified, a random name will be generated.
             comment: Comment associated with the model version. Defaults to None.
             metrics: A JSON serializable dictionary containing metrics linked to the model version. Defaults to None.
-            signatures: Model data signatures for inputs and outputs for various target methods. If it is None,
-                sample_input_data would be used to infer the signatures for those models that cannot automatically
-                infer the signature. If not None, sample_input_data should not be specified. Defaults to None.
-            sample_input_data: Sample input data to infer model signatures from.
-                It would also be used as background data in explanation and to capture data lineage. Defaults to None.
             conda_dependencies: List of Conda package specifications. Use "[channel::]package [operator version]" syntax
                 to specify a dependency. It is a recommended way to specify your dependencies using conda. When channel
                 is not specified, Snowflake Anaconda Channel will be used. Defaults to None.
             pip_requirements: List of Pip package specifications. Defaults to None.
-                Currently it is not supported since Model can only executed in Snowflake Warehouse where all
-                dependencies are required to be retrieved from Snowflake Anaconda Channel.
+                Models with pip requirements are currently only runnable in Snowpark Container Services.
+                See https://docs.snowflake.com/en/developer-guide/snowflake-ml/model-registry/container for more.
+                Models with pip requirements specified will not be executable in Snowflake Warehouse where all
+                dependencies must be retrieved from Snowflake Anaconda Channel.
             target_platforms: List of target platforms to run the model. The only acceptable inputs are a combination of
                 {"WAREHOUSE", "SNOWPARK_CONTAINER_SERVICES"}. Defaults to None.
             python_version: Python version in which the model is run. Defaults to None.
+            signatures: Model data signatures for inputs and outputs for various target methods. If it is None,
+                sample_input_data would be used to infer the signatures for those models that cannot automatically
+                infer the signature. If not None, sample_input_data should not be specified. Defaults to None.
+            sample_input_data: Sample input data to infer model signatures from.
+                It would also be used as background data in explanation and to capture data lineage. Defaults to None.
+            user_files: Dictionary where the keys are subdirectories, and values are lists of local file name
+                strings. The local file name strings can include wildcards (? or *) for matching multiple files.
             code_paths: List of directories containing code to import. Defaults to None.
             ext_modules: List of external modules to pickle with the model object.
                 Only supported when logging the following types of model:
@@ -261,7 +279,7 @@ class Registry:
                 - embed_local_ml_library: Embed local Snowpark ML into the code directory or folder.
                     Override to True if the local Snowpark ML version is not available in the Snowflake Anaconda
                     Channel. Otherwise, defaults to False
-                - relax_version: Whether or not relax the version constraints of the dependencies when running in the
+                - relax_version: Whether to relax the version constraints of the dependencies when running in the
                     Warehouse. It detects any ==x.y.z in specifiers and replaced with >=x.y, <(x+1). Defaults to True.
                 - function_type: Set the method function type globally. To set method function types individually see
                   function_type in model_options.
@@ -301,6 +319,7 @@ class Registry:
             python_version=python_version,
             signatures=signatures,
             sample_input_data=sample_input_data,
+            user_files=user_files,
             code_paths=code_paths,
             ext_modules=ext_modules,
             task=task,
@@ -388,15 +407,15 @@ class Registry:
         source_config: model_monitor_config.ModelMonitorSourceConfig,
         model_monitor_config: model_monitor_config.ModelMonitorConfig,
     ) -> model_monitor.ModelMonitor:
-        """Add a Model Monitor to the Registry
+        """Add a Model Monitor to the Registry.
 
         Args:
-            name: Name of Model Monitor to create
-            source_config: Configuration options of table for ModelMonitor.
-            model_monitor_config: Configuration options of ModelMonitor.
+            name: Name of Model Monitor to create.
+            source_config: Configuration options of table for Model Monitor.
+            model_monitor_config: Configuration options of Model Monitor.
 
         Returns:
-            The newly added ModelMonitor object.
+            The newly added Model Monitor object.
 
         Raises:
             ValueError: If monitoring is not enabled in the Registry.
@@ -407,16 +426,16 @@ class Registry:
 
     @overload
     def get_monitor(self, model_version: model_version_impl.ModelVersion) -> model_monitor.ModelMonitor:
-        """Get a Model Monitor on a ModelVersion from the Registry
+        """Get a Model Monitor on a Model Version from the Registry.
 
         Args:
-            model_version: ModelVersion for which to retrieve the ModelMonitor.
+            model_version: Model Version for which to retrieve the Model Monitor.
         """
         ...
 
     @overload
     def get_monitor(self, name: str) -> model_monitor.ModelMonitor:
-        """Get a Model Monitor from the Registry
+        """Get a Model Monitor by name from the Registry.
 
         Args:
             name: Name of Model Monitor to retrieve.
@@ -431,14 +450,14 @@ class Registry:
     def get_monitor(
         self, *, name: Optional[str] = None, model_version: Optional[model_version_impl.ModelVersion] = None
     ) -> model_monitor.ModelMonitor:
-        """Get a Model Monitor from the Registry
+        """Get a Model Monitor from the Registry.
 
         Args:
             name: Name of Model Monitor to retrieve.
-            model_version: ModelVersion for which to retrieve the ModelMonitor.
+            model_version: Model Version for which to retrieve the Model Monitor.
 
         Returns:
-            The fetched ModelMonitor.
+            The fetched Model Monitor.
 
         Raises:
             ValueError: If monitoring is not enabled in the Registry.
@@ -476,7 +495,7 @@ class Registry:
     )
     @snowpark._internal.utils.private_preview(version=model_monitor_version.SNOWFLAKE_ML_MONITORING_MIN_VERSION)
     def delete_monitor(self, name: str) -> None:
-        """Delete a Model Monitor from the Registry
+        """Delete a Model Monitor by name from the Registry.
 
         Args:
             name: Name of the Model Monitor to delete.

snowflake/ml/utils/authentication.py

@@ -0,0 +1,75 @@
+import http
+import logging
+from datetime import timedelta
+from typing import Dict, Optional
+
+import requests
+from cryptography.hazmat.primitives.asymmetric import types
+from requests import auth
+
+from snowflake.ml._internal.utils import jwt_generator
+
+logger = logging.getLogger(__name__)
+_JWT_TOKEN_CACHE: Dict[str, Dict[int, str]] = {}
+
+
+def get_jwt_token_generator(
+    account: str,
+    user: str,
+    private_key: types.PRIVATE_KEY_TYPES,
+    lifetime: Optional[timedelta] = None,
+    renewal_delay: Optional[timedelta] = None,
+) -> jwt_generator.JWTGenerator:
+    return jwt_generator.JWTGenerator(account, user, private_key, lifetime=lifetime, renewal_delay=renewal_delay)
+
+
+def _get_snowflake_token_by_jwt(
+    jwt_token_generator: jwt_generator.JWTGenerator,
+    account: Optional[str] = None,
+    role: Optional[str] = None,
+    endpoint: Optional[str] = None,
+    snowflake_account_url: Optional[str] = None,
+) -> str:
+    scope_role = f"session:role:{role}" if role is not None else None
+    scope = " ".join(filter(None, [scope_role, endpoint]))
+    data = {
+        "grant_type": "urn:ietf:params:oauth:grant-type:jwt-bearer",
+        "scope": scope or None,
+        "assertion": jwt_token_generator.get_token(),
+    }
+    account = account or jwt_token_generator.account
+    url = f"https://{account}.snowflakecomputing.com/oauth/token"
+    if snowflake_account_url:
+        url = f"{snowflake_account_url}/oauth/token"
+
+    cache_key = hash(frozenset(data.items()))
+    if url in _JWT_TOKEN_CACHE:
+        if cache_key in _JWT_TOKEN_CACHE[url]:
+            return _JWT_TOKEN_CACHE[url][cache_key]
+    else:
+        _JWT_TOKEN_CACHE[url] = {}
+
+    response = requests.post(url, data=data)
+    if response.status_code != http.HTTPStatus.OK:
+        raise RuntimeError(f"Failed to get snowflake token: {response.status_code} {response.content!r}")
+    auth_token = response.text
+    _JWT_TOKEN_CACHE[url][cache_key] = auth_token
+    return auth_token
+
+
+class SnowflakeJWTTokenAuth(auth.AuthBase):
+    def __init__(
+        self,
+        jwt_token_generator: jwt_generator.JWTGenerator,
+        account: Optional[str] = None,
+        role: Optional[str] = None,
+        endpoint: Optional[str] = None,
+        snowflake_account_url: Optional[str] = None,
+    ) -> None:
+        self.snowflake_token = _get_snowflake_token_by_jwt(
+            jwt_token_generator, account, role, endpoint, snowflake_account_url
+        )
+
+    def __call__(self, r: requests.PreparedRequest) -> requests.PreparedRequest:
+        r.headers["Authorization"] = f'Snowflake Token="{self.snowflake_token}"'
+        return r

snowflake/ml/version.py

@@ -1 +1 @@
-VERSION="1.7.1"
+VERSION="1.7.3"