snowflake-ml-python 1.1.1__py3-none-any.whl → 1.1.2__py3-none-any.whl

snowflake/ml/model/_deploy_client/image_builds/templates/image_build_job_spec_template

@@ -1,22 +1,38 @@
 spec:
   container:
-    - name: $container_name
-      image: $base_image
+    - name: "${container_name}"
+      image: "${base_image}"
       command:
         - sh
       args:
         - -c
-        - >-
-          while [ ! -f "$script_path" ]; do echo "File not found: $script_path"; sleep 1; done;
-          chmod +x $script_path;
-          sh $script_path;
+        - |
+          wait_for_file() {
+              file_path="$1"
+              timeout="$2"
+              elapsed_time=0
+              while [ ! -f "${file_path}" ]; do
+                if [ "${elapsed_time}" -ge "${timeout}" ]; then
+                  echo "Error: ${file_path} not found within ${timeout} seconds. Exiting."
+                  exit 1
+                fi
+                elapsed_time=$((elapsed_time + 1))
+                remaining_time=$((timeout - elapsed_time))
+                echo "Awaiting the mounting of ${file_path}. Wait time remaining: ${remaining_time} seconds"
+                sleep 1
+              done
+           }
+           wait_for_file "${script_path}" 300
+           wait_for_file "${mounted_token_path}" 300
+           chmod +x "${script_path}"
+           sh "${script_path}"
       volumeMounts:
         - name: vol1
           mountPath: /local/user/vol1
         - name: stagemount
-          mountPath: /$stage
+          mountPath: "/${stage}"
   volume:
     - name: vol1
       source: local  # only local emptyDir volume is supported
     - name: stagemount
-      source: "@$stage"
+      source: "@${stage}"

snowflake/ml/model/_deploy_client/image_builds/templates/kaniko_shell_script_template

@@ -11,18 +11,41 @@ cleanup() {
   kill -- -$$$ # Kill the entire process group. Extra $ to escape, the generated shell script should have two $.
 }
 
+# SNOW-990976, This is an additional safety check to ensure token file exists, on top of the token file check upon
+# launching SPCS job. This additional check could provide value in cases things go wrong with token refresh that result
+# in token file to disappear.
+wait_till_token_file_exists() {
+  timeout=60  # 1 minute timeout
+  elapsed_time=0
+
+  while [ ! -f "${SESSION_TOKEN_PATH}" ] && [ "$elapsed_time" -lt "$timeout" ]; do
+    sleep 1
+    elapsed_time=$((elapsed_time + 1))
+    remaining_time=$((timeout - elapsed_time))
+    echo "Waiting for token file to exist. Wait time remaining: ${remaining_time} seconds."
+  done
+
+  if [ ! -f "${SESSION_TOKEN_PATH}" ]; then
+    echo "Error: Token file '${SESSION_TOKEN_PATH}' does not show up within the ${timeout} seconds timeout period."
+    exit 1
+  fi
+}
+
 generate_registry_cred() {
+  wait_till_token_file_exists
   AUTH_TOKEN=$(printf '0auth2accesstoken:%s' "$(cat ${SESSION_TOKEN_PATH})" | base64);
   echo '{"auths":{"$image_repo":{"auth":"'"$AUTH_TOKEN"'"}}}' | tr -d '\n' > $REGISTRY_CRED_PATH;
 }
 
 on_session_token_change() {
+  wait_till_token_file_exists
   # Get the initial checksum of the file
   CHECKSUM=$(md5sum "${SESSION_TOKEN_PATH}" | awk '{ print $1 }')
   # Run the command once before the loop
   echo "Monitoring session token changes in the background..."
   (
     while true; do
+      wait_till_token_file_exists
       # Get the current checksum of the file
       CURRENT_CHECKSUM=$(md5sum "${SESSION_TOKEN_PATH}" | awk '{ print $1 }')
       if [ "${CURRENT_CHECKSUM}" != "${CHECKSUM}" ]; then

snowflake/ml/model/_deploy_client/snowservice/deploy.py

@@ -10,14 +10,19 @@ from typing import Any, Dict, Generator, Optional, cast
 
 import importlib_resources
 import yaml
+from packaging import requirements
 from typing_extensions import Unpack
 
-from snowflake.ml._internal import file_utils
+from snowflake.ml._internal import env_utils, file_utils
 from snowflake.ml._internal.exceptions import (
     error_codes,
     exceptions as snowml_exceptions,
 )
-from snowflake.ml._internal.utils import identifier, query_result_checker
+from snowflake.ml._internal.utils import (
+    identifier,
+    query_result_checker,
+    spcs_attribution_utils,
+)
 from snowflake.ml.model import type_hints
 from snowflake.ml.model._deploy_client import snowservice
 from snowflake.ml.model._deploy_client.image_builds import (
@@ -161,6 +166,11 @@ def _deploy(
         # Set conda-forge as backup channel for SPCS deployment
         if "conda-forge" not in model_meta_deploy.env._conda_dependencies:
             model_meta_deploy.env._conda_dependencies["conda-forge"] = []
+        # Snowflake connector needs pyarrow to work correctly.
+        env_utils.append_conda_dependency(
+            model_meta_deploy.env._conda_dependencies,
+            (env_utils.DEFAULT_CHANNEL_NAME, requirements.Requirement("pyarrow")),
+        )
         if options.use_gpu:
             # Make mypy happy
             assert options.num_gpus is not None
@@ -585,6 +595,8 @@ class SnowServiceDeployment:
         )
         logger.info(f"Service {self._service_name} is ready. Creating service function...")
 
+        spcs_attribution_utils.record_service_start(self.session, self._service_name)
+
         service_function_sql = client.create_or_replace_service_function(
             service_func_name=self.service_func_name,
             service_name=self._service_name,

snowflake/ml/model/_deploy_client/utils/constants.py

@@ -50,3 +50,4 @@ KANIKO_SHELL_SCRIPT_NAME = "kaniko_shell_script_fixture.sh"
 KANIKO_CONTAINER_NAME = "kaniko"
 LATEST_IMAGE_TAG = "latest"
 KANIKO_IMAGE = "kaniko-project/executor:v1.16.0-debug"
+SPCS_MOUNTED_TOKEN_PATH = "/snowflake/session/token"

snowflake/ml/model/_deploy_client/warehouse/deploy.py

@@ -173,7 +173,7 @@ def _get_model_final_packages(
     else:
         required_packages = meta.env._conda_dependencies[env_utils.DEFAULT_CHANNEL_NAME]
 
-    final_packages = env_utils.validate_requirements_in_snowflake_conda_channel(
+    final_packages = env_utils.validate_requirements_in_information_schema(
         session, required_packages, python_version=meta.env.python_version
     )
 
@@ -182,7 +182,7 @@ def _get_model_final_packages(
         raise snowml_exceptions.SnowflakeMLException(
             error_code=error_codes.DEPENDENCY_VERSION_ERROR,
             original_exception=RuntimeError(
-                "The model's dependencyies are not available in Snowflake Anaconda Channel. "
+                "The model's dependencies are not available in Snowflake Anaconda Channel. "
                 + relax_version_info_str
                 + "Required packages are:\n"
                 + " ".join(map(lambda x: f'"{x}"', required_packages))

snowflake/ml/model/_model_composer/model_composer.py

@@ -3,7 +3,7 @@ import pathlib
 import tempfile
 import zipfile
 from types import ModuleType
-from typing import Dict, List, Optional
+from typing import Any, Dict, List, Optional
 
 from absl import logging
 from packaging import requirements
@@ -32,8 +32,15 @@ class ModelComposer:
     """
 
     MODEL_FILE_REL_PATH = "model.zip"
+    MODEL_DIR_REL_PATH = "model"
 
-    def __init__(self, session: Session, stage_path: str) -> None:
+    def __init__(
+        self,
+        session: Session,
+        stage_path: str,
+        *,
+        statement_params: Optional[Dict[str, Any]] = None,
+    ) -> None:
         self.session = session
         self.stage_path = pathlib.PurePosixPath(stage_path)
 
@@ -43,6 +50,8 @@ class ModelComposer:
         self.packager = model_packager.ModelPackager(local_dir_path=str(self._packager_workspace_path))
         self.manifest = model_manifest.ModelManifest(workspace_path=self.workspace_path)
 
+        self._statement_params = statement_params
+
     def __del__(self) -> None:
         self._workspace.cleanup()
         self._packager_workspace.cleanup()
@@ -82,13 +91,11 @@ class ModelComposer:
             options = model_types.BaseModelSaveOption()
 
         if not snowpark_utils.is_in_stored_procedure():  # type: ignore[no-untyped-call]
-            snowml_server_availability = env_utils.validate_requirements_in_snowflake_conda_channel(
-                session=self.session,
-                reqs=[requirements.Requirement(f"snowflake-ml-python=={snowml_env.VERSION}")],
-                python_version=snowml_env.PYTHON_VERSION,
+            snowml_matched_versions = env_utils.get_matched_package_versions_in_snowflake_conda_channel(
+                req=requirements.Requirement(f"snowflake-ml-python=={snowml_env.VERSION}")
             )
 
-            if snowml_server_availability is None and options.get("embed_local_ml_library", False) is False:
+            if len(snowml_matched_versions) < 1 and options.get("embed_local_ml_library", False) is False:
                 logging.info(
                     f"Local snowflake-ml-python library has version {snowml_env.VERSION},"
                     " which is not available in the Snowflake server, embedding local ML library automatically."
@@ -111,6 +118,13 @@ class ModelComposer:
 
         assert self.packager.meta is not None
 
+        if not options.get("_legacy_save", False):
+            # Keep both loose files and zipped file.
+            # TODO(SNOW-726678): Remove once import a directory is possible.
+            file_utils.copytree(
+                str(self._packager_workspace_path), str(self.workspace_path / ModelComposer.MODEL_DIR_REL_PATH)
+            )
+
         file_utils.make_archive(self.model_local_path, str(self._packager_workspace_path))
 
         self.manifest.save(
@@ -120,7 +134,12 @@ class ModelComposer:
             options=options,
         )
 
-        file_utils.upload_directory_to_stage(self.session, local_path=self.workspace_path, stage_path=self.stage_path)
+        file_utils.upload_directory_to_stage(
+            self.session,
+            local_path=self.workspace_path,
+            stage_path=self.stage_path,
+            statement_params=self._statement_params,
+        )
 
     def load(
         self,
@@ -129,7 +148,10 @@ class ModelComposer:
         options: Optional[model_types.ModelLoadOption] = None,
     ) -> None:
         file_utils.download_directory_from_stage(
-            self.session, stage_path=self.stage_path, local_path=self.workspace_path
+            self.session,
+            stage_path=self.stage_path,
+            local_path=self.workspace_path,
+            statement_params=self._statement_params,
         )
 
         # TODO (Server-side Model Rollout): Remove this section.

snowflake/ml/model/_model_composer/model_manifest/model_manifest.py

@@ -1,5 +1,6 @@
+import collections
 import pathlib
-from typing import List, Optional
+from typing import List, Optional, cast
 
 import yaml
 
@@ -48,7 +49,6 @@ class ModelManifest:
         ]
         self.function_generator = function_generator.FunctionGenerator(model_file_rel_path=model_file_rel_path)
         self.methods: List[model_method.ModelMethod] = []
-        _seen_method_names: List[str] = []
         for target_method in model_meta.signatures.keys():
             method = model_method.ModelMethod(
                 model_meta=model_meta,
@@ -57,17 +57,18 @@ class ModelManifest:
                 function_generator=self.function_generator,
                 options=model_method.get_model_method_options_from_options(options, target_method),
             )
-            if method.method_name in _seen_method_names:
-                raise ValueError(
-                    f"Found duplicate method named resolved as {method.method_name} in the model. "
-                    "This might because you have methods with same letters but different cases. "
-                    "In this case, set case_sensitive as True for those methods to distinguish them"
-                )
-            else:
-                _seen_method_names.append(method.method_name)
 
             self.methods.append(method)
 
+        method_name_counter = collections.Counter([method.method_name for method in self.methods])
+        dup_method_names = [k for k, v in method_name_counter.items() if v > 1]
+        if dup_method_names:
+            raise ValueError(
+                f"Found duplicate method named resolved as {', '.join(dup_method_names)} in the model. "
+                "This might because you have methods with same letters but different cases. "
+                "In this case, set case_sensitive as True for those methods to distinguish them."
+            )
+
         manifest_dict = model_manifest_schema.ModelManifestDict(
             manifest_version=model_manifest_schema.MODEL_MANIFEST_VERSION,
             runtimes={runtime.name: runtime.save(self.workspace_path) for runtime in self.runtimes},
@@ -84,3 +85,17 @@ class ModelManifest:
 
         with (self.workspace_path / ModelManifest.MANIFEST_FILE_REL_PATH).open("w", encoding="utf-8") as f:
             yaml.safe_dump(manifest_dict, f)
+
+    def load(self) -> model_manifest_schema.ModelManifestDict:
+        with (self.workspace_path / ModelManifest.MANIFEST_FILE_REL_PATH).open("r", encoding="utf-8") as f:
+            raw_input = yaml.safe_load(f)
+        if not isinstance(raw_input, dict):
+            raise ValueError(f"Read ill-formatted model MANIFEST, should be a dict, received {type(raw_input)}")
+
+        original_loaded_manifest_version = raw_input.get("manifest_version", None)
+        if not original_loaded_manifest_version:
+            raise ValueError("Unable to get the version of the MANIFEST file.")
+
+        res = cast(model_manifest_schema.ModelManifestDict, raw_input)
+
+        return res

snowflake/ml/model/_model_composer/model_manifest/model_manifest_schema.py

@@ -1,6 +1,6 @@
 # This files contains schema definition of what will be written into MANIFEST.yml
 
-from typing import Dict, List, Literal, TypedDict
+from typing import Any, Dict, List, Literal, TypedDict
 
 from typing_extensions import NotRequired, Required
 
@@ -42,4 +42,4 @@ class ModelManifestDict(TypedDict):
     manifest_version: Required[str]
     runtimes: Required[Dict[str, ModelRuntimeDict]]
     methods: Required[List[ModelMethodDict]]
-    user_data: NotRequired[Dict[str, str]]
+    user_data: NotRequired[Dict[str, Any]]

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

@@ -73,6 +73,7 @@ dtype_map = {{feature.name: feature.as_dtype() for feature in features}}
 # Actual function
 @vectorized(input=pd.DataFrame, max_batch_size=MAX_BATCH_SIZE)
 def {function_name}(df: pd.DataFrame) -> dict:
-    input_df = pd.json_normalize(df[0]).astype(dtype=dtype_map)
+    df.columns = input_cols
+    input_df = df.astype(dtype=dtype_map)
     predictions_df = runner(input_df[input_cols])
     return predictions_df.to_dict("records")

snowflake/ml/model/_model_composer/model_method/model_method.py

@@ -1,13 +1,15 @@
+import collections
 import pathlib
 from typing import Optional, TypedDict
 
 from typing_extensions import NotRequired
 
 from snowflake.ml._internal.utils import sql_identifier
-from snowflake.ml.model import type_hints
+from snowflake.ml.model import model_signature, type_hints
 from snowflake.ml.model._model_composer.model_manifest import model_manifest_schema
 from snowflake.ml.model._model_composer.model_method import function_generator
 from snowflake.ml.model._packager.model_meta import model_meta as model_meta_api
+from snowflake.snowpark._internal import type_utils
 
 
 class ModelMethodOptions(TypedDict):
@@ -69,6 +71,22 @@ class ModelMethod:
         if self.target_method not in self.model_meta.signatures.keys():
             raise ValueError(f"Target method {self.target_method} is not available in the signatures of the model.")
 
+    @staticmethod
+    def _get_method_arg_from_feature(
+        feature: model_signature.BaseFeatureSpec, case_sensitive: bool = False
+    ) -> model_manifest_schema.ModelMethodSignatureFieldWithName:
+        assert isinstance(feature, model_signature.FeatureSpec), "FeatureGroupSpec is not supported."
+        try:
+            feature_name = sql_identifier.SqlIdentifier(feature.name, case_sensitive=case_sensitive)
+        except ValueError as e:
+            raise ValueError(
+                f"Your feature {feature.name} cannot be resolved as valid SQL identifier. "
+                "Try specify `case_sensitive` as True."
+            ) from e
+        return model_manifest_schema.ModelMethodSignatureFieldWithName(
+            name=feature_name.resolved(), type=type_utils.convert_sp_to_sf_type(feature.as_snowpark_type())
+        )
+
     def save(
         self, workspace_path: pathlib.Path, options: Optional[function_generator.FunctionGenerateOptions] = None
     ) -> model_manifest_schema.ModelMethodDict:
@@ -78,13 +96,26 @@ class ModelMethod:
             self.target_method,
             options=options,
         )
+        input_list = [
+            ModelMethod._get_method_arg_from_feature(ft, case_sensitive=self.options.get("case_sensitive", False))
+            for ft in self.model_meta.signatures[self.target_method].inputs
+        ]
+        input_name_counter = collections.Counter([input_info["name"] for input_info in input_list])
+        dup_input_names = [k for k, v in input_name_counter.items() if v > 1]
+        if dup_input_names:
+            raise ValueError(
+                f"Found duplicate input feature named resolved as {', '.join(dup_input_names)} in the method"
+                f" {self.target_method} This might because you have methods with same letters but different cases. "
+                "In this case, set case_sensitive as True for those methods to distinguish them."
+            )
+
         return model_manifest_schema.ModelFunctionMethodDict(
-            name=self.method_name.identifier(),
+            name=self.method_name.resolved(),
             runtime=self.runtime_name,
             type="FUNCTION",
             handler=".".join(
                 [ModelMethod.FUNCTIONS_DIR_REL_PATH, self.target_method, self.function_generator.FUNCTION_NAME]
             ),
-            inputs=[model_manifest_schema.ModelMethodSignatureFieldWithName(name="tmp_input", type="OBJECT")],
+            inputs=input_list,
             outputs=[model_manifest_schema.ModelMethodSignatureField(type="OBJECT")],
         )

snowflake/ml/model/_model_composer/model_runtime/model_runtime.py

@@ -44,7 +44,7 @@ class ModelRuntime:
         if self.runtime_env._snowpark_ml_version.local:
             self.embed_local_ml_library = True
         else:
-            snowml_server_availability = env_utils.validate_requirements_in_snowflake_conda_channel(
+            snowml_server_availability = env_utils.validate_requirements_in_information_schema(
                 session=session,
                 reqs=[requirements.Requirement(snowml_pkg_spec)],
                 python_version=snowml_env.PYTHON_VERSION,

snowflake/ml/model/_packager/model_handlers/huggingface_pipeline.py

@@ -59,7 +59,7 @@ def get_requirements_from_task(task: str, spcs_only: bool = False) -> List[model
         return (
             [model_env.ModelDependency(requirement="tokenizers>=0.13.3", pip_name="tokenizers")]
             if spcs_only
-            else [model_env.ModelDependency(requirement="tokenizers", pip_name="tokenizers")]
+            else [model_env.ModelDependency(requirement="tokenizers<=0.13.2", pip_name="tokenizers")]
         )
 
     return []
@@ -170,6 +170,7 @@ class HuggingFacePipelineHandler(
                         " `snowflake.ml.model.models.huggingface_pipeline.HuggingFacePipelineModel` object. "
                         "Please make sure you are providing correct model signatures.",
                         UserWarning,
+                        stacklevel=2,
                     )
             else:
                 handlers_utils.validate_target_methods(model, target_methods)
@@ -179,6 +180,7 @@ class HuggingFacePipelineHandler(
                         + "Model signature will automatically be inferred from pipeline task. "
                         + "Or, you could specify model signature manually.",
                         UserWarning,
+                        stacklevel=2,
                     )
                 if inferred_pipe_sig is None:
                     raise NotImplementedError(f"Cannot auto infer the signature of pipeline for task {task}")

snowflake/ml/model/_packager/model_handlers/snowmlmodel.py

@@ -1,4 +1,5 @@
 import os
+import warnings
 from typing import TYPE_CHECKING, Callable, Dict, Optional, Type, cast, final
 
 import cloudpickle
@@ -9,7 +10,7 @@ from typing_extensions import TypeGuard, Unpack
 from snowflake.ml._internal import type_utils
 from snowflake.ml.model import custom_model, model_signature, type_hints as model_types
 from snowflake.ml.model._packager.model_env import model_env
-from snowflake.ml.model._packager.model_handlers import _base, _utils as handlers_utils
+from snowflake.ml.model._packager.model_handlers import _base
 from snowflake.ml.model._packager.model_handlers_migrator import base_migrator
 from snowflake.ml.model._packager.model_meta import (
     model_blob_meta,
@@ -78,34 +79,15 @@ class SnowMLModelHandler(_base.BaseModelHandler["BaseEstimator"]):
         # Pipeline is inherited from BaseEstimator, so no need to add one more check
 
         if not is_sub_model:
-            if (not model_meta.signatures) and sample_input is None:
-                assert hasattr(model, "model_signatures")
-                model_meta.signatures = getattr(model, "model_signatures", {})
-            else:
-                target_methods = handlers_utils.get_target_methods(
-                    model=model,
-                    target_methods=kwargs.pop("target_methods", None),
-                    default_target_methods=cls.DEFAULT_TARGET_METHODS,
-                )
-
-                def get_prediction(
-                    target_method_name: str, sample_input: model_types.SupportedLocalDataType
-                ) -> model_types.SupportedLocalDataType:
-                    if not isinstance(sample_input, (pd.DataFrame,)):
-                        sample_input = model_signature._convert_local_data_to_df(sample_input)
-
-                    target_method = getattr(model, target_method_name, None)
-                    assert callable(target_method)
-                    predictions_df = target_method(sample_input)
-                    return predictions_df
-
-                model_meta = handlers_utils.validate_signature(
-                    model=model,
-                    model_meta=model_meta,
-                    target_methods=target_methods,
-                    sample_input=sample_input,
-                    get_prediction_fn=get_prediction,
+            if sample_input is not None or model_meta.signatures:
+                warnings.warn(
+                    "Inferring model signature from sample input or providing model signature for Snowpark ML "
+                    + "Modeling model is not required. Model signature will automatically be inferred during fitting. ",
+                    UserWarning,
+                    stacklevel=2,
                 )
+            assert hasattr(model, "model_signatures"), "Model does not have model signatures as expected."
+            model_meta.signatures = getattr(model, "model_signatures", {})
 
         model_blob_path = os.path.join(model_blobs_dir_path, name)
         os.makedirs(model_blob_path, exist_ok=True)

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

@@ -72,20 +72,22 @@ def create_model_metadata(
     """
     model_dir_path = os.path.normpath(model_dir_path)
     embed_local_ml_library = kwargs.pop("embed_local_ml_library", False)
-    # Use the last one which is loaded first, that is mean, it is loaded from site-packages.
-    # We could make sure that user does not overwrite our library with their code follow the same naming.
-    snowml_path, snowml_start_path = file_utils.get_package_path(_SNOWFLAKE_ML_PKG_NAME, strategy="last")
-    if os.path.isdir(snowml_start_path):
-        path_to_copy = snowml_path
-    # If the package is zip-imported, then the path will be `../path_to_zip.zip/snowflake/ml`
-    # It is not a valid path in fact and we need to get the path to the zip file to verify it.
-    elif os.path.isfile(snowml_start_path):
-        extract_root = tempfile.mkdtemp()
-        with zipfile.ZipFile(os.path.abspath(snowml_start_path), mode="r", compression=zipfile.ZIP_DEFLATED) as zf:
-            zf.extractall(path=extract_root)
-        path_to_copy = os.path.join(extract_root, *(_SNOWFLAKE_ML_PKG_NAME.split(".")))
-    else:
-        raise ValueError("`snowflake.ml` is imported via a way that embedding local ML library is not supported.")
+    legacy_save = kwargs.pop("_legacy_save", False)
+    if embed_local_ml_library:
+        # Use the last one which is loaded first, that is mean, it is loaded from site-packages.
+        # We could make sure that user does not overwrite our library with their code follow the same naming.
+        snowml_path, snowml_start_path = file_utils.get_package_path(_SNOWFLAKE_ML_PKG_NAME, strategy="last")
+        if os.path.isdir(snowml_start_path):
+            path_to_copy = snowml_path
+        # If the package is zip-imported, then the path will be `../path_to_zip.zip/snowflake/ml`
+        # It is not a valid path in fact and we need to get the path to the zip file to verify it.
+        elif os.path.isfile(snowml_start_path):
+            extract_root = tempfile.mkdtemp()
+            with zipfile.ZipFile(os.path.abspath(snowml_start_path), mode="r", compression=zipfile.ZIP_DEFLATED) as zf:
+                zf.extractall(path=extract_root)
+            path_to_copy = os.path.join(extract_root, *(_SNOWFLAKE_ML_PKG_NAME.split(".")))
+        else:
+            raise ValueError("`snowflake.ml` is imported via a way that embedding local ML library is not supported.")
 
     env = _create_env_for_model_metadata(
         conda_dependencies=conda_dependencies,
@@ -106,10 +108,10 @@ def create_model_metadata(
     )
 
     code_dir_path = os.path.join(model_dir_path, MODEL_CODE_DIR)
-    if embed_local_ml_library or code_paths:
+    if (embed_local_ml_library and legacy_save) or code_paths:
         os.makedirs(code_dir_path, exist_ok=True)
 
-    if embed_local_ml_library:
+    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)

snowflake/ml/model/_signatures/snowpark_handler.py

@@ -51,7 +51,7 @@ class SnowparkDataFrameHandler(base_handler.BaseDataHandler[snowflake.snowpark.D
         data: snowflake.snowpark.DataFrame, role: Literal["input", "output"]
     ) -> Sequence[core.BaseFeatureSpec]:
         return pandas_handler.PandasDataFrameHandler.infer_signature(
-            SnowparkDataFrameHandler.convert_to_df(data), role=role
+            SnowparkDataFrameHandler.convert_to_df(data.limit(n=1)), role=role
         )
 
     @staticmethod