snowflake-ml-python 1.8.2__py3-none-any.whl → 1.8.4__py3-none-any.whl

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

@@ -1,5 +1,6 @@
+import json
 import pathlib
-from typing import List, Optional, Union
+from typing import Any, Optional, Union
 
 import yaml
 
@@ -18,99 +19,392 @@ class ModelDeploymentSpec:
 
     def __init__(self, workspace_path: Optional[pathlib.Path] = None) -> None:
         self.workspace_path = workspace_path
+        self._models: list[model_deployment_spec_schema.Model] = []
+        self._image_build: Optional[model_deployment_spec_schema.ImageBuild] = None
+        self._service: Optional[model_deployment_spec_schema.Service] = None
+        self._job: Optional[model_deployment_spec_schema.Job] = None
+        self._model_loggings: Optional[list[model_deployment_spec_schema.ModelLogging]] = None
+        self._inference_spec: dict[str, Any] = {}  # Common inference spec for service/job
 
-    def save(
+        self.database: Optional[sql_identifier.SqlIdentifier] = None
+        self.schema: Optional[sql_identifier.SqlIdentifier] = None
+
+    def add_model_spec(
         self,
-        *,
         database_name: sql_identifier.SqlIdentifier,
         schema_name: sql_identifier.SqlIdentifier,
         model_name: sql_identifier.SqlIdentifier,
         version_name: sql_identifier.SqlIdentifier,
-        service_database_name: Optional[sql_identifier.SqlIdentifier],
-        service_schema_name: Optional[sql_identifier.SqlIdentifier],
-        service_name: sql_identifier.SqlIdentifier,
+    ) -> "ModelDeploymentSpec":
+        """Add model specification to the deployment spec.
+
+        Args:
+            database_name: Database name containing the model.
+            schema_name: Schema name containing the model.
+            model_name: Name of the model.
+            version_name: Version of the model.
+
+        Returns:
+            Self for chaining.
+        """
+        fq_model_name = identifier.get_schema_level_object_identifier(
+            database_name.identifier(), schema_name.identifier(), model_name.identifier()
+        )
+        if not self.database:
+            self.database = database_name
+        if not self.schema:
+            self.schema = schema_name
+        model = model_deployment_spec_schema.Model(name=fq_model_name, version=version_name.identifier())
+        self._models.append(model)
+        return self
+
+    def add_image_build_spec(
+        self,
         image_build_compute_pool_name: sql_identifier.SqlIdentifier,
-        service_compute_pool_name: sql_identifier.SqlIdentifier,
-        image_repo_database_name: Optional[sql_identifier.SqlIdentifier],
-        image_repo_schema_name: Optional[sql_identifier.SqlIdentifier],
         image_repo_name: sql_identifier.SqlIdentifier,
-        ingress_enabled: bool,
-        max_instances: int,
+        image_repo_database_name: Optional[sql_identifier.SqlIdentifier] = None,
+        image_repo_schema_name: Optional[sql_identifier.SqlIdentifier] = None,
+        force_rebuild: bool = False,
+        external_access_integrations: Optional[list[sql_identifier.SqlIdentifier]] = None,
+    ) -> "ModelDeploymentSpec":
+        """Add image build specification to the deployment spec.
+
+        Args:
+            image_build_compute_pool_name: Compute pool for image building.
+            image_repo_name: Name of the image repository.
+            image_repo_database_name: Database name for the image repository.
+            image_repo_schema_name: Schema name for the image repository.
+            force_rebuild: Whether to force rebuilding the image.
+            external_access_integrations: List of external access integrations.
+
+        Returns:
+            Self for chaining.
+        """
+        saved_image_repo_database = image_repo_database_name or self.database
+        saved_image_repo_schema = image_repo_schema_name or self.schema
+        assert saved_image_repo_database is not None
+        assert saved_image_repo_schema is not None
+        fq_image_repo_name = identifier.get_schema_level_object_identifier(
+            db=saved_image_repo_database.identifier(),
+            schema=saved_image_repo_schema.identifier(),
+            object_name=image_repo_name.identifier(),
+        )
+
+        self._image_build = model_deployment_spec_schema.ImageBuild(
+            compute_pool=image_build_compute_pool_name.identifier(),
+            image_repo=fq_image_repo_name,
+            force_rebuild=force_rebuild,
+            external_access_integrations=(
+                [eai.identifier() for eai in external_access_integrations] if external_access_integrations else None
+            ),
+        )
+        return self
+
+    def _add_inference_spec(
+        self,
         cpu: Optional[str],
         memory: Optional[str],
         gpu: Optional[Union[str, int]],
         num_workers: Optional[int],
         max_batch_rows: Optional[int],
-        force_rebuild: bool,
-        external_access_integrations: Optional[List[sql_identifier.SqlIdentifier]],
-    ) -> str:
-        # create the deployment spec
-        # models spec
-        fq_model_name = identifier.get_schema_level_object_identifier(
-            database_name.identifier(), schema_name.identifier(), model_name.identifier()
-        )
-        model_dict = model_deployment_spec_schema.ModelDict(name=fq_model_name, version=version_name.identifier())
+    ) -> None:
+        """Internal helper to store common inference specs."""
+        if cpu:
+            self._inference_spec["cpu"] = cpu
+        if memory:
+            self._inference_spec["memory"] = memory
+        if gpu:
+            if isinstance(gpu, int):
+                gpu_str = str(gpu)
+            else:
+                gpu_str = gpu
+            self._inference_spec["gpu"] = gpu_str
+        if num_workers:
+            self._inference_spec["num_workers"] = num_workers
+        if max_batch_rows:
+            self._inference_spec["max_batch_rows"] = max_batch_rows
 
-        # image_build spec
-        saved_image_repo_database = image_repo_database_name or database_name
-        saved_image_repo_schema = image_repo_schema_name or schema_name
-        fq_image_repo_name = identifier.get_schema_level_object_identifier(
-            saved_image_repo_database.identifier(), saved_image_repo_schema.identifier(), image_repo_name.identifier()
-        )
-        image_build_dict: model_deployment_spec_schema.ImageBuildDict = {
-            "compute_pool": image_build_compute_pool_name.identifier(),
-            "image_repo": fq_image_repo_name,
-            "force_rebuild": force_rebuild,
-        }
-        if external_access_integrations is not None:
-            image_build_dict["external_access_integrations"] = [
-                eai.identifier() for eai in external_access_integrations
-            ]
-
-        # service spec
-        saved_service_database = service_database_name or database_name
-        saved_service_schema = service_schema_name or schema_name
+    def add_service_spec(
+        self,
+        service_name: sql_identifier.SqlIdentifier,
+        inference_compute_pool_name: sql_identifier.SqlIdentifier,
+        service_database_name: Optional[sql_identifier.SqlIdentifier] = None,
+        service_schema_name: Optional[sql_identifier.SqlIdentifier] = None,
+        ingress_enabled: bool = True,
+        max_instances: int = 1,
+        cpu: Optional[str] = None,
+        memory: Optional[str] = None,
+        gpu: Optional[Union[str, int]] = None,
+        num_workers: Optional[int] = None,
+        max_batch_rows: Optional[int] = None,
+    ) -> "ModelDeploymentSpec":
+        """Add service specification to the deployment spec.
+
+        Args:
+            service_name: Name of the service.
+            inference_compute_pool_name: Compute pool for inference.
+            service_database_name: Database name for the service.
+            service_schema_name: Schema name for the service.
+            ingress_enabled: Whether ingress is enabled.
+            max_instances: Maximum number of service instances.
+            cpu: CPU requirement.
+            memory: Memory requirement.
+            gpu: GPU requirement.
+            num_workers: Number of workers.
+            max_batch_rows: Maximum batch rows for inference.
+
+        Raises:
+            ValueError: If a job spec already exists.
+
+        Returns:
+            Self for chaining.
+        """
+        if self._job:
+            raise ValueError("Cannot add a service spec when a job spec already exists.")
+
+        saved_service_database = service_database_name or self.database
+        saved_service_schema = service_schema_name or self.schema
+        assert saved_service_database is not None
+        assert saved_service_schema is not None
         fq_service_name = identifier.get_schema_level_object_identifier(
             saved_service_database.identifier(), saved_service_schema.identifier(), service_name.identifier()
         )
-        service_dict = model_deployment_spec_schema.ServiceDict(
+
+        self._add_inference_spec(cpu, memory, gpu, num_workers, max_batch_rows)
+
+        self._service = model_deployment_spec_schema.Service(
             name=fq_service_name,
-            compute_pool=service_compute_pool_name.identifier(),
+            compute_pool=inference_compute_pool_name.identifier(),
             ingress_enabled=ingress_enabled,
             max_instances=max_instances,
+            **self._inference_spec,
         )
-        if cpu:
-            service_dict["cpu"] = cpu
+        return self
 
-        if memory:
-            service_dict["memory"] = memory
+    def add_job_spec(
+        self,
+        job_name: sql_identifier.SqlIdentifier,
+        inference_compute_pool_name: sql_identifier.SqlIdentifier,
+        warehouse: sql_identifier.SqlIdentifier,
+        target_method: str,
+        input_table_name: sql_identifier.SqlIdentifier,
+        output_table_name: sql_identifier.SqlIdentifier,
+        job_database_name: Optional[sql_identifier.SqlIdentifier] = None,
+        job_schema_name: Optional[sql_identifier.SqlIdentifier] = None,
+        input_table_database_name: Optional[sql_identifier.SqlIdentifier] = None,
+        input_table_schema_name: Optional[sql_identifier.SqlIdentifier] = None,
+        output_table_database_name: Optional[sql_identifier.SqlIdentifier] = None,
+        output_table_schema_name: Optional[sql_identifier.SqlIdentifier] = None,
+        cpu: Optional[str] = None,
+        memory: Optional[str] = None,
+        gpu: Optional[Union[str, int]] = None,
+        num_workers: Optional[int] = None,
+        max_batch_rows: Optional[int] = None,
+    ) -> "ModelDeploymentSpec":
+        """Add job specification to the deployment spec.
 
-        if gpu:
-            if isinstance(gpu, int):
-                gpu_str = str(gpu)
-            else:
-                gpu_str = gpu
-            service_dict["gpu"] = gpu_str
+        Args:
+            job_name: Name of the job.
+            inference_compute_pool_name: Compute pool for inference.
+            job_database_name: Database name for the job.
+            job_schema_name: Schema name for the job.
+            warehouse: Warehouse for the job.
+            target_method: Target method for inference.
+            input_table_name: Input table name.
+            output_table_name: Output table name.
+            input_table_database_name: Database for input table.
+            input_table_schema_name: Schema for input table.
+            output_table_database_name: Database for output table.
+            output_table_schema_name: Schema for output table.
+            cpu: CPU requirement.
+            memory: Memory requirement.
+            gpu: GPU requirement.
+            num_workers: Number of workers.
+            max_batch_rows: Maximum batch rows for inference.
 
-        if num_workers:
-            service_dict["num_workers"] = num_workers
+        Raises:
+            ValueError: If a service spec already exists.
 
-        if max_batch_rows:
-            service_dict["max_batch_rows"] = max_batch_rows
+        Returns:
+            Self for chaining.
+        """
+        if self._service:
+            raise ValueError("Cannot add a job spec when a service spec already exists.")
+
+        saved_job_database = job_database_name or self.database
+        saved_job_schema = job_schema_name or self.schema
+        input_table_database_name = input_table_database_name or self.database
+        input_table_schema_name = input_table_schema_name or self.schema
+        output_table_database_name = output_table_database_name or self.database
+        output_table_schema_name = output_table_schema_name or self.schema
 
-        # model deployment spec
-        model_deployment_spec_dict = model_deployment_spec_schema.ModelDeploymentSpecDict(
-            models=[model_dict],
-            image_build=image_build_dict,
-            service=service_dict,
+        assert saved_job_database is not None
+        assert saved_job_schema is not None
+        assert input_table_database_name is not None
+        assert input_table_schema_name is not None
+        assert output_table_database_name is not None
+        assert output_table_schema_name is not None
+
+        fq_job_name = identifier.get_schema_level_object_identifier(
+            saved_job_database.identifier(), saved_job_schema.identifier(), job_name.identifier()
+        )
+        fq_input_table_name = identifier.get_schema_level_object_identifier(
+            input_table_database_name.identifier(),
+            input_table_schema_name.identifier(),
+            input_table_name.identifier(),
         )
+        fq_output_table_name = identifier.get_schema_level_object_identifier(
+            output_table_database_name.identifier(),
+            output_table_schema_name.identifier(),
+            output_table_name.identifier(),
+        )
+
+        self._add_inference_spec(cpu, memory, gpu, num_workers, max_batch_rows)
+
+        self._job = model_deployment_spec_schema.Job(
+            name=fq_job_name,
+            compute_pool=inference_compute_pool_name.identifier(),
+            warehouse=warehouse.identifier(),
+            target_method=target_method,
+            input_table_name=fq_input_table_name,
+            output_table_name=fq_output_table_name,
+            **self._inference_spec,
+        )
+        return self
+
+    def add_hf_logger_spec(
+        self,
+        hf_model_name: str,
+        hf_task: Optional[str] = None,
+        hf_token: Optional[str] = None,
+        hf_tokenizer: Optional[str] = None,
+        hf_revision: Optional[str] = None,
+        hf_trust_remote_code: Optional[bool] = False,
+        pip_requirements: Optional[list[str]] = None,
+        conda_dependencies: Optional[list[str]] = None,
+        target_platforms: Optional[list[str]] = None,
+        comment: Optional[str] = None,
+        warehouse: Optional[str] = None,
+        **kwargs: Any,
+    ) -> "ModelDeploymentSpec":
+        """Add Hugging Face logger specification.
+
+        Args:
+            hf_model_name: Hugging Face model name.
+            hf_task: Hugging Face task.
+            hf_token: Hugging Face token.
+            hf_tokenizer: Hugging Face tokenizer.
+            hf_revision: Hugging Face model revision.
+            hf_trust_remote_code: Whether to trust remote code.
+            pip_requirements: List of pip requirements.
+            conda_dependencies: List of conda dependencies.
+            target_platforms: List of target platforms.
+            comment: Comment for the model.
+            warehouse: Warehouse used to log the model.
+            **kwargs: Additional Hugging Face model arguments.
+
+        Raises:
+            ValueError: If Hugging Face model name is missing when other HF parameters are provided.
+
+        Returns:
+            Self for chaining.
+        """
+        # Validation moved here from save
+        if (
+            any(
+                [
+                    hf_task,
+                    hf_token,
+                    hf_tokenizer,
+                    hf_revision,
+                    hf_trust_remote_code,
+                    pip_requirements,
+                ]
+            )
+            and not hf_model_name
+        ):
+            # This condition might be redundant now as hf_model_name is mandatory
+            raise ValueError("Hugging Face model name is required when using Hugging Face model deployment.")
+
+        log_model_args = model_deployment_spec_schema.LogModelArgs(
+            pip_requirements=pip_requirements,
+            conda_dependencies=conda_dependencies,
+            target_platforms=target_platforms,
+            comment=comment,
+            warehouse=warehouse,
+        )
+        hf_model = model_deployment_spec_schema.HuggingFaceModel(
+            hf_model_name=hf_model_name,
+            task=hf_task,
+            hf_token=hf_token,
+            tokenizer=hf_tokenizer,
+            trust_remote_code=hf_trust_remote_code,
+            revision=hf_revision,
+            hf_model_kwargs=json.dumps(kwargs),
+        )
+        model_logging = model_deployment_spec_schema.ModelLogging(
+            log_model_args=log_model_args,
+            hf_model=hf_model,
+        )
+        if self._model_loggings is None:
+            self._model_loggings = [model_logging]
+        else:
+            self._model_loggings.append(model_logging)
+        return self
+
+    def save(self) -> str:
+        """Constructs the final deployment spec from added components and saves it.
+
+        Raises:
+            ValueError: If required components are missing or conflicting specs are added.
+            RuntimeError: If no service or job spec is found despite validation.
+
+        Returns:
+            The path to the saved YAML file as a string, or the YAML content as a string
+            if workspace_path was not provided.
+        """
+        # Validations
+        if not self._models:
+            raise ValueError("Model specification is required. Call add_model_spec().")
+        if not self._image_build:
+            raise ValueError("Image build specification is required. Call add_image_build_spec().")
+        if not self._service and not self._job:
+            raise ValueError(
+                "Either service or job specification is required. Call add_service_spec() or add_job_spec()."
+            )
+        if self._service and self._job:
+            # This case should be prevented by checks in add_service_spec/add_job_spec, but double-check
+            raise ValueError("Cannot have both service and job specifications.")
+
+        # Construct the final spec object
+        if self._service:
+            model_deployment_spec: Union[
+                model_deployment_spec_schema.ModelServiceDeploymentSpec,
+                model_deployment_spec_schema.ModelJobDeploymentSpec,
+            ] = model_deployment_spec_schema.ModelServiceDeploymentSpec(
+                models=self._models,
+                image_build=self._image_build,
+                service=self._service,
+                model_loggings=self._model_loggings,
+            )
+        elif self._job:
+            model_deployment_spec = model_deployment_spec_schema.ModelJobDeploymentSpec(
+                models=self._models,
+                image_build=self._image_build,
+                job=self._job,
+                model_loggings=self._model_loggings,
+            )
+        else:
+            # Should not happen due to earlier validation
+            raise RuntimeError("Internal error: No service or job spec found despite validation.")
+
+        # Serialize and save/return
+        yaml_content = model_deployment_spec.model_dump(exclude_none=True)
 
-        # Anchors are not supported in the server, avoid that.
-        yaml.SafeDumper.ignore_aliases = lambda *args: True  # type: ignore[method-assign]
         if self.workspace_path is None:
-            return yaml.safe_dump(model_deployment_spec_dict)
-        # save the yaml
+            return yaml.safe_dump(yaml_content)
+
         file_path = self.workspace_path / self.DEPLOY_SPEC_FILE_REL_PATH
         with file_path.open("w", encoding="utf-8") as f:
-            yaml.safe_dump(model_deployment_spec_dict, f)
+            yaml.safe_dump(yaml_content, f)
         return str(file_path.resolve())

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

@@ -1,33 +1,78 @@
-from typing import List, TypedDict
+from typing import Optional
 
-from typing_extensions import NotRequired, Required
+from pydantic import BaseModel
 
 
-class ModelDict(TypedDict):
-    name: Required[str]
-    version: Required[str]
+class Model(BaseModel):
+    name: str
+    version: str
 
 
-class ImageBuildDict(TypedDict):
-    compute_pool: Required[str]
-    image_repo: Required[str]
-    force_rebuild: Required[bool]
-    external_access_integrations: NotRequired[List[str]]
+class ImageBuild(BaseModel):
+    compute_pool: str
+    image_repo: str
+    force_rebuild: bool
+    external_access_integrations: Optional[list[str]] = None
 
 
-class ServiceDict(TypedDict):
-    name: Required[str]
-    compute_pool: Required[str]
-    ingress_enabled: Required[bool]
-    max_instances: Required[int]
-    cpu: NotRequired[str]
-    memory: NotRequired[str]
-    gpu: NotRequired[str]
-    num_workers: NotRequired[int]
-    max_batch_rows: NotRequired[int]
+class Service(BaseModel):
+    name: str
+    compute_pool: str
+    ingress_enabled: bool
+    max_instances: int
+    cpu: Optional[str] = None
+    memory: Optional[str] = None
+    gpu: Optional[str] = None
+    num_workers: Optional[int] = None
+    max_batch_rows: Optional[int] = None
 
 
-class ModelDeploymentSpecDict(TypedDict):
-    models: Required[List[ModelDict]]
-    image_build: Required[ImageBuildDict]
-    service: Required[ServiceDict]
+class Job(BaseModel):
+    name: str
+    compute_pool: str
+    cpu: Optional[str] = None
+    memory: Optional[str] = None
+    gpu: Optional[str] = None
+    num_workers: Optional[int] = None
+    max_batch_rows: Optional[int] = None
+    warehouse: str
+    target_method: str
+    input_table_name: str
+    output_table_name: str
+
+
+class LogModelArgs(BaseModel):
+    pip_requirements: Optional[list[str]] = None
+    conda_dependencies: Optional[list[str]] = None
+    target_platforms: Optional[list[str]] = None
+    comment: Optional[str] = None
+    warehouse: Optional[str] = None
+
+
+class HuggingFaceModel(BaseModel):
+    hf_model_name: str
+    task: Optional[str] = None
+    tokenizer: Optional[str] = None
+    hf_token: Optional[str] = None
+    trust_remote_code: Optional[bool] = False
+    revision: Optional[str] = None
+    hf_model_kwargs: Optional[str] = "{}"
+
+
+class ModelLogging(BaseModel):
+    log_model_args: Optional[LogModelArgs] = None
+    hf_model: Optional[HuggingFaceModel] = None
+
+
+class ModelServiceDeploymentSpec(BaseModel):
+    models: list[Model]
+    image_build: ImageBuild
+    service: Service
+    model_loggings: Optional[list[ModelLogging]] = None
+
+
+class ModelJobDeploymentSpec(BaseModel):
+    models: list[Model]
+    image_build: ImageBuild
+    job: Job
+    model_loggings: Optional[list[ModelLogging]] = None

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

@@ -1,4 +1,4 @@
-from typing import Any, Dict, List, Optional
+from typing import Any, Optional
 
 from snowflake.ml._internal.utils import query_result_checker, sql_identifier
 from snowflake.ml.model._client.sql import _base
@@ -24,8 +24,8 @@ class ModelSQLClient(_base._BaseSQLClient):
         schema_name: Optional[sql_identifier.SqlIdentifier],
         model_name: Optional[sql_identifier.SqlIdentifier] = None,
         validate_result: bool = True,
-        statement_params: Optional[Dict[str, Any]] = None,
-    ) -> List[row.Row]:
+        statement_params: Optional[dict[str, Any]] = None,
+    ) -> list[row.Row]:
         actual_database_name = database_name or self._database_name
         actual_schema_name = schema_name or self._schema_name
         fully_qualified_schema_name = ".".join([actual_database_name.identifier(), actual_schema_name.identifier()])
@@ -57,8 +57,8 @@ class ModelSQLClient(_base._BaseSQLClient):
         version_name: Optional[sql_identifier.SqlIdentifier] = None,
         validate_result: bool = True,
         check_model_details: bool = False,
-        statement_params: Optional[Dict[str, Any]] = None,
-    ) -> List[row.Row]:
+        statement_params: Optional[dict[str, Any]] = None,
+    ) -> list[row.Row]:
         like_sql = ""
         if version_name:
             like_sql = f" LIKE '{version_name.resolved()}'"
@@ -90,7 +90,7 @@ class ModelSQLClient(_base._BaseSQLClient):
         schema_name: Optional[sql_identifier.SqlIdentifier],
         model_name: sql_identifier.SqlIdentifier,
         comment: str,
-        statement_params: Optional[Dict[str, Any]] = None,
+        statement_params: Optional[dict[str, Any]] = None,
     ) -> None:
         query_result_checker.SqlResultValidator(
             self._session,
@@ -107,7 +107,7 @@ class ModelSQLClient(_base._BaseSQLClient):
         database_name: Optional[sql_identifier.SqlIdentifier],
         schema_name: Optional[sql_identifier.SqlIdentifier],
         model_name: sql_identifier.SqlIdentifier,
-        statement_params: Optional[Dict[str, Any]] = None,
+        statement_params: Optional[dict[str, Any]] = None,
     ) -> None:
         query_result_checker.SqlResultValidator(
             self._session,
@@ -124,7 +124,7 @@ class ModelSQLClient(_base._BaseSQLClient):
         new_model_db: Optional[sql_identifier.SqlIdentifier],
         new_model_schema: Optional[sql_identifier.SqlIdentifier],
         new_model_name: sql_identifier.SqlIdentifier,
-        statement_params: Optional[Dict[str, Any]] = None,
+        statement_params: Optional[dict[str, Any]] = None,
     ) -> None:
         # Use registry's database and schema if a non fully qualified new model name is provided.
         new_fully_qualified_name = self.fully_qualified_object_name(new_model_db, new_model_schema, new_model_name)