snowflake-ml-python 1.5.1__py3-none-any.whl → 1.5.2__py3-none-any.whl

snowflake/cortex/_sentiment.py

@@ -11,7 +11,7 @@ from snowflake.ml._internal import telemetry
 )
 def Sentiment(
     text: Union[str, snowpark.Column], session: Optional[snowpark.Session] = None
-) -> Union[str, snowpark.Column]:
+) -> Union[float, snowpark.Column]:
     """Sentiment calls into the LLM inference service to perform sentiment analysis on the input text.
 
     Args:
@@ -21,11 +21,14 @@ def Sentiment(
     Returns:
         A column of floats. 1 represents positive sentiment, -1 represents negative sentiment.
     """
-
     return _sentiment_impl("snowflake.cortex.sentiment", text, session=session)
 
 
 def _sentiment_impl(
     function: str, text: Union[str, snowpark.Column], session: Optional[snowpark.Session] = None
-) -> Union[str, snowpark.Column]:
-    return call_sql_function(function, session, text)
+) -> Union[float, snowpark.Column]:
+
+    output = call_sql_function(function, session, text)
+    if isinstance(output, snowpark.Column):
+        return output
+    return float(output)

snowflake/ml/_internal/utils/temp_file_utils.py

@@ -8,14 +8,17 @@ from absl.logging import logging
 logger = logging.getLogger(__name__)
 
 
-def get_temp_file_path() -> str:
+def get_temp_file_path(prefix: str = "") -> str:
     """Returns a new random temp file path.
 
+    Args:
+        prefix: A prefix to the temp file path, this can help add stored file information. Defaults to None.
+
     Returns:
         A new temp file path.
     """
     # TODO(snandamuri): Use in-memory filesystem for temp files.
-    local_file = tempfile.NamedTemporaryFile(delete=True)
+    local_file = tempfile.NamedTemporaryFile(prefix=prefix, delete=True)
     local_file_name = local_file.name
     local_file.close()
     return local_file_name

snowflake/ml/feature_store/access_manager.py

@@ -42,6 +42,8 @@ class _SessionInfo:
 # Lists of permissions as tuples of (OBJECT_TYPE, [PRIVILEGES, ...])
 _PRE_INIT_PRIVILEGES: Dict[_FeatureStoreRole, List[_Privilege]] = {
     _FeatureStoreRole.PRODUCER: [
+        _Privilege("DATABASE", "{database}", ["USAGE"]),
+        _Privilege("SCHEMA", "{database}.{schema}", ["USAGE"]),
         _Privilege(
             "SCHEMA",
             "{database}.{schema}",
@@ -69,8 +71,7 @@ _PRE_INIT_PRIVILEGES: Dict[_FeatureStoreRole, List[_Privilege]] = {
         _Privilege("DYNAMIC TABLE", _ALL_OBJECTS, ["SELECT", "MONITOR"], "SCHEMA {database}.{schema}"),
         _Privilege("VIEW", _ALL_OBJECTS, ["SELECT", "REFERENCES"], "SCHEMA {database}.{schema}"),
         _Privilege("TABLE", _ALL_OBJECTS, ["SELECT", "REFERENCES"], "SCHEMA {database}.{schema}"),
-        # FIXME(dhung): FUTURE DATASETS not supported until 8.19
-        # _Privilege("DATASET", _ALL_OBJECTS, ["USAGE"], "SCHEMA {database}.{schema}"),
+        _Privilege("DATASET", _ALL_OBJECTS, ["USAGE"], "SCHEMA {database}.{schema}"),
         # User should decide whether they want to grant warehouse usage to CONSUMER
         # _Privilege("WAREHOUSE", "{warehouse}", ["USAGE"]),
     ],
@@ -128,8 +129,7 @@ def _grant_privileges(
 def _configure_pre_init_privileges(
     session: Session,
     session_info: _SessionInfo,
-    producer_role: str = "SNOWML_FEATURE_STORE_PRODUCER_RL",
-    consumer_role: str = "SNOWML_FEATURE_STORE_CONSUMER_RL",
+    roles_to_create: Dict[_FeatureStoreRole, str],
 ) -> None:
     """
     Configure Feature Store role privileges. Must be run with ACCOUNTADMIN
@@ -141,8 +141,7 @@ def _configure_pre_init_privileges(
     Args:
         session: Snowpark Session to interact with Snowflake backend.
         session_info: Session info like database and schema for the FeatureStore instance.
-        producer_role: Name of producer role to be configured.
-        consumer_role: Name of consumer role to be configured.
+        roles_to_create: Producer and optional consumer roles to create.
     """
 
     # Create schema if not already exists
@@ -159,29 +158,30 @@ def _configure_pre_init_privileges(
 
     # Pass schema ownership from admin to PRODUCER
     if schema_created:
+        # TODO: we are missing a test case for this code path
         session.sql(
-            f"GRANT OWNERSHIP ON SCHEMA {session_info.database}.{session_info.schema} TO ROLE {producer_role}"
+            f"GRANT OWNERSHIP ON SCHEMA {session_info.database}.{session_info.schema} "
+            f"TO ROLE {roles_to_create[_FeatureStoreRole.PRODUCER]}"
         ).collect()
 
     # Grant privileges to roles
-    _grant_privileges(session, producer_role, _PRE_INIT_PRIVILEGES[_FeatureStoreRole.PRODUCER], session_info)
-    _grant_privileges(session, consumer_role, _PRE_INIT_PRIVILEGES[_FeatureStoreRole.CONSUMER], session_info)
+    for role_type, role in roles_to_create.items():
+        _grant_privileges(session, role, _PRE_INIT_PRIVILEGES[role_type], session_info)
 
 
 def _configure_post_init_privileges(
     session: Session,
     session_info: _SessionInfo,
-    producer_role: str = "FS_PRODUCER",
-    consumer_role: str = "FS_CONSUMER",
+    roles_to_create: Dict[_FeatureStoreRole, str],
 ) -> None:
-    _grant_privileges(session, producer_role, _POST_INIT_PRIVILEGES[_FeatureStoreRole.PRODUCER], session_info)
-    _grant_privileges(session, consumer_role, _POST_INIT_PRIVILEGES[_FeatureStoreRole.CONSUMER], session_info)
+    for role_type, role in roles_to_create.items():
+        _grant_privileges(session, role, _POST_INIT_PRIVILEGES[role_type], session_info)
 
 
 def _configure_role_hierarchy(
     session: Session,
     producer_role: str,
-    consumer_role: str,
+    consumer_role: Optional[str],
 ) -> None:
     """
     Create Feature Store roles and configure role hierarchy hierarchy. Must be run with
@@ -195,18 +195,17 @@ def _configure_role_hierarchy(
         producer_role: Name of producer role to be configured.
         consumer_role: Name of consumer role to be configured.
     """
+    # Create the necessary roles and build role hierarchy
     producer_role = SqlIdentifier(producer_role)
-    consumer_role = SqlIdentifier(consumer_role)
-
-    # Create the necessary roles
     session.sql(f"CREATE ROLE IF NOT EXISTS {producer_role}").collect()
-    session.sql(f"CREATE ROLE IF NOT EXISTS {consumer_role}").collect()
-
-    # Build role hierarchy
-    session.sql(f"GRANT ROLE {consumer_role} TO ROLE {producer_role}").collect()
     session.sql(f"GRANT ROLE {producer_role} TO ROLE SYSADMIN").collect()
     session.sql(f"GRANT ROLE {producer_role} TO ROLE {session.get_current_role()}").collect()
 
+    if consumer_role is not None:
+        consumer_role = SqlIdentifier(consumer_role)
+        session.sql(f"CREATE ROLE IF NOT EXISTS {consumer_role}").collect()
+        session.sql(f"GRANT ROLE {consumer_role} TO ROLE {producer_role}").collect()
+
 
 @telemetry.send_api_usage_telemetry(project=_PROJECT)
 def setup_feature_store(
@@ -215,7 +214,7 @@ def setup_feature_store(
     schema: str,
     warehouse: str,
     producer_role: str = "FS_PRODUCER",
-    consumer_role: str = "FS_CONSUMER",
+    consumer_role: Optional[str] = None,
 ) -> FeatureStore:
     """
     Sets up a new Feature Store including role/privilege setup. Must be run with ACCOUNTADMIN
@@ -230,7 +229,7 @@ def setup_feature_store(
         schema: Schema to create the FeatureStore instance.
         warehouse: Default warehouse for Feature Store compute.
         producer_role: Name of producer role to be configured.
-        consumer_role: Name of consumer role to be configured.
+        consumer_role: Name of consumer role to be configured. If not specified, consumer role won't be created.
 
     Returns:
         Feature Store instance.
@@ -249,20 +248,25 @@ def setup_feature_store(
     )
 
     try:
+        roles_to_create = {_FeatureStoreRole.PRODUCER: producer_role}
+        if consumer_role is not None:
+            roles_to_create.update({_FeatureStoreRole.CONSUMER: consumer_role})
         _configure_role_hierarchy(session, producer_role=producer_role, consumer_role=consumer_role)
     except exceptions.SnowparkSQLException:
         # Error can be safely ignored if roles already exist and hierarchy is already built
-        for role in (producer_role, consumer_role):
+        for _, role in roles_to_create.items():
             # Ensure roles already exist
             if session.sql(f"SHOW ROLES LIKE '{role}' STARTS WITH '{role}'").count() == 0:
                 raise
-        # Ensure hierarchy already configured
-        consumer_grants = session.sql(f"SHOW GRANTS ON ROLE {consumer_role}").collect()
-        if not any(r["granted_to"] == "ROLE" and r["grantee_name"] == producer_role for r in consumer_grants):
-            raise
+
+        if consumer_role is not None:
+            # Ensure hierarchy already configured
+            consumer_grants = session.sql(f"SHOW GRANTS ON ROLE {consumer_role}").collect()
+            if not any(r["granted_to"] == "ROLE" and r["grantee_name"] == producer_role for r in consumer_grants):
+                raise
 
     # Do any pre-FeatureStore.__init__() privilege setup
-    _configure_pre_init_privileges(session, session_info, producer_role, consumer_role)
+    _configure_pre_init_privileges(session, session_info, roles_to_create)
 
     # Use PRODUCER role to create and operate new Feature Store
     current_role = session.get_current_role()
@@ -274,6 +278,6 @@ def setup_feature_store(
         session.use_role(current_role)
 
     # Do any post-FeatureStore.__init__() privilege setup
-    _configure_post_init_privileges(session, session_info, producer_role, consumer_role)
+    _configure_post_init_privileges(session, session_info, roles_to_create)
 
     return fs

snowflake/ml/feature_store/feature_store.py

@@ -1761,7 +1761,7 @@ class FeatureStore:
             self._session.sql(
                 f"""
                     SELECT * FROM TABLE(
-                        INFORMATION_SCHEMA.TAG_REFERENCES_INTERNAL(
+                        {self._config.database}.INFORMATION_SCHEMA.TAG_REFERENCES_INTERNAL(
                             TAG_NAME => '{_FEATURE_STORE_OBJECT_TAG}'
                         )
                     ) LIMIT 1;

snowflake/ml/feature_store/feature_view.py

@@ -7,10 +7,6 @@ from dataclasses import asdict, dataclass
 from enum import Enum
 from typing import Any, Dict, List, Optional
 
-from snowflake.ml._internal.exceptions import (
-    error_codes,
-    exceptions as snowml_exceptions,
-)
 from snowflake.ml._internal.utils.identifier import concat_names
 from snowflake.ml._internal.utils.sql_identifier import (
     SqlIdentifier,
@@ -34,6 +30,11 @@ _FEATURE_OBJ_TYPE = "FEATURE_OBJ_TYPE"
 _FEATURE_VIEW_VERSION_RE = re.compile(r"^[a-zA-Z0-9][a-zA-Z0-9_.\-]*$")
 _FEATURE_VIEW_VERSION_MAX_LENGTH = 128
 
+_RESULT_SCAN_QUERY_PATTERN = re.compile(
+    r".*FROM\s*TABLE\s*\(\s*RESULT_SCAN\s*\(.*",
+    flags=re.DOTALL | re.IGNORECASE | re.X,
+)
+
 
 @dataclass(frozen=True)
 class _FeatureViewMetadata:
@@ -54,13 +55,10 @@ class _FeatureViewMetadata:
 class FeatureViewVersion(str):
     def __new__(cls, version: str) -> FeatureViewVersion:
         if not _FEATURE_VIEW_VERSION_RE.match(version) or len(version) > _FEATURE_VIEW_VERSION_MAX_LENGTH:
-            raise snowml_exceptions.SnowflakeMLException(
-                error_code=error_codes.INVALID_ARGUMENT,
-                original_exception=ValueError(
-                    f"`{version}` is not a valid feature view version. "
-                    "It must start with letter or digit, and followed by letter, digit, '_', '-' or '.'. "
-                    f"The length limit is {_FEATURE_VIEW_VERSION_MAX_LENGTH}."
-                ),
+            raise ValueError(
+                f"`{version}` is not a valid feature view version. "
+                "It must start with letter or digit, and followed by letter, digit, '_', '-' or '.'. "
+                f"The length limit is {_FEATURE_VIEW_VERSION_MAX_LENGTH}."
             )
         return super().__new__(cls, version)
 
@@ -352,6 +350,9 @@ Got {len(self._feature_df.queries['queries'])}: {self._feature_df.queries['queri
             if not isinstance(col_type, (DateType, TimeType, TimestampType, _NumericType)):
                 raise ValueError(f"Invalid data type for timestamp_col {ts_col}: {col_type}.")
 
+        if re.match(_RESULT_SCAN_QUERY_PATTERN, self._query) is not None:
+            raise ValueError(f"feature_df should not be reading from RESULT_SCAN. Invalid query: {self._query}")
+
     def _get_feature_names(self) -> List[SqlIdentifier]:
         join_keys = [k for e in self._entities for k in e.join_keys]
         ts_col = [self._timestamp_col] if self._timestamp_col is not None else []

snowflake/ml/fileset/snowfs.py

@@ -1,10 +1,9 @@
 import collections
 import logging
 import re
-from typing import Any, Dict, Optional
+from typing import Any, Optional
 
 import fsspec
-import packaging.version as pkg_version
 
 from snowflake import snowpark
 from snowflake.connector import connection
@@ -12,7 +11,7 @@ from snowflake.ml._internal.exceptions import (
     error_codes,
     exceptions as snowml_exceptions,
 )
-from snowflake.ml._internal.utils import identifier, snowflake_env
+from snowflake.ml._internal.utils import identifier
 from snowflake.ml.fileset import embedded_stage_fs, sfcfs
 
 PROTOCOL_NAME = "snow"
@@ -28,10 +27,6 @@ _SNOWURL_PATTERN = re.compile(
     r"(?P<path>versions/(?:(?P<version>[^/]+)(?:/(?P<relpath>.*))?)?)"
 )
 
-# FIXME(dhung): Temporary fix for bug in GS version 8.17
-_BUG_VERSION_MIN = pkg_version.Version("8.17")  # Inclusive minimum version with bugged behavior
-_BUG_VERSION_MAX = pkg_version.Version("8.18")  # Exclusive maximum version with bugged behavior
-
 
 class SnowFileSystem(sfcfs.SFFileSystem):
     """A filesystem that allows user to access Snowflake embedded stage files with valid Snowflake locations.
@@ -54,21 +49,6 @@ class SnowFileSystem(sfcfs.SFFileSystem):
     ) -> None:
         super().__init__(sf_connection=sf_connection, snowpark_session=snowpark_session, **kwargs)
 
-        # FIXME(dhung): Temporary fix for bug in GS version 8.17
-        if SnowFileSystem._IS_BUGGED_VERSION is None:
-            try:
-                sf_version = snowflake_env.get_current_snowflake_version(self._session)
-                SnowFileSystem._IS_BUGGED_VERSION = _BUG_VERSION_MIN <= sf_version < _BUG_VERSION_MAX
-            except Exception:
-                SnowFileSystem._IS_BUGGED_VERSION = False
-
-    def info(self, path: str, **kwargs: Any) -> Dict[str, Any]:
-        # FIXME(dhung): Temporary fix for bug in GS version 8.17
-        res: Dict[str, Any] = super().info(path, **kwargs)
-        if res.get("type") == "directory" and not res["name"].endswith("/"):
-            res["name"] += "/"
-        return res
-
     def _get_stage_fs(
         self, sf_file_path: _SFFileEntityPath  # type: ignore[override]
     ) -> embedded_stage_fs.SFEmbeddedStageFileSystem:
@@ -100,12 +80,6 @@ class SnowFileSystem(sfcfs.SFFileSystem):
         if stage_name.startswith(protocol):
             stage_name = stage_name[len(protocol) :]
         abs_path = stage_name + "/" + path
-        # FIXME(dhung): Temporary fix for bug in GS version 8.17
-        if self._IS_BUGGED_VERSION:
-            match = _SNOWURL_PATTERN.fullmatch(abs_path)
-            assert match is not None
-            if match.group("relpath"):
-                abs_path = abs_path.replace(match.group("relpath"), match.group("relpath").lstrip("/"))
         return abs_path
 
     @classmethod
@@ -144,9 +118,6 @@ class SnowFileSystem(sfcfs.SFFileSystem):
             version = snowurl_match.group("version")
             relative_path = snowurl_match.group("relpath") or ""
             logging.debug(f"Parsed snow URL: {snowurl_match.groups()}")
-            # FIXME(dhung): Temporary fix for bug in GS version 8.17
-            if cls._IS_BUGGED_VERSION:
-                filepath = f"versions/{version}//{relative_path}"
             return _SFFileEntityPath(
                 domain=domain, name=name, version=version, relative_path=relative_path, filepath=filepath
             )

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

@@ -140,6 +140,49 @@ class ModelOperator:
                 statement_params=statement_params,
             )
 
+    def create_from_model_version(
+        self,
+        *,
+        source_database_name: Optional[sql_identifier.SqlIdentifier],
+        source_schema_name: Optional[sql_identifier.SqlIdentifier],
+        source_model_name: sql_identifier.SqlIdentifier,
+        source_version_name: sql_identifier.SqlIdentifier,
+        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:
+        if self.validate_existence(
+            database_name=database_name,
+            schema_name=schema_name,
+            model_name=model_name,
+            statement_params=statement_params,
+        ):
+            return self._model_version_client.add_version_from_model_version(
+                source_database_name=source_database_name,
+                source_schema_name=source_schema_name,
+                source_model_name=source_model_name,
+                source_version_name=source_version_name,
+                database_name=database_name,
+                schema_name=schema_name,
+                model_name=model_name,
+                version_name=version_name,
+                statement_params=statement_params,
+            )
+        else:
+            return self._model_version_client.create_from_model_version(
+                source_database_name=source_database_name,
+                source_schema_name=source_schema_name,
+                source_model_name=source_model_name,
+                source_version_name=source_version_name,
+                database_name=database_name,
+                schema_name=schema_name,
+                model_name=model_name,
+                version_name=version_name,
+                statement_params=statement_params,
+            )
+
     def show_models_or_versions(
         self,
         *,

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

@@ -44,6 +44,32 @@ class ModelVersionSQLClient(_base._BaseSQLClient):
             statement_params=statement_params,
         ).has_dimensions(expected_rows=1, expected_cols=1).validate()
 
+    def create_from_model_version(
+        self,
+        *,
+        source_database_name: Optional[sql_identifier.SqlIdentifier],
+        source_schema_name: Optional[sql_identifier.SqlIdentifier],
+        source_model_name: sql_identifier.SqlIdentifier,
+        source_version_name: sql_identifier.SqlIdentifier,
+        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:
+        fq_source_model_name = self.fully_qualified_object_name(
+            source_database_name, source_schema_name, source_model_name
+        )
+        fq_model_name = self.fully_qualified_object_name(database_name, schema_name, model_name)
+        query_result_checker.SqlResultValidator(
+            self._session,
+            (
+                f"CREATE MODEL {fq_model_name} WITH VERSION {version_name} FROM MODEL {fq_source_model_name}"
+                f" VERSION {source_version_name}"
+            ),
+            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,
@@ -64,6 +90,32 @@ class ModelVersionSQLClient(_base._BaseSQLClient):
             statement_params=statement_params,
         ).has_dimensions(expected_rows=1, expected_cols=1).validate()
 
+    def add_version_from_model_version(
+        self,
+        *,
+        source_database_name: Optional[sql_identifier.SqlIdentifier],
+        source_schema_name: Optional[sql_identifier.SqlIdentifier],
+        source_model_name: sql_identifier.SqlIdentifier,
+        source_version_name: sql_identifier.SqlIdentifier,
+        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:
+        fq_source_model_name = self.fully_qualified_object_name(
+            source_database_name, source_schema_name, source_model_name
+        )
+        fq_model_name = self.fully_qualified_object_name(database_name, schema_name, model_name)
+        query_result_checker.SqlResultValidator(
+            self._session,
+            (
+                f"ALTER MODEL {fq_model_name} ADD VERSION {version_name} FROM MODEL {fq_source_model_name}"
+                f" VERSION {source_version_name}"
+            ),
+            statement_params=statement_params,
+        ).has_dimensions(expected_rows=1, expected_cols=1).validate()
+
     def set_default_version(
         self,
         *,
@@ -145,7 +197,7 @@ class ModelVersionSQLClient(_base._BaseSQLClient):
         if snowpark_utils.is_in_stored_procedure():  # type: ignore[no-untyped-call]
             options = {"parallel": 10}
             cursor = self._session._conn._cursor
-            cursor._download(stage_location_url, str(target_path), options)  # type: ignore[attr-defined]
+            cursor._download(stage_location_url, str(target_path), options)  # type: ignore[union-attr]
             cursor.fetchall()
         else:
             query_result_checker.SqlResultValidator(

snowflake/ml/model/_model_composer/model_composer.py

@@ -136,7 +136,7 @@ class ModelComposer:
             model_meta=self.packager.meta,
             model_file_rel_path=pathlib.PurePosixPath(self.model_file_rel_path),
             options=options,
-            data_sources=self._get_data_sources(model),
+            data_sources=self._get_data_sources(model, sample_input_data),
         )
 
         file_utils.upload_directory_to_stage(
@@ -179,8 +179,12 @@ class ModelComposer:
         mp.load(meta_only=meta_only, options=options)
         return mp
 
-    def _get_data_sources(self, model: model_types.SupportedModelType) -> Optional[List[data_source.DataSource]]:
+    def _get_data_sources(
+        self, model: model_types.SupportedModelType, sample_input_data: Optional[model_types.SupportedDataType] = None
+    ) -> Optional[List[data_source.DataSource]]:
         data_sources = getattr(model, lineage_utils.DATA_SOURCES_ATTR, None)
+        if not data_sources and sample_input_data is not None:
+            data_sources = getattr(sample_input_data, lineage_utils.DATA_SOURCES_ATTR, None)
         if isinstance(data_sources, list) and all(isinstance(item, data_source.DataSource) for item in data_sources):
             return data_sources
         return None

snowflake/ml/model/_packager/model_meta/model_meta.py

@@ -281,9 +281,7 @@ class ModelMetadata:
             "cpu": model_runtime.ModelRuntime("cpu", self.env),
         }
         if self.env.cuda_version:
-            runtimes.update(
-                {"gpu": model_runtime.ModelRuntime("gpu", self.env, is_gpu=True, server_availability_source="conda")}
-            )
+            runtimes.update({"gpu": model_runtime.ModelRuntime("gpu", self.env, is_gpu=True)})
         return runtimes
 
     def save(self, model_dir_path: str) -> None:

snowflake/ml/model/_packager/model_runtime/model_runtime.py

@@ -1,11 +1,11 @@
 import copy
 import pathlib
 import warnings
-from typing import List, Literal, Optional
+from typing import List, Optional
 
 from packaging import requirements
 
-from snowflake.ml._internal import env as snowml_env, env_utils, file_utils
+from snowflake.ml._internal import env_utils, file_utils
 from snowflake.ml.model._packager.model_env import model_env
 from snowflake.ml.model._packager.model_meta import model_meta_schema
 from snowflake.ml.model._packager.model_runtime import (
@@ -37,7 +37,6 @@ class ModelRuntime:
         env: model_env.ModelEnv,
         imports: Optional[List[pathlib.PurePosixPath]] = None,
         is_gpu: bool = False,
-        server_availability_source: Literal["snowflake", "conda"] = "snowflake",
         loading_from_file: bool = False,
     ) -> None:
         self.name = name
@@ -48,30 +47,7 @@ class ModelRuntime:
             return
 
         snowml_pkg_spec = f"{env_utils.SNOWPARK_ML_PKG_NAME}=={self.runtime_env.snowpark_ml_version}"
-        if self.runtime_env._snowpark_ml_version.local:
-            self.embed_local_ml_library = True
-        else:
-            if server_availability_source == "snowflake":
-                snowml_server_availability = (
-                    len(
-                        env_utils.get_matched_package_versions_in_information_schema_with_active_session(
-                            reqs=[requirements.Requirement(snowml_pkg_spec)],
-                            python_version=snowml_env.PYTHON_VERSION,
-                        ).get(env_utils.SNOWPARK_ML_PKG_NAME, [])
-                    )
-                    >= 1
-                )
-            else:
-                snowml_server_availability = (
-                    len(
-                        env_utils.get_matched_package_versions_in_snowflake_conda_channel(
-                            req=requirements.Requirement(snowml_pkg_spec),
-                            python_version=snowml_env.PYTHON_VERSION,
-                        )
-                    )
-                    >= 1
-                )
-            self.embed_local_ml_library = not snowml_server_availability
+        self.embed_local_ml_library = self.runtime_env._snowpark_ml_version.local
 
         additional_package = (
             _SNOWML_INFERENCE_ALTERNATIVE_DEPENDENCIES if self.embed_local_ml_library else [snowml_pkg_spec]