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

snowflake/ml/jobs/manager.py

@@ -0,0 +1,298 @@
+import pathlib
+import textwrap
+from typing import Any, Callable, Dict, List, Literal, Optional, Union
+from uuid import uuid4
+
+import yaml
+
+from snowflake import snowpark
+from snowflake.ml._internal import telemetry
+from snowflake.ml._internal.utils import identifier
+from snowflake.ml.jobs import job as jb
+from snowflake.ml.jobs._utils import payload_utils, spec_utils
+from snowflake.snowpark.context import get_active_session
+from snowflake.snowpark.exceptions import SnowparkSQLException
+
+_PROJECT = "MLJob"
+JOB_ID_PREFIX = "MLJOB_"
+
+
+@snowpark._internal.utils.private_preview(version="1.7.4")
+@telemetry.send_api_usage_telemetry(project=_PROJECT, func_params_to_log=["limit", "scope"])
+def list_jobs(
+    limit: int = 10,
+    scope: Union[Literal["account", "database", "schema"], str, None] = None,
+    session: Optional[snowpark.Session] = None,
+) -> snowpark.DataFrame:
+    """
+    Returns a Snowpark DataFrame with the list of jobs in the current session.
+
+    Args:
+        limit: The maximum number of jobs to return. Non-positive values are treated as no limit.
+        scope: The scope to list jobs from, such as "schema" or "compute pool <pool_name>".
+        session: The Snowpark session to use. If none specified, uses active session.
+
+    Returns:
+        A DataFrame with the list of jobs.
+
+    Examples:
+        >>> from snowflake.ml.jobs import list_jobs
+        >>> list_jobs(limit=5).show()
+    """
+    session = session or get_active_session()
+    query = "SHOW JOB SERVICES"
+    query += f" LIKE '{JOB_ID_PREFIX}%'"
+    if scope:
+        query += f" IN {scope}"
+    if limit > 0:
+        query += f" LIMIT {limit}"
+    df = session.sql(query)
+    df = df.select(
+        df['"name"'].alias('"id"'),
+        df['"owner"'],
+        df['"status"'],
+        df['"created_on"'],
+        df['"compute_pool"'],
+    ).order_by('"created_on"', ascending=False)
+    return df
+
+
+@snowpark._internal.utils.private_preview(version="1.7.4")
+@telemetry.send_api_usage_telemetry(project=_PROJECT)
+def get_job(job_id: str, session: Optional[snowpark.Session] = None) -> jb.MLJob:
+    """Retrieve a job service from the backend."""
+    session = session or get_active_session()
+
+    try:
+        # Validate job_id
+        job_id = identifier.resolve_identifier(job_id)
+    except ValueError as e:
+        raise ValueError(f"Invalid job ID: {job_id}") from e
+
+    try:
+        # Validate that job exists by doing a status check
+        job = jb.MLJob(job_id, session=session)
+        _ = job.status
+        return job
+    except SnowparkSQLException as e:
+        if "does not exist" in e.message:
+            raise ValueError(f"Job does not exist: {job_id}") from e
+        raise
+
+
+@snowpark._internal.utils.private_preview(version="1.7.4")
+@telemetry.send_api_usage_telemetry(project=_PROJECT)
+def delete_job(job: Union[str, jb.MLJob], session: Optional[snowpark.Session] = None) -> None:
+    """Delete a job service from the backend. Status and logs will be lost."""
+    if isinstance(job, jb.MLJob):
+        job_id = job.id
+        session = job._session or session
+    else:
+        job_id = job
+    session = session or get_active_session()
+    session.sql("DROP SERVICE IDENTIFIER(?)", params=(job_id,)).collect()
+
+
+@snowpark._internal.utils.private_preview(version="1.7.4")
+@telemetry.send_api_usage_telemetry(project=_PROJECT)
+def submit_file(
+    file_path: str,
+    compute_pool: str,
+    *,
+    stage_name: str,
+    args: Optional[List[str]] = None,
+    env_vars: Optional[Dict[str, str]] = None,
+    pip_requirements: Optional[List[str]] = None,
+    external_access_integrations: Optional[List[str]] = None,
+    query_warehouse: Optional[str] = None,
+    spec_overrides: Optional[Dict[str, Any]] = None,
+    session: Optional[snowpark.Session] = None,
+) -> jb.MLJob:
+    """
+    Submit a Python file as a job to the compute pool.
+
+    Args:
+        file_path: The path to the file containing the source code for the job.
+        compute_pool: The compute pool to use for the job.
+        stage_name: The name of the stage where the job payload will be uploaded.
+        args: A list of arguments to pass to the job.
+        env_vars: Environment variables to set in container
+        pip_requirements: A list of pip requirements for the job.
+        external_access_integrations: A list of external access integrations.
+        query_warehouse: The query warehouse to use. Defaults to session warehouse.
+        spec_overrides: Custom service specification overrides to apply.
+        session: The Snowpark session to use. If none specified, uses active session.
+
+    Returns:
+        An object representing the submitted job.
+    """
+    return _submit_job(
+        source=file_path,
+        args=args,
+        compute_pool=compute_pool,
+        stage_name=stage_name,
+        env_vars=env_vars,
+        pip_requirements=pip_requirements,
+        external_access_integrations=external_access_integrations,
+        query_warehouse=query_warehouse,
+        spec_overrides=spec_overrides,
+        session=session,
+    )
+
+
+@snowpark._internal.utils.private_preview(version="1.7.4")
+@telemetry.send_api_usage_telemetry(project=_PROJECT)
+def submit_directory(
+    dir_path: str,
+    compute_pool: str,
+    *,
+    entrypoint: str,
+    stage_name: str,
+    args: Optional[List[str]] = None,
+    env_vars: Optional[Dict[str, str]] = None,
+    pip_requirements: Optional[List[str]] = None,
+    external_access_integrations: Optional[List[str]] = None,
+    query_warehouse: Optional[str] = None,
+    spec_overrides: Optional[Dict[str, Any]] = None,
+    session: Optional[snowpark.Session] = None,
+) -> jb.MLJob:
+    """
+    Submit a directory containing Python script(s) as a job to the compute pool.
+
+    Args:
+        dir_path: The path to the directory containing the job payload.
+        compute_pool: The compute pool to use for the job.
+        entrypoint: The relative path to the entry point script inside the source directory.
+        stage_name: The name of the stage where the job payload will be uploaded.
+        args: A list of arguments to pass to the job.
+        env_vars: Environment variables to set in container
+        pip_requirements: A list of pip requirements for the job.
+        external_access_integrations: A list of external access integrations.
+        query_warehouse: The query warehouse to use. Defaults to session warehouse.
+        spec_overrides: Custom service specification overrides to apply.
+        session: The Snowpark session to use. If none specified, uses active session.
+
+    Returns:
+        An object representing the submitted job.
+    """
+    return _submit_job(
+        source=dir_path,
+        entrypoint=entrypoint,
+        args=args,
+        compute_pool=compute_pool,
+        stage_name=stage_name,
+        env_vars=env_vars,
+        pip_requirements=pip_requirements,
+        external_access_integrations=external_access_integrations,
+        query_warehouse=query_warehouse,
+        spec_overrides=spec_overrides,
+        session=session,
+    )
+
+
+@telemetry.send_api_usage_telemetry(
+    project=_PROJECT,
+    func_params_to_log=[
+        # TODO: Log the source type (callable, file, directory, etc)
+        # TODO: Log instance type of compute pool used
+        # TODO: Log lengths of args, env_vars, and spec_overrides values
+        "pip_requirements",
+        "external_access_integrations",
+    ],
+)
+def _submit_job(
+    source: Union[str, Callable[..., Any]],
+    compute_pool: str,
+    *,
+    stage_name: str,
+    entrypoint: Optional[str] = None,
+    args: Optional[List[str]] = None,
+    env_vars: Optional[Dict[str, str]] = None,
+    pip_requirements: Optional[List[str]] = None,
+    external_access_integrations: Optional[List[str]] = None,
+    query_warehouse: Optional[str] = None,
+    spec_overrides: Optional[Dict[str, Any]] = None,
+    session: Optional[snowpark.Session] = None,
+) -> jb.MLJob:
+    """
+    Submit a job to the compute pool.
+
+    Args:
+        source: The file/directory path containing payload source code or a serializable Python callable.
+        compute_pool: The compute pool to use for the job.
+        stage_name: The name of the stage where the job payload will be uploaded.
+        entrypoint: The entry point for the job execution. Required if source is a directory.
+        args: A list of arguments to pass to the job.
+        env_vars: Environment variables to set in container
+        pip_requirements: A list of pip requirements for the job.
+        external_access_integrations: A list of external access integrations.
+        query_warehouse: The query warehouse to use. Defaults to session warehouse.
+        spec_overrides: Custom service specification overrides to apply.
+        session: The Snowpark session to use. If none specified, uses active session.
+
+    Returns:
+        An object representing the submitted job.
+
+    Raises:
+        RuntimeError: If required Snowflake features are not enabled.
+    """
+    session = session or get_active_session()
+    job_id = f"{JOB_ID_PREFIX}{str(uuid4()).replace('-', '_').upper()}"
+    stage_name = "@" + stage_name.lstrip("@").rstrip("/")
+    stage_path = pathlib.PurePosixPath(f"{stage_name}/{job_id}")
+
+    # Upload payload
+    uploaded_payload = payload_utils.JobPayload(
+        source,
+        entrypoint=entrypoint,
+        pip_requirements=pip_requirements,
+    ).upload(session, stage_path)
+
+    # Generate service spec
+    spec = spec_utils.generate_service_spec(
+        session,
+        compute_pool=compute_pool,
+        payload=uploaded_payload,
+        args=args,
+    )
+    spec_overrides = spec_utils.generate_spec_overrides(
+        environment_vars=env_vars,
+        custom_overrides=spec_overrides,
+    )
+    if spec_overrides:
+        spec = spec_utils.merge_patch(spec, spec_overrides, display_name="spec_overrides")
+
+    # Generate SQL command for job submission
+    query_template = textwrap.dedent(
+        f"""\
+        EXECUTE JOB SERVICE
+        IN COMPUTE POOL {compute_pool}
+        FROM SPECIFICATION $$
+        {{}}
+        $$
+        NAME = {job_id}
+        ASYNC = TRUE
+        """
+    )
+    query = query_template.format(yaml.dump(spec)).splitlines()
+    if external_access_integrations:
+        external_access_integration_list = ",".join(f"{e}" for e in external_access_integrations)
+        query.append(f"EXTERNAL_ACCESS_INTEGRATIONS = ({external_access_integration_list})")
+    query_warehouse = query_warehouse or session.get_current_warehouse()
+    if query_warehouse:
+        query.append(f"QUERY_WAREHOUSE = {query_warehouse}")
+
+    # Submit job
+    query_text = "\n".join(line for line in query if line)
+
+    try:
+        _ = session.sql(query_text).collect()
+    except SnowparkSQLException as e:
+        if "invalid property 'ASYNC'" in e.message:
+            raise RuntimeError(
+                "SPCS Async Jobs not enabled. Set parameter `ENABLE_SNOWSERVICES_ASYNC_JOBS = TRUE` to enable."
+            ) from e
+        raise
+
+    # TODO: Wrap snowflake.core.service.JobService object
+    return jb.MLJob(job_id, session=session)

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

@@ -33,6 +33,7 @@ from snowflake.snowpark._internal import utils as snowpark_utils
 
 class ServiceInfo(TypedDict):
     name: str
+    status: str
     inference_endpoint: Optional[str]
 
 
@@ -550,9 +551,13 @@ class ModelOperator:
         fully_qualified_service_names = [str(service) for service in json_array if "MODEL_BUILD_" not in service]
 
         result = []
-        ingress_url: Optional[str] = None
+
         for fully_qualified_service_name in fully_qualified_service_names:
+            ingress_url: Optional[str] = None
             db, schema, service_name = sql_identifier.parse_fully_qualified_name(fully_qualified_service_name)
+            service_status, _ = self._service_client.get_service_status(
+                database_name=db, schema_name=schema, service_name=service_name, statement_params=statement_params
+            )
             for res_row in self._service_client.show_endpoints(
                 database_name=db, schema_name=schema, service_name=service_name, statement_params=statement_params
             ):
@@ -566,7 +571,11 @@ class ModelOperator:
                     )
                     if not ingress_url.endswith(ModelOperator.INGRESS_ENDPOINT_URL_SUFFIX):
                         ingress_url = None
-            result.append(ServiceInfo(name=fully_qualified_service_name, inference_endpoint=ingress_url))
+            result.append(
+                ServiceInfo(
+                    name=fully_qualified_service_name, status=service_status.value, inference_endpoint=ingress_url
+                )
+            )
 
         return result
 

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

@@ -8,11 +8,9 @@ import threading
 import time
 from typing import Any, Dict, List, Optional, Tuple, Union, cast
 
-from packaging import version
-
 from snowflake import snowpark
 from snowflake.ml._internal import file_utils
-from snowflake.ml._internal.utils import service_logger, snowflake_env, sql_identifier
+from snowflake.ml._internal.utils import service_logger, sql_identifier
 from snowflake.ml.model._client.service import model_deployment_spec
 from snowflake.ml.model._client.sql import service as service_sql, stage as stage_sql
 from snowflake.snowpark import async_job, exceptions, row, session
@@ -133,14 +131,6 @@ class ServiceOperator:
         )
         stage_path = self._stage_client.fully_qualified_object_name(database_name, schema_name, stage_name)
 
-        # TODO(hayu): Remove the version check after Snowflake 8.40.0 release
-        if (
-            snowflake_env.get_current_snowflake_version(self._session, statement_params=statement_params)
-            < version.parse("8.40.0")
-            and build_external_access_integrations is None
-        ):
-            raise ValueError("External access integrations are required in Snowflake < 8.40.0.")
-
         self._model_deployment_spec.save(
             database_name=database_name,
             schema_name=schema_name,

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

@@ -4,6 +4,7 @@ import textwrap
 from typing import Any, Dict, List, Optional, Tuple
 
 from snowflake import snowpark
+from snowflake.ml._internal import platform_capabilities
 from snowflake.ml._internal.utils import (
     identifier,
     query_result_checker,
@@ -120,12 +121,18 @@ class ServiceSQLClient(_base._BaseSQLClient):
             args_sql_list.append(input_arg_value)
         args_sql = ", ".join(args_sql_list)
 
-        function_name = identifier.concat_names([service_name.identifier(), "_", method_name.identifier()])
-        fully_qualified_function_name = identifier.get_schema_level_object_identifier(
-            actual_database_name.identifier(),
-            actual_schema_name.identifier(),
-            function_name,
-        )
+        if platform_capabilities.PlatformCapabilities.get_instance().is_nested_function_enabled():
+            fully_qualified_service_name = self.fully_qualified_object_name(
+                actual_database_name, actual_schema_name, service_name
+            )
+            fully_qualified_function_name = f"{fully_qualified_service_name}!{method_name.identifier()}"
+        else:
+            function_name = identifier.concat_names([service_name.identifier(), "_", method_name.identifier()])
+            fully_qualified_function_name = identifier.get_schema_level_object_identifier(
+                actual_database_name.identifier(),
+                actual_schema_name.identifier(),
+                function_name,
+            )
 
         sql = textwrap.dedent(
             f"""{with_sql}

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

@@ -38,8 +38,10 @@ def _is_callable(model: model_types.SupportedModelType, method_name: str) -> boo
     return callable(getattr(model, method_name, None))
 
 
-def get_truncated_sample_data(sample_input_data: model_types.SupportedDataType) -> model_types.SupportedLocalDataType:
-    trunc_sample_input = model_signature._truncate_data(sample_input_data)
+def get_truncated_sample_data(
+    sample_input_data: model_types.SupportedDataType, length: int = 100
+) -> model_types.SupportedLocalDataType:
+    trunc_sample_input = model_signature._truncate_data(sample_input_data, length=length)
     local_sample_input: model_types.SupportedLocalDataType = None
     if isinstance(sample_input_data, SnowparkDataFrame):
         # Added because of Any from missing stubs.
@@ -78,7 +80,14 @@ def validate_signature(
     local_sample_input = get_truncated_sample_data(sample_input_data)
     for target_method in target_methods:
         predictions_df = get_prediction_fn(target_method, local_sample_input)
-        sig = model_signature.infer_signature(local_sample_input, predictions_df)
+        sig = model_signature.infer_signature(
+            sample_input_data,
+            predictions_df,
+            input_feature_names=None,
+            output_feature_names=None,
+            input_data_limit=100,
+            output_data_limit=100,
+        )
         model_meta.signatures[target_method] = sig
 
     return model_meta

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

@@ -66,7 +66,7 @@ class CustomModelHandler(_base.BaseModelHandler["custom_model.CustomModel"]):
                 sample_input_data = model_signature._convert_local_data_to_df(sample_input_data)
 
             if inspect.iscoroutinefunction(target_method):
-                with anyio.start_blocking_portal() as portal:
+                with anyio.from_thread.start_blocking_portal() as portal:
                     predictions_df = portal.call(target_method, model, sample_input_data)
             else:
                 predictions_df = target_method(model, sample_input_data)
@@ -98,7 +98,6 @@ class CustomModelHandler(_base.BaseModelHandler["custom_model.CustomModel"]):
         if model.context.model_refs:
             for sub_name, model_ref in model.context.model_refs.items():
                 handler = model_handler.find_handler(model_ref.model)
-                assert handler is not None
                 if handler is None:
                     raise TypeError("Your input type to custom model is not currently supported")
                 sub_model = handler.cast_model(model_ref.model)

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

@@ -323,6 +323,7 @@ class HuggingFacePipelineHandler(
                 model_blob_options["task"],
                 model=model_blob_file_or_dir_path,
                 trust_remote_code=True,
+                torch_dtype="auto",
                 **device_config,
             )
 

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

@@ -1,2 +1,2 @@
-REQUIREMENTS = ['absl-py>=0.15,<2', 'aiohttp!=4.0.0a0, !=4.0.0a1', 'anyio>=3.5.0,<4', 'cachetools>=3.1.1,<6', 'cloudpickle>=2.0.0', 'cryptography', 'fsspec>=2024.6.1,<2026', 'importlib_resources>=6.1.1, <7', 'numpy>=1.23,<2', 'packaging>=20.9,<25', 'pandas>=1.0.0,<3', 'pyarrow', 'pyjwt>=2.0.0, <3', 'pytimeparse>=1.1.8,<2', 'pyyaml>=6.0,<7', 'requests', 'retrying>=1.3.3,<2', 's3fs>=2024.6.1,<2026', 'scikit-learn>=1.4,<1.6', 'scipy>=1.9,<2', 'snowflake-connector-python>=3.5.0,<4', 'snowflake-snowpark-python>=1.17.0,<2', 'sqlparse>=0.4,<1', 'typing-extensions>=4.1.0,<5', 'xgboost>=1.7.3,<3']
-ALL_REQUIREMENTS=['absl-py>=0.15,<2', 'aiohttp!=4.0.0a0, !=4.0.0a1', 'anyio>=3.5.0,<4', 'cachetools>=3.1.1,<6', 'catboost>=1.2.0, <2', 'cloudpickle>=2.0.0', 'cryptography', 'fsspec>=2024.6.1,<2026', 'huggingface_hub<0.26', 'importlib_resources>=6.1.1, <7', 'lightgbm>=4.1.0, <5', 'mlflow>=2.16.0, <3', 'numpy>=1.23,<2', 'packaging>=20.9,<25', 'pandas>=1.0.0,<3', 'pyarrow', 'pyjwt>=2.0.0, <3', 'pytimeparse>=1.1.8,<2', 'pytorch>=2.0.1,<2.3.0', 'pyyaml>=6.0,<7', 'requests', 'retrying>=1.3.3,<2', 's3fs>=2024.6.1,<2026', 'scikit-learn>=1.4,<1.6', 'scipy>=1.9,<2', 'sentence-transformers>=2.2.2,<3', 'sentencepiece>=0.1.95,<1', 'shap>=0.46.0,<1', 'snowflake-connector-python>=3.5.0,<4', 'snowflake-snowpark-python>=1.17.0,<2', 'sqlparse>=0.4,<1', 'tensorflow>=2.12.0,<3', 'tokenizers>=0.10,<1', 'torchdata>=0.4,<1', 'transformers>=4.32.1,<5', 'typing-extensions>=4.1.0,<5', 'xgboost>=1.7.3,<3']
+REQUIREMENTS = ['absl-py>=0.15,<2', 'aiohttp!=4.0.0a0, !=4.0.0a1', 'anyio>=3.5.0,<5', 'cachetools>=3.1.1,<6', 'cloudpickle>=2.0.0', 'cryptography', 'fsspec>=2024.6.1,<2026', 'importlib_resources>=6.1.1, <7', 'numpy>=1.23,<2', 'packaging>=20.9,<25', 'pandas>=1.0.0,<3', 'pyarrow', 'pyjwt>=2.0.0, <3', 'pytimeparse>=1.1.8,<2', 'pyyaml>=6.0,<7', 'requests', 'retrying>=1.3.3,<2', 's3fs>=2024.6.1,<2026', 'scikit-learn>=1.4,<1.6', 'scipy>=1.9,<2', 'snowflake-connector-python>=3.5.0,<4', 'snowflake-snowpark-python>=1.17.0,<2,!=1.26.0', 'sqlparse>=0.4,<1', 'typing-extensions>=4.1.0,<5', 'xgboost>=1.7.3,<3']
+ALL_REQUIREMENTS=['absl-py>=0.15,<2', 'aiohttp!=4.0.0a0, !=4.0.0a1', 'anyio>=3.5.0,<5', 'cachetools>=3.1.1,<6', 'catboost>=1.2.0, <2', 'cloudpickle>=2.0.0', 'cryptography', 'fsspec>=2024.6.1,<2026', 'huggingface_hub<0.26', 'importlib_resources>=6.1.1, <7', 'lightgbm>=4.1.0, <5', 'mlflow>=2.16.0, <3', 'numpy>=1.23,<2', 'packaging>=20.9,<25', 'pandas>=1.0.0,<3', 'pyarrow', 'pyjwt>=2.0.0, <3', 'pytimeparse>=1.1.8,<2', 'pytorch>=2.0.1,<2.3.0', 'pyyaml>=6.0,<7', 'requests', 'retrying>=1.3.3,<2', 's3fs>=2024.6.1,<2026', 'scikit-learn>=1.4,<1.6', 'scipy>=1.9,<2', 'sentence-transformers>=2.2.2,<3', 'sentencepiece>=0.1.95,<1', 'shap>=0.46.0,<1', 'snowflake-connector-python>=3.5.0,<4', 'snowflake-snowpark-python>=1.17.0,<2,!=1.26.0', 'sqlparse>=0.4,<1', 'tensorflow>=2.12.0,<3', 'tokenizers>=0.10,<1', 'torchdata>=0.4,<1', 'transformers>=4.32.1,<5', 'typing-extensions>=4.1.0,<5', 'xgboost>=1.7.3,<3']

snowflake/ml/model/_signatures/base_handler.py

@@ -12,7 +12,6 @@ class BaseDataHandler(ABC, Generic[model_types._DataType]):
     FEATURE_PREFIX: Final[str] = "feature"
     INPUT_PREFIX: Final[str] = "input"
     OUTPUT_PREFIX: Final[str] = "output"
-    SIG_INFER_ROWS_COUNT_LIMIT: Final[int] = 10
 
     @staticmethod
     @abstractmethod
@@ -26,7 +25,7 @@ class BaseDataHandler(ABC, Generic[model_types._DataType]):
 
     @staticmethod
     @abstractmethod
-    def truncate(data: model_types._DataType) -> model_types._DataType:
+    def truncate(data: model_types._DataType, length: int) -> model_types._DataType:
         ...
 
     @staticmethod

snowflake/ml/model/_signatures/builtins_handler.py

@@ -35,8 +35,8 @@ class ListOfBuiltinHandler(base_handler.BaseDataHandler[model_types._SupportedBu
         return len(data)
 
     @staticmethod
-    def truncate(data: model_types._SupportedBuiltinsList) -> model_types._SupportedBuiltinsList:
-        return data[: min(ListOfBuiltinHandler.count(data), ListOfBuiltinHandler.SIG_INFER_ROWS_COUNT_LIMIT)]
+    def truncate(data: model_types._SupportedBuiltinsList, length: int) -> model_types._SupportedBuiltinsList:
+        return data[: min(ListOfBuiltinHandler.count(data), length)]
 
     @staticmethod
     def validate(data: model_types._SupportedBuiltinsList) -> None:

snowflake/ml/model/_signatures/numpy_handler.py

@@ -23,8 +23,8 @@ class NumpyArrayHandler(base_handler.BaseDataHandler[model_types._SupportedNumpy
         return data.shape[0]
 
     @staticmethod
-    def truncate(data: model_types._SupportedNumpyArray) -> model_types._SupportedNumpyArray:
-        return data[: min(NumpyArrayHandler.count(data), NumpyArrayHandler.SIG_INFER_ROWS_COUNT_LIMIT)]
+    def truncate(data: model_types._SupportedNumpyArray, length: int) -> model_types._SupportedNumpyArray:
+        return data[: min(NumpyArrayHandler.count(data), length)]
 
     @staticmethod
     def validate(data: model_types._SupportedNumpyArray) -> None:
@@ -94,11 +94,10 @@ class SeqOfNumpyArrayHandler(base_handler.BaseDataHandler[Sequence[model_types._
         return min(NumpyArrayHandler.count(data_col) for data_col in data)
 
     @staticmethod
-    def truncate(data: Sequence[model_types._SupportedNumpyArray]) -> Sequence[model_types._SupportedNumpyArray]:
-        return [
-            data_col[: min(SeqOfNumpyArrayHandler.count(data), SeqOfNumpyArrayHandler.SIG_INFER_ROWS_COUNT_LIMIT)]
-            for data_col in data
-        ]
+    def truncate(
+        data: Sequence[model_types._SupportedNumpyArray], length: int
+    ) -> Sequence[model_types._SupportedNumpyArray]:
+        return [data_col[: min(SeqOfNumpyArrayHandler.count(data), length)] for data_col in data]
 
     @staticmethod
     def validate(data: Sequence[model_types._SupportedNumpyArray]) -> None:

snowflake/ml/model/_signatures/pandas_handler.py

@@ -23,8 +23,8 @@ class PandasDataFrameHandler(base_handler.BaseDataHandler[pd.DataFrame]):
         return len(data.index)
 
     @staticmethod
-    def truncate(data: pd.DataFrame) -> pd.DataFrame:
-        return data.head(min(PandasDataFrameHandler.count(data), PandasDataFrameHandler.SIG_INFER_ROWS_COUNT_LIMIT))
+    def truncate(data: pd.DataFrame, length: int) -> pd.DataFrame:
+        return data.head(min(PandasDataFrameHandler.count(data), length))
 
     @staticmethod
     def validate(data: Union[pd.DataFrame, pd.Series]) -> None:

snowflake/ml/model/_signatures/pytorch_handler.py

@@ -33,11 +33,8 @@ class SeqOfPyTorchTensorHandler(base_handler.BaseDataHandler[Sequence["torch.Ten
         return min(data_col.shape[0] for data_col in data)  # type: ignore[no-any-return]
 
     @staticmethod
-    def truncate(data: Sequence["torch.Tensor"]) -> Sequence["torch.Tensor"]:
-        return [
-            data_col[: min(SeqOfPyTorchTensorHandler.count(data), SeqOfPyTorchTensorHandler.SIG_INFER_ROWS_COUNT_LIMIT)]
-            for data_col in data
-        ]
+    def truncate(data: Sequence["torch.Tensor"], length: int) -> Sequence["torch.Tensor"]:
+        return [data_col[: min(SeqOfPyTorchTensorHandler.count(data), 10)] for data_col in data]
 
     @staticmethod
     def validate(data: Sequence["torch.Tensor"]) -> None:

snowflake/ml/model/_signatures/snowpark_handler.py

@@ -29,8 +29,8 @@ class SnowparkDataFrameHandler(base_handler.BaseDataHandler[snowflake.snowpark.D
         return data.count()
 
     @staticmethod
-    def truncate(data: snowflake.snowpark.DataFrame) -> snowflake.snowpark.DataFrame:
-        return cast(snowflake.snowpark.DataFrame, data.limit(SnowparkDataFrameHandler.SIG_INFER_ROWS_COUNT_LIMIT))
+    def truncate(data: snowflake.snowpark.DataFrame, length: int) -> snowflake.snowpark.DataFrame:
+        return cast(snowflake.snowpark.DataFrame, data.limit(length))
 
     @staticmethod
     def validate(data: snowflake.snowpark.DataFrame) -> None:
@@ -52,7 +52,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.limit(n=1)), role=role
+            SnowparkDataFrameHandler.convert_to_df(data), role=role
         )
 
     @staticmethod

snowflake/ml/model/_signatures/tensorflow_handler.py

@@ -60,14 +60,9 @@ class SeqOfTensorflowTensorHandler(
 
     @staticmethod
     def truncate(
-        data: Sequence[Union["tensorflow.Tensor", "tensorflow.Variable"]]
+        data: Sequence[Union["tensorflow.Tensor", "tensorflow.Variable"]], length: int
     ) -> Sequence[Union["tensorflow.Tensor", "tensorflow.Variable"]]:
-        return [
-            data_col[
-                : min(SeqOfTensorflowTensorHandler.count(data), SeqOfTensorflowTensorHandler.SIG_INFER_ROWS_COUNT_LIMIT)
-            ]
-            for data_col in data
-        ]
+        return [data_col[: min(SeqOfTensorflowTensorHandler.count(data), length)] for data_col in data]
 
     @staticmethod
     def validate(data: Sequence[Union["tensorflow.Tensor", "tensorflow.Variable"]]) -> None:

snowflake/ml/model/model_signature.py

@@ -59,11 +59,16 @@ _ALL_DATA_HANDLERS = _LOCAL_DATA_HANDLERS + [snowpark_handler.SnowparkDataFrameH
 
 def _truncate_data(
     data: model_types.SupportedDataType,
+    length: Optional[int] = 100,
 ) -> model_types.SupportedDataType:
     for handler in _ALL_DATA_HANDLERS:
         if handler.can_handle(data):
+            # If length is None, return the original data
+            if length is None:
+                return data
+
             row_count = handler.count(data)
-            if row_count <= handler.SIG_INFER_ROWS_COUNT_LIMIT:
+            if row_count <= length:
                 return data
 
             warnings.warn(
@@ -77,7 +82,7 @@ def _truncate_data(
                 category=UserWarning,
                 stacklevel=1,
             )
-            return handler.truncate(data)
+            return handler.truncate(data, length)
     raise snowml_exceptions.SnowflakeMLException(
         error_code=error_codes.NOT_IMPLEMENTED,
         original_exception=NotImplementedError(
@@ -687,6 +692,8 @@ def infer_signature(
     output_data: model_types.SupportedLocalDataType,
     input_feature_names: Optional[List[str]] = None,
     output_feature_names: Optional[List[str]] = None,
+    input_data_limit: Optional[int] = 100,
+    output_data_limit: Optional[int] = 100,
 ) -> core.ModelSignature:
     """
     Infer model signature from given input and output sample data.
@@ -710,12 +717,18 @@ def infer_signature(
         output_data: Sample output data for the model.
         input_feature_names: Names for input features. Defaults to None.
         output_feature_names: Names for output features. Defaults to None.
+        input_data_limit: Limit the number of rows to be used in signature inference in the input data. Defaults to 100.
+            If None, all rows are used. If the number of rows in the input data is less than the limit, all rows are
+            used.
+        output_data_limit: Limit the number of rows to be used in signature inference in the output data. Defaults to
+            100. If None, all rows are used. If the number of rows in the output data is less than the limit, all rows
+            are used.
 
     Returns:
         A model signature inferred from the given input and output sample data.
     """
-    inputs = _infer_signature(input_data, role="input")
+    inputs = _infer_signature(_truncate_data(input_data, input_data_limit), role="input")
     inputs = utils.rename_features(inputs, input_feature_names)
-    outputs = _infer_signature(output_data, role="output")
+    outputs = _infer_signature(_truncate_data(output_data, output_data_limit), role="output")
     outputs = utils.rename_features(outputs, output_feature_names)
     return core.ModelSignature(inputs, outputs)