snowflake-ml-python 1.8.1__py3-none-any.whl → 1.8.3__py3-none-any.whl

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

@@ -1,33 +1,53 @@
-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 ModelServiceDeploymentSpec(BaseModel):
+    models: list[Model]
+    image_build: ImageBuild
+    service: Service
+
+
+class ModelJobDeploymentSpec(BaseModel):
+    models: list[Model]
+    image_build: ImageBuild
+    job: Job

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)

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

@@ -1,7 +1,7 @@
 import json
 import pathlib
 import textwrap
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any, Optional
 from urllib.parse import ParseResult
 
 from snowflake.ml._internal.utils import (
@@ -34,7 +34,7 @@ class ModelVersionSQLClient(_base._BaseSQLClient):
         model_name: sql_identifier.SqlIdentifier,
         version_name: sql_identifier.SqlIdentifier,
         stage_path: str,
-        statement_params: Optional[Dict[str, Any]] = None,
+        statement_params: Optional[dict[str, Any]] = None,
     ) -> None:
         query_result_checker.SqlResultValidator(
             self._session,
@@ -56,7 +56,7 @@ class ModelVersionSQLClient(_base._BaseSQLClient):
         schema_name: Optional[sql_identifier.SqlIdentifier],
         model_name: sql_identifier.SqlIdentifier,
         version_name: sql_identifier.SqlIdentifier,
-        statement_params: Optional[Dict[str, Any]] = None,
+        statement_params: Optional[dict[str, Any]] = None,
     ) -> None:
         fq_source_model_name = self.fully_qualified_object_name(
             source_database_name, source_schema_name, source_model_name
@@ -78,7 +78,7 @@ class ModelVersionSQLClient(_base._BaseSQLClient):
         schema_name: Optional[sql_identifier.SqlIdentifier],
         model_name: sql_identifier.SqlIdentifier,
         version_name: sql_identifier.SqlIdentifier,
-        statement_params: Optional[Dict[str, Any]] = None,
+        statement_params: Optional[dict[str, Any]] = None,
     ) -> None:
         sql = (
             f"CREATE MODEL {self.fully_qualified_object_name(database_name, schema_name, model_name)}"
@@ -97,7 +97,7 @@ class ModelVersionSQLClient(_base._BaseSQLClient):
         schema_name: Optional[sql_identifier.SqlIdentifier],
         model_name: sql_identifier.SqlIdentifier,
         version_name: sql_identifier.SqlIdentifier,
-        statement_params: Optional[Dict[str, Any]] = None,
+        statement_params: Optional[dict[str, Any]] = None,
     ) -> None:
         sql = (
             f"ALTER MODEL {self.fully_qualified_object_name(database_name, schema_name, model_name)}"
@@ -116,7 +116,7 @@ class ModelVersionSQLClient(_base._BaseSQLClient):
         schema_name: Optional[sql_identifier.SqlIdentifier],
         model_name: sql_identifier.SqlIdentifier,
         version_name: sql_identifier.SqlIdentifier,
-        statement_params: Optional[Dict[str, Any]] = None,
+        statement_params: Optional[dict[str, Any]] = None,
     ) -> None:
         sql = (
             f"ALTER MODEL {self.fully_qualified_object_name(database_name, schema_name, model_name)}"
@@ -138,7 +138,7 @@ class ModelVersionSQLClient(_base._BaseSQLClient):
         model_name: sql_identifier.SqlIdentifier,
         version_name: sql_identifier.SqlIdentifier,
         stage_path: str,
-        statement_params: Optional[Dict[str, Any]] = None,
+        statement_params: Optional[dict[str, Any]] = None,
     ) -> None:
         query_result_checker.SqlResultValidator(
             self._session,
@@ -160,7 +160,7 @@ class ModelVersionSQLClient(_base._BaseSQLClient):
         schema_name: Optional[sql_identifier.SqlIdentifier],
         model_name: sql_identifier.SqlIdentifier,
         version_name: sql_identifier.SqlIdentifier,
-        statement_params: Optional[Dict[str, Any]] = None,
+        statement_params: Optional[dict[str, Any]] = None,
     ) -> None:
         fq_source_model_name = self.fully_qualified_object_name(
             source_database_name, source_schema_name, source_model_name
@@ -182,7 +182,7 @@ class ModelVersionSQLClient(_base._BaseSQLClient):
         schema_name: Optional[sql_identifier.SqlIdentifier],
         model_name: sql_identifier.SqlIdentifier,
         version_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,
@@ -201,7 +201,7 @@ class ModelVersionSQLClient(_base._BaseSQLClient):
         model_name: sql_identifier.SqlIdentifier,
         version_name: sql_identifier.SqlIdentifier,
         alias_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,
@@ -219,7 +219,7 @@ class ModelVersionSQLClient(_base._BaseSQLClient):
         schema_name: Optional[sql_identifier.SqlIdentifier],
         model_name: sql_identifier.SqlIdentifier,
         version_or_alias_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,
@@ -239,8 +239,8 @@ class ModelVersionSQLClient(_base._BaseSQLClient):
         version_name: sql_identifier.SqlIdentifier,
         file_path: pathlib.PurePosixPath,
         is_dir: bool = False,
-        statement_params: Optional[Dict[str, Any]] = None,
-    ) -> List[row.Row]:
+        statement_params: Optional[dict[str, Any]] = None,
+    ) -> list[row.Row]:
         # Workaround for snowURL bug.
         trailing_slash = "/" if is_dir else ""
 
@@ -276,7 +276,7 @@ class ModelVersionSQLClient(_base._BaseSQLClient):
         version_name: sql_identifier.SqlIdentifier,
         file_path: pathlib.PurePosixPath,
         target_path: pathlib.Path,
-        statement_params: Optional[Dict[str, Any]] = None,
+        statement_params: Optional[dict[str, Any]] = None,
     ) -> pathlib.Path:
         stage_location = pathlib.PurePosixPath(
             self.fully_qualified_object_name(database_name, schema_name, model_name),
@@ -310,8 +310,8 @@ class ModelVersionSQLClient(_base._BaseSQLClient):
         schema_name: Optional[sql_identifier.SqlIdentifier],
         model_name: sql_identifier.SqlIdentifier,
         version_name: sql_identifier.SqlIdentifier,
-        statement_params: Optional[Dict[str, Any]] = None,
-    ) -> List[row.Row]:
+        statement_params: Optional[dict[str, Any]] = None,
+    ) -> list[row.Row]:
         res = query_result_checker.SqlResultValidator(
             self._session,
             (
@@ -331,7 +331,7 @@ class ModelVersionSQLClient(_base._BaseSQLClient):
         model_name: sql_identifier.SqlIdentifier,
         version_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,
@@ -351,9 +351,9 @@ class ModelVersionSQLClient(_base._BaseSQLClient):
         version_name: sql_identifier.SqlIdentifier,
         method_name: sql_identifier.SqlIdentifier,
         input_df: dataframe.DataFrame,
-        input_args: List[sql_identifier.SqlIdentifier],
-        returns: List[Tuple[str, spt.DataType, sql_identifier.SqlIdentifier]],
-        statement_params: Optional[Dict[str, Any]] = None,
+        input_args: list[sql_identifier.SqlIdentifier],
+        returns: list[tuple[str, spt.DataType, sql_identifier.SqlIdentifier]],
+        statement_params: Optional[dict[str, Any]] = None,
     ) -> dataframe.DataFrame:
         with_statements = []
         if len(input_df.queries["queries"]) == 1 and len(input_df.queries["post_actions"]) == 0:
@@ -433,10 +433,10 @@ class ModelVersionSQLClient(_base._BaseSQLClient):
         version_name: sql_identifier.SqlIdentifier,
         method_name: sql_identifier.SqlIdentifier,
         input_df: dataframe.DataFrame,
-        input_args: List[sql_identifier.SqlIdentifier],
-        returns: List[Tuple[str, spt.DataType, sql_identifier.SqlIdentifier]],
+        input_args: list[sql_identifier.SqlIdentifier],
+        returns: list[tuple[str, spt.DataType, sql_identifier.SqlIdentifier]],
         partition_column: Optional[sql_identifier.SqlIdentifier],
-        statement_params: Optional[Dict[str, Any]] = None,
+        statement_params: Optional[dict[str, Any]] = None,
         is_partitioned: bool = True,
     ) -> dataframe.DataFrame:
         with_statements = []
@@ -529,13 +529,13 @@ class ModelVersionSQLClient(_base._BaseSQLClient):
 
     def set_metadata(
         self,
-        metadata_dict: Dict[str, Any],
+        metadata_dict: dict[str, Any],
         *,
         database_name: Optional[sql_identifier.SqlIdentifier],
         schema_name: Optional[sql_identifier.SqlIdentifier],
         model_name: sql_identifier.SqlIdentifier,
         version_name: sql_identifier.SqlIdentifier,
-        statement_params: Optional[Dict[str, Any]] = None,
+        statement_params: Optional[dict[str, Any]] = None,
     ) -> None:
         json_metadata = json.dumps(metadata_dict)
         query_result_checker.SqlResultValidator(
@@ -554,7 +554,7 @@ class ModelVersionSQLClient(_base._BaseSQLClient):
         schema_name: Optional[sql_identifier.SqlIdentifier],
         model_name: sql_identifier.SqlIdentifier,
         version_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,

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

@@ -1,7 +1,7 @@
 import enum
 import json
 import textwrap
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any, Optional, Union
 
 from snowflake import snowpark
 from snowflake.ml._internal import platform_capabilities
@@ -47,7 +47,7 @@ class ServiceSQLClient(_base._BaseSQLClient):
         gpu: Optional[Union[str, int]],
         force_rebuild: bool,
         external_access_integration: sql_identifier.SqlIdentifier,
-        statement_params: Optional[Dict[str, Any]] = None,
+        statement_params: Optional[dict[str, Any]] = None,
     ) -> None:
         actual_image_repo_database = image_repo_database_name or self._database_name
         actual_image_repo_schema = image_repo_schema_name or self._schema_name
@@ -73,13 +73,17 @@ class ServiceSQLClient(_base._BaseSQLClient):
     def deploy_model(
         self,
         *,
-        stage_path: str,
-        model_deployment_spec_file_rel_path: str,
-        statement_params: Optional[Dict[str, Any]] = None,
-    ) -> Tuple[str, snowpark.AsyncJob]:
-        async_job = self._session.sql(
-            f"CALL SYSTEM$DEPLOY_MODEL('@{stage_path}/{model_deployment_spec_file_rel_path}')"
-        ).collect(block=False, statement_params=statement_params)
+        stage_path: Optional[str] = None,
+        model_deployment_spec_yaml_str: Optional[str] = None,
+        model_deployment_spec_file_rel_path: Optional[str] = None,
+        statement_params: Optional[dict[str, Any]] = None,
+    ) -> tuple[str, snowpark.AsyncJob]:
+        assert model_deployment_spec_yaml_str or model_deployment_spec_file_rel_path
+        if model_deployment_spec_yaml_str:
+            sql_str = f"CALL SYSTEM$DEPLOY_MODEL('{model_deployment_spec_yaml_str}')"
+        else:
+            sql_str = f"CALL SYSTEM$DEPLOY_MODEL('@{stage_path}/{model_deployment_spec_file_rel_path}')"
+        async_job = self._session.sql(sql_str).collect(block=False, statement_params=statement_params)
         assert isinstance(async_job, snowpark.AsyncJob)
         return async_job.query_id, async_job
 
@@ -91,9 +95,9 @@ class ServiceSQLClient(_base._BaseSQLClient):
         service_name: sql_identifier.SqlIdentifier,
         method_name: sql_identifier.SqlIdentifier,
         input_df: dataframe.DataFrame,
-        input_args: List[sql_identifier.SqlIdentifier],
-        returns: List[Tuple[str, spt.DataType, sql_identifier.SqlIdentifier]],
-        statement_params: Optional[Dict[str, Any]] = None,
+        input_args: list[sql_identifier.SqlIdentifier],
+        returns: list[tuple[str, spt.DataType, sql_identifier.SqlIdentifier]],
+        statement_params: Optional[dict[str, Any]] = None,
     ) -> dataframe.DataFrame:
         with_statements = []
         actual_database_name = database_name or self._database_name
@@ -177,7 +181,7 @@ class ServiceSQLClient(_base._BaseSQLClient):
         service_name: sql_identifier.SqlIdentifier,
         instance_id: str = "0",
         container_name: str,
-        statement_params: Optional[Dict[str, Any]] = None,
+        statement_params: Optional[dict[str, Any]] = None,
     ) -> str:
         system_func = "SYSTEM$GET_SERVICE_LOGS"
         rows = (
@@ -202,8 +206,8 @@ class ServiceSQLClient(_base._BaseSQLClient):
         schema_name: Optional[sql_identifier.SqlIdentifier],
         service_name: sql_identifier.SqlIdentifier,
         include_message: bool = False,
-        statement_params: Optional[Dict[str, Any]] = None,
-    ) -> Tuple[ServiceStatus, Optional[str]]:
+        statement_params: Optional[dict[str, Any]] = None,
+    ) -> tuple[ServiceStatus, Optional[str]]:
         system_func = "SYSTEM$GET_SERVICE_STATUS"
         rows = (
             query_result_checker.SqlResultValidator(
@@ -227,7 +231,7 @@ class ServiceSQLClient(_base._BaseSQLClient):
         database_name: Optional[sql_identifier.SqlIdentifier],
         schema_name: Optional[sql_identifier.SqlIdentifier],
         service_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,
@@ -241,8 +245,8 @@ class ServiceSQLClient(_base._BaseSQLClient):
         database_name: Optional[sql_identifier.SqlIdentifier],
         schema_name: Optional[sql_identifier.SqlIdentifier],
         service_name: sql_identifier.SqlIdentifier,
-        statement_params: Optional[Dict[str, Any]] = None,
-    ) -> List[row.Row]:
+        statement_params: Optional[dict[str, Any]] = None,
+    ) -> list[row.Row]:
         fully_qualified_service_name = self.fully_qualified_object_name(database_name, schema_name, service_name)
         res = (
             query_result_checker.SqlResultValidator(

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

@@ -1,4 +1,4 @@
-from typing import Any, Dict, 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
@@ -11,7 +11,7 @@ class StageSQLClient(_base._BaseSQLClient):
         database_name: Optional[sql_identifier.SqlIdentifier],
         schema_name: Optional[sql_identifier.SqlIdentifier],
         stage_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,

snowflake/ml/model/_client/sql/tag.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
@@ -16,7 +16,7 @@ class ModuleTagSQLClient(_base._BaseSQLClient):
         tag_schema_name: Optional[sql_identifier.SqlIdentifier],
         tag_name: sql_identifier.SqlIdentifier,
         tag_value: str,
-        statement_params: Optional[Dict[str, Any]] = None,
+        statement_params: Optional[dict[str, Any]] = None,
     ) -> None:
         fq_model_name = self.fully_qualified_object_name(database_name, schema_name, model_name)
         fq_tag_name = self.fully_qualified_object_name(tag_database_name, tag_schema_name, tag_name)
@@ -35,7 +35,7 @@ class ModuleTagSQLClient(_base._BaseSQLClient):
         tag_database_name: Optional[sql_identifier.SqlIdentifier],
         tag_schema_name: Optional[sql_identifier.SqlIdentifier],
         tag_name: sql_identifier.SqlIdentifier,
-        statement_params: Optional[Dict[str, Any]] = None,
+        statement_params: Optional[dict[str, Any]] = None,
     ) -> None:
         fq_model_name = self.fully_qualified_object_name(database_name, schema_name, model_name)
         fq_tag_name = self.fully_qualified_object_name(tag_database_name, tag_schema_name, tag_name)
@@ -54,7 +54,7 @@ class ModuleTagSQLClient(_base._BaseSQLClient):
         tag_database_name: Optional[sql_identifier.SqlIdentifier],
         tag_schema_name: Optional[sql_identifier.SqlIdentifier],
         tag_name: sql_identifier.SqlIdentifier,
-        statement_params: Optional[Dict[str, Any]] = None,
+        statement_params: Optional[dict[str, Any]] = None,
     ) -> row.Row:
         fq_model_name = self.fully_qualified_object_name(database_name, schema_name, model_name)
         fq_tag_name = self.fully_qualified_object_name(tag_database_name, tag_schema_name, tag_name)
@@ -75,8 +75,8 @@ class ModuleTagSQLClient(_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,
-    ) -> List[row.Row]:
+        statement_params: Optional[dict[str, Any]] = None,
+    ) -> list[row.Row]:
         fq_model_name = self.fully_qualified_object_name(database_name, schema_name, model_name)
         actual_database_name = database_name or self._database_name
         return (

snowflake/ml/model/_model_composer/model_composer.py

@@ -3,13 +3,14 @@ import tempfile
 import uuid
 import warnings
 from types import ModuleType
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Optional, Union
 from urllib import parse
 
 from absl import logging
 from packaging import requirements
 
 from snowflake import snowpark
+from snowflake.ml import version as snowml_version
 from snowflake.ml._internal import env as snowml_env, env_utils, file_utils
 from snowflake.ml._internal.lineage import lineage_utils
 from snowflake.ml.data import data_source
@@ -43,7 +44,8 @@ class ModelComposer:
         session: Session,
         stage_path: str,
         *,
-        statement_params: Optional[Dict[str, Any]] = None,
+        statement_params: Optional[dict[str, Any]] = None,
+        save_location: Optional[str] = None,
     ) -> None:
         self.session = session
         self.stage_path: Union[pathlib.PurePosixPath, parse.ParseResult] = None  # type: ignore[assignment]
@@ -54,10 +56,29 @@ class ModelComposer:
             # The stage path is a user stage path
             self.stage_path = pathlib.PurePosixPath(stage_path)
 
-        self._workspace = tempfile.TemporaryDirectory()
-        self._packager_workspace = tempfile.TemporaryDirectory()
+        # Set up workspace based on save_location if provided, otherwise use temporary directory
+        self.save_location = save_location
+        if save_location:
+            # Use the save_location directory directly
+            self._workspace_path = pathlib.Path(save_location)
+            self._workspace_path.mkdir(exist_ok=True)
+            # ensure that the directory is empty
+            if any(self._workspace_path.iterdir()):
+                raise ValueError(f"The directory {self._workspace_path} is not empty.")
+            self._workspace = None
+
+            self._packager_workspace_path = self._workspace_path / ModelComposer.MODEL_DIR_REL_PATH
+            self._packager_workspace_path.mkdir(exist_ok=True)
+            self._packager_workspace = None
+        else:
+            # Use a temporary directory
+            self._workspace = tempfile.TemporaryDirectory()
+            self._workspace_path = pathlib.Path(self._workspace.name)
+
+            self._packager_workspace_path = self._workspace_path / ModelComposer.MODEL_DIR_REL_PATH
+            self._packager_workspace_path.mkdir(exist_ok=True)
 
-        self.packager = model_packager.ModelPackager(local_dir_path=str(self._packager_workspace_path))
+        self.packager = model_packager.ModelPackager(local_dir_path=str(self.packager_workspace_path))
         self.manifest = model_manifest.ModelManifest(workspace_path=self.workspace_path)
 
         self.model_file_rel_path = f"model-{uuid.uuid4().hex}.zip"
@@ -65,16 +86,16 @@ class ModelComposer:
         self._statement_params = statement_params
 
     def __del__(self) -> None:
-        self._workspace.cleanup()
-        self._packager_workspace.cleanup()
+        if self._workspace:
+            self._workspace.cleanup()
 
     @property
     def workspace_path(self) -> pathlib.Path:
-        return pathlib.Path(self._workspace.name)
+        return self._workspace_path
 
     @property
-    def _packager_workspace_path(self) -> pathlib.Path:
-        return pathlib.Path(self._packager_workspace.name)
+    def packager_workspace_path(self) -> pathlib.Path:
+        return self._packager_workspace_path
 
     @property
     def model_stage_path(self) -> str:
@@ -102,17 +123,18 @@ class ModelComposer:
         *,
         name: str,
         model: model_types.SupportedModelType,
-        signatures: Optional[Dict[str, model_signature.ModelSignature]] = None,
+        signatures: Optional[dict[str, model_signature.ModelSignature]] = None,
         sample_input_data: Optional[model_types.SupportedDataType] = None,
-        metadata: Optional[Dict[str, str]] = None,
-        conda_dependencies: Optional[List[str]] = None,
-        pip_requirements: Optional[List[str]] = None,
-        artifact_repository_map: Optional[Dict[str, str]] = None,
-        target_platforms: Optional[List[model_types.TargetPlatform]] = None,
+        metadata: Optional[dict[str, str]] = None,
+        conda_dependencies: Optional[list[str]] = None,
+        pip_requirements: Optional[list[str]] = None,
+        artifact_repository_map: Optional[dict[str, str]] = None,
+        resource_constraint: Optional[dict[str, str]] = None,
+        target_platforms: Optional[list[model_types.TargetPlatform]] = None,
         python_version: Optional[str] = None,
-        user_files: Optional[Dict[str, List[str]]] = None,
-        ext_modules: Optional[List[ModuleType]] = None,
-        code_paths: Optional[List[str]] = None,
+        user_files: Optional[dict[str, list[str]]] = None,
+        ext_modules: Optional[list[ModuleType]] = None,
+        code_paths: Optional[list[str]] = None,
         task: model_types.Task = model_types.Task.UNKNOWN,
         options: Optional[model_types.ModelSaveOption] = None,
     ) -> model_meta.ModelMetadata:
@@ -146,14 +168,14 @@ class ModelComposer:
         if not snowpark_utils.is_in_stored_procedure():  # type: ignore[no-untyped-call]
             snowml_matched_versions = env_utils.get_matched_package_versions_in_information_schema(
                 self.session,
-                reqs=[requirements.Requirement(f"{env_utils.SNOWPARK_ML_PKG_NAME}=={snowml_env.VERSION}")],
+                reqs=[requirements.Requirement(f"{env_utils.SNOWPARK_ML_PKG_NAME}=={snowml_version.VERSION}")],
                 python_version=python_version or snowml_env.PYTHON_VERSION,
                 statement_params=self._statement_params,
             ).get(env_utils.SNOWPARK_ML_PKG_NAME, [])
 
             if len(snowml_matched_versions) < 1 and options.get("embed_local_ml_library", False) is False:
                 logging.info(
-                    f"Local snowflake-ml-python library has version {snowml_env.VERSION},"
+                    f"Local snowflake-ml-python library has version {snowml_version.VERSION},"
                     " which is not available in the Snowflake server, embedding local ML library automatically."
                 )
                 options["embed_local_ml_library"] = True
@@ -167,6 +189,8 @@ class ModelComposer:
             conda_dependencies=conda_dependencies,
             pip_requirements=pip_requirements,
             artifact_repository_map=artifact_repository_map,
+            resource_constraint=resource_constraint,
+            target_platforms=target_platforms,
             python_version=python_version,
             ext_modules=ext_modules,
             code_paths=code_paths,
@@ -175,9 +199,6 @@ class ModelComposer:
         )
         assert self.packager.meta is not None
 
-        file_utils.copytree(
-            str(self._packager_workspace_path), str(self.workspace_path / ModelComposer.MODEL_DIR_REL_PATH)
-        )
         self.manifest.save(
             model_meta=self.packager.meta,
             model_rel_path=pathlib.PurePosixPath(ModelComposer.MODEL_DIR_REL_PATH),
@@ -208,7 +229,7 @@ class ModelComposer:
 
     def _get_data_sources(
         self, model: model_types.SupportedModelType, sample_input_data: Optional[model_types.SupportedDataType] = None
-    ) -> Optional[List[data_source.DataSource]]:
+    ) -> Optional[list[data_source.DataSource]]:
         data_sources = lineage_utils.get_data_sources(model)
         if not data_sources and sample_input_data is not None:
             data_sources = lineage_utils.get_data_sources(sample_input_data)