snowflake-ml-python 1.22.0__py3-none-any.whl → 1.23.0__py3-none-any.whl

snowflake/ml/jobs/manager.py

@@ -1,11 +1,5 @@
-import json
 import logging
-import os
-import pathlib
-import sys
-from pathlib import PurePath
 from typing import Any, Callable, Optional, TypeVar, Union, cast, overload
-from uuid import uuid4
 
 import pandas as pd
 
@@ -13,13 +7,8 @@ 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 (
-    constants,
-    feature_flags,
-    payload_utils,
-    query_helper,
-    types,
-)
+from snowflake.ml.jobs._utils import query_helper
+from snowflake.ml.jobs.job_definition import MLJobDefinition
 from snowflake.snowpark.context import get_active_session
 from snowflake.snowpark.exceptions import SnowparkSQLException
 from snowflake.snowpark.functions import coalesce, col, lit, when
@@ -457,7 +446,6 @@ def _submit_job(
         An object representing the submitted job.
 
     Raises:
-        ValueError: If database or schema value(s) are invalid
         RuntimeError: If schema is not specified in session context or job submission
     """
     session = _ensure_session(session)
@@ -469,94 +457,30 @@ def _submit_job(
         )
         target_instances = max(target_instances, kwargs.pop("num_instances"))
 
-    imports = None
     if "additional_payloads" in kwargs:
         logger.warning(
             "'additional_payloads' is deprecated and will be removed in a future release. Use 'imports' instead."
         )
-        imports = kwargs.pop("additional_payloads")
+        if "imports" not in kwargs:
+            imports = kwargs.pop("additional_payloads", None)
+            kwargs.update({"imports": imports})
 
     if "runtime_environment" in kwargs:
         logger.warning("'runtime_environment' is in private preview since 1.15.0, do not use it in production.")
 
-    # Use kwargs for less common optional parameters
-    database = kwargs.pop("database", None)
-    schema = kwargs.pop("schema", None)
-    min_instances = kwargs.pop("min_instances", target_instances)
-    pip_requirements = kwargs.pop("pip_requirements", None)
-    external_access_integrations = kwargs.pop("external_access_integrations", None)
-    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", session.get_current_warehouse())
-    imports = kwargs.pop("imports", None) or imports
-    # if the mljob is submitted from a notebook, we use the same image tag as the notebook
-    runtime_environment = kwargs.pop("runtime_environment", os.environ.get(constants.RUNTIME_IMAGE_TAG_ENV_VAR, None))
-
-    # Warn if there are unknown kwargs
-    if kwargs:
-        logger.warning(f"Ignoring unknown kwargs: {kwargs.keys()}")
-
-    # Validate parameters
-    if database and not schema:
-        raise ValueError("Schema must be specified if database is specified.")
-    if target_instances < 1:
-        raise ValueError("target_instances must be greater than 0.")
-    if not (0 < min_instances <= target_instances):
-        raise ValueError("min_instances must be greater than 0 and less than or equal to target_instances.")
-    if min_instances > 1:
-        # Validate min_instances against compute pool max_nodes
-        pool_info = jb._get_compute_pool_info(session, compute_pool)
-        max_nodes = int(pool_info["max_nodes"])
-        if min_instances > max_nodes:
-            raise ValueError(
-                f"The requested min_instances ({min_instances}) exceeds the max_nodes ({max_nodes}) "
-                f"of compute pool '{compute_pool}'. Reduce min_instances or increase max_nodes."
-            )
-
-    job_name = f"{JOB_ID_PREFIX}{str(uuid4()).replace('-', '_').upper()}"
-    job_id = identifier.get_schema_level_object_identifier(database, schema, job_name)
-    stage_path_parts = identifier.parse_snowflake_stage_path(stage_name.lstrip("@"))
-    stage_name = f"@{'.'.join(filter(None, stage_path_parts[:3]))}"
-    stage_path = pathlib.PurePosixPath(f"{stage_name}{stage_path_parts[-1].rstrip('/')}/{job_name}")
-
-    try:
-        # Upload payload
-        uploaded_payload = payload_utils.JobPayload(
-            source, entrypoint=entrypoint, pip_requirements=pip_requirements, imports=imports
-        ).upload(session, stage_path)
-    except 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"
-            )
-        elif e.sql_error_code == 3001 and "schema" in str(e).lower():
-            raise RuntimeError(
-                "please grant privileges on schema before submitting a job, see",
-                "https://docs.snowflake.com/en/developer-guide/snowflake-ml/ml-jobs/access-control-requirements",
-                " for more details",
-            ) from e
-        raise
-
-    combined_env_vars = {**uploaded_payload.env_vars, **(env_vars or {})}
+    job_definition = MLJobDefinition.register(
+        source,
+        compute_pool,
+        stage_name,
+        session or get_active_session(),
+        entrypoint,
+        target_instances,
+        generate_suffix=True,
+        **kwargs,
+    )
 
     try:
-        return _do_submit_job(
-            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,
-            runtime_environment=runtime_environment,
-        )
+        return job_definition(*(args or []))
     except SnowparkSQLException as e:
         if e.sql_error_code == 3001 and "schema" in str(e).lower():
             raise RuntimeError(
@@ -567,91 +491,6 @@ def _submit_job(
         raise
 
 
-def _do_submit_job(
-    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,
-    runtime_environment: Optional[str] = None,
-) -> 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.
-        runtime_environment: image tag or full image URL to use for the job.
-
-    Returns:
-        The job object.
-    """
-    args = [(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,
-    }
-    if runtime_environment:
-        # for the image tag or full image URL, we use that directly
-        spec_options["RUNTIME"] = runtime_environment
-    elif feature_flags.FeatureFlags.ENABLE_RUNTIME_VERSIONS.is_enabled():
-        # when feature flag is enabled, we get the local python version and wrap it in a dict
-        # in system function, we can know whether it is python version or image tag or full image URL through the format
-        spec_options["RUNTIME"] = json.dumps({"pythonVersion": f"{sys.version_info.major}.{sys.version_info.minor}"})
-
-    job_options = {
-        "EXTERNAL_ACCESS_INTEGRATIONS": external_access_integrations,
-        "QUERY_WAREHOUSE": query_warehouse,
-        "TARGET_INSTANCES": target_instances,
-        "MIN_INSTANCES": min_instances,
-        "ASYNC": use_async,
-    }
-
-    if feature_flags.FeatureFlags.ENABLE_STAGE_MOUNT_V2.is_enabled(default=True):
-        spec_options["ENABLE_STAGE_MOUNT_V2"] = True
-    if payload.payload_name:
-        job_options["GENERATE_SUFFIX"] = True
-    job_options = {k: v for k, v in job_options.items() if v is not None}
-
-    query_template = "CALL SYSTEM$EXECUTE_ML_JOB(?, ?, ?, ?)"
-    if job_id:
-        database, schema, _ = identifier.parse_schema_level_object_identifier(job_id)
-    params = [
-        job_id
-        if payload.payload_name is None
-        else identifier.get_schema_level_object_identifier(database, schema, payload.payload_name) + "_",
-        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()

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

@@ -1,6 +1,4 @@
-import base64
 import enum
-import json
 import pathlib
 import tempfile
 import uuid
@@ -8,7 +6,6 @@ import warnings
 from typing import Any, Callable, Optional, Union, overload
 
 import pandas as pd
-from pydantic import TypeAdapter
 
 from snowflake import snowpark
 from snowflake.ml._internal import telemetry
@@ -33,7 +30,10 @@ _TELEMETRY_PROJECT = "MLOps"
 _TELEMETRY_SUBPROJECT = "ModelManagement"
 _BATCH_INFERENCE_JOB_ID_PREFIX = "BATCH_INFERENCE_"
 _BATCH_INFERENCE_TEMPORARY_FOLDER = "_temporary"
-_UTF8_ENCODING = "utf-8"
+VLLM_SUPPORTED_TASKS = [
+    "text-generation",
+    "image-text-to-text",
+]
 
 
 class ExportMode(enum.Enum):
@@ -649,41 +649,6 @@ class ModelVersion(lineage_node.LineageNode):
             method_options, target_function_info["name"]
         )
 
-    @staticmethod
-    def _encode_column_handling(
-        column_handling: Optional[dict[str, batch_inference_specs.ColumnHandlingOptions]],
-    ) -> Optional[str]:
-        """Validate and encode column_handling to a base64 string.
-
-        Args:
-            column_handling: Optional dictionary mapping column names to file encoding options.
-
-        Returns:
-            Base64 encoded JSON string of the column handling options, or None if input is None.
-        """
-        # TODO: validation for column names
-        if column_handling is None:
-            return None
-        adapter = TypeAdapter(dict[str, batch_inference_specs.ColumnHandlingOptions])
-        # TODO: throw error if the validate_python function fails
-        validated_input = adapter.validate_python(column_handling)
-        return base64.b64encode(adapter.dump_json(validated_input)).decode(_UTF8_ENCODING)
-
-    @staticmethod
-    def _encode_params(params: Optional[dict[str, Any]]) -> Optional[str]:
-        """Encode params dictionary to a base64 string.
-
-        Args:
-            params: Optional dictionary of model inference parameters.
-
-        Returns:
-            Base64 encoded JSON string of the params, or None if input is None.
-        """
-        if params is None:
-            return None
-        # TODO: validation for param names, types
-        return base64.b64encode(json.dumps(params).encode(_UTF8_ENCODING)).decode(_UTF8_ENCODING)
-
     @telemetry.send_api_usage_telemetry(
         project=_TELEMETRY_PROJECT,
         subproject=_TELEMETRY_SUBPROJECT,
@@ -703,6 +668,7 @@ class ModelVersion(lineage_node.LineageNode):
         job_spec: Optional[batch_inference_specs.JobSpec] = None,
         params: Optional[dict[str, Any]] = None,
         column_handling: Optional[dict[str, batch_inference_specs.ColumnHandlingOptions]] = None,
+        inference_engine_options: Optional[dict[str, Any]] = None,
     ) -> job.MLJob[Any]:
         """Execute batch inference on datasets as an SPCS job.
 
@@ -722,6 +688,10 @@ class ModelVersion(lineage_node.LineageNode):
             column_handling (Optional[dict[str, batch_inference_specs.FileEncoding]]): Optional dictionary
                 specifying how to handle specific columns during file I/O. Maps column names to their
                 file encoding configuration.
+            inference_engine_options: Options for the service creation with custom inference engine.
+                Supports `engine` and `engine_args_override`.
+                `engine` is the type of the inference engine to use.
+                `engine_args_override` is a list of string arguments to pass to the inference engine.
 
         Returns:
             job.MLJob[Any]: A batch inference job object that can be used to monitor progress and manage the job
@@ -777,12 +747,18 @@ class ModelVersion(lineage_node.LineageNode):
             subproject=_TELEMETRY_SUBPROJECT,
         )
 
-        column_handling_as_string = self._encode_column_handling(column_handling)
-        params_as_string = self._encode_params(params)
-
         if job_spec is None:
             job_spec = batch_inference_specs.JobSpec()
 
+        # Validate GPU support if GPU resources are requested
+        self._throw_error_if_gpu_is_not_supported(job_spec.gpu_requests, statement_params)
+
+        inference_engine_args = self._prepare_inference_engine_args(
+            inference_engine_options,
+            job_spec.gpu_requests,
+            statement_params,
+        )
+
         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.")
@@ -807,12 +783,14 @@ class ModelVersion(lineage_node.LineageNode):
         else:
             job_name = job_spec.job_name
 
+        target_function_info = self._get_function_info(function_name=job_spec.function_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"],
+            function_name=target_function_info["target_method"],
             compute_pool_name=sql_identifier.SqlIdentifier(compute_pool),
             force_rebuild=job_spec.force_rebuild,
             image_repo_name=job_spec.image_repo,
@@ -827,12 +805,14 @@ class ModelVersion(lineage_node.LineageNode):
             # input and output
             input_stage_location=input_stage_location,
             input_file_pattern="*",
-            column_handling=column_handling_as_string,
-            params=params_as_string,
+            column_handling=column_handling,
+            params=params,
+            signature_params=target_function_info["signature"].params,
             output_stage_location=output_stage_location,
             completion_filename="_SUCCESS",
             # misc
             statement_params=statement_params,
+            inference_engine_args=inference_engine_args,
         )
 
     def _get_function_info(self, function_name: Optional[str]) -> model_manifest_schema.ModelFunctionInfo:
@@ -1048,20 +1028,55 @@ class ModelVersion(lineage_node.LineageNode):
                 " the `log_model` function."
             )
 
-    def _check_huggingface_text_generation_model(
+    def _prepare_inference_engine_args(
+        self,
+        inference_engine_options: Optional[dict[str, Any]],
+        gpu_requests: Optional[Union[str, int]],
+        statement_params: Optional[dict[str, Any]] = None,
+    ) -> Optional[service_ops.InferenceEngineArgs]:
+        """Prepare and validate inference engine arguments.
+
+        This method handles the common logic for processing inference engine options:
+        1. Parse inference engine options into InferenceEngineArgs
+        2. Validate that the model is a HuggingFace text-generation model (if inference engine is specified)
+        3. Enrich inference engine args
+
+        Args:
+            inference_engine_options: Optional dictionary containing inference engine configuration.
+            gpu_requests: GPU resource request string (e.g., "4").
+            statement_params: Optional dictionary of statement parameters for SQL commands.
+
+        Returns:
+            Prepared InferenceEngineArgs or None if no inference engine is specified.
+        """
+        inference_engine_args = inference_engine_utils._get_inference_engine_args(inference_engine_options)
+
+        if inference_engine_args is not None:
+            # Validate that model is HuggingFace vLLM supported model and is logged with
+            # OpenAI compatible signature.
+            self._check_huggingface_vllm_supported_model(statement_params)
+            # Enrich with GPU configuration
+            inference_engine_args = inference_engine_utils._enrich_inference_engine_args(
+                inference_engine_args,
+                gpu_requests,
+            )
+
+        return inference_engine_args
+
+    def _check_huggingface_vllm_supported_model(
         self,
         statement_params: Optional[dict[str, Any]] = None,
     ) -> None:
-        """Check if the model is a HuggingFace pipeline with text-generation task
-        and is logged with OPENAI_CHAT_SIGNATURE.
+        """Check if the model is a HuggingFace pipeline with vLLM supported task
+        and is logged with OpenAI compatible signature.
 
         Args:
             statement_params: Optional dictionary of statement parameters to include
                 in the SQL command to fetch model spec.
 
         Raises:
-            ValueError: If the model is not a HuggingFace text-generation model or
-                if the model is not logged with OPENAI_CHAT_SIGNATURE.
+            ValueError: If the model is not a HuggingFace vLLM supported model or
+                if the model is not logged with OpenAI compatible signature.
         """
         # Fetch model spec
         model_spec = self._get_model_spec(statement_params)
@@ -1070,34 +1085,37 @@ class ModelVersion(lineage_node.LineageNode):
         model_type = model_spec.get("model_type")
         if model_type != "huggingface_pipeline":
             raise ValueError(
-                f"Inference engine is only supported for HuggingFace text-generation models. "
+                f"Inference engine is only supported for HuggingFace vLLM supported models. "
                 f"Found model_type: {model_type}"
             )
 
-        # Check if model supports text-generation task
+        # Check if model supports vLLM supported task
         # There should only be one model in the list because we don't support multiple models in a single model spec
         models = model_spec.get("models", {})
-        is_text_generation = False
+        is_vllm_supported_task = False
         found_tasks: list[str] = []
 
-        # As long as the model supports text-generation task, we can use it
+        # As long as the model supports vLLM supported task, we can use it
         for _, model_info in models.items():
             options = model_info.get("options", {})
             task = options.get("task")
             if task:
                 found_tasks.append(str(task))
-                if task == "text-generation":
-                    is_text_generation = True
+                if task in VLLM_SUPPORTED_TASKS:
+                    is_vllm_supported_task = True
                     break
 
-        if not is_text_generation:
+        if not is_vllm_supported_task:
             tasks_str = ", ".join(found_tasks)
             found_tasks_str = (
                 f"Found task(s): {tasks_str} in model spec." if found_tasks else "No task found in model spec."
             )
-            raise ValueError(f"Inference engine is only supported for task 'text-generation'. {found_tasks_str}")
+            supported_tasks_str = ", ".join(VLLM_SUPPORTED_TASKS)
+            raise ValueError(
+                f"Inference engine is only supported for vLLM supported tasks. {supported_tasks_str}. {found_tasks_str}"
+            )
 
-        # Check if the model is logged with OPENAI_CHAT_SIGNATURE
+        # Check if the model is logged with OpenAI compatible signature.
         signatures_dict = model_spec.get("signatures", {})
 
         # Deserialize signatures from model spec to ModelSignature objects for proper semantic comparison.
@@ -1105,11 +1123,16 @@ class ModelVersion(lineage_node.LineageNode):
             func_name: core.ModelSignature.from_dict(sig_dict) for func_name, sig_dict in signatures_dict.items()
         }
 
-        if deserialized_signatures != openai_signatures.OPENAI_CHAT_SIGNATURE:
+        if deserialized_signatures not in [
+            openai_signatures.OPENAI_CHAT_SIGNATURE,
+            openai_signatures.OPENAI_CHAT_SIGNATURE_WITH_CONTENT_FORMAT_STRING,
+        ]:
             raise ValueError(
-                "Inference engine requires the model to be logged with OPENAI_CHAT_SIGNATURE. "
+                "Inference engine requires the model to be logged with openai_signatures.OPENAI_CHAT_SIGNATURE or "
+                "openai_signatures.OPENAI_CHAT_SIGNATURE_WITH_CONTENT_FORMAT_STRING. "
                 f"Found signatures: {signatures_dict}. "
-                "Please log the model with: signatures=openai_signatures.OPENAI_CHAT_SIGNATURE"
+                "Please log the model again with: signatures=openai_signatures.OPENAI_CHAT_SIGNATURE or "
+                "signatures=openai_signatures.OPENAI_CHAT_SIGNATURE_WITH_CONTENT_FORMAT_STRING"
             )
 
     @overload
@@ -1350,20 +1373,11 @@ class ModelVersion(lineage_node.LineageNode):
         # Validate GPU support if GPU resources are requested
         self._throw_error_if_gpu_is_not_supported(gpu_requests, statement_params)
 
-        inference_engine_args = inference_engine_utils._get_inference_engine_args(inference_engine_options)
-
-        # Check if model is HuggingFace text-generation and is logged with
-        # OPENAI_CHAT_SIGNATURE before doing inference engine checks
-        # Only validate if inference engine is actually specified
-        if inference_engine_args is not None:
-            self._check_huggingface_text_generation_model(statement_params)
-
-        # Enrich inference engine args if inference engine is specified
-        if inference_engine_args is not None:
-            inference_engine_args = inference_engine_utils._enrich_inference_engine_args(
-                inference_engine_args,
-                gpu_requests,
-            )
+        inference_engine_args = self._prepare_inference_engine_args(
+            inference_engine_options,
+            gpu_requests,
+            statement_params,
+        )
 
         from snowflake.ml.model import event_handler
         from snowflake.snowpark import exceptions

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

@@ -14,7 +14,7 @@ from snowflake.ml._internal import platform_capabilities
 from snowflake.ml._internal.exceptions import error_codes, exceptions
 from snowflake.ml._internal.utils import formatting, identifier, sql_identifier, url
 from snowflake.ml.model import model_signature, type_hints
-from snowflake.ml.model._client.ops import deployment_step, metadata_ops
+from snowflake.ml.model._client.ops import deployment_step, metadata_ops, param_utils
 from snowflake.ml.model._client.sql import (
     model as model_sql,
     model_version as model_version_sql,
@@ -1063,23 +1063,7 @@ class ModelOperator:
                 col_name = sql_identifier.SqlIdentifier(input_feature.name.upper(), case_sensitive=True)
             input_args.append(col_name)
 
-        method_parameters: Optional[list[tuple[sql_identifier.SqlIdentifier, Any]]] = None
-        if signature.params:
-            # Start with defaults from signature
-            final_params = {}
-            for param_spec in signature.params:
-                if hasattr(param_spec, "default_value"):
-                    final_params[param_spec.name] = param_spec.default_value
-
-            # Override with provided runtime parameters
-            if params:
-                final_params.update(params)
-
-            # Convert to list of tuples with SqlIdentifier for parameter names
-            method_parameters = [
-                (sql_identifier.SqlIdentifier(param_name), param_value)
-                for param_name, param_value in final_params.items()
-            ]
+        method_parameters = param_utils.validate_and_resolve_params(params, signature.params)
 
         returns = []
         for output_feature in signature.outputs: