snowflake-ml-python 1.5.1__py3-none-any.whl → 1.5.3__py3-none-any.whl

snowflake/ml/feature_store/feature_view.py

@@ -7,10 +7,6 @@ from dataclasses import asdict, dataclass
 from enum import Enum
 from typing import Any, Dict, List, Optional
 
-from snowflake.ml._internal.exceptions import (
-    error_codes,
-    exceptions as snowml_exceptions,
-)
 from snowflake.ml._internal.utils.identifier import concat_names
 from snowflake.ml._internal.utils.sql_identifier import (
     SqlIdentifier,
@@ -34,6 +30,11 @@ _FEATURE_OBJ_TYPE = "FEATURE_OBJ_TYPE"
 _FEATURE_VIEW_VERSION_RE = re.compile(r"^[a-zA-Z0-9][a-zA-Z0-9_.\-]*$")
 _FEATURE_VIEW_VERSION_MAX_LENGTH = 128
 
+_RESULT_SCAN_QUERY_PATTERN = re.compile(
+    r".*FROM\s*TABLE\s*\(\s*RESULT_SCAN\s*\(.*",
+    flags=re.DOTALL | re.IGNORECASE | re.X,
+)
+
 
 @dataclass(frozen=True)
 class _FeatureViewMetadata:
@@ -54,13 +55,10 @@ class _FeatureViewMetadata:
 class FeatureViewVersion(str):
     def __new__(cls, version: str) -> FeatureViewVersion:
         if not _FEATURE_VIEW_VERSION_RE.match(version) or len(version) > _FEATURE_VIEW_VERSION_MAX_LENGTH:
-            raise snowml_exceptions.SnowflakeMLException(
-                error_code=error_codes.INVALID_ARGUMENT,
-                original_exception=ValueError(
-                    f"`{version}` is not a valid feature view version. "
-                    "It must start with letter or digit, and followed by letter, digit, '_', '-' or '.'. "
-                    f"The length limit is {_FEATURE_VIEW_VERSION_MAX_LENGTH}."
-                ),
+            raise ValueError(
+                f"`{version}` is not a valid feature view version. "
+                "It must start with letter or digit, and followed by letter, digit, '_', '-' or '.'. "
+                f"The length limit is {_FEATURE_VIEW_VERSION_MAX_LENGTH}."
             )
         return super().__new__(cls, version)
 
@@ -352,6 +350,9 @@ Got {len(self._feature_df.queries['queries'])}: {self._feature_df.queries['queri
             if not isinstance(col_type, (DateType, TimeType, TimestampType, _NumericType)):
                 raise ValueError(f"Invalid data type for timestamp_col {ts_col}: {col_type}.")
 
+        if re.match(_RESULT_SCAN_QUERY_PATTERN, self._query) is not None:
+            raise ValueError(f"feature_df should not be reading from RESULT_SCAN. Invalid query: {self._query}")
+
     def _get_feature_names(self) -> List[SqlIdentifier]:
         join_keys = [k for e in self._entities for k in e.join_keys]
         ts_col = [self._timestamp_col] if self._timestamp_col is not None else []

snowflake/ml/fileset/snowfs.py

@@ -1,10 +1,9 @@
 import collections
 import logging
 import re
-from typing import Any, Dict, Optional
+from typing import Any, Optional
 
 import fsspec
-import packaging.version as pkg_version
 
 from snowflake import snowpark
 from snowflake.connector import connection
@@ -12,7 +11,7 @@ from snowflake.ml._internal.exceptions import (
     error_codes,
     exceptions as snowml_exceptions,
 )
-from snowflake.ml._internal.utils import identifier, snowflake_env
+from snowflake.ml._internal.utils import identifier
 from snowflake.ml.fileset import embedded_stage_fs, sfcfs
 
 PROTOCOL_NAME = "snow"
@@ -28,10 +27,6 @@ _SNOWURL_PATTERN = re.compile(
     r"(?P<path>versions/(?:(?P<version>[^/]+)(?:/(?P<relpath>.*))?)?)"
 )
 
-# FIXME(dhung): Temporary fix for bug in GS version 8.17
-_BUG_VERSION_MIN = pkg_version.Version("8.17")  # Inclusive minimum version with bugged behavior
-_BUG_VERSION_MAX = pkg_version.Version("8.18")  # Exclusive maximum version with bugged behavior
-
 
 class SnowFileSystem(sfcfs.SFFileSystem):
     """A filesystem that allows user to access Snowflake embedded stage files with valid Snowflake locations.
@@ -54,21 +49,6 @@ class SnowFileSystem(sfcfs.SFFileSystem):
     ) -> None:
         super().__init__(sf_connection=sf_connection, snowpark_session=snowpark_session, **kwargs)
 
-        # FIXME(dhung): Temporary fix for bug in GS version 8.17
-        if SnowFileSystem._IS_BUGGED_VERSION is None:
-            try:
-                sf_version = snowflake_env.get_current_snowflake_version(self._session)
-                SnowFileSystem._IS_BUGGED_VERSION = _BUG_VERSION_MIN <= sf_version < _BUG_VERSION_MAX
-            except Exception:
-                SnowFileSystem._IS_BUGGED_VERSION = False
-
-    def info(self, path: str, **kwargs: Any) -> Dict[str, Any]:
-        # FIXME(dhung): Temporary fix for bug in GS version 8.17
-        res: Dict[str, Any] = super().info(path, **kwargs)
-        if res.get("type") == "directory" and not res["name"].endswith("/"):
-            res["name"] += "/"
-        return res
-
     def _get_stage_fs(
         self, sf_file_path: _SFFileEntityPath  # type: ignore[override]
     ) -> embedded_stage_fs.SFEmbeddedStageFileSystem:
@@ -100,12 +80,6 @@ class SnowFileSystem(sfcfs.SFFileSystem):
         if stage_name.startswith(protocol):
             stage_name = stage_name[len(protocol) :]
         abs_path = stage_name + "/" + path
-        # FIXME(dhung): Temporary fix for bug in GS version 8.17
-        if self._IS_BUGGED_VERSION:
-            match = _SNOWURL_PATTERN.fullmatch(abs_path)
-            assert match is not None
-            if match.group("relpath"):
-                abs_path = abs_path.replace(match.group("relpath"), match.group("relpath").lstrip("/"))
         return abs_path
 
     @classmethod
@@ -144,9 +118,6 @@ class SnowFileSystem(sfcfs.SFFileSystem):
             version = snowurl_match.group("version")
             relative_path = snowurl_match.group("relpath") or ""
             logging.debug(f"Parsed snow URL: {snowurl_match.groups()}")
-            # FIXME(dhung): Temporary fix for bug in GS version 8.17
-            if cls._IS_BUGGED_VERSION:
-                filepath = f"versions/{version}//{relative_path}"
             return _SFFileEntityPath(
                 domain=domain, name=name, version=version, relative_path=relative_path, filepath=filepath
             )

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

@@ -140,6 +140,49 @@ class ModelOperator:
                 statement_params=statement_params,
             )
 
+    def create_from_model_version(
+        self,
+        *,
+        source_database_name: Optional[sql_identifier.SqlIdentifier],
+        source_schema_name: Optional[sql_identifier.SqlIdentifier],
+        source_model_name: sql_identifier.SqlIdentifier,
+        source_version_name: sql_identifier.SqlIdentifier,
+        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,
+    ) -> None:
+        if self.validate_existence(
+            database_name=database_name,
+            schema_name=schema_name,
+            model_name=model_name,
+            statement_params=statement_params,
+        ):
+            return self._model_version_client.add_version_from_model_version(
+                source_database_name=source_database_name,
+                source_schema_name=source_schema_name,
+                source_model_name=source_model_name,
+                source_version_name=source_version_name,
+                database_name=database_name,
+                schema_name=schema_name,
+                model_name=model_name,
+                version_name=version_name,
+                statement_params=statement_params,
+            )
+        else:
+            return self._model_version_client.create_from_model_version(
+                source_database_name=source_database_name,
+                source_schema_name=source_schema_name,
+                source_model_name=source_model_name,
+                source_version_name=source_version_name,
+                database_name=database_name,
+                schema_name=schema_name,
+                model_name=model_name,
+                version_name=version_name,
+                statement_params=statement_params,
+            )
+
     def show_models_or_versions(
         self,
         *,

snowflake/ml/model/_client/sql/model_version.py

@@ -44,6 +44,32 @@ class ModelVersionSQLClient(_base._BaseSQLClient):
             statement_params=statement_params,
         ).has_dimensions(expected_rows=1, expected_cols=1).validate()
 
+    def create_from_model_version(
+        self,
+        *,
+        source_database_name: Optional[sql_identifier.SqlIdentifier],
+        source_schema_name: Optional[sql_identifier.SqlIdentifier],
+        source_model_name: sql_identifier.SqlIdentifier,
+        source_version_name: sql_identifier.SqlIdentifier,
+        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,
+    ) -> None:
+        fq_source_model_name = self.fully_qualified_object_name(
+            source_database_name, source_schema_name, source_model_name
+        )
+        fq_model_name = self.fully_qualified_object_name(database_name, schema_name, model_name)
+        query_result_checker.SqlResultValidator(
+            self._session,
+            (
+                f"CREATE MODEL {fq_model_name} WITH VERSION {version_name} FROM MODEL {fq_source_model_name}"
+                f" VERSION {source_version_name}"
+            ),
+            statement_params=statement_params,
+        ).has_dimensions(expected_rows=1, expected_cols=1).validate()
+
     # TODO(SNOW-987381): Merge with above when we have `create or alter module m [with] version v1 ...`
     def add_version_from_stage(
         self,
@@ -64,6 +90,32 @@ class ModelVersionSQLClient(_base._BaseSQLClient):
             statement_params=statement_params,
         ).has_dimensions(expected_rows=1, expected_cols=1).validate()
 
+    def add_version_from_model_version(
+        self,
+        *,
+        source_database_name: Optional[sql_identifier.SqlIdentifier],
+        source_schema_name: Optional[sql_identifier.SqlIdentifier],
+        source_model_name: sql_identifier.SqlIdentifier,
+        source_version_name: sql_identifier.SqlIdentifier,
+        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,
+    ) -> None:
+        fq_source_model_name = self.fully_qualified_object_name(
+            source_database_name, source_schema_name, source_model_name
+        )
+        fq_model_name = self.fully_qualified_object_name(database_name, schema_name, model_name)
+        query_result_checker.SqlResultValidator(
+            self._session,
+            (
+                f"ALTER MODEL {fq_model_name} ADD VERSION {version_name} FROM MODEL {fq_source_model_name}"
+                f" VERSION {source_version_name}"
+            ),
+            statement_params=statement_params,
+        ).has_dimensions(expected_rows=1, expected_cols=1).validate()
+
     def set_default_version(
         self,
         *,
@@ -145,7 +197,7 @@ class ModelVersionSQLClient(_base._BaseSQLClient):
         if snowpark_utils.is_in_stored_procedure():  # type: ignore[no-untyped-call]
             options = {"parallel": 10}
             cursor = self._session._conn._cursor
-            cursor._download(stage_location_url, str(target_path), options)  # type: ignore[attr-defined]
+            cursor._download(stage_location_url, str(target_path), options)  # type: ignore[union-attr]
             cursor.fetchall()
         else:
             query_result_checker.SqlResultValidator(
@@ -220,7 +272,7 @@ class ModelVersionSQLClient(_base._BaseSQLClient):
                 actual_schema_name.identifier(),
                 tmp_table_name,
             )
-            input_df.write.save_as_table(  # type: ignore[call-overload]
+            input_df.write.save_as_table(
                 table_name=INTERMEDIATE_TABLE_NAME,
                 mode="errorifexists",
                 table_type="temporary",
@@ -296,7 +348,7 @@ class ModelVersionSQLClient(_base._BaseSQLClient):
                 actual_schema_name.identifier(),
                 tmp_table_name,
             )
-            input_df.write.save_as_table(  # type: ignore[call-overload]
+            input_df.write.save_as_table(
                 table_name=INTERMEDIATE_TABLE_NAME,
                 mode="errorifexists",
                 table_type="temporary",

snowflake/ml/model/_model_composer/model_composer.py

@@ -136,7 +136,7 @@ class ModelComposer:
             model_meta=self.packager.meta,
             model_file_rel_path=pathlib.PurePosixPath(self.model_file_rel_path),
             options=options,
-            data_sources=self._get_data_sources(model),
+            data_sources=self._get_data_sources(model, sample_input_data),
         )
 
         file_utils.upload_directory_to_stage(
@@ -179,8 +179,12 @@ class ModelComposer:
         mp.load(meta_only=meta_only, options=options)
         return mp
 
-    def _get_data_sources(self, model: model_types.SupportedModelType) -> Optional[List[data_source.DataSource]]:
-        data_sources = getattr(model, lineage_utils.DATA_SOURCES_ATTR, None)
+    def _get_data_sources(
+        self, model: model_types.SupportedModelType, sample_input_data: Optional[model_types.SupportedDataType] = None
+    ) -> Optional[List[data_source.DataSource]]:
+        data_sources = lineage_utils.get_data_sources(model)
+        if not data_sources and sample_input_data is not None:
+            data_sources = lineage_utils.get_data_sources(sample_input_data)
         if isinstance(data_sources, list) and all(isinstance(item, data_source.DataSource) for item in data_sources):
             return data_sources
         return None

snowflake/ml/model/_model_composer/model_method/infer_table_function.py_template

@@ -74,4 +74,6 @@ dtype_map = {{feature.name: feature.as_dtype() for feature in features}}
 class {function_name}:
     @vectorized(input=pd.DataFrame)
     def end_partition(self, df: pd.DataFrame) -> pd.DataFrame:
-        return runner(df)
+        df.columns = input_cols
+        input_df = df.astype(dtype=dtype_map)
+        return runner(input_df[input_cols])

snowflake/ml/model/_packager/model_meta/_core_requirements.py

@@ -6,6 +6,6 @@ REQUIREMENTS = [
     "packaging>=20.9,<24",
     "pandas>=1.0.0,<3",
     "pyyaml>=6.0,<7",
-    "snowflake-snowpark-python>=1.11.1,<2,!=1.12.0",
+    "snowflake-snowpark-python>=1.15.0,<2",
     "typing-extensions>=4.1.0,<5"
 ]

snowflake/ml/model/_packager/model_meta/model_meta.py

@@ -281,9 +281,7 @@ class ModelMetadata:
             "cpu": model_runtime.ModelRuntime("cpu", self.env),
         }
         if self.env.cuda_version:
-            runtimes.update(
-                {"gpu": model_runtime.ModelRuntime("gpu", self.env, is_gpu=True, server_availability_source="conda")}
-            )
+            runtimes.update({"gpu": model_runtime.ModelRuntime("gpu", self.env, is_gpu=True)})
         return runtimes
 
     def save(self, model_dir_path: str) -> None:

snowflake/ml/model/_packager/model_runtime/_snowml_inference_alternative_requirements.py

@@ -5,6 +5,6 @@ REQUIREMENTS = [
     "packaging>=20.9,<24",
     "pandas>=1.0.0,<3",
     "pyyaml>=6.0,<7",
-    "snowflake-snowpark-python>=1.11.1,<2,!=1.12.0",
+    "snowflake-snowpark-python>=1.15.0,<2",
     "typing-extensions>=4.1.0,<5"
 ]

snowflake/ml/model/_packager/model_runtime/model_runtime.py

@@ -1,11 +1,11 @@
 import copy
 import pathlib
 import warnings
-from typing import List, Literal, Optional
+from typing import List, Optional
 
 from packaging import requirements
 
-from snowflake.ml._internal import env as snowml_env, env_utils, file_utils
+from snowflake.ml._internal import env_utils, file_utils
 from snowflake.ml.model._packager.model_env import model_env
 from snowflake.ml.model._packager.model_meta import model_meta_schema
 from snowflake.ml.model._packager.model_runtime import (
@@ -37,7 +37,6 @@ class ModelRuntime:
         env: model_env.ModelEnv,
         imports: Optional[List[pathlib.PurePosixPath]] = None,
         is_gpu: bool = False,
-        server_availability_source: Literal["snowflake", "conda"] = "snowflake",
         loading_from_file: bool = False,
     ) -> None:
         self.name = name
@@ -48,30 +47,7 @@ class ModelRuntime:
             return
 
         snowml_pkg_spec = f"{env_utils.SNOWPARK_ML_PKG_NAME}=={self.runtime_env.snowpark_ml_version}"
-        if self.runtime_env._snowpark_ml_version.local:
-            self.embed_local_ml_library = True
-        else:
-            if server_availability_source == "snowflake":
-                snowml_server_availability = (
-                    len(
-                        env_utils.get_matched_package_versions_in_information_schema_with_active_session(
-                            reqs=[requirements.Requirement(snowml_pkg_spec)],
-                            python_version=snowml_env.PYTHON_VERSION,
-                        ).get(env_utils.SNOWPARK_ML_PKG_NAME, [])
-                    )
-                    >= 1
-                )
-            else:
-                snowml_server_availability = (
-                    len(
-                        env_utils.get_matched_package_versions_in_snowflake_conda_channel(
-                            req=requirements.Requirement(snowml_pkg_spec),
-                            python_version=snowml_env.PYTHON_VERSION,
-                        )
-                    )
-                    >= 1
-                )
-            self.embed_local_ml_library = not snowml_server_availability
+        self.embed_local_ml_library = self.runtime_env._snowpark_ml_version.local
 
         additional_package = (
             _SNOWML_INFERENCE_ALTERNATIVE_DEPENDENCIES if self.embed_local_ml_library else [snowml_pkg_spec]

snowflake/ml/model/_signatures/builtins_handler.py

@@ -1,3 +1,4 @@
+import datetime
 from collections import abc
 from typing import Literal, Sequence
 
@@ -24,7 +25,7 @@ class ListOfBuiltinHandler(base_handler.BaseDataHandler[model_types._SupportedBu
             # String is a Sequence but we take them as an whole
             if isinstance(element, abc.Sequence) and not isinstance(element, str):
                 can_handle = ListOfBuiltinHandler.can_handle(element)
-            elif not isinstance(element, (int, float, bool, str)):
+            elif not isinstance(element, (int, float, bool, str, datetime.datetime)):
                 can_handle = False
                 break
         return can_handle

snowflake/ml/model/_signatures/core.py

@@ -53,6 +53,8 @@ class DataType(Enum):
     STRING = ("string", spt.StringType, np.str_)
     BYTES = ("bytes", spt.BinaryType, np.bytes_)
 
+    TIMESTAMP_NTZ = ("datetime64[ns]", spt.TimestampType, "datetime64[ns]")
+
     def as_snowpark_type(self) -> spt.DataType:
         """Convert to corresponding Snowpark Type.
 
@@ -78,6 +80,13 @@ class DataType(Enum):
             Corresponding DataType.
         """
         np_to_snowml_type_mapping = {i._numpy_type: i for i in DataType}
+
+        # Add datetime types:
+        datetime_res = ["Y", "M", "W", "D", "h", "m", "s", "ms", "us", "ns"]
+
+        for res in datetime_res:
+            np_to_snowml_type_mapping[f"datetime64[{res}]"] = DataType.TIMESTAMP_NTZ
+
         for potential_type in np_to_snowml_type_mapping.keys():
             if np.can_cast(np_type, potential_type, casting="no"):
                 # This is used since the same dtype might represented in different ways.
@@ -247,9 +256,12 @@ class FeatureSpec(BaseFeatureSpec):
             result_type = spt.ArrayType(result_type)
         return result_type
 
-    def as_dtype(self) -> npt.DTypeLike:
+    def as_dtype(self) -> Union[npt.DTypeLike, str]:
         """Convert to corresponding local Type."""
         if not self._shape:
+            # scalar dtype: use keys from `np.sctypeDict` to prevent unit-less dtype 'datetime64'
+            if "datetime64" in self._dtype._value:
+                return self._dtype._value
             return self._dtype._numpy_type
         return np.object_
 

snowflake/ml/model/_signatures/pandas_handler.py

@@ -147,6 +147,8 @@ class PandasDataFrameHandler(base_handler.BaseDataHandler[pd.DataFrame]):
                     specs.append(core.FeatureSpec(dtype=core.DataType.STRING, name=ft_name))
                 elif isinstance(data[df_col].iloc[0], bytes):
                     specs.append(core.FeatureSpec(dtype=core.DataType.BYTES, name=ft_name))
+            elif isinstance(data[df_col].iloc[0], np.datetime64):
+                specs.append(core.FeatureSpec(dtype=core.DataType.TIMESTAMP_NTZ, name=ft_name))
             else:
                 specs.append(core.FeatureSpec(dtype=core.DataType.from_numpy_type(df_col_dtype), name=ft_name))
         return specs

snowflake/ml/model/_signatures/snowpark_handler.py

@@ -107,6 +107,9 @@ class SnowparkDataFrameHandler(base_handler.BaseDataHandler[snowflake.snowpark.D
         if not features:
             features = pandas_handler.PandasDataFrameHandler.infer_signature(df, role="input")
         # Role will be no effect on the column index. That is to say, the feature name is the actual column name.
+        if keep_order:
+            df = df.reset_index(drop=True)
+            df[infer_template._KEEP_ORDER_COL_NAME] = df.index
         sp_df = session.create_dataframe(df)
         column_names = []
         columns = []
@@ -122,7 +125,4 @@ class SnowparkDataFrameHandler(base_handler.BaseDataHandler[snowflake.snowpark.D
 
         sp_df = sp_df.with_columns(column_names, columns)
 
-        if keep_order:
-            sp_df = sp_df.with_column(infer_template._KEEP_ORDER_COL_NAME, F.monotonically_increasing_id())
-
         return sp_df

snowflake/ml/model/model_signature.py

@@ -168,6 +168,8 @@ def _validate_numpy_array(
             max_v <= np.finfo(feature_type._numpy_type).max  # type: ignore[arg-type]
             and min_v >= np.finfo(feature_type._numpy_type).min  # type: ignore[arg-type]
         )
+    elif feature_type in [core.DataType.TIMESTAMP_NTZ]:
+        return np.issubdtype(arr.dtype, np.datetime64)
     else:
         return np.can_cast(arr.dtype, feature_type._numpy_type, casting="no")
 

snowflake/ml/model/type_hints.py

@@ -54,6 +54,7 @@ _SupportedNumpyDtype = Union[
     "np.bool_",
     "np.str_",
     "np.bytes_",
+    "np.datetime64",
 ]
 _SupportedNumpyArray = npt.NDArray[_SupportedNumpyDtype]
 _SupportedBuiltinsList = Sequence[_SupportedBuiltins]

snowflake/ml/modeling/_internal/estimator_utils.py

@@ -1,15 +1,19 @@
 import inspect
 import numbers
+import os
 from typing import Any, Callable, Dict, List, Set, Tuple
 
+import cloudpickle as cp
 import numpy as np
 from numpy import typing as npt
-from typing_extensions import TypeGuard
 
 from snowflake.ml._internal.exceptions import error_codes, exceptions
+from snowflake.ml._internal.utils import temp_file_utils
+from snowflake.ml._internal.utils.query_result_checker import SqlResultValidator
 from snowflake.ml.modeling.framework._utils import to_native_format
 from snowflake.ml.modeling.framework.base import BaseTransformer
 from snowflake.snowpark import Session
+from snowflake.snowpark._internal import utils as snowpark_utils
 
 
 def validate_sklearn_args(args: Dict[str, Tuple[Any, Any, bool]], klass: type) -> Dict[str, Any]:
@@ -97,6 +101,7 @@ def original_estimator_has_callable(attr: str) -> Callable[[Any], bool]:
     Returns:
         A function which checks for the existence of callable `attr` on the given object.
     """
+    from typing_extensions import TypeGuard
 
     def check(self: BaseTransformer) -> TypeGuard[Callable[..., object]]:
         """Check for the existence of callable `attr` in self.
@@ -218,3 +223,55 @@ def handle_inference_result(
                     )
 
     return transformed_numpy_array, output_cols
+
+
+def create_temp_stage(session: Session) -> str:
+    """Creates temporary stage.
+
+    Args:
+        session: Session
+
+    Returns:
+        Temp stage name.
+    """
+    # Create temp stage to upload pickled model file.
+    transform_stage_name = snowpark_utils.random_name_for_temp_object(snowpark_utils.TempObjectType.STAGE)
+    stage_creation_query = f"CREATE OR REPLACE TEMPORARY STAGE {transform_stage_name};"
+    SqlResultValidator(session=session, query=stage_creation_query).has_dimensions(
+        expected_rows=1, expected_cols=1
+    ).validate()
+    return transform_stage_name
+
+
+def upload_model_to_stage(
+    stage_name: str, estimator: object, session: Session, statement_params: Dict[str, str]
+) -> str:
+    """Util method to pickle and upload the model to a temp Snowflake stage.
+
+
+    Args:
+        stage_name: Stage name to save model.
+        estimator: Estimator object to upload to stage (sklearn model object)
+        session: The snowpark session to use.
+        statement_params: Statement parameters for query telemetry.
+
+    Returns:
+        a tuple containing stage file paths for pickled input model for training and location to store trained
+        models(response from training sproc).
+    """
+    # Create a temp file and dump the transform to that file.
+    local_transform_file_name = temp_file_utils.get_temp_file_path()
+    with open(local_transform_file_name, mode="w+b") as local_transform_file:
+        cp.dump(estimator, local_transform_file)
+
+    # Put locally serialized transform on stage.
+    session.file.put(
+        local_file_name=local_transform_file_name,
+        stage_location=stage_name,
+        auto_compress=False,
+        overwrite=True,
+        statement_params=statement_params,
+    )
+
+    temp_file_utils.cleanup_temp_files([local_transform_file_name])
+    return os.path.basename(local_transform_file_name)