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

snowflake/ml/modeling/_internal/snowpark_implementations/snowpark_handlers.py

@@ -2,6 +2,7 @@ import importlib
 import inspect
 import os
 import posixpath
+import sys
 from typing import Any, Dict, List, Optional
 from uuid import uuid4
 
@@ -13,12 +14,10 @@ from snowflake.ml._internal.utils import (
     identifier,
     pkg_version_utils,
     snowpark_dataframe_utils,
+    temp_file_utils,
 )
 from snowflake.ml._internal.utils.query_result_checker import SqlResultValidator
-from snowflake.ml._internal.utils.temp_file_utils import (
-    cleanup_temp_files,
-    get_temp_file_path,
-)
+from snowflake.ml.modeling._internal import estimator_utils
 from snowflake.ml.modeling._internal.estimator_utils import handle_inference_result
 from snowflake.snowpark import DataFrame, Session, functions as F, types as T
 from snowflake.snowpark._internal.utils import (
@@ -26,7 +25,7 @@ from snowflake.snowpark._internal.utils import (
     random_name_for_temp_object,
 )
 
-cp.register_pickle_by_value(inspect.getmodule(get_temp_file_path))
+cp.register_pickle_by_value(inspect.getmodule(temp_file_utils.get_temp_file_path))
 cp.register_pickle_by_value(inspect.getmodule(identifier.get_inferred_name))
 cp.register_pickle_by_value(inspect.getmodule(handle_inference_result))
 
@@ -97,7 +96,25 @@ class SnowparkTransformHandlers:
 
         dependencies = self._get_validated_snowpark_dependencies(session, dependencies)
         dataset = self.dataset
-        estimator = self.estimator
+
+        statement_params = telemetry.get_function_usage_statement_params(
+            project=_PROJECT,
+            subproject=self._subproject,
+            function_name=telemetry.get_statement_params_full_func_name(inspect.currentframe(), self._class_name),
+            api_calls=[F.pandas_udf],
+            custom_tags={"autogen": True} if self._autogenerated else None,
+        )
+
+        temp_stage_name = estimator_utils.create_temp_stage(session)
+
+        estimator_file_name = estimator_utils.upload_model_to_stage(
+            stage_name=temp_stage_name,
+            estimator=self.estimator,
+            session=session,
+            statement_params=statement_params,
+        )
+        imports = [f"@{temp_stage_name}/{estimator_file_name}"]
+
         # Register vectorized UDF for batch inference
         batch_inference_udf_name = random_name_for_temp_object(TempObjectType.FUNCTION)
 
@@ -113,13 +130,13 @@ class SnowparkTransformHandlers:
         for field in fields:
             input_datatypes.append(field.datatype)
 
-        statement_params = telemetry.get_function_usage_statement_params(
-            project=_PROJECT,
-            subproject=self._subproject,
-            function_name=telemetry.get_statement_params_full_func_name(inspect.currentframe(), self._class_name),
-            api_calls=[F.pandas_udf],
-            custom_tags=dict([("autogen", True)]) if self._autogenerated else None,
-        )
+        # TODO(xjiang): for optimization, use register_from_file to reduce duplicate loading estimator object
+        # or use cachetools here
+        def load_estimator() -> object:
+            estimator_file_path = os.path.join(sys._xoptions["snowflake_import_directory"], f"{estimator_file_name}")
+            with open(estimator_file_path, mode="rb") as local_estimator_file_obj:
+                estimator_object = cp.load(local_estimator_file_obj)
+            return estimator_object
 
         @F.pandas_udf(  # type: ignore[arg-type, misc]
             is_permanent=False,
@@ -129,6 +146,7 @@ class SnowparkTransformHandlers:
             session=session,
             statement_params=statement_params,
             input_types=[T.PandasDataFrameType(input_datatypes)],
+            imports=imports,  # type: ignore[arg-type]
         )
         def vec_batch_infer(input_df: pd.DataFrame) -> T.PandasSeries[dict]:  # type: ignore[type-arg]
             import numpy as np  # noqa: F401
@@ -136,6 +154,8 @@ class SnowparkTransformHandlers:
 
             input_df.columns = snowpark_cols
 
+            estimator = load_estimator()
+
             if hasattr(estimator, "n_jobs"):
                 # Vectorized UDF cannot handle joblib multiprocessing right now, deactivate the n_jobs
                 estimator.n_jobs = 1
@@ -225,7 +245,7 @@ class SnowparkTransformHandlers:
         queries = dataset.queries["queries"]
 
         # Create a temp file and dump the score to that file.
-        local_score_file_name = get_temp_file_path()
+        local_score_file_name = temp_file_utils.get_temp_file_path()
         with open(local_score_file_name, mode="w+b") as local_score_file:
             cp.dump(estimator, local_score_file)
 
@@ -247,7 +267,7 @@ class SnowparkTransformHandlers:
                 inspect.currentframe(), self.__class__.__name__
             ),
             api_calls=[F.sproc],
-            custom_tags=dict([("autogen", True)]) if self._autogenerated else None,
+            custom_tags={"autogen": True} if self._autogenerated else None,
         )
         # Put locally serialized score on stage.
         session.file.put(
@@ -290,7 +310,7 @@ class SnowparkTransformHandlers:
             df: pd.DataFrame = sp_df.to_pandas(statement_params=score_statement_params)
             df.columns = sp_df.columns
 
-            local_score_file_name = get_temp_file_path()
+            local_score_file_name = temp_file_utils.get_temp_file_path()
             session.file.get(stage_score_file_name, local_score_file_name, statement_params=score_statement_params)
 
             local_score_file_name_path = os.path.join(local_score_file_name, os.listdir(local_score_file_name)[0])
@@ -323,7 +343,7 @@ class SnowparkTransformHandlers:
                 inspect.currentframe(), self.__class__.__name__
             ),
             api_calls=[Session.call],
-            custom_tags=dict([("autogen", True)]) if self._autogenerated else None,
+            custom_tags={"autogen": True} if self._autogenerated else None,
         )
 
         kwargs = telemetry.get_sproc_statement_params_kwargs(score_wrapper_sproc, score_statement_params)
@@ -338,7 +358,7 @@ class SnowparkTransformHandlers:
             **kwargs,
         )
 
-        cleanup_temp_files([local_score_file_name])
+        temp_file_utils.cleanup_temp_files([local_score_file_name])
 
         return score
 

snowflake/ml/modeling/_internal/snowpark_implementations/snowpark_trainer.py

@@ -17,30 +17,19 @@ from snowflake.ml._internal.utils import (
     identifier,
     pkg_version_utils,
     snowpark_dataframe_utils,
+    temp_file_utils,
 )
-from snowflake.ml._internal.utils.query_result_checker import SqlResultValidator
-from snowflake.ml._internal.utils.temp_file_utils import (
-    cleanup_temp_files,
-    get_temp_file_path,
-)
+from snowflake.ml.modeling._internal import estimator_utils
 from snowflake.ml.modeling._internal.estimator_utils import handle_inference_result
 from snowflake.ml.modeling._internal.model_specifications import (
     ModelSpecifications,
     ModelSpecificationsBuilder,
 )
-from snowflake.snowpark import (
-    DataFrame,
-    Session,
-    exceptions as snowpark_exceptions,
-    functions as F,
-)
-from snowflake.snowpark._internal.utils import (
-    TempObjectType,
-    random_name_for_temp_object,
-)
+from snowflake.snowpark import DataFrame, Session, exceptions as snowpark_exceptions
+from snowflake.snowpark._internal import utils as snowpark_utils
 from snowflake.snowpark.stored_procedure import StoredProcedure
 
-cp.register_pickle_by_value(inspect.getmodule(get_temp_file_path))
+cp.register_pickle_by_value(inspect.getmodule(temp_file_utils.get_temp_file_path))
 cp.register_pickle_by_value(inspect.getmodule(identifier.get_inferred_name))
 cp.register_pickle_by_value(inspect.getmodule(handle_inference_result))
 
@@ -90,60 +79,6 @@ class SnowparkModelTrainer:
         self._subproject = subproject
         self._class_name = estimator.__class__.__name__
 
-    def _create_temp_stage(self) -> str:
-        """
-        Creates temporary stage.
-
-        Returns:
-            Temp stage name.
-        """
-        # Create temp stage to upload pickled model file.
-        transform_stage_name = random_name_for_temp_object(TempObjectType.STAGE)
-        stage_creation_query = f"CREATE OR REPLACE TEMPORARY STAGE {transform_stage_name};"
-        SqlResultValidator(session=self.session, query=stage_creation_query).has_dimensions(
-            expected_rows=1, expected_cols=1
-        ).validate()
-        return transform_stage_name
-
-    def _upload_model_to_stage(self, stage_name: str) -> Tuple[str, str]:
-        """
-        Util method to pickle and upload the model to a temp Snowflake stage.
-
-        Args:
-            stage_name: Stage name to save model.
-
-        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 = get_temp_file_path()
-        with open(local_transform_file_name, mode="w+b") as local_transform_file:
-            cp.dump(self.estimator, local_transform_file)
-
-        # Use posixpath to construct stage paths
-        stage_transform_file_name = posixpath.join(stage_name, os.path.basename(local_transform_file_name))
-        stage_result_file_name = posixpath.join(stage_name, os.path.basename(local_transform_file_name))
-
-        statement_params = telemetry.get_function_usage_statement_params(
-            project=_PROJECT,
-            subproject=self._subproject,
-            function_name=telemetry.get_statement_params_full_func_name(inspect.currentframe(), self._class_name),
-            api_calls=[F.sproc],
-            custom_tags=dict([("autogen", True)]) if self._autogenerated else None,
-        )
-        # Put locally serialized transform on stage.
-        self.session.file.put(
-            local_transform_file_name,
-            stage_transform_file_name,
-            auto_compress=False,
-            overwrite=True,
-            statement_params=statement_params,
-        )
-
-        cleanup_temp_files([local_transform_file_name])
-        return (stage_transform_file_name, stage_result_file_name)
-
     def _fetch_model_from_stage(self, dir_path: str, file_name: str, statement_params: Dict[str, str]) -> object:
         """
         Downloads the serialized model from a stage location and unpickles it.
@@ -156,7 +91,7 @@ class SnowparkModelTrainer:
         Returns:
             Deserialized model object.
         """
-        local_result_file_name = get_temp_file_path()
+        local_result_file_name = temp_file_utils.get_temp_file_path()
         self.session.file.get(
             posixpath.join(dir_path, file_name),
             local_result_file_name,
@@ -166,13 +101,13 @@ class SnowparkModelTrainer:
         with open(os.path.join(local_result_file_name, file_name), mode="r+b") as result_file_obj:
             fit_estimator = cp.load(result_file_obj)
 
-        cleanup_temp_files([local_result_file_name])
+        temp_file_utils.cleanup_temp_files([local_result_file_name])
         return fit_estimator
 
     def _build_fit_wrapper_sproc(
         self,
         model_spec: ModelSpecifications,
-    ) -> Callable[[Any, List[str], str, str, List[str], List[str], Optional[str], Dict[str, str]], str]:
+    ) -> Callable[[Any, List[str], str, List[str], List[str], Optional[str], Dict[str, str]], str]:
         """
         Constructs and returns a python stored procedure function to be used for training model.
 
@@ -188,8 +123,7 @@ class SnowparkModelTrainer:
         def fit_wrapper_function(
             session: Session,
             sql_queries: List[str],
-            stage_transform_file_name: str,
-            stage_result_file_name: str,
+            temp_stage_name: str,
             input_cols: List[str],
             label_cols: List[str],
             sample_weight_col: Optional[str],
@@ -212,9 +146,13 @@ class SnowparkModelTrainer:
             df: pd.DataFrame = sp_df.to_pandas(statement_params=statement_params)
             df.columns = sp_df.columns
 
-            local_transform_file_name = get_temp_file_path()
+            local_transform_file_name = temp_file_utils.get_temp_file_path()
 
-            session.file.get(stage_transform_file_name, local_transform_file_name, statement_params=statement_params)
+            session.file.get(
+                stage_location=temp_stage_name,
+                target_directory=local_transform_file_name,
+                statement_params=statement_params,
+            )
 
             local_transform_file_path = os.path.join(
                 local_transform_file_name, os.listdir(local_transform_file_name)[0]
@@ -233,14 +171,14 @@ class SnowparkModelTrainer:
 
             estimator.fit(**args)
 
-            local_result_file_name = get_temp_file_path()
+            local_result_file_name = temp_file_utils.get_temp_file_path()
 
             with open(local_result_file_name, mode="w+b") as local_result_file_obj:
                 cp.dump(estimator, local_result_file_obj)
 
             session.file.put(
-                local_result_file_name,
-                stage_result_file_name,
+                local_file_name=local_result_file_name,
+                stage_location=temp_stage_name,
                 auto_compress=False,
                 overwrite=True,
                 statement_params=statement_params,
@@ -254,7 +192,7 @@ class SnowparkModelTrainer:
 
     def _get_fit_wrapper_sproc_anonymous(self, statement_params: Dict[str, str]) -> StoredProcedure:
         model_spec = ModelSpecificationsBuilder.build(model=self.estimator)
-        fit_sproc_name = random_name_for_temp_object(TempObjectType.PROCEDURE)
+        fit_sproc_name = snowpark_utils.random_name_for_temp_object(snowpark_utils.TempObjectType.PROCEDURE)
 
         relaxed_dependencies = pkg_version_utils.get_valid_pkg_versions_supported_in_snowflake_conda_channel(
             pkg_versions=model_spec.pkgDependencies, session=self.session
@@ -284,7 +222,7 @@ class SnowparkModelTrainer:
             fit_sproc: StoredProcedure = self.session._FIT_WRAPPER_SPROCS[fit_sproc_key]  # type: ignore[attr-defined]
             return fit_sproc
 
-        fit_sproc_name = random_name_for_temp_object(TempObjectType.PROCEDURE)
+        fit_sproc_name = snowpark_utils.random_name_for_temp_object(snowpark_utils.TempObjectType.PROCEDURE)
 
         relaxed_dependencies = pkg_version_utils.get_valid_pkg_versions_supported_in_snowflake_conda_channel(
             pkg_versions=model_spec.pkgDependencies, session=self.session
@@ -307,7 +245,7 @@ class SnowparkModelTrainer:
     def _build_fit_predict_wrapper_sproc(
         self,
         model_spec: ModelSpecifications,
-    ) -> Callable[[Session, List[str], str, str, List[str], Dict[str, str], bool, List[str], str], str]:
+    ) -> Callable[[Session, List[str], str, List[str], Dict[str, str], bool, List[str], str], str]:
         """
         Constructs and returns a python stored procedure function to be used for training model.
 
@@ -323,8 +261,7 @@ class SnowparkModelTrainer:
         def fit_predict_wrapper_function(
             session: Session,
             sql_queries: List[str],
-            stage_transform_file_name: str,
-            stage_result_file_name: str,
+            temp_stage_name: str,
             input_cols: List[str],
             statement_params: Dict[str, str],
             drop_input_cols: bool,
@@ -347,9 +284,13 @@ class SnowparkModelTrainer:
             df: pd.DataFrame = sp_df.to_pandas(statement_params=statement_params)
             df.columns = sp_df.columns
 
-            local_transform_file_name = get_temp_file_path()
+            local_transform_file_name = temp_file_utils.get_temp_file_path()
 
-            session.file.get(stage_transform_file_name, local_transform_file_name, statement_params=statement_params)
+            session.file.get(
+                stage_location=temp_stage_name,
+                target_directory=local_transform_file_name,
+                statement_params=statement_params,
+            )
 
             local_transform_file_path = os.path.join(
                 local_transform_file_name, os.listdir(local_transform_file_name)[0]
@@ -359,14 +300,14 @@ class SnowparkModelTrainer:
 
             fit_predict_result = estimator.fit_predict(X=df[input_cols])
 
-            local_result_file_name = get_temp_file_path()
+            local_result_file_name = temp_file_utils.get_temp_file_path()
 
             with open(local_result_file_name, mode="w+b") as local_result_file_obj:
                 cp.dump(estimator, local_result_file_obj)
 
             session.file.put(
-                local_result_file_name,
-                stage_result_file_name,
+                local_file_name=local_result_file_name,
+                stage_location=temp_stage_name,
                 auto_compress=False,
                 overwrite=True,
                 statement_params=statement_params,
@@ -407,7 +348,6 @@ class SnowparkModelTrainer:
             Session,
             List[str],
             str,
-            str,
             List[str],
             Optional[List[str]],
             Optional[str],
@@ -433,8 +373,7 @@ class SnowparkModelTrainer:
         def fit_transform_wrapper_function(
             session: Session,
             sql_queries: List[str],
-            stage_transform_file_name: str,
-            stage_result_file_name: str,
+            temp_stage_name: str,
             input_cols: List[str],
             label_cols: Optional[List[str]],
             sample_weight_col: Optional[str],
@@ -459,9 +398,13 @@ class SnowparkModelTrainer:
             df: pd.DataFrame = sp_df.to_pandas(statement_params=statement_params)
             df.columns = sp_df.columns
 
-            local_transform_file_name = get_temp_file_path()
+            local_transform_file_name = temp_file_utils.get_temp_file_path()
 
-            session.file.get(stage_transform_file_name, local_transform_file_name, statement_params=statement_params)
+            session.file.get(
+                stage_location=temp_stage_name,
+                target_directory=local_transform_file_name,
+                statement_params=statement_params,
+            )
 
             local_transform_file_path = os.path.join(
                 local_transform_file_name, os.listdir(local_transform_file_name)[0]
@@ -480,14 +423,14 @@ class SnowparkModelTrainer:
 
             fit_transform_result = estimator.fit_transform(**args)
 
-            local_result_file_name = get_temp_file_path()
+            local_result_file_name = temp_file_utils.get_temp_file_path()
 
             with open(local_result_file_name, mode="w+b") as local_result_file_obj:
                 cp.dump(estimator, local_result_file_obj)
 
             session.file.put(
-                local_result_file_name,
-                stage_result_file_name,
+                local_file_name=local_result_file_name,
+                stage_location=temp_stage_name,
                 auto_compress=False,
                 overwrite=True,
                 statement_params=statement_params,
@@ -535,7 +478,7 @@ class SnowparkModelTrainer:
     def _get_fit_predict_wrapper_sproc_anonymous(self, statement_params: Dict[str, str]) -> StoredProcedure:
         model_spec = ModelSpecificationsBuilder.build(model=self.estimator)
 
-        fit_predict_sproc_name = random_name_for_temp_object(TempObjectType.PROCEDURE)
+        fit_predict_sproc_name = snowpark_utils.random_name_for_temp_object(snowpark_utils.TempObjectType.PROCEDURE)
 
         relaxed_dependencies = pkg_version_utils.get_valid_pkg_versions_supported_in_snowflake_conda_channel(
             pkg_versions=model_spec.pkgDependencies, session=self.session
@@ -567,7 +510,7 @@ class SnowparkModelTrainer:
             ]
             return fit_sproc
 
-        fit_predict_sproc_name = random_name_for_temp_object(TempObjectType.PROCEDURE)
+        fit_predict_sproc_name = snowpark_utils.random_name_for_temp_object(snowpark_utils.TempObjectType.PROCEDURE)
 
         relaxed_dependencies = pkg_version_utils.get_valid_pkg_versions_supported_in_snowflake_conda_channel(
             pkg_versions=model_spec.pkgDependencies, session=self.session
@@ -592,7 +535,7 @@ class SnowparkModelTrainer:
     def _get_fit_transform_wrapper_sproc_anonymous(self, statement_params: Dict[str, str]) -> StoredProcedure:
         model_spec = ModelSpecificationsBuilder.build(model=self.estimator)
 
-        fit_transform_sproc_name = random_name_for_temp_object(TempObjectType.PROCEDURE)
+        fit_transform_sproc_name = snowpark_utils.random_name_for_temp_object(snowpark_utils.TempObjectType.PROCEDURE)
 
         relaxed_dependencies = pkg_version_utils.get_valid_pkg_versions_supported_in_snowflake_conda_channel(
             pkg_versions=model_spec.pkgDependencies, session=self.session
@@ -623,7 +566,7 @@ class SnowparkModelTrainer:
             ]
             return fit_sproc
 
-        fit_transform_sproc_name = random_name_for_temp_object(TempObjectType.PROCEDURE)
+        fit_transform_sproc_name = snowpark_utils.random_name_for_temp_object(snowpark_utils.TempObjectType.PROCEDURE)
 
         relaxed_dependencies = pkg_version_utils.get_valid_pkg_versions_supported_in_snowflake_conda_channel(
             pkg_versions=model_spec.pkgDependencies, session=self.session
@@ -663,19 +606,21 @@ class SnowparkModelTrainer:
         # Extract query that generated the dataframe. We will need to pass it to the fit procedure.
         queries = dataset.queries["queries"]
 
-        transform_stage_name = self._create_temp_stage()
-        (stage_transform_file_name, stage_result_file_name) = self._upload_model_to_stage(
-            stage_name=transform_stage_name
-        )
-
-        # Call fit sproc
+        temp_stage_name = estimator_utils.create_temp_stage(self.session)
         statement_params = telemetry.get_function_usage_statement_params(
             project=_PROJECT,
             subproject=self._subproject,
             function_name=telemetry.get_statement_params_full_func_name(inspect.currentframe(), self._class_name),
             api_calls=[Session.call],
-            custom_tags=dict([("autogen", True)]) if self._autogenerated else None,
+            custom_tags={"autogen": True} if self._autogenerated else None,
         )
+        estimator_utils.upload_model_to_stage(
+            stage_name=temp_stage_name,
+            estimator=self.estimator,
+            session=self.session,
+            statement_params=statement_params,
+        )
+        # Call fit sproc
 
         if _ENABLE_ANONYMOUS_SPROC:
             fit_wrapper_sproc = self._get_fit_wrapper_sproc_anonymous(statement_params=statement_params)
@@ -686,8 +631,7 @@ class SnowparkModelTrainer:
             sproc_export_file_name: str = fit_wrapper_sproc(
                 self.session,
                 queries,
-                stage_transform_file_name,
-                stage_result_file_name,
+                temp_stage_name,
                 self.input_cols,
                 self.label_cols,
                 self.sample_weight_col,
@@ -706,7 +650,7 @@ class SnowparkModelTrainer:
             sproc_export_file_name = fields[0]
 
         return self._fetch_model_from_stage(
-            dir_path=stage_result_file_name,
+            dir_path=temp_stage_name,
             file_name=sproc_export_file_name,
             statement_params=statement_params,
         )
@@ -734,32 +678,34 @@ class SnowparkModelTrainer:
         # Extract query that generated the dataframe. We will need to pass it to the fit procedure.
         queries = dataset.queries["queries"]
 
-        transform_stage_name = self._create_temp_stage()
-        (stage_transform_file_name, stage_result_file_name) = self._upload_model_to_stage(
-            stage_name=transform_stage_name
-        )
-
-        # Call fit sproc
         statement_params = telemetry.get_function_usage_statement_params(
             project=_PROJECT,
             subproject=self._subproject,
             function_name=telemetry.get_statement_params_full_func_name(inspect.currentframe(), self._class_name),
             api_calls=[Session.call],
-            custom_tags=dict([("autogen", True)]) if self._autogenerated else None,
+            custom_tags={"autogen": True} if self._autogenerated else None,
         )
 
+        temp_stage_name = estimator_utils.create_temp_stage(self.session)
+        estimator_utils.upload_model_to_stage(
+            stage_name=temp_stage_name,
+            estimator=self.estimator,
+            session=self.session,
+            statement_params=statement_params,
+        )
+
+        # Call fit sproc
         if _ENABLE_ANONYMOUS_SPROC:
             fit_predict_wrapper_sproc = self._get_fit_predict_wrapper_sproc_anonymous(statement_params=statement_params)
         else:
             fit_predict_wrapper_sproc = self._get_fit_predict_wrapper_sproc(statement_params=statement_params)
 
-        fit_predict_result_name = random_name_for_temp_object(TempObjectType.TABLE)
+        fit_predict_result_name = snowpark_utils.random_name_for_temp_object(snowpark_utils.TempObjectType.TABLE)
 
         sproc_export_file_name: str = fit_predict_wrapper_sproc(
             self.session,
             queries,
-            stage_transform_file_name,
-            stage_result_file_name,
+            temp_stage_name,
             self.input_cols,
             statement_params,
             drop_input_cols,
@@ -769,7 +715,7 @@ class SnowparkModelTrainer:
 
         output_result_sp = self.session.table(fit_predict_result_name)
         fitted_estimator = self._fetch_model_from_stage(
-            dir_path=stage_result_file_name,
+            dir_path=temp_stage_name,
             file_name=sproc_export_file_name,
             statement_params=statement_params,
         )
@@ -799,20 +745,23 @@ class SnowparkModelTrainer:
         # Extract query that generated the dataframe. We will need to pass it to the fit procedure.
         queries = dataset.queries["queries"]
 
-        transform_stage_name = self._create_temp_stage()
-        (stage_transform_file_name, stage_result_file_name) = self._upload_model_to_stage(
-            stage_name=transform_stage_name
-        )
-
-        # Call fit sproc
         statement_params = telemetry.get_function_usage_statement_params(
             project=_PROJECT,
             subproject=self._subproject,
             function_name=telemetry.get_statement_params_full_func_name(inspect.currentframe(), self._class_name),
             api_calls=[Session.call],
-            custom_tags=dict([("autogen", True)]) if self._autogenerated else None,
+            custom_tags={"autogen": True} if self._autogenerated else None,
+        )
+
+        temp_stage_name = estimator_utils.create_temp_stage(self.session)
+        estimator_utils.upload_model_to_stage(
+            stage_name=temp_stage_name,
+            estimator=self.estimator,
+            session=self.session,
+            statement_params=statement_params,
         )
 
+        # Call fit sproc
         if _ENABLE_ANONYMOUS_SPROC:
             fit_transform_wrapper_sproc = self._get_fit_transform_wrapper_sproc_anonymous(
                 statement_params=statement_params
@@ -820,13 +769,12 @@ class SnowparkModelTrainer:
         else:
             fit_transform_wrapper_sproc = self._get_fit_transform_wrapper_sproc(statement_params=statement_params)
 
-        fit_transform_result_name = random_name_for_temp_object(TempObjectType.TABLE)
+        fit_transform_result_name = snowpark_utils.random_name_for_temp_object(snowpark_utils.TempObjectType.TABLE)
 
         sproc_export_file_name: str = fit_transform_wrapper_sproc(
             self.session,
             queries,
-            stage_transform_file_name,
-            stage_result_file_name,
+            temp_stage_name,
             self.input_cols,
             self.label_cols,
             self.sample_weight_col,
@@ -838,7 +786,7 @@ class SnowparkModelTrainer:
 
         output_result_sp = self.session.table(fit_transform_result_name)
         fitted_estimator = self._fetch_model_from_stage(
-            dir_path=stage_result_file_name,
+            dir_path=temp_stage_name,
             file_name=sproc_export_file_name,
             statement_params=statement_params,
         )