snowflake-ml-python 1.6.2__py3-none-any.whl → 1.6.3__py3-none-any.whl

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

@@ -2,7 +2,6 @@ import os
 import pathlib
 import sys
 import tempfile
-import warnings
 import zipfile
 from contextlib import contextmanager
 from datetime import datetime
@@ -18,7 +17,6 @@ from snowflake.ml._internal import env as snowml_env, env_utils, file_utils
 from snowflake.ml.model import model_signature, type_hints as model_types
 from snowflake.ml.model._packager.model_env import model_env
 from snowflake.ml.model._packager.model_meta import (
-    _core_requirements,
     _packaging_requirements,
     model_blob_meta,
     model_meta_schema,
@@ -29,14 +27,10 @@ from snowflake.ml.model._packager.model_runtime import model_runtime
 MODEL_METADATA_FILE = "model.yaml"
 MODEL_CODE_DIR = "code"
 
-_PACKAGING_CORE_DEPENDENCIES = [
-    str(env_utils.get_package_spec_with_supported_ops_only(requirements.Requirement(r)))
-    for r in _core_requirements.REQUIREMENTS
-]  # Legacy Model only
 _PACKAGING_REQUIREMENTS = [
     str(env_utils.get_package_spec_with_supported_ops_only(requirements.Requirement(r)))
     for r in _packaging_requirements.REQUIREMENTS
-]  # New Model only
+]
 _SNOWFLAKE_PKG_NAME = "snowflake"
 _SNOWFLAKE_ML_PKG_NAME = f"{_SNOWFLAKE_PKG_NAME}.ml"
 
@@ -55,7 +49,7 @@ def create_model_metadata(
     conda_dependencies: Optional[List[str]] = None,
     pip_requirements: Optional[List[str]] = None,
     python_version: Optional[str] = None,
-    model_objective: model_types.ModelObjective = model_types.ModelObjective.UNKNOWN,
+    task: model_types.Task = model_types.Task.UNKNOWN,
     **kwargs: Any,
 ) -> Generator["ModelMetadata", None, None]:
     """Create a generator for model metadata object. Use generator to ensure correct register and unregister for
@@ -75,9 +69,9 @@ def create_model_metadata(
         pip_requirements: List of pip Python packages requirements for running the model. Defaults to None.
         python_version: A string of python version where model is run. Used for user override. If specified as None,
             current version would be captured. Defaults to None.
-        model_objective: The objective of the Model Version. It is an enum class ModelObjective with values REGRESSION,
-            BINARY_CLASSIFICATION, MULTI_CLASSIFICATION, RANKING, or UNKNOWN. By default it is set to
-            ModelObjective.UNKNOWN and may be overridden by inferring from the Model Object.
+        task: The task of the Model Version. It is an enum class Task with values TABULAR_REGRESSION,
+            TABULAR_BINARY_CLASSIFICATION, TABULAR_MULTI_CLASSIFICATION, TABULAR_RANKING, or UNKNOWN. By default,
+            it is set to Task.UNKNOWN and may be overridden by inferring from the Model Object.
         **kwargs: Dict of attributes and values of the metadata. Used when loading from file.
 
     Raises:
@@ -88,18 +82,6 @@ def create_model_metadata(
     """
     model_dir_path = os.path.normpath(model_dir_path)
     embed_local_ml_library = kwargs.pop("embed_local_ml_library", False)
-    legacy_save = kwargs.pop("_legacy_save", False)
-    if "relax_version" not in kwargs:
-        warnings.warn(
-            (
-                "`relax_version` is not set and therefore defaulted to True. Dependency version constraints relaxed "
-                "from ==x.y.z to >=x.y, <(x+1). To use specific dependency versions for compatibility, "
-                "reproducibility, etc., set `options={'relax_version': False}` when logging the model."
-            ),
-            category=UserWarning,
-            stacklevel=2,
-        )
-    relax_version = kwargs.pop("relax_version", True)
 
     if embed_local_ml_library:
         # Use the last one which is loaded first, that is mean, it is loaded from site-packages.
@@ -122,7 +104,6 @@ def create_model_metadata(
         pip_requirements=pip_requirements,
         python_version=python_version,
         embed_local_ml_library=embed_local_ml_library,
-        legacy_save=legacy_save,
     )
 
     if embed_local_ml_library:
@@ -135,18 +116,13 @@ def create_model_metadata(
         model_type=model_type,
         signatures=signatures,
         function_properties=function_properties,
-        model_objective=model_objective,
+        task=task,
     )
 
     code_dir_path = os.path.join(model_dir_path, MODEL_CODE_DIR)
-    if (embed_local_ml_library and legacy_save) or code_paths:
+    if code_paths:
         os.makedirs(code_dir_path, exist_ok=True)
 
-    if embed_local_ml_library and legacy_save:
-        snowml_path_in_code = os.path.join(code_dir_path, _SNOWFLAKE_PKG_NAME)
-        os.makedirs(snowml_path_in_code, exist_ok=True)
-        file_utils.copy_file_or_tree(path_to_copy, snowml_path_in_code)
-
     if code_paths:
         for code_path in code_paths:
             # This part is to prevent users from providing code following our naming and overwrite our code.
@@ -165,8 +141,6 @@ def create_model_metadata(
                     cloudpickle.register_pickle_by_value(mod)
                     imported_modules.append(mod)
         yield model_meta
-        if relax_version:
-            model_meta.env.relax_version()
         model_meta.save(model_dir_path)
     finally:
         for mod in imported_modules:
@@ -179,7 +153,6 @@ def _create_env_for_model_metadata(
     pip_requirements: Optional[List[str]] = None,
     python_version: Optional[str] = None,
     embed_local_ml_library: bool = False,
-    legacy_save: bool = False,
 ) -> model_env.ModelEnv:
     env = model_env.ModelEnv()
 
@@ -189,7 +162,7 @@ def _create_env_for_model_metadata(
     env.python_version = python_version  # type: ignore[assignment]
     env.snowpark_ml_version = snowml_env.VERSION
 
-    requirements_to_add = _PACKAGING_CORE_DEPENDENCIES if legacy_save else _PACKAGING_REQUIREMENTS
+    requirements_to_add = _PACKAGING_REQUIREMENTS
 
     if embed_local_ml_library:
         env.include_if_absent(
@@ -242,7 +215,7 @@ class ModelMetadata:
         function_properties: A dict mapping function names to dict mapping function property key to value.
         metadata: User provided key-value metadata of the model. Defaults to None.
         creation_timestamp: Unix timestamp when the model metadata is created.
-        model_objective: Model objective like regression, classification etc.
+        task: Model task like TABULAR_REGRESSION, tabular_classification, timeseries_forecasting etc.
     """
 
     def telemetry_metadata(self) -> ModelMetadataTelemetryDict:
@@ -266,7 +239,7 @@ class ModelMetadata:
         min_snowpark_ml_version: Optional[str] = None,
         models: Optional[Dict[str, model_blob_meta.ModelBlobMeta]] = None,
         original_metadata_version: Optional[str] = model_meta_schema.MODEL_METADATA_VERSION,
-        model_objective: model_types.ModelObjective = model_types.ModelObjective.UNKNOWN,
+        task: model_types.Task = model_types.Task.UNKNOWN,
         explain_algorithm: Optional[model_meta_schema.ModelExplainAlgorithm] = None,
     ) -> None:
         self.name = name
@@ -292,7 +265,7 @@ class ModelMetadata:
 
         self.original_metadata_version = original_metadata_version
 
-        self.model_objective: model_types.ModelObjective = model_objective
+        self.task: model_types.Task = task
         self.explain_algorithm: Optional[model_meta_schema.ModelExplainAlgorithm] = explain_algorithm
 
     @property
@@ -309,10 +282,10 @@ class ModelMetadata:
         if self._runtimes and "cpu" in self._runtimes:
             return self._runtimes
         runtimes = {
-            "cpu": model_runtime.ModelRuntime("cpu", self.env),
+            "cpu": model_runtime.ModelRuntime("cpu", self.env, is_warehouse=False),
         }
         if self.env.cuda_version:
-            runtimes.update({"gpu": model_runtime.ModelRuntime("gpu", self.env, is_gpu=True)})
+            runtimes.update({"gpu": model_runtime.ModelRuntime("gpu", self.env, is_warehouse=False, is_gpu=True)})
         return runtimes
 
     def save(self, model_dir_path: str) -> None:
@@ -346,7 +319,7 @@ class ModelMetadata:
                 "signatures": {func_name: sig.to_dict() for func_name, sig in self.signatures.items()},
                 "version": model_meta_schema.MODEL_METADATA_VERSION,
                 "min_snowpark_ml_version": self.min_snowpark_ml_version,
-                "model_objective": self.model_objective.value,
+                "task": self.task.value,
                 "explainability": (
                     model_meta_schema.ExplainabilityMetadataDict(algorithm=self.explain_algorithm.value)
                     if self.explain_algorithm
@@ -390,7 +363,7 @@ class ModelMetadata:
             signatures=loaded_meta["signatures"],
             version=original_loaded_meta_version,
             min_snowpark_ml_version=loaded_meta_min_snowpark_ml_version,
-            model_objective=loaded_meta.get("model_objective", model_types.ModelObjective.UNKNOWN.value),
+            task=loaded_meta.get("task", model_types.Task.UNKNOWN.value),
             explainability=loaded_meta.get("explainability", None),
             function_properties=loaded_meta.get("function_properties", {}),
         )
@@ -445,9 +418,7 @@ class ModelMetadata:
             min_snowpark_ml_version=model_dict["min_snowpark_ml_version"],
             models=models,
             original_metadata_version=model_dict["version"],
-            model_objective=model_types.ModelObjective(
-                model_dict.get("model_objective", model_types.ModelObjective.UNKNOWN.value)
-            ),
+            task=model_types.Task(model_dict.get("task", model_types.Task.UNKNOWN.value)),
             explain_algorithm=explanation_algorithm,
             function_properties=model_dict.get("function_properties", {}),
         )

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

@@ -50,10 +50,6 @@ class LightGBMModelBlobOptions(BaseModelBlobOptions):
     lightgbm_estimator_type: Required[str]
 
 
-class LLMModelBlobOptions(BaseModelBlobOptions):
-    batch_size: Required[int]
-
-
 class MLFlowModelBlobOptions(BaseModelBlobOptions):
     artifact_path: Required[str]
 
@@ -65,7 +61,6 @@ class XgboostModelBlobOptions(BaseModelBlobOptions):
 ModelBlobOptions = Union[
     BaseModelBlobOptions,
     HuggingFacePipelineModelBlobOptions,
-    LLMModelBlobOptions,
     MLFlowModelBlobOptions,
     XgboostModelBlobOptions,
 ]
@@ -96,7 +91,7 @@ class ModelMetadataDict(TypedDict):
     signatures: Required[Dict[str, Dict[str, Any]]]
     version: Required[str]
     min_snowpark_ml_version: Required[str]
-    model_objective: Required[str]
+    task: Required[str]
     explainability: NotRequired[Optional[ExplainabilityMetadataDict]]
     function_properties: NotRequired[Dict[str, Dict[str, Any]]]
 

snowflake/ml/model/_packager/model_packager.py

@@ -47,8 +47,9 @@ class ModelPackager:
         ext_modules: Optional[List[ModuleType]] = None,
         code_paths: Optional[List[str]] = None,
         options: Optional[model_types.ModelSaveOption] = None,
-        model_objective: model_types.ModelObjective = model_types.ModelObjective.UNKNOWN,
+        task: model_types.Task = model_types.Task.UNKNOWN,
     ) -> model_meta.ModelMetadata:
+
         if (signatures is None) and (sample_input_data is None) and not model_handler.is_auto_signature_model(model):
             raise snowml_exceptions.SnowflakeMLException(
                 error_code=error_codes.INVALID_ARGUMENT,
@@ -57,17 +58,20 @@ class ModelPackager:
                 ),
             )
 
-        if (signatures is not None) and (sample_input_data is not None):
-            raise snowml_exceptions.SnowflakeMLException(
-                error_code=error_codes.INVALID_ARGUMENT,
-                original_exception=ValueError(
-                    "Signatures and sample_input_data both cannot be specified at the same time."
-                ),
-            )
-
         if not options:
             options = model_types.BaseModelSaveOption()
 
+        # here handling the case of enable_explainability is False/None
+        enable_explainability = options.get("enable_explainability", None)
+        if enable_explainability is False or enable_explainability is None:
+            if (signatures is not None) and (sample_input_data is not None):
+                raise snowml_exceptions.SnowflakeMLException(
+                    error_code=error_codes.INVALID_ARGUMENT,
+                    original_exception=ValueError(
+                        "Signatures and sample_input_data both cannot be specified at the same time."
+                    ),
+                )
+
         handler = model_handler.find_handler(model)
         if handler is None:
             raise snowml_exceptions.SnowflakeMLException(
@@ -85,7 +89,7 @@ class ModelPackager:
             conda_dependencies=conda_dependencies,
             pip_requirements=pip_requirements,
             python_version=python_version,
-            model_objective=model_objective,
+            task=task,
             **options,
         ) as meta:
             model_blobs_path = os.path.join(self.local_dir_path, ModelPackager.MODEL_BLOBS_DIR)

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

@@ -36,6 +36,7 @@ class ModelRuntime:
         name: str,
         env: model_env.ModelEnv,
         imports: Optional[List[str]] = None,
+        is_warehouse: bool = False,
         is_gpu: bool = False,
         loading_from_file: bool = False,
     ) -> None:
@@ -60,6 +61,16 @@ class ModelRuntime:
             ],
         )
 
+        if not is_warehouse and self.embed_local_ml_library:
+            self.runtime_env.include_if_absent(
+                [
+                    model_env.ModelDependency(
+                        requirement="pyarrow",
+                        pip_name="pyarrow",
+                    )
+                ],
+            )
+
         if is_gpu:
             self.runtime_env.generate_env_for_cuda()
 

snowflake/ml/model/_signatures/snowpark_handler.py

@@ -14,9 +14,10 @@ from snowflake.ml._internal.exceptions import (
 )
 from snowflake.ml._internal.utils import identifier
 from snowflake.ml.model import type_hints as model_types
-from snowflake.ml.model._deploy_client.warehouse import infer_template
 from snowflake.ml.model._signatures import base_handler, core, pandas_handler
 
+_KEEP_ORDER_COL_NAME = "_ID"
+
 
 class SnowparkDataFrameHandler(base_handler.BaseDataHandler[snowflake.snowpark.DataFrame]):
     @staticmethod
@@ -109,7 +110,7 @@ class SnowparkDataFrameHandler(base_handler.BaseDataHandler[snowflake.snowpark.D
         # 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
+            df[_KEEP_ORDER_COL_NAME] = df.index
         sp_df = session.create_dataframe(df)
         column_names = []
         columns = []

snowflake/ml/model/type_hints.py

@@ -1,23 +1,9 @@
 # mypy: disable-error-code="import"
 from enum import Enum
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    Dict,
-    List,
-    Literal,
-    Optional,
-    Sequence,
-    TypedDict,
-    TypeVar,
-    Union,
-)
+from typing import TYPE_CHECKING, Dict, Literal, Sequence, TypedDict, TypeVar, Union
 
 import numpy.typing as npt
-from typing_extensions import NotRequired, Required
-
-from snowflake.ml.model import deploy_platforms
-from snowflake.ml.model._signatures import core
+from typing_extensions import NotRequired
 
 if TYPE_CHECKING:
     import catboost
@@ -35,7 +21,6 @@ if TYPE_CHECKING:
 
     import snowflake.ml.model.custom_model
     import snowflake.ml.model.models.huggingface_pipeline
-    import snowflake.ml.model.models.llm
     import snowflake.snowpark
     from snowflake.ml.modeling.framework import base  # noqa: F401
 
@@ -91,7 +76,6 @@ SupportedNoSignatureRequirementsModelType = Union[
     "transformers.Pipeline",
     "sentence_transformers.SentenceTransformer",
     "snowflake.ml.model.models.huggingface_pipeline.HuggingFacePipelineModel",
-    "snowflake.ml.model.models.llm.LLM",
 ]
 
 SupportedModelType = Union[
@@ -134,86 +118,11 @@ SupportedModelHandlerType = Literal[
     "tensorflow",
     "torchscript",
     "xgboost",
-    "llm",
 ]
 
 _ModelType = TypeVar("_ModelType", bound=SupportedModelType)
 
 
-class DeployOptions(TypedDict):
-    """Common Options for deploying to Snowflake."""
-
-    ...
-
-
-class WarehouseDeployOptions(DeployOptions):
-    """Options for deploying to the Snowflake Warehouse.
-
-
-    permanent_udf_stage_location: A Snowflake stage option where the UDF should be persisted. If specified, the model
-        will be deployed as a permanent UDF, otherwise temporary.
-    relax_version: Whether or not relax the version constraints of the dependencies if unresolvable. It detects any
-        ==x.y.z in specifiers and replaced with >=x.y, <(x+1). Defaults to False.
-    replace_udf: Flag to indicate when deploying model as permanent UDF, whether overwriting existed UDF is allowed.
-        Default to False.
-    """
-
-    permanent_udf_stage_location: NotRequired[str]
-    relax_version: NotRequired[bool]
-    replace_udf: NotRequired[bool]
-
-
-class SnowparkContainerServiceDeployOptions(DeployOptions):
-    """Deployment options for deploying to SnowService.
-    When type hint is updated, please ensure the concrete class is updated accordingly at:
-    //snowflake/ml/model/_deploy_client/snowservice/_deploy_options
-
-    compute_pool[REQUIRED]: SnowService compute pool name. Please refer to official doc for how to create a
-        compute pool: https://docs.snowflake.com/LIMITEDACCESS/snowpark-containers/reference/compute-pool
-    image_repo: SnowService image repo path. e.g. "<image_registry>/<db>/<schema>/<repo>". Default to auto
-        inferred based on session information.
-    min_instances: Minimum number of service replicas. Default to 1.
-    max_instances: Maximum number of service replicas. Default to 1.
-    prebuilt_snowflake_image: When provided, the image-building step is skipped, and the pre-built image from
-        Snowflake is used as is. This option is for users who consistently use the same image for multiple use
-        cases, allowing faster deployment. The snowflake image used for deployment is logged to the console for
-        future use. Default to None.
-    num_gpus: Number of GPUs to be used for the service. Default to 0.
-    num_workers: Number of workers used for model inference. Please ensure that the number of workers is set lower than
-        the total available memory divided by the size of model to prevent memory-related issues. Default is number of
-        CPU cores * 2 + 1.
-    enable_remote_image_build: When set to True, will enable image build on a remote SnowService job. Default is True.
-    force_image_build: When set to True, an image rebuild will occur. The default is False, which means the system
-        will automatically check whether a previously built image can be reused
-    model_in_image: When set to True, image would container full model weights. The default if False, which
-                means image without model weights and we do stage mount to access weights.
-    debug_mode: When set to True, deployment artifacts will be persisted in a local temp directory.
-    enable_ingress: When set to True, will expose HTTP endpoint for access to the predict method of the created
-        service.
-    external_access_integrations: External Access Integrations name used to build image and deploy the model.
-        Please refer to the doc for how to create an External Access Integrations: https://docs.snowflake.com/
-        developer-guide/snowpark-container-services/additional-considerations-services-jobs
-        #configuring-network-capabilities .
-        To make sure your image could be built, access to the following endpoint must be allowed.
-        docker.com:80, docker.com:443, anaconda.com:80, anaconda.com:443, anaconda.org:80, anaconda.org:443,
-        pypi.org:80, pypi.org:443
-    """
-
-    compute_pool: str
-    image_repo: NotRequired[str]
-    min_instances: NotRequired[int]
-    max_instances: NotRequired[int]
-    prebuilt_snowflake_image: NotRequired[str]
-    num_gpus: NotRequired[int]
-    num_workers: NotRequired[int]
-    enable_remote_image_build: NotRequired[bool]
-    force_image_build: NotRequired[bool]
-    model_in_image: NotRequired[bool]
-    debug_mode: NotRequired[bool]
-    enable_ingress: NotRequired[bool]
-    external_access_integrations: List[str]
-
-
 class ModelMethodSaveOptions(TypedDict):
     case_sensitive: NotRequired[bool]
     max_batch_size: NotRequired[int]
@@ -224,13 +133,12 @@ class BaseModelSaveOption(TypedDict):
     """Options for saving the model.
 
     embed_local_ml_library: Embedding local SnowML into the code directory of the folder.
-    relax_version: Whether or not relax the version constraints of the dependencies if unresolvable. It detects any
-        ==x.y.z in specifiers and replaced with >=x.y, <(x+1). Defaults to False.
+    relax_version: Whether or not relax the version constraints of the dependencies if unresolvable in Warehouse.
+        It detects any ==x.y.z in specifiers and replaced with >=x.y, <(x+1). Defaults to True.
     """
 
     embed_local_ml_library: NotRequired[bool]
     relax_version: NotRequired[bool]
-    _legacy_save: NotRequired[bool]
     function_type: NotRequired[Literal["FUNCTION", "TABLE_FUNCTION"]]
     method_options: NotRequired[Dict[str, ModelMethodSaveOptions]]
     enable_explainability: NotRequired[bool]
@@ -293,10 +201,6 @@ class SentenceTransformersSaveOptions(BaseModelSaveOption):
     cuda_version: NotRequired[str]
 
 
-class LLMSaveOptions(BaseModelSaveOption):
-    cuda_version: NotRequired[str]
-
-
 ModelSaveOption = Union[
     BaseModelSaveOption,
     CatBoostModelSaveOptions,
@@ -311,7 +215,6 @@ ModelSaveOption = Union[
     MLFlowSaveOptions,
     HuggingFaceSaveOptions,
     SentenceTransformersSaveOptions,
-    LLMSaveOptions,
 ]
 
 
@@ -369,10 +272,7 @@ class HuggingFaceLoadOptions(BaseModelLoadOption):
 
 class SentenceTransformersLoadOptions(BaseModelLoadOption):
     use_gpu: NotRequired[bool]
-
-
-class LLMLoadOptions(BaseModelLoadOption):
-    ...
+    device: NotRequired[str]
 
 
 ModelLoadOption = Union[
@@ -389,53 +289,12 @@ ModelLoadOption = Union[
     MLFlowLoadOptions,
     HuggingFaceLoadOptions,
     SentenceTransformersLoadOptions,
-    LLMLoadOptions,
-]
-
-
-class SnowparkContainerServiceDeployDetails(TypedDict):
-    """
-    Attributes:
-        service_info: A snowpark row containing the result of "describe service"
-        service_function_sql: SQL for service function creation.
-    """
-
-    service_info: Optional[Dict[str, Any]]
-    service_function_sql: str
-
-
-class WarehouseDeployDetails(TypedDict):
-    ...
-
-
-DeployDetails = Union[
-    SnowparkContainerServiceDeployDetails,
-    WarehouseDeployDetails,
 ]
 
 
-class Deployment(TypedDict):
-    """Deployment information.
-
-    Attributes:
-        name: Name of the deployment.
-        platform: Target platform to deploy the model.
-        target_method: Target method name.
-        signature: The signature of the model method.
-        options: Additional options when deploying the model.
-    """
-
-    name: Required[str]
-    platform: Required[deploy_platforms.TargetPlatform]
-    target_method: Required[str]
-    signature: core.ModelSignature
-    options: Required[DeployOptions]
-    details: NotRequired[DeployDetails]
-
-
-class ModelObjective(Enum):
-    UNKNOWN = "unknown"
-    BINARY_CLASSIFICATION = "binary_classification"
-    MULTI_CLASSIFICATION = "multi_classification"
-    REGRESSION = "regression"
-    RANKING = "ranking"
+class Task(Enum):
+    UNKNOWN = "UNKNOWN"
+    TABULAR_BINARY_CLASSIFICATION = "TABULAR_BINARY_CLASSIFICATION"
+    TABULAR_MULTI_CLASSIFICATION = "TABULAR_MULTI_CLASSIFICATION"
+    TABULAR_REGRESSION = "TABULAR_REGRESSION"
+    TABULAR_RANKING = "TABULAR_RANKING"

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

@@ -377,7 +377,6 @@ class DistributedHPOTrainer(SnowparkModelTrainer):
             anonymous=True,
             imports=imports,  # type: ignore[arg-type]
             statement_params=sproc_statement_params,
-            execute_as="caller",
         )
         def _distributed_search(
             session: Session,
@@ -783,7 +782,6 @@ class DistributedHPOTrainer(SnowparkModelTrainer):
             anonymous=True,
             imports=imports,  # type: ignore[arg-type]
             statement_params=sproc_statement_params,
-            execute_as="caller",
         )
         def _distributed_search(
             session: Session,

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

@@ -230,7 +230,6 @@ class SnowparkModelTrainer:
             replace=True,
             session=self.session,
             statement_params=statement_params,
-            execute_as="caller",
             anonymous=anonymous,
         )
         return fit_wrapper_sproc
@@ -461,9 +460,7 @@ class SnowparkModelTrainer:
             session.write_pandas(
                 transformed_pandas_df,
                 fit_transform_result_name,
-                auto_create_table=True,
-                table_type="temp",
-                quote_identifiers=False,
+                overwrite=True,
             )
 
             return str(os.path.basename(local_result_file_name))
@@ -488,7 +485,6 @@ class SnowparkModelTrainer:
             session=self.session,
             statement_params=statement_params,
             anonymous=anonymous,
-            execute_as="caller",
         )
 
         return fit_predict_wrapper_sproc
@@ -510,7 +506,6 @@ class SnowparkModelTrainer:
             replace=True,
             session=self.session,
             statement_params=statement_params,
-            execute_as="caller",
             anonymous=anonymous,
         )
 
@@ -730,6 +725,22 @@ class SnowparkModelTrainer:
 
         fit_transform_result_name = snowpark_utils.random_name_for_temp_object(snowpark_utils.TempObjectType.TABLE)
 
+        # Create a temp table in advance to store the output
+        # This would allow us to use the same table outside the stored procedure
+        df_one_line = dataset.limit(1).to_pandas(statement_params=statement_params)
+        df_one_line[
+            expected_output_cols_list[0]
+        ] = "[0]"  # Add one column as the output_col; this is a dummy value to represent the OBJECT type
+        if drop_input_cols:
+            self.session.write_pandas(
+                df_one_line[expected_output_cols_list[0]],
+                fit_transform_result_name,
+                auto_create_table=True,
+                table_type="temp",
+            )
+        else:
+            self.session.write_pandas(df_one_line, fit_transform_result_name, auto_create_table=True, table_type="temp")
+
         sproc_export_file_name: str = fit_transform_wrapper_sproc(
             self.session,
             queries,

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

@@ -303,7 +303,6 @@ class XGBoostExternalMemoryTrainer(SnowparkModelTrainer):
             statement_params=statement_params,
             anonymous=True,
             imports=list(import_file_paths),
-            execute_as="caller",
         )  # type: ignore[misc]
         def fit_wrapper_sproc(
             session: Session,

snowflake/ml/modeling/calibration/calibrated_classifier_cv.py

@@ -576,6 +576,7 @@ class CalibratedClassifierCV(BaseTransformer):
         """
         self._infer_input_output_cols(dataset)
         super()._check_dataset_type(dataset)
+
         model_trainer = ModelTrainerBuilder.build_fit_transform(
             estimator=self._sklearn_object,
             dataset=dataset,

snowflake/ml/modeling/cluster/affinity_propagation.py

@@ -555,6 +555,7 @@ class AffinityPropagation(BaseTransformer):
         """
         self._infer_input_output_cols(dataset)
         super()._check_dataset_type(dataset)
+
         model_trainer = ModelTrainerBuilder.build_fit_transform(
             estimator=self._sklearn_object,
             dataset=dataset,

snowflake/ml/modeling/cluster/agglomerative_clustering.py

@@ -586,6 +586,7 @@ class AgglomerativeClustering(BaseTransformer):
         """
         self._infer_input_output_cols(dataset)
         super()._check_dataset_type(dataset)
+
         model_trainer = ModelTrainerBuilder.build_fit_transform(
             estimator=self._sklearn_object,
             dataset=dataset,

snowflake/ml/modeling/cluster/birch.py

@@ -550,6 +550,7 @@ class Birch(BaseTransformer):
         """
         self._infer_input_output_cols(dataset)
         super()._check_dataset_type(dataset)
+
         model_trainer = ModelTrainerBuilder.build_fit_transform(
             estimator=self._sklearn_object,
             dataset=dataset,

snowflake/ml/modeling/cluster/bisecting_k_means.py

@@ -599,6 +599,7 @@ class BisectingKMeans(BaseTransformer):
         """
         self._infer_input_output_cols(dataset)
         super()._check_dataset_type(dataset)
+
         model_trainer = ModelTrainerBuilder.build_fit_transform(
             estimator=self._sklearn_object,
             dataset=dataset,