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

snowflake/ml/jobs/decorators.py

@@ -0,0 +1,91 @@
+import copy
+import functools
+import inspect
+from typing import Callable, Dict, List, Optional, TypeVar
+
+from typing_extensions import ParamSpec
+
+from snowflake import snowpark
+from snowflake.ml._internal import telemetry
+from snowflake.ml.jobs import job as jb, manager as jm
+from snowflake.ml.jobs._utils import payload_utils
+
+_PROJECT = "MLJob"
+
+_Args = ParamSpec("_Args")
+_ReturnValue = TypeVar("_ReturnValue")
+
+
+@snowpark._internal.utils.private_preview(version="1.7.4")
+@telemetry.send_api_usage_telemetry(project=_PROJECT)
+def remote(
+    compute_pool: str,
+    stage_name: str,
+    pip_requirements: Optional[List[str]] = None,
+    external_access_integrations: Optional[List[str]] = None,
+    query_warehouse: Optional[str] = None,
+    env_vars: Optional[Dict[str, str]] = None,
+    session: Optional[snowpark.Session] = None,
+) -> Callable[[Callable[_Args, _ReturnValue]], Callable[_Args, jb.MLJob]]:
+    """
+    Submit a job to the compute pool.
+
+    Args:
+        compute_pool: The compute pool to use for the job.
+        stage_name: The name of the stage where the job payload will be uploaded.
+        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.
+        env_vars: Environment variables to set in container
+        session: The Snowpark session to use. If none specified, uses active session.
+
+    Returns:
+        Decorator that dispatches invocations of the decorated function as remote jobs.
+    """
+
+    def decorator(func: Callable[_Args, _ReturnValue]) -> Callable[_Args, jb.MLJob]:
+        # Copy the function to avoid modifying the original
+        # We need to modify the line number of the function to exclude the
+        # decorator from the copied source code
+        wrapped_func = copy.copy(func)
+        wrapped_func.__code__ = wrapped_func.__code__.replace(co_firstlineno=func.__code__.co_firstlineno + 1)
+
+        # Validate function arguments based on signature
+        signature = inspect.signature(func)
+        pos_arg_names = []
+        for name, param in signature.parameters.items():
+            param_type = payload_utils.get_parameter_type(param)
+            if param_type is not None:
+                payload_utils.validate_parameter_type(param_type, name)
+            if param.kind in (param.POSITIONAL_ONLY, param.POSITIONAL_OR_KEYWORD):
+                pos_arg_names.append(name)
+
+        @functools.wraps(func)
+        def wrapper(*args: _Args.args, **kwargs: _Args.kwargs) -> jb.MLJob:
+            # Validate positional args
+            for i, arg in enumerate(args):
+                arg_name = pos_arg_names[i] if i < len(pos_arg_names) else f"args[{i}]"
+                payload_utils.validate_parameter_type(type(arg), arg_name)
+
+            # Validate keyword args
+            for k, v in kwargs.items():
+                payload_utils.validate_parameter_type(type(v), k)
+
+            arg_list = [str(v) for v in args] + [x for k, v in kwargs.items() for x in (f"--{k}", str(v))]
+            job = jm._submit_job(
+                source=wrapped_func,
+                args=arg_list,
+                stage_name=stage_name,
+                compute_pool=compute_pool,
+                pip_requirements=pip_requirements,
+                external_access_integrations=external_access_integrations,
+                query_warehouse=query_warehouse,
+                env_vars=env_vars,
+                session=session,
+            )
+            assert isinstance(job, jb.MLJob)
+            return job
+
+        return wrapper
+
+    return decorator

snowflake/ml/jobs/job.py

@@ -0,0 +1,113 @@
+import time
+from typing import Any, List, Optional, cast
+
+from snowflake import snowpark
+from snowflake.ml._internal import telemetry
+from snowflake.ml.jobs._utils import constants, types
+from snowflake.snowpark.context import get_active_session
+
+_PROJECT = "MLJob"
+TERMINAL_JOB_STATUSES = {"FAILED", "DONE", "INTERNAL_ERROR"}
+
+
+class MLJob:
+    def __init__(self, id: str, session: Optional[snowpark.Session] = None) -> None:
+        self._id = id
+        self._session = session or get_active_session()
+        self._status: types.JOB_STATUS = "PENDING"
+
+    @property
+    def id(self) -> str:
+        """Get the unique job ID"""
+        return self._id
+
+    @property
+    def status(self) -> types.JOB_STATUS:
+        """Get the job's execution status."""
+        if self._status not in TERMINAL_JOB_STATUSES:
+            # Query backend for job status if not in terminal state
+            self._status = _get_status(self._session, self.id)
+        return self._status
+
+    @snowpark._internal.utils.private_preview(version="1.7.4")
+    def get_logs(self, limit: int = -1) -> str:
+        """
+        Return the job's execution logs.
+
+        Args:
+            limit: The maximum number of lines to return. Negative values are treated as no limit.
+
+        Returns:
+            The job's execution logs.
+        """
+        logs = _get_logs(self._session, self.id, limit)
+        assert isinstance(logs, str)  # mypy
+        return logs
+
+    @snowpark._internal.utils.private_preview(version="1.7.4")
+    def show_logs(self, limit: int = -1) -> None:
+        """
+        Display the job's execution logs.
+
+        Args:
+            limit: The maximum number of lines to display. Negative values are treated as no limit.
+        """
+        print(self.get_logs(limit))  # noqa: T201: we need to print here.
+
+    @snowpark._internal.utils.private_preview(version="1.7.4")
+    @telemetry.send_api_usage_telemetry(project=_PROJECT)
+    def wait(self, timeout: float = -1) -> types.JOB_STATUS:
+        """
+        Block until completion. Returns completion status.
+
+        Args:
+            timeout: The maximum time to wait in seconds. Negative values are treated as no timeout.
+
+        Returns:
+            The job's completion status.
+
+        Raises:
+            TimeoutError: If the job does not complete within the specified timeout.
+        """
+        delay = constants.JOB_POLL_INITIAL_DELAY_SECONDS  # Start with 100ms delay
+        start_time = time.monotonic()
+        while self.status not in TERMINAL_JOB_STATUSES:
+            if timeout >= 0 and (elapsed := time.monotonic() - start_time) >= timeout:
+                raise TimeoutError(f"Job {self.id} did not complete within {elapsed} seconds")
+            time.sleep(delay)
+            delay = min(delay * 2, constants.JOB_POLL_MAX_DELAY_SECONDS)  # Exponential backoff
+        return self.status
+
+
+@telemetry.send_api_usage_telemetry(project=_PROJECT)
+def _get_status(session: snowpark.Session, job_id: str) -> types.JOB_STATUS:
+    """Retrieve job execution status."""
+    # TODO: snowflake-snowpark-python<1.24.0 shows spurious error messages on
+    #       `DESCRIBE` queries with bind variables
+    #       Switch to use bind variables instead of client side formatting after
+    #       updating to snowflake-snowpark-python>=1.24.0
+    (row,) = session.sql(f"DESCRIBE SERVICE {job_id}").collect()
+    return cast(types.JOB_STATUS, row["status"])
+
+
+@telemetry.send_api_usage_telemetry(project=_PROJECT)
+def _get_logs(session: snowpark.Session, job_id: str, limit: int = -1) -> str:
+    """
+    Retrieve the job's execution logs.
+
+    Args:
+        job_id: The job ID.
+        limit: The maximum number of lines to return. Negative values are treated as no limit.
+        session: The Snowpark session to use. If none specified, uses active session.
+
+    Returns:
+        The job's execution logs.
+    """
+    params: List[Any] = [job_id]
+    if limit > 0:
+        params.append(limit)
+    (row,) = session.sql(
+        f"SELECT SYSTEM$GET_SERVICE_LOGS(?, 0, '{constants.DEFAULT_CONTAINER_NAME}'{f', ?' if limit > 0 else ''})",
+        params=params,
+    ).collect()
+    return str(row[0])

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/model/model_version_impl.py

@@ -447,13 +447,15 @@ class ModelVersion(lineage_node.LineageNode):
             target_function_info = functions[0]
 
         if service_name:
+            database_name_id, schema_name_id, service_name_id = sql_identifier.parse_fully_qualified_name(service_name)
+
             return self._model_ops.invoke_method(
                 method_name=sql_identifier.SqlIdentifier(target_function_info["name"]),
                 signature=target_function_info["signature"],
                 X=X,
-                database_name=None,
-                schema_name=None,
-                service_name=sql_identifier.SqlIdentifier(service_name),
+                database_name=database_name_id,
+                schema_name=schema_name_id,
+                service_name=service_name_id,
                 strict_input_validation=strict_input_validation,
                 statement_params=statement_params,
             )

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]
 
 
@@ -168,14 +169,10 @@ class ModelOperator:
         schema_name: Optional[sql_identifier.SqlIdentifier],
         model_name: sql_identifier.SqlIdentifier,
         version_name: sql_identifier.SqlIdentifier,
+        model_exists: bool,
         statement_params: Optional[Dict[str, Any]] = None,
     ) -> None:
-        if self.validate_existence(
-            database_name=database_name,
-            schema_name=schema_name,
-            model_name=model_name,
-            statement_params=statement_params,
-        ):
+        if model_exists:
             return self._model_version_client.add_version_from_model_version(
                 source_database_name=source_database_name,
                 source_schema_name=source_schema_name,
@@ -554,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
             ):
@@ -570,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/model_version.py

@@ -10,6 +10,7 @@ from snowflake.ml._internal.utils import (
     sql_identifier,
 )
 from snowflake.ml.model._client.sql import _base
+from snowflake.ml.model._model_composer.model_method import constants
 from snowflake.snowpark import dataframe, functions as F, row, types as spt
 from snowflake.snowpark._internal import utils as snowpark_utils
 
@@ -333,6 +334,11 @@ class ModelVersionSQLClient(_base._BaseSQLClient):
 
         args_sql = ", ".join(args_sql_list)
 
+        wide_input = len(input_args) > constants.SNOWPARK_UDF_INPUT_COL_LIMIT
+        if wide_input:
+            input_args_sql = ", ".join(f"'{arg}', {arg.identifier()}" for arg in input_args)
+            args_sql = f"object_construct_keep_null({input_args_sql})"
+
         sql = textwrap.dedent(
             f"""WITH {','.join(with_statements)}
                 SELECT *,
@@ -412,6 +418,11 @@ class ModelVersionSQLClient(_base._BaseSQLClient):
 
         args_sql = ", ".join(args_sql_list)
 
+        wide_input = len(input_args) > constants.SNOWPARK_UDF_INPUT_COL_LIMIT
+        if wide_input:
+            input_args_sql = ", ".join(f"'{arg}', {arg.identifier()}" for arg in input_args)
+            args_sql = f"object_construct_keep_null({input_args_sql})"
+
         sql = textwrap.dedent(
             f"""WITH {','.join(with_statements)}
                 SELECT *,

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}