snowflake-ml-python 1.7.4__py3-none-any.whl → 1.8.0__py3-none-any.whl

snowflake/ml/jobs/decorators.py

@@ -1,6 +1,5 @@
 import copy
 import functools
-import inspect
 from typing import Callable, Dict, List, Optional, TypeVar
 
 from typing_extensions import ParamSpec
@@ -8,7 +7,7 @@ from typing_extensions import ParamSpec
 from snowflake import snowpark
 from snowflake.ml._internal import telemetry
 from snowflake.ml.jobs import job as jb, manager as jm
-from snowflake.ml.jobs._utils import payload_utils
+from snowflake.ml.jobs._utils import constants
 
 _PROJECT = "MLJob"
 
@@ -26,6 +25,7 @@ def remote(
     query_warehouse: Optional[str] = None,
     env_vars: Optional[Dict[str, str]] = None,
     session: Optional[snowpark.Session] = None,
+    num_instances: Optional[int] = None,
 ) -> Callable[[Callable[_Args, _ReturnValue]], Callable[_Args, jb.MLJob]]:
     """
     Submit a job to the compute pool.
@@ -38,6 +38,7 @@ def remote(
         query_warehouse: The query warehouse to use. Defaults to session warehouse.
         env_vars: Environment variables to set in container
         session: The Snowpark session to use. If none specified, uses active session.
+        num_instances: The number of nodes in the job. If none specified, create a single node job.
 
     Returns:
         Decorator that dispatches invocations of the decorated function as remote jobs.
@@ -50,31 +51,12 @@ def remote(
         wrapped_func = copy.copy(func)
         wrapped_func.__code__ = wrapped_func.__code__.replace(co_firstlineno=func.__code__.co_firstlineno + 1)
 
-        # Validate function arguments based on signature
-        signature = inspect.signature(func)
-        pos_arg_names = []
-        for name, param in signature.parameters.items():
-            param_type = payload_utils.get_parameter_type(param)
-            if param_type is not None:
-                payload_utils.validate_parameter_type(param_type, name)
-            if param.kind in (param.POSITIONAL_ONLY, param.POSITIONAL_OR_KEYWORD):
-                pos_arg_names.append(name)
-
         @functools.wraps(func)
         def wrapper(*args: _Args.args, **kwargs: _Args.kwargs) -> jb.MLJob:
-            # Validate positional args
-            for i, arg in enumerate(args):
-                arg_name = pos_arg_names[i] if i < len(pos_arg_names) else f"args[{i}]"
-                payload_utils.validate_parameter_type(type(arg), arg_name)
-
-            # Validate keyword args
-            for k, v in kwargs.items():
-                payload_utils.validate_parameter_type(type(v), k)
-
-            arg_list = [str(v) for v in args] + [x for k, v in kwargs.items() for x in (f"--{k}", str(v))]
+            payload = functools.partial(func, *args, **kwargs)
+            setattr(payload, constants.IS_MLJOB_REMOTE_ATTR, True)
             job = jm._submit_job(
-                source=wrapped_func,
-                args=arg_list,
+                source=payload,
                 stage_name=stage_name,
                 compute_pool=compute_pool,
                 pip_requirements=pip_requirements,
@@ -82,8 +64,9 @@ def remote(
                 query_warehouse=query_warehouse,
                 env_vars=env_vars,
                 session=session,
+                num_instances=num_instances,
             )
-            assert isinstance(job, jb.MLJob)
+            assert isinstance(job, jb.MLJob), f"Unexpected job type: {type(job)}"
             return job
 
         return wrapper

snowflake/ml/jobs/job.py

@@ -4,7 +4,7 @@ from typing import Any, List, Optional, cast
 from snowflake import snowpark
 from snowflake.ml._internal import telemetry
 from snowflake.ml.jobs._utils import constants, types
-from snowflake.snowpark.context import get_active_session
+from snowflake.snowpark import context as sp_context
 
 _PROJECT = "MLJob"
 TERMINAL_JOB_STATUSES = {"FAILED", "DONE", "INTERNAL_ERROR"}
@@ -13,7 +13,7 @@ TERMINAL_JOB_STATUSES = {"FAILED", "DONE", "INTERNAL_ERROR"}
 class MLJob:
     def __init__(self, id: str, session: Optional[snowpark.Session] = None) -> None:
         self._id = id
-        self._session = session or get_active_session()
+        self._session = session or sp_context.get_active_session()
         self._status: types.JOB_STATUS = "PENDING"
 
     @property
@@ -79,7 +79,7 @@ class MLJob:
         return self.status
 
 
-@telemetry.send_api_usage_telemetry(project=_PROJECT)
+@telemetry.send_api_usage_telemetry(project=_PROJECT, func_params_to_log=["job_id"])
 def _get_status(session: snowpark.Session, job_id: str) -> types.JOB_STATUS:
     """Retrieve job execution status."""
     # TODO: snowflake-snowpark-python<1.24.0 shows spurious error messages on
@@ -90,7 +90,7 @@ def _get_status(session: snowpark.Session, job_id: str) -> types.JOB_STATUS:
     return cast(types.JOB_STATUS, row["status"])
 
 
-@telemetry.send_api_usage_telemetry(project=_PROJECT)
+@telemetry.send_api_usage_telemetry(project=_PROJECT, func_params_to_log=["job_id", "limit"])
 def _get_logs(session: snowpark.Session, job_id: str, limit: int = -1) -> str:
     """
     Retrieve the job's execution logs.

snowflake/ml/jobs/manager.py

@@ -213,6 +213,7 @@ def _submit_job(
     query_warehouse: Optional[str] = None,
     spec_overrides: Optional[Dict[str, Any]] = None,
     session: Optional[snowpark.Session] = None,
+    num_instances: Optional[int] = None,
 ) -> jb.MLJob:
     """
     Submit a job to the compute pool.
@@ -229,6 +230,7 @@ def _submit_job(
         query_warehouse: The query warehouse to use. Defaults to session warehouse.
         spec_overrides: Custom service specification overrides to apply.
         session: The Snowpark session to use. If none specified, uses active session.
+        num_instances: The number of instances to use for the job. If none specified, single node job is created.
 
     Returns:
         An object representing the submitted job.
@@ -254,6 +256,7 @@ def _submit_job(
         compute_pool=compute_pool,
         payload=uploaded_payload,
         args=args,
+        num_instances=num_instances,
     )
     spec_overrides = spec_utils.generate_spec_overrides(
         environment_vars=env_vars,
@@ -281,6 +284,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,

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,

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
 
@@ -41,7 +42,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 +122,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

@@ -78,6 +78,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 +125,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={

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: