snowflake-ml-python 1.11.0__py3-none-any.whl → 1.13.0__py3-none-any.whl

snowflake/ml/jobs/manager.py

@@ -1,6 +1,8 @@
+import json
 import logging
 import pathlib
 import textwrap
+from pathlib import PurePath
 from typing import Any, Callable, Optional, TypeVar, Union, cast, overload
 from uuid import uuid4
 
@@ -11,7 +13,13 @@ 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, query_helper, spec_utils
+from snowflake.ml.jobs._utils import (
+    feature_flags,
+    payload_utils,
+    query_helper,
+    spec_utils,
+    types,
+)
 from snowflake.snowpark.context import get_active_session
 from snowflake.snowpark.exceptions import SnowparkSQLException
 from snowflake.snowpark.functions import coalesce, col, lit, when
@@ -50,7 +58,8 @@ def list_jobs(
         >>> from snowflake.ml.jobs import list_jobs
         >>> list_jobs(limit=5)
     """
-    session = session or get_active_session()
+
+    session = _ensure_session(session)
     try:
         df = _get_job_history_spcs(
             session,
@@ -154,7 +163,7 @@ def _get_job_history_spcs(
 @telemetry.send_api_usage_telemetry(project=_PROJECT)
 def get_job(job_id: str, session: Optional[snowpark.Session] = None) -> jb.MLJob[Any]:
     """Retrieve a job service from the backend."""
-    session = session or get_active_session()
+    session = _ensure_session(session)
     try:
         database, schema, job_name = identifier.parse_schema_level_object_identifier(job_id)
         database = identifier.resolve_identifier(cast(str, database or session.get_current_database()))
@@ -426,8 +435,10 @@ def _submit_job(
 
     Raises:
         ValueError: If database or schema value(s) are invalid
+        RuntimeError: If schema is not specified in session context or job submission
+        snowpark.exceptions.SnowparkSQLException: if failed to upload payload
     """
-    session = session or get_active_session()
+    session = _ensure_session(session)
 
     # Check for deprecated args
     if "num_instances" in kwargs:
@@ -445,7 +456,7 @@ def _submit_job(
     env_vars = kwargs.pop("env_vars", None)
     spec_overrides = kwargs.pop("spec_overrides", None)
     enable_metrics = kwargs.pop("enable_metrics", True)
-    query_warehouse = kwargs.pop("query_warehouse", None)
+    query_warehouse = kwargs.pop("query_warehouse", session.get_current_warehouse())
     additional_payloads = kwargs.pop("additional_payloads", None)
 
     if additional_payloads:
@@ -478,11 +489,39 @@ def _submit_job(
     stage_name = f"@{'.'.join(filter(None, stage_path_parts[:3]))}"
     stage_path = pathlib.PurePosixPath(f"{stage_name}{stage_path_parts[-1].rstrip('/')}/{job_name}")
 
-    # Upload payload
-    uploaded_payload = payload_utils.JobPayload(
-        source, entrypoint=entrypoint, pip_requirements=pip_requirements, additional_payloads=additional_payloads
-    ).upload(session, stage_path)
+    try:
+        # Upload payload
+        uploaded_payload = payload_utils.JobPayload(
+            source, entrypoint=entrypoint, pip_requirements=pip_requirements, additional_payloads=additional_payloads
+        ).upload(session, stage_path)
+    except snowpark.exceptions.SnowparkSQLException as e:
+        if e.sql_error_code == 90106:
+            raise RuntimeError(
+                "Please specify a schema, either in the session context or as a parameter in the job submission"
+            )
+        raise
 
+    if feature_flags.FeatureFlags.USE_SUBMIT_JOB_V2.is_enabled():
+        # Add default env vars (extracted from spec_utils.generate_service_spec)
+        combined_env_vars = {**uploaded_payload.env_vars, **(env_vars or {})}
+
+        return _do_submit_job_v2(
+            session=session,
+            payload=uploaded_payload,
+            args=args,
+            env_vars=combined_env_vars,
+            spec_overrides=spec_overrides,
+            compute_pool=compute_pool,
+            job_id=job_id,
+            external_access_integrations=external_access_integrations,
+            query_warehouse=query_warehouse,
+            target_instances=target_instances,
+            min_instances=min_instances,
+            enable_metrics=enable_metrics,
+            use_async=True,
+        )
+
+    # Fall back to v1
     # Generate service spec
     spec = spec_utils.generate_service_spec(
         session,
@@ -493,6 +532,8 @@ def _submit_job(
         min_instances=min_instances,
         enable_metrics=enable_metrics,
     )
+
+    # Generate spec overrides
     spec_overrides = spec_utils.generate_spec_overrides(
         environment_vars=env_vars,
         custom_overrides=spec_overrides,
@@ -500,26 +541,25 @@ def _submit_job(
     if spec_overrides:
         spec = spec_utils.merge_patch(spec, spec_overrides, display_name="spec_overrides")
 
-    query_text, params = _generate_submission_query(
-        spec, external_access_integrations, query_warehouse, target_instances, session, compute_pool, job_id
+    return _do_submit_job_v1(
+        session, spec, external_access_integrations, query_warehouse, target_instances, compute_pool, job_id
     )
-    _ = query_helper.run_query(session, query_text, params=params)
-    return get_job(job_id, session=session)
 
 
-def _generate_submission_query(
+def _do_submit_job_v1(
+    session: snowpark.Session,
     spec: dict[str, Any],
     external_access_integrations: list[str],
     query_warehouse: Optional[str],
     target_instances: int,
-    session: snowpark.Session,
     compute_pool: str,
     job_id: str,
-) -> tuple[str, list[Any]]:
+) -> jb.MLJob[Any]:
     """
     Generate the SQL query for job submission.
 
     Args:
+        session: The Snowpark session to use.
         spec: The service spec for the job.
         external_access_integrations: The external access integrations for the job.
         query_warehouse: The query warehouse for the job.
@@ -529,7 +569,7 @@ def _generate_submission_query(
         job_id: The ID of the job.
 
     Returns:
-        A tuple containing the SQL query text and the parameters for the query.
+        The job object.
     """
     query_template = textwrap.dedent(
         """\
@@ -547,12 +587,89 @@ def _generate_submission_query(
     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("QUERY_WAREHOUSE = IDENTIFIER(?)")
         params.append(query_warehouse)
     if target_instances > 1:
         query.append("REPLICAS = ?")
         params.append(target_instances)
+
     query_text = "\n".join(line for line in query if line)
-    return query_text, params
+    _ = query_helper.run_query(session, query_text, params=params)
+
+    return get_job(job_id, session=session)
+
+
+def _do_submit_job_v2(
+    session: snowpark.Session,
+    payload: types.UploadedPayload,
+    args: Optional[list[str]],
+    env_vars: dict[str, str],
+    spec_overrides: dict[str, Any],
+    compute_pool: str,
+    job_id: Optional[str] = None,
+    external_access_integrations: Optional[list[str]] = None,
+    query_warehouse: Optional[str] = None,
+    target_instances: int = 1,
+    min_instances: int = 1,
+    enable_metrics: bool = True,
+    use_async: bool = True,
+) -> jb.MLJob[Any]:
+    """
+    Generate the SQL query for job submission.
+
+    Args:
+        session: The Snowpark session to use.
+        payload: The uploaded job payload.
+        args: Arguments to pass to the entrypoint script.
+        env_vars: Environment variables to set in the job container.
+        spec_overrides: Custom service specification overrides.
+        compute_pool: The compute pool to use for job execution.
+        job_id: The ID of the job.
+        external_access_integrations: Optional list of external access integrations.
+        query_warehouse: Optional query warehouse to use.
+        target_instances: Number of instances for multi-node job.
+        min_instances: Minimum number of instances required to start the job.
+        enable_metrics: Whether to enable platform metrics for the job.
+        use_async: Whether to run the job asynchronously.
+
+    Returns:
+        The job object.
+    """
+    args = [
+        (payload.stage_path.joinpath(v).as_posix() if isinstance(v, PurePath) else v) for v in payload.entrypoint
+    ] + (args or [])
+    spec_options = {
+        "STAGE_PATH": payload.stage_path.as_posix(),
+        "ENTRYPOINT": ["/usr/local/bin/_entrypoint.sh"],
+        "ARGS": args,
+        "ENV_VARS": env_vars,
+        "ENABLE_METRICS": enable_metrics,
+        "SPEC_OVERRIDES": spec_overrides,
+    }
+    job_options = {
+        "EXTERNAL_ACCESS_INTEGRATIONS": external_access_integrations,
+        "QUERY_WAREHOUSE": query_warehouse,
+        "TARGET_INSTANCES": target_instances,
+        "MIN_INSTANCES": min_instances,
+        "ASYNC": use_async,
+    }
+    job_options = {k: v for k, v in job_options.items() if v is not None}
+
+    query_template = "CALL SYSTEM$EXECUTE_ML_JOB(?, ?, ?, ?)"
+    params = [job_id, compute_pool, json.dumps(spec_options), json.dumps(job_options)]
+    actual_job_id = query_helper.run_query(session, query_template, params=params)[0][0]
+
+    return get_job(actual_job_id, session=session)
+
+
+def _ensure_session(session: Optional[snowpark.Session]) -> snowpark.Session:
+    try:
+        session = session or get_active_session()
+    except snowpark.exceptions.SnowparkSessionException as e:
+        if "More than one active session" in e.message:
+            raise RuntimeError("Please specify the session as a parameter in API call")
+        if "No default Session is found" in e.message:
+            raise RuntimeError("Please create a session before API call")
+        raise
+    return session

snowflake/ml/model/__init__.py

@@ -1,5 +1,10 @@
+from snowflake.ml.model._client.model.batch_inference_specs import (
+    InputSpec,
+    JobSpec,
+    OutputSpec,
+)
 from snowflake.ml.model._client.model.model_impl import Model
 from snowflake.ml.model._client.model.model_version_impl import ExportMode, ModelVersion
 from snowflake.ml.model.models.huggingface_pipeline import HuggingFacePipelineModel
 
-__all__ = ["Model", "ModelVersion", "ExportMode", "HuggingFacePipelineModel"]
+__all__ = ["Model", "ModelVersion", "ExportMode", "HuggingFacePipelineModel", "InputSpec", "JobSpec", "OutputSpec"]

snowflake/ml/model/_client/model/batch_inference_specs.py

@@ -0,0 +1,25 @@
+from typing import Optional, Union
+
+from pydantic import BaseModel
+
+
+class InputSpec(BaseModel):
+    stage_location: str
+
+
+class OutputSpec(BaseModel):
+    stage_location: str
+
+
+class JobSpec(BaseModel):
+    image_repo: Optional[str] = None
+    job_name: Optional[str] = None
+    num_workers: Optional[int] = None
+    function_name: Optional[str] = None
+    gpu: Optional[Union[str, int]] = None
+    force_rebuild: bool = False
+    max_batch_rows: int = 1024
+    warehouse: Optional[str] = None
+    cpu_requests: Optional[str] = None
+    memory_requests: Optional[str] = None
+    replicas: Optional[int] = None

snowflake/ml/model/_client/model/model_version_impl.py

@@ -1,16 +1,18 @@
 import enum
 import pathlib
 import tempfile
+import uuid
 import warnings
 from typing import Any, Callable, Optional, Union, overload
 
 import pandas as pd
 
-from snowflake import snowpark
+from snowflake.ml import jobs
 from snowflake.ml._internal import telemetry
 from snowflake.ml._internal.utils import sql_identifier
 from snowflake.ml.lineage import lineage_node
 from snowflake.ml.model import task, type_hints
+from snowflake.ml.model._client.model import batch_inference_specs
 from snowflake.ml.model._client.ops import metadata_ops, model_ops, service_ops
 from snowflake.ml.model._model_composer import model_composer
 from snowflake.ml.model._model_composer.model_manifest import model_manifest_schema
@@ -19,6 +21,7 @@ from snowflake.snowpark import Session, async_job, dataframe
 
 _TELEMETRY_PROJECT = "MLOps"
 _TELEMETRY_SUBPROJECT = "ModelManagement"
+_BATCH_INFERENCE_JOB_ID_PREFIX = "BATCH_INFERENCE_"
 
 
 class ExportMode(enum.Enum):
@@ -539,6 +542,64 @@ class ModelVersion(lineage_node.LineageNode):
                 is_partitioned=target_function_info["is_partitioned"],
             )
 
+    @telemetry.send_api_usage_telemetry(
+        project=_TELEMETRY_PROJECT,
+        subproject=_TELEMETRY_SUBPROJECT,
+        func_params_to_log=[
+            "compute_pool",
+        ],
+    )
+    def _run_batch(
+        self,
+        *,
+        compute_pool: str,
+        input_spec: batch_inference_specs.InputSpec,
+        output_spec: batch_inference_specs.OutputSpec,
+        job_spec: Optional[batch_inference_specs.JobSpec] = None,
+    ) -> jobs.MLJob[Any]:
+        statement_params = telemetry.get_statement_params(
+            project=_TELEMETRY_PROJECT,
+            subproject=_TELEMETRY_SUBPROJECT,
+        )
+
+        if job_spec is None:
+            job_spec = batch_inference_specs.JobSpec()
+
+        warehouse = job_spec.warehouse or self._service_ops._session.get_current_warehouse()
+        if warehouse is None:
+            raise ValueError("Warehouse is not set. Please set the warehouse field in the JobSpec.")
+
+        if job_spec.job_name is None:
+            # Same as the MLJob ID generation logic with a different prefix
+            job_name = f"{_BATCH_INFERENCE_JOB_ID_PREFIX}{str(uuid.uuid4()).replace('-', '_').upper()}"
+        else:
+            job_name = job_spec.job_name
+
+        return self._service_ops.invoke_batch_job_method(
+            # model version info
+            model_name=self._model_name,
+            version_name=self._version_name,
+            # job spec
+            function_name=self._get_function_info(function_name=job_spec.function_name)["target_method"],
+            compute_pool_name=sql_identifier.SqlIdentifier(compute_pool),
+            force_rebuild=job_spec.force_rebuild,
+            image_repo_name=job_spec.image_repo,
+            num_workers=job_spec.num_workers,
+            max_batch_rows=job_spec.max_batch_rows,
+            warehouse=sql_identifier.SqlIdentifier(warehouse),
+            cpu_requests=job_spec.cpu_requests,
+            memory_requests=job_spec.memory_requests,
+            job_name=job_name,
+            replicas=job_spec.replicas,
+            # input and output
+            input_stage_location=input_spec.stage_location,
+            input_file_pattern="*",
+            output_stage_location=output_spec.stage_location,
+            completion_filename="_SUCCESS",
+            # misc
+            statement_params=statement_params,
+        )
+
     def _get_function_info(self, function_name: Optional[str]) -> model_manifest_schema.ModelFunctionInfo:
         functions: list[model_manifest_schema.ModelFunctionInfo] = self._functions
 
@@ -1184,69 +1245,5 @@ class ModelVersion(lineage_node.LineageNode):
             statement_params=statement_params,
         )
 
-    @snowpark._internal.utils.private_preview(version="1.8.3")
-    @telemetry.send_api_usage_telemetry(
-        project=_TELEMETRY_PROJECT,
-        subproject=_TELEMETRY_SUBPROJECT,
-    )
-    def _run_job(
-        self,
-        X: Union[pd.DataFrame, "dataframe.DataFrame"],
-        *,
-        job_name: str,
-        compute_pool: str,
-        image_repo: Optional[str] = None,
-        output_table_name: str,
-        function_name: Optional[str] = None,
-        cpu_requests: Optional[str] = None,
-        memory_requests: Optional[str] = None,
-        gpu_requests: Optional[Union[str, int]] = None,
-        num_workers: Optional[int] = None,
-        max_batch_rows: Optional[int] = None,
-        force_rebuild: bool = False,
-        build_external_access_integrations: Optional[list[str]] = None,
-    ) -> Union[pd.DataFrame, dataframe.DataFrame]:
-        statement_params = telemetry.get_statement_params(
-            project=_TELEMETRY_PROJECT,
-            subproject=_TELEMETRY_SUBPROJECT,
-        )
-        target_function_info = self._get_function_info(function_name=function_name)
-        job_db_id, job_schema_id, job_id = sql_identifier.parse_fully_qualified_name(job_name)
-        output_table_db_id, output_table_schema_id, output_table_id = sql_identifier.parse_fully_qualified_name(
-            output_table_name
-        )
-        warehouse = self._service_ops._session.get_current_warehouse()
-        assert warehouse, "No active warehouse selected in the current session."
-        return self._service_ops.invoke_job_method(
-            target_method=target_function_info["target_method"],
-            signature=target_function_info["signature"],
-            X=X,
-            database_name=None,
-            schema_name=None,
-            model_name=self._model_name,
-            version_name=self._version_name,
-            job_database_name=job_db_id,
-            job_schema_name=job_schema_id,
-            job_name=job_id,
-            compute_pool_name=sql_identifier.SqlIdentifier(compute_pool),
-            warehouse_name=sql_identifier.SqlIdentifier(warehouse),
-            image_repo_name=image_repo,
-            output_table_database_name=output_table_db_id,
-            output_table_schema_name=output_table_schema_id,
-            output_table_name=output_table_id,
-            cpu_requests=cpu_requests,
-            memory_requests=memory_requests,
-            gpu_requests=gpu_requests,
-            num_workers=num_workers,
-            max_batch_rows=max_batch_rows,
-            force_rebuild=force_rebuild,
-            build_external_access_integrations=(
-                None
-                if build_external_access_integrations is None
-                else [sql_identifier.SqlIdentifier(eai) for eai in build_external_access_integrations]
-            ),
-            statement_params=statement_params,
-        )
-
 
 lineage_node.DOMAIN_LINEAGE_REGISTRY["model"] = ModelVersion

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

@@ -47,6 +47,7 @@ class ServiceInfo(TypedDict):
 class ModelOperator:
     INFERENCE_SERVICE_ENDPOINT_NAME = "inference"
     INGRESS_ENDPOINT_URL_SUFFIX = "snowflakecomputing.app"
+    PRIVATELINK_INGRESS_ENDPOINT_URL_SUBSTRING = "privatelink.snowflakecomputing"
 
     def __init__(
         self,
@@ -612,6 +613,30 @@ class ModelOperator:
             statement_params=statement_params,
         )
 
+    def _is_privatelink_connection(self) -> bool:
+        """Detect if the current session is using a privatelink connection."""
+        try:
+            host = self._session.connection.host
+            return ModelOperator.PRIVATELINK_INGRESS_ENDPOINT_URL_SUBSTRING in host
+        except AttributeError:
+            return False
+
+    def _extract_and_validate_ingress_url(self, res_row: "row.Row") -> Optional[str]:
+        """Extract and validate ingress URL from endpoint row."""
+        url_value = res_row[self._service_client.MODEL_INFERENCE_SERVICE_ENDPOINT_INGRESS_URL_COL_NAME]
+        if url_value is None:
+            return None
+        url_str = str(url_value)
+        return url_str if url_str.endswith(ModelOperator.INGRESS_ENDPOINT_URL_SUFFIX) else None
+
+    def _extract_and_validate_privatelink_url(self, res_row: "row.Row") -> Optional[str]:
+        """Extract and validate privatelink ingress URL from endpoint row."""
+        url_value = res_row[self._service_client.MODEL_INFERENCE_SERVICE_ENDPOINT_PRIVATELINK_INGRESS_URL_COL_NAME]
+        if url_value is None:
+            return None
+        url_str = str(url_value)
+        return url_str if ModelOperator.PRIVATELINK_INGRESS_ENDPOINT_URL_SUBSTRING in url_str else None
+
     def show_services(
         self,
         *,
@@ -644,8 +669,10 @@ class ModelOperator:
         fully_qualified_service_names = [str(service) for service in json_array if "MODEL_BUILD_" not in service]
 
         result: list[ServiceInfo] = []
+        is_privatelink_connection = self._is_privatelink_connection()
+
         for fully_qualified_service_name in fully_qualified_service_names:
-            ingress_url: Optional[str] = None
+            inference_endpoint: Optional[str] = None
             db, schema, service_name = sql_identifier.parse_fully_qualified_name(fully_qualified_service_name)
             statuses = self._service_client.get_service_container_statuses(
                 database_name=db, schema_name=schema, service_name=service_name, statement_params=statement_params
@@ -659,17 +686,23 @@ class ModelOperator:
             ):
                 if (
                     res_row[self._service_client.MODEL_INFERENCE_SERVICE_ENDPOINT_NAME_COL_NAME]
-                    == self.INFERENCE_SERVICE_ENDPOINT_NAME
-                    and res_row[self._service_client.MODEL_INFERENCE_SERVICE_ENDPOINT_INGRESS_URL_COL_NAME] is not None
+                    != self.INFERENCE_SERVICE_ENDPOINT_NAME
                 ):
-                    ingress_url = str(
-                        res_row[self._service_client.MODEL_INFERENCE_SERVICE_ENDPOINT_INGRESS_URL_COL_NAME]
-                    )
-                    if not ingress_url.endswith(ModelOperator.INGRESS_ENDPOINT_URL_SUFFIX):
-                        ingress_url = None
+                    continue
+
+                ingress_url = self._extract_and_validate_ingress_url(res_row)
+                privatelink_ingress_url = self._extract_and_validate_privatelink_url(res_row)
+
+                if is_privatelink_connection and privatelink_ingress_url is not None:
+                    inference_endpoint = privatelink_ingress_url
+                else:
+                    inference_endpoint = ingress_url
+
             result.append(
                 ServiceInfo(
-                    name=fully_qualified_service_name, status=service_status.value, inference_endpoint=ingress_url
+                    name=fully_qualified_service_name,
+                    status=service_status.value,
+                    inference_endpoint=inference_endpoint,
                 )
             )