snowflake-ml-python 1.7.3__py3-none-any.whl → 1.7.5__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_env/model_env.py

@@ -113,7 +113,33 @@ class ModelEnv:
             self._snowpark_ml_version = version.parse(snowpark_ml_version)
 
     def include_if_absent(self, pkgs: List[ModelDependency], check_local_version: bool = False) -> None:
-        """Append requirements into model env if absent.
+        """Append requirements into model env if absent. Depending on the environment, requirements may be added
+        to either the pip requirements or conda dependencies.
+
+        Args:
+            pkgs: A list of ModelDependency namedtuple to be appended.
+            check_local_version: Flag to indicate if it is required to pin to local version. Defaults to False.
+        """
+        if self.pip_requirements and not self.conda_dependencies and pkgs:
+            pip_pkg_reqs: List[str] = []
+            warnings.warn(
+                (
+                    "Dependencies specified from pip requirements."
+                    " This may prevent model deploying to Snowflake Warehouse."
+                ),
+                category=UserWarning,
+                stacklevel=2,
+            )
+            for conda_req_str, pip_name in pkgs:
+                _, conda_req = env_utils._validate_conda_dependency_string(conda_req_str)
+                pip_req = requirements.Requirement(f"{pip_name}{conda_req.specifier}")
+                pip_pkg_reqs.append(str(pip_req))
+            self._include_if_absent_pip(pip_pkg_reqs, check_local_version)
+        else:
+            self._include_if_absent_conda(pkgs, check_local_version)
+
+    def _include_if_absent_conda(self, pkgs: List[ModelDependency], check_local_version: bool = False) -> None:
+        """Append requirements into model env conda dependencies if absent.
 
         Args:
             pkgs: A list of ModelDependency namedtuple to be appended.
@@ -134,8 +160,8 @@ class ModelEnv:
                 if show_warning_message:
                     warnings.warn(
                         (
-                            f"Basic dependency {req_to_add.name} specified from PIP requirements."
-                            + " This may prevent model deploying to Snowflake Warehouse."
+                            f"Basic dependency {req_to_add.name} specified from pip requirements."
+                            " This may prevent model deploying to Snowflake Warehouse."
                         ),
                         category=UserWarning,
                         stacklevel=2,
@@ -157,11 +183,11 @@ class ModelEnv:
                         stacklevel=2,
                     )
 
-    def include_if_absent_pip(self, pkgs: List[str], check_local_version: bool = False) -> None:
-        """Append pip requirements into model env if absent.
+    def _include_if_absent_pip(self, pkgs: List[str], check_local_version: bool = False) -> None:
+        """Append pip requirements into model env pip requirements if absent.
 
         Args:
-            pkgs: A list of string to be appended in pip requirement.
+            pkgs: A list of strings to be appended to pip environment.
             check_local_version: Flag to indicate if it is required to pin to local version. Defaults to False.
         """
 
@@ -187,25 +213,6 @@ class ModelEnv:
                 self._conda_dependencies[channel].remove(spec)
 
     def generate_env_for_cuda(self) -> None:
-        if self.cuda_version is None:
-            return
-
-        cuda_spec = env_utils.find_dep_spec(
-            self._conda_dependencies, self._pip_requirements, conda_pkg_name="cuda", remove_spec=False
-        )
-        if cuda_spec and not cuda_spec.specifier.contains(self.cuda_version):
-            raise ValueError(
-                "The CUDA requirement you specified in your conda dependencies or pip requirements is"
-                " conflicting with CUDA version required. Please do not specify CUDA dependency using conda"
-                " dependencies or pip requirements."
-            )
-
-        if not cuda_spec:
-            self.include_if_absent(
-                [ModelDependency(requirement=f"nvidia::cuda=={self.cuda_version}.*", pip_name="cuda")],
-                check_local_version=False,
-            )
-
         xgboost_spec = env_utils.find_dep_spec(
             self._conda_dependencies, self._pip_requirements, conda_pkg_name="xgboost", remove_spec=True
         )
@@ -236,7 +243,7 @@ class ModelEnv:
                 check_local_version=False,
             )
 
-            self.include_if_absent_pip(["bitsandbytes>=0.41.0"], check_local_version=False)
+            self._include_if_absent_pip(["bitsandbytes>=0.41.0"], check_local_version=False)
 
     def relax_version(self) -> None:
         """Relax the version requirements for both conda dependencies and pip requirements.
@@ -252,7 +259,9 @@ class ModelEnv:
         self._pip_requirements = list(map(env_utils.relax_requirement_version, self._pip_requirements))
 
     def load_from_conda_file(self, conda_env_path: pathlib.Path) -> None:
-        conda_dependencies_dict, pip_requirements_list, python_version = env_utils.load_conda_env_file(conda_env_path)
+        conda_dependencies_dict, pip_requirements_list, python_version, cuda_version = env_utils.load_conda_env_file(
+            conda_env_path
+        )
 
         for channel, channel_dependencies in conda_dependencies_dict.items():
             if channel != env_utils.DEFAULT_CHANNEL_NAME:
@@ -310,6 +319,9 @@ class ModelEnv:
         if python_version:
             self.python_version = python_version
 
+        if cuda_version:
+            self.cuda_version = cuda_version
+
     def load_from_pip_file(self, pip_requirements_path: pathlib.Path) -> None:
         pip_requirements_list = env_utils.load_requirements_file(pip_requirements_path)
 
@@ -342,12 +354,17 @@ class ModelEnv:
         self.snowpark_ml_version = env_dict["snowpark_ml_version"]
 
     def save_as_dict(
-        self, base_dir: pathlib.Path, default_channel_override: str = env_utils.SNOWFLAKE_CONDA_CHANNEL_URL
+        self,
+        base_dir: pathlib.Path,
+        default_channel_override: str = env_utils.SNOWFLAKE_CONDA_CHANNEL_URL,
+        is_gpu: Optional[bool] = False,
     ) -> model_meta_schema.ModelEnvDict:
+        cuda_version = self.cuda_version if is_gpu else None
         env_utils.save_conda_env_file(
             pathlib.Path(base_dir / self.conda_env_rel_path),
             self._conda_dependencies,
             self.python_version,
+            cuda_version,
             default_channel_override=default_channel_override,
         )
         env_utils.save_requirements_file(

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

@@ -38,13 +38,17 @@ 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, is_for_modeling_model: bool = False
+) -> 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.
         trunc_sample_input = cast(SnowparkDataFrame, trunc_sample_input)
         local_sample_input = snowpark_handler.SnowparkDataFrameHandler.convert_to_df(trunc_sample_input)
+        if is_for_modeling_model:
+            local_sample_input.columns = trunc_sample_input.columns
     else:
         local_sample_input = trunc_sample_input
     return local_sample_input
@@ -56,13 +60,15 @@ def validate_signature(
     target_methods: Iterable[str],
     sample_input_data: Optional[model_types.SupportedDataType],
     get_prediction_fn: Callable[[str, model_types.SupportedLocalDataType], model_types.SupportedLocalDataType],
+    is_for_modeling_model: bool = False,
 ) -> model_meta.ModelMetadata:
     if model_meta.signatures:
         validate_target_methods(model, list(model_meta.signatures.keys()))
         if sample_input_data is not None:
-            local_sample_input = get_truncated_sample_data(sample_input_data)
+            local_sample_input = get_truncated_sample_data(
+                sample_input_data, is_for_modeling_model=is_for_modeling_model
+            )
             for target_method in model_meta.signatures.keys():
-
                 model_signature_inst = model_meta.signatures.get(target_method)
                 if model_signature_inst is not None:
                     # strict validation the input signature
@@ -75,10 +81,17 @@ def validate_signature(
     assert (
         sample_input_data is not None
     ), "Model signature and sample input are None at the same time. This should not happen with local model."
-    local_sample_input = get_truncated_sample_data(sample_input_data)
+    local_sample_input = get_truncated_sample_data(sample_input_data, is_for_modeling_model=is_for_modeling_model)
     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

@@ -146,6 +146,10 @@ class HuggingFacePipelineHandler(
             framework = getattr(model, "framework", None)
             batch_size = getattr(model, "batch_size", None)
 
+        has_tokenizer = getattr(model, "tokenizer", None) is not None
+        has_feature_extractor = getattr(model, "feature_extractor", None) is not None
+        has_image_preprocessor = getattr(model, "image_preprocessor", None) is not None
+
         if type_utils.LazyType("transformers.Pipeline").isinstance(model):
             params = {
                 **model._preprocess_params,  # type:ignore[attr-defined]
@@ -234,6 +238,9 @@ class HuggingFacePipelineHandler(
                 {
                     "task": task,
                     "batch_size": batch_size if batch_size is not None else 1,
+                    "has_tokenizer": has_tokenizer,
+                    "has_feature_extractor": has_feature_extractor,
+                    "has_image_preprocessor": has_image_preprocessor,
                 }
             ),
         )
@@ -308,6 +315,14 @@ class HuggingFacePipelineHandler(
         if os.path.isdir(model_blob_file_or_dir_path):
             import transformers
 
+            additional_pipeline_params = {}
+            if model_blob_options.get("has_tokenizer", False):
+                additional_pipeline_params["tokenizer"] = model_blob_file_or_dir_path
+            if model_blob_options.get("has_feature_extractor", False):
+                additional_pipeline_params["feature_extractor"] = model_blob_file_or_dir_path
+            if model_blob_options.get("has_image_preprocessor", False):
+                additional_pipeline_params["image_preprocessor"] = model_blob_file_or_dir_path
+
             with open(
                 os.path.join(
                     model_blob_file_or_dir_path,
@@ -323,6 +338,8 @@ class HuggingFacePipelineHandler(
                 model_blob_options["task"],
                 model=model_blob_file_or_dir_path,
                 trust_remote_code=True,
+                torch_dtype="auto",
+                **additional_pipeline_params,
                 **device_config,
             )