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

snowflake/ml/jobs/manager.py

@@ -106,6 +106,8 @@ def submit_file(
     external_access_integrations: Optional[List[str]] = None,
     query_warehouse: Optional[str] = None,
     spec_overrides: Optional[Dict[str, Any]] = None,
+    num_instances: Optional[int] = None,
+    enable_metrics: bool = False,
     session: Optional[snowpark.Session] = None,
 ) -> jb.MLJob:
     """
@@ -121,6 +123,8 @@ def submit_file(
         external_access_integrations: A list of external access integrations.
         query_warehouse: The query warehouse to use. Defaults to session warehouse.
         spec_overrides: Custom service specification overrides to apply.
+        num_instances: The number of instances to use for the job. If none specified, single node job is created.
+        enable_metrics: Whether to enable metrics publishing for the job.
         session: The Snowpark session to use. If none specified, uses active session.
 
     Returns:
@@ -136,6 +140,8 @@ def submit_file(
         external_access_integrations=external_access_integrations,
         query_warehouse=query_warehouse,
         spec_overrides=spec_overrides,
+        num_instances=num_instances,
+        enable_metrics=enable_metrics,
         session=session,
     )
 
@@ -154,6 +160,8 @@ def submit_directory(
     external_access_integrations: Optional[List[str]] = None,
     query_warehouse: Optional[str] = None,
     spec_overrides: Optional[Dict[str, Any]] = None,
+    num_instances: Optional[int] = None,
+    enable_metrics: bool = False,
     session: Optional[snowpark.Session] = None,
 ) -> jb.MLJob:
     """
@@ -170,6 +178,8 @@ def submit_directory(
         external_access_integrations: A list of external access integrations.
         query_warehouse: The query warehouse to use. Defaults to session warehouse.
         spec_overrides: Custom service specification overrides to apply.
+        num_instances: The number of instances to use for the job. If none specified, single node job is created.
+        enable_metrics: Whether to enable metrics publishing for the job.
         session: The Snowpark session to use. If none specified, uses active session.
 
     Returns:
@@ -186,6 +196,8 @@ def submit_directory(
         external_access_integrations=external_access_integrations,
         query_warehouse=query_warehouse,
         spec_overrides=spec_overrides,
+        num_instances=num_instances,
+        enable_metrics=enable_metrics,
         session=session,
     )
 
@@ -212,6 +224,8 @@ def _submit_job(
     external_access_integrations: Optional[List[str]] = None,
     query_warehouse: Optional[str] = None,
     spec_overrides: Optional[Dict[str, Any]] = None,
+    num_instances: Optional[int] = None,
+    enable_metrics: bool = False,
     session: Optional[snowpark.Session] = None,
 ) -> jb.MLJob:
     """
@@ -228,6 +242,8 @@ def _submit_job(
         external_access_integrations: A list of external access integrations.
         query_warehouse: The query warehouse to use. Defaults to session warehouse.
         spec_overrides: Custom service specification overrides to apply.
+        num_instances: The number of instances to use for the job. If none specified, single node job is created.
+        enable_metrics: Whether to enable metrics publishing for the job.
         session: The Snowpark session to use. If none specified, uses active session.
 
     Returns:
@@ -254,6 +270,8 @@ def _submit_job(
         compute_pool=compute_pool,
         payload=uploaded_payload,
         args=args,
+        num_instances=num_instances,
+        enable_metrics=enable_metrics,
     )
     spec_overrides = spec_utils.generate_spec_overrides(
         environment_vars=env_vars,
@@ -281,6 +299,8 @@ def _submit_job(
     query_warehouse = query_warehouse or session.get_current_warehouse()
     if query_warehouse:
         query.append(f"QUERY_WAREHOUSE = {query_warehouse}")
+    if num_instances:
+        query.append(f"REPLICAS = {num_instances}")
 
     # Submit job
     query_text = "\n".join(line for line in query if line)

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

@@ -746,7 +746,7 @@ class ModelVersion(lineage_node.LineageNode):
         max_instances: int = 1,
         cpu_requests: Optional[str] = None,
         memory_requests: Optional[str] = None,
-        gpu_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,

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

@@ -1,3 +1,4 @@
+import enum
 import json
 import os
 import pathlib
@@ -31,6 +32,12 @@ from snowflake.snowpark import dataframe, row, session
 from snowflake.snowpark._internal import utils as snowpark_utils
 
 
+# An enum class to represent Create Or Alter Model SQL command.
+class ModelAction(enum.Enum):
+    CREATE = "CREATE"
+    ALTER = "ALTER"
+
+
 class ServiceInfo(TypedDict):
     name: str
     status: str
@@ -92,7 +99,7 @@ class ModelOperator:
             and self._model_version_client == __value._model_version_client
         )
 
-    def prepare_model_stage_path(
+    def prepare_model_temp_stage_path(
         self,
         *,
         database_name: Optional[sql_identifier.SqlIdentifier],
@@ -110,17 +117,28 @@ class ModelOperator:
         )
         return f"@{self._stage_client.fully_qualified_object_name(database_name, schema_name, stage_name)}/model"
 
-    def create_from_stage(
+    def get_model_version_stage_path(
+        self,
+        *,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
+        model_name: sql_identifier.SqlIdentifier,
+        version_name: sql_identifier.SqlIdentifier,
+    ) -> str:
+        return (
+            f"snow://model/{self._stage_client.fully_qualified_object_name(database_name, schema_name, model_name)}"
+            f"/versions/{version_name}/"
+        )
+
+    def get_model_action_from_model_name_and_version(
         self,
-        composed_model: model_composer.ModelComposer,
         *,
         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,
-    ) -> None:
-        stage_path = str(composed_model.stage_path)
+    ) -> ModelAction:
         if self.validate_existence(
             database_name=database_name,
             schema_name=schema_name,
@@ -140,6 +158,79 @@ class ModelOperator:
                     f" version {version_name} already existed."
                 )
             else:
+                return ModelAction.ALTER
+        else:
+            return ModelAction.CREATE
+
+    def add_or_create_live_version(
+        self,
+        *,
+        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,
+    ) -> None:
+        model_action = self.get_model_action_from_model_name_and_version(
+            database_name=database_name,
+            schema_name=schema_name,
+            model_name=model_name,
+            version_name=version_name,
+            statement_params=statement_params,
+        )
+        if model_action == ModelAction.CREATE:
+            self._model_version_client.create_live_version(
+                database_name=database_name,
+                schema_name=schema_name,
+                model_name=model_name,
+                version_name=version_name,
+                statement_params=statement_params,
+            )
+        elif model_action == ModelAction.ALTER:
+            self._model_version_client.add_live_version(
+                database_name=database_name,
+                schema_name=schema_name,
+                model_name=model_name,
+                version_name=version_name,
+                statement_params=statement_params,
+            )
+        else:
+            raise AssertionError(f"The model_action is {model_action}. Expected CREATE or ALTER.")
+
+    def create_from_stage(
+        self,
+        composed_model: model_composer.ModelComposer,
+        *,
+        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,
+        use_live_commit: Optional[bool] = False,
+    ) -> None:
+
+        if use_live_commit:
+            # if the model version is live, we can only commit the version
+            self._model_version_client.commit_version(
+                database_name=database_name,
+                schema_name=schema_name,
+                model_name=model_name,
+                version_name=version_name,
+                statement_params=statement_params,
+            )
+        else:
+            stage_path = str(composed_model.stage_path)
+            # if the model version is not live,
+            # find whether the model exists and whether the version exists
+            # and then decide whether to create or alter the model
+            model_action = self.get_model_action_from_model_name_and_version(
+                database_name=database_name,
+                schema_name=schema_name,
+                model_name=model_name,
+                version_name=version_name,
+                statement_params=statement_params,
+            )
+            if model_action == ModelAction.ALTER:
                 self._model_version_client.add_version_from_stage(
                     database_name=database_name,
                     schema_name=schema_name,
@@ -148,15 +239,17 @@ class ModelOperator:
                     version_name=version_name,
                     statement_params=statement_params,
                 )
-        else:
-            self._model_version_client.create_from_stage(
-                database_name=database_name,
-                schema_name=schema_name,
-                stage_path=stage_path,
-                model_name=model_name,
-                version_name=version_name,
-                statement_params=statement_params,
-            )
+            elif model_action == ModelAction.CREATE:
+                self._model_version_client.create_from_stage(
+                    database_name=database_name,
+                    schema_name=schema_name,
+                    stage_path=stage_path,
+                    model_name=model_name,
+                    version_name=version_name,
+                    statement_params=statement_params,
+                )
+            else:
+                raise AssertionError(f"The model_action is {model_action}. Expected CREATE or ALTER.")
 
     def create_from_model_version(
         self,
@@ -696,14 +789,17 @@ class ModelOperator:
         version_name: sql_identifier.SqlIdentifier,
         statement_params: Optional[Dict[str, Any]] = None,
     ) -> type_hints.Task:
-        model_spec = self._fetch_model_spec(
+        model_version = self._model_client.show_versions(
             database_name=database_name,
             schema_name=schema_name,
             model_name=model_name,
             version_name=version_name,
+            validate_result=True,
             statement_params=statement_params,
-        )
-        task_val = model_spec.get("task", type_hints.Task.UNKNOWN.value)
+        )[0]
+
+        model_attributes = json.loads(model_version.model_attributes)
+        task_val = model_attributes.get("task", type_hints.Task.UNKNOWN.value)
         return type_hints.Task(task_val)
 
     def get_functions(

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

@@ -100,7 +100,7 @@ class ServiceOperator:
         max_instances: int,
         cpu_requests: Optional[str],
         memory_requests: Optional[str],
-        gpu_requests: Optional[str],
+        gpu_requests: Optional[Union[int, str]],
         num_workers: Optional[int],
         max_batch_rows: Optional[int],
         force_rebuild: bool,
@@ -161,12 +161,16 @@ class ServiceOperator:
             statement_params=statement_params,
         )
 
-        # check if the inference service is already running
+        # check if the inference service is already running/suspended
         model_inference_service_exists = self._check_if_service_exists(
             database_name=service_database_name,
             schema_name=service_schema_name,
             service_name=service_name,
-            service_status_list_if_exists=[service_sql.ServiceStatus.READY],
+            service_status_list_if_exists=[
+                service_sql.ServiceStatus.READY,
+                service_sql.ServiceStatus.SUSPENDING,
+                service_sql.ServiceStatus.SUSPENDED,
+            ],
             statement_params=statement_params,
         )
 
@@ -309,7 +313,10 @@ class ServiceOperator:
                         set_service_log_metadata_to_model_inference(
                             service_log_meta,
                             model_inference_service,
-                            "Model Inference image build is not rebuilding the image and using previously built image.",
+                            (
+                                "Model Inference image build is not rebuilding the image, but using a previously built "
+                                "image."
+                            ),
                         )
                         continue
 
@@ -366,7 +373,9 @@ class ServiceOperator:
             time.sleep(5)
 
         if model_inference_service_exists:
-            module_logger.info(f"Inference service {model_inference_service.display_service_name} is already RUNNING.")
+            module_logger.info(
+                f"Inference service {model_inference_service.display_service_name} has already been deployed."
+            )
         else:
             self._finalize_logs(
                 service_log_meta.service_logger, service_log_meta.service, service_log_meta.log_offset, statement_params
@@ -416,6 +425,8 @@ class ServiceOperator:
             service_status_list_if_exists = [
                 service_sql.ServiceStatus.PENDING,
                 service_sql.ServiceStatus.READY,
+                service_sql.ServiceStatus.SUSPENDING,
+                service_sql.ServiceStatus.SUSPENDED,
                 service_sql.ServiceStatus.DONE,
                 service_sql.ServiceStatus.FAILED,
             ]

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

@@ -1,5 +1,5 @@
 import pathlib
-from typing import List, Optional
+from typing import List, Optional, Union
 
 import yaml
 
@@ -38,7 +38,7 @@ class ModelDeploymentSpec:
         max_instances: int,
         cpu: Optional[str],
         memory: Optional[str],
-        gpu: Optional[str],
+        gpu: Optional[Union[str, int]],
         num_workers: Optional[int],
         max_batch_rows: Optional[int],
         force_rebuild: bool,
@@ -86,7 +86,11 @@ class ModelDeploymentSpec:
             service_dict["memory"] = memory
 
         if gpu:
-            service_dict["gpu"] = gpu
+            if isinstance(gpu, int):
+                gpu_str = str(gpu)
+            else:
+                gpu_str = gpu
+            service_dict["gpu"] = gpu_str
 
         if num_workers:
             service_dict["num_workers"] = num_workers

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

@@ -71,6 +71,64 @@ class ModelVersionSQLClient(_base._BaseSQLClient):
             statement_params=statement_params,
         ).has_dimensions(expected_rows=1, expected_cols=1).validate()
 
+    def create_live_version(
+        self,
+        *,
+        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,
+    ) -> None:
+        sql = (
+            f"CREATE MODEL {self.fully_qualified_object_name(database_name, schema_name, model_name)}"
+            f" WITH LIVE VERSION {version_name.identifier()}"
+        )
+        query_result_checker.SqlResultValidator(
+            self._session,
+            sql,
+            statement_params=statement_params,
+        ).has_dimensions(expected_rows=1, expected_cols=1).validate()
+
+    def add_live_version(
+        self,
+        *,
+        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,
+    ) -> None:
+        sql = (
+            f"ALTER MODEL {self.fully_qualified_object_name(database_name, schema_name, model_name)}"
+            f" ADD LIVE VERSION {version_name.identifier()}"
+        )
+        query_result_checker.SqlResultValidator(
+            self._session,
+            sql,
+            statement_params=statement_params,
+        ).has_dimensions(expected_rows=1, expected_cols=1).validate()
+
+    def commit_version(
+        self,
+        *,
+        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,
+    ) -> None:
+        sql = (
+            f"ALTER MODEL {self.fully_qualified_object_name(database_name, schema_name, model_name)}"
+            f" COMMIT VERSION {version_name.identifier()}"
+        )
+
+        query_result_checker.SqlResultValidator(
+            self._session,
+            sql,
+            statement_params=statement_params,
+        ).has_dimensions(expected_rows=1, expected_cols=1).validate()
+
     # TODO(SNOW-987381): Merge with above when we have `create or alter module m [with] version v1 ...`
     def add_version_from_stage(
         self,

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
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 from snowflake import snowpark
 from snowflake.ml._internal import platform_capabilities
@@ -11,6 +11,7 @@ from snowflake.ml._internal.utils import (
     sql_identifier,
 )
 from snowflake.ml.model._client.sql import _base
+from snowflake.ml.model._model_composer.model_method import constants
 from snowflake.snowpark import dataframe, functions as F, row, types as spt
 from snowflake.snowpark._internal import utils as snowpark_utils
 
@@ -19,6 +20,8 @@ class ServiceStatus(enum.Enum):
     UNKNOWN = "UNKNOWN"  # status is unknown because we have not received enough data from K8s yet.
     PENDING = "PENDING"  # resource set is being created, can't be used yet
     READY = "READY"  # resource set has been deployed.
+    SUSPENDING = "SUSPENDING"  # the service is set to suspended but the resource set is still in deleting state
+    SUSPENDED = "SUSPENDED"  # the service is suspended and the resource set is deleted
     DELETING = "DELETING"  # resource set is being deleted
     FAILED = "FAILED"  # resource set has failed and cannot be used anymore
     DONE = "DONE"  # resource set has finished running
@@ -41,7 +44,7 @@ class ServiceSQLClient(_base._BaseSQLClient):
         image_repo_database_name: Optional[sql_identifier.SqlIdentifier],
         image_repo_schema_name: Optional[sql_identifier.SqlIdentifier],
         image_repo_name: sql_identifier.SqlIdentifier,
-        gpu: Optional[str],
+        gpu: Optional[Union[str, int]],
         force_rebuild: bool,
         external_access_integration: sql_identifier.SqlIdentifier,
         statement_params: Optional[Dict[str, Any]] = None,
@@ -121,6 +124,11 @@ class ServiceSQLClient(_base._BaseSQLClient):
             args_sql_list.append(input_arg_value)
         args_sql = ", ".join(args_sql_list)
 
+        wide_input = len(input_args) > constants.SNOWPARK_UDF_INPUT_COL_LIMIT
+        if wide_input:
+            input_args_sql = ", ".join(f"'{arg}', {arg.identifier()}" for arg in input_args)
+            args_sql = f"object_construct_keep_null({input_args_sql})"
+
         if platform_capabilities.PlatformCapabilities.get_instance().is_nested_function_enabled():
             fully_qualified_service_name = self.fully_qualified_object_name(
                 actual_database_name, actual_schema_name, service_name

snowflake/ml/model/_model_composer/model_composer.py

@@ -1,8 +1,10 @@
 import pathlib
 import tempfile
 import uuid
+import warnings
 from types import ModuleType
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Union
+from urllib import parse
 
 from absl import logging
 from packaging import requirements
@@ -44,7 +46,13 @@ class ModelComposer:
         statement_params: Optional[Dict[str, Any]] = None,
     ) -> None:
         self.session = session
-        self.stage_path = pathlib.PurePosixPath(stage_path)
+        self.stage_path: Union[pathlib.PurePosixPath, parse.ParseResult] = None  # type: ignore[assignment]
+        if stage_path.startswith("snow://"):
+            # The stage path is a snowflake internal stage path
+            self.stage_path = parse.urlparse(stage_path)
+        else:
+            # The stage path is a user stage path
+            self.stage_path = pathlib.PurePosixPath(stage_path)
 
         self._workspace = tempfile.TemporaryDirectory()
         self._packager_workspace = tempfile.TemporaryDirectory()
@@ -70,7 +78,20 @@ class ModelComposer:
 
     @property
     def model_stage_path(self) -> str:
-        return (self.stage_path / self.model_file_rel_path).as_posix()
+        if isinstance(self.stage_path, parse.ParseResult):
+            model_file_path = (pathlib.PosixPath(self.stage_path.path) / self.model_file_rel_path).as_posix()
+            new_url = parse.ParseResult(
+                scheme=self.stage_path.scheme,
+                netloc=self.stage_path.netloc,
+                path=str(model_file_path),
+                params=self.stage_path.params,
+                query=self.stage_path.query,
+                fragment=self.stage_path.fragment,
+            )
+            return str(parse.urlunparse(new_url))
+        else:
+            assert isinstance(self.stage_path, pathlib.PurePosixPath)
+            return (self.stage_path / self.model_file_rel_path).as_posix()
 
     @property
     def model_local_path(self) -> str:
@@ -86,6 +107,7 @@ class ModelComposer:
         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,
         python_version: Optional[str] = None,
         user_files: Optional[Dict[str, List[str]]] = None,
@@ -94,8 +116,32 @@ class ModelComposer:
         task: model_types.Task = model_types.Task.UNKNOWN,
         options: Optional[model_types.ModelSaveOption] = None,
     ) -> model_meta.ModelMetadata:
+        # set enable_explainability=False if the model is not runnable in WH or the target platforms include SPCS
+        conda_dep_dict = env_utils.validate_conda_dependency_string_list(
+            conda_dependencies if conda_dependencies else []
+        )
+        is_warehouse_runnable = (
+            not conda_dep_dict
+            or all(
+                chan == env_utils.DEFAULT_CHANNEL_NAME or chan == env_utils.SNOWFLAKE_CONDA_CHANNEL_URL
+                for chan in conda_dep_dict
+            )
+        ) and (not pip_requirements)
+        disable_explainability = (
+            target_platforms and model_types.TargetPlatform.SNOWPARK_CONTAINER_SERVICES in target_platforms
+        ) or (not is_warehouse_runnable)
+
+        if disable_explainability and options and options.get("enable_explainability", False):
+            warnings.warn(
+                ("The model can be deployed to Snowpark Container Services only if `enable_explainability=False`."),
+                category=UserWarning,
+                stacklevel=2,
+            )
+
         if not options:
             options = model_types.BaseModelSaveOption()
+            if disable_explainability:
+                options["enable_explainability"] = False
 
         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(
@@ -120,6 +166,7 @@ class ModelComposer:
             metadata=metadata,
             conda_dependencies=conda_dependencies,
             pip_requirements=pip_requirements,
+            artifact_repository_map=artifact_repository_map,
             python_version=python_version,
             ext_modules=ext_modules,
             code_paths=code_paths,

snowflake/ml/model/_model_composer/model_manifest/model_manifest.py

@@ -36,7 +36,6 @@ class ModelManifest:
     """
 
     MANIFEST_FILE_REL_PATH = "MANIFEST.yml"
-    _ENABLE_USER_FILES = False
     _DEFAULT_RUNTIME_NAME = "python_runtime"
 
     def __init__(self, workspace_path: pathlib.Path) -> None:
@@ -78,6 +77,7 @@ class ModelManifest:
             logger.info("Relaxing version constraints for dependencies in the model.")
             logger.info(f"Conda dependencies: {runtime_to_use.runtime_env.conda_dependencies}")
             logger.info(f"Pip requirements: {runtime_to_use.runtime_env.pip_requirements}")
+            logger.info(f"artifact_repository_map: {runtime_to_use.runtime_env.artifact_repository_map}")
         runtime_dict = runtime_to_use.save(
             self.workspace_path, default_channel_override=env_utils.SNOWFLAKE_CONDA_CHANNEL_URL
         )
@@ -124,6 +124,9 @@ class ModelManifest:
         if len(model_meta.env.pip_requirements) > 0:
             dependencies["pip"] = runtime_dict["dependencies"]["pip"]
 
+        if model_meta.env.artifact_repository_map:
+            dependencies["artifact_repository_map"] = runtime_dict["dependencies"]["artifact_repository_map"]
+
         manifest_dict = model_manifest_schema.ModelManifestDict(
             manifest_version=model_manifest_schema.MODEL_MANIFEST_VERSION,
             runtimes={
@@ -145,7 +148,7 @@ class ModelManifest:
             ],
         )
 
-        if self._ENABLE_USER_FILES:
+        if self.user_files:
             manifest_dict["user_files"] = [user_file.save(self.workspace_path) for user_file in self.user_files]
 
         lineage_sources = self._extract_lineage_info(data_sources)

snowflake/ml/model/_model_composer/model_manifest/model_manifest_schema.py

@@ -1,6 +1,6 @@
 # This files contains schema definition of what will be written into MANIFEST.yml
 import enum
-from typing import Any, Dict, List, Literal, TypedDict, Union
+from typing import Any, Dict, List, Literal, Optional, TypedDict, Union
 
 from typing_extensions import NotRequired, Required
 
@@ -20,6 +20,7 @@ class ModelMethodFunctionTypes(enum.Enum):
 class ModelRuntimeDependenciesDict(TypedDict):
     conda: NotRequired[str]
     pip: NotRequired[str]
+    artifact_repository_map: NotRequired[Optional[Dict[str, str]]]
 
 
 class ModelRuntimeDict(TypedDict):

snowflake/ml/model/_model_composer/model_method/model_method.py

@@ -98,7 +98,6 @@ class ModelMethod:
     def _get_method_arg_from_feature(
         feature: model_signature.BaseFeatureSpec, case_sensitive: bool = False
     ) -> model_manifest_schema.ModelMethodSignatureFieldWithName:
-        assert isinstance(feature, model_signature.FeatureSpec), "FeatureGroupSpec is not supported."
         try:
             feature_name = sql_identifier.SqlIdentifier(feature.name, case_sensitive=case_sensitive)
         except ValueError as e:

snowflake/ml/model/_packager/model_env/model_env.py

@@ -3,7 +3,7 @@ import itertools
 import os
 import pathlib
 import warnings
-from typing import DefaultDict, List, Optional
+from typing import DefaultDict, Dict, List, Optional
 
 from packaging import requirements, version
 
@@ -36,6 +36,7 @@ class ModelEnv:
             pip_requirements_rel_path = os.path.join(_DEFAULT_ENV_DIR, _DEFAULT_PIP_REQUIREMENTS_FILENAME)
         self.conda_env_rel_path = pathlib.PurePosixPath(pathlib.Path(conda_env_rel_path).as_posix())
         self.pip_requirements_rel_path = pathlib.PurePosixPath(pathlib.Path(pip_requirements_rel_path).as_posix())
+        self.artifact_repository_map: Optional[Dict[str, str]] = None
         self._conda_dependencies: DefaultDict[str, List[requirements.Requirement]] = collections.defaultdict(list)
         self._pip_requirements: List[requirements.Requirement] = []
         self._python_version: version.Version = version.parse(snowml_env.PYTHON_VERSION)
@@ -345,6 +346,7 @@ class ModelEnv:
     def load_from_dict(self, base_dir: pathlib.Path, env_dict: model_meta_schema.ModelEnvDict) -> None:
         self.conda_env_rel_path = pathlib.PurePosixPath(env_dict["conda"])
         self.pip_requirements_rel_path = pathlib.PurePosixPath(env_dict["pip"])
+        self.artifact_repository_map = env_dict.get("artifact_repository_map", None)
 
         self.load_from_conda_file(base_dir / self.conda_env_rel_path)
         self.load_from_pip_file(base_dir / self.pip_requirements_rel_path)
@@ -373,6 +375,7 @@ class ModelEnv:
         return {
             "conda": self.conda_env_rel_path.as_posix(),
             "pip": self.pip_requirements_rel_path.as_posix(),
+            "artifact_repository_map": self.artifact_repository_map if self.artifact_repository_map is not None else {},
             "python_version": self.python_version,
             "cuda_version": self.cuda_version,
             "snowpark_ml_version": self.snowpark_ml_version,