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

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

@@ -9,7 +9,8 @@ from snowflake.ml._internal.utils import (
     query_result_checker,
     sql_identifier,
 )
-from snowflake.snowpark import dataframe, functions as F, row, session, types as spt
+from snowflake.ml.model._client.sql import _base
+from snowflake.snowpark import dataframe, functions as F, row, types as spt
 from snowflake.snowpark._internal import utils as snowpark_utils
 
 
@@ -20,44 +21,51 @@ def _normalize_url_for_sql(url: str) -> str:
     return f"'{url}'"
 
 
-class ModelVersionSQLClient:
+class ModelVersionSQLClient(_base._BaseSQLClient):
     FUNCTION_NAME_COL_NAME = "name"
     FUNCTION_RETURN_TYPE_COL_NAME = "return_type"
 
-    def __init__(
+    def create_from_stage(
         self,
-        session: session.Session,
         *,
-        database_name: sql_identifier.SqlIdentifier,
-        schema_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,
+        stage_path: str,
+        statement_params: Optional[Dict[str, Any]] = None,
     ) -> None:
-        self._session = session
-        self._database_name = database_name
-        self._schema_name = schema_name
-
-    def __eq__(self, __value: object) -> bool:
-        if not isinstance(__value, ModelVersionSQLClient):
-            return False
-        return self._database_name == __value._database_name and self._schema_name == __value._schema_name
-
-    def fully_qualified_model_name(self, model_name: sql_identifier.SqlIdentifier) -> str:
-        return identifier.get_schema_level_object_identifier(
-            self._database_name.identifier(), self._schema_name.identifier(), model_name.identifier()
-        )
+        query_result_checker.SqlResultValidator(
+            self._session,
+            (
+                f"CREATE MODEL {self.fully_qualified_object_name(database_name, schema_name, model_name)}"
+                f" WITH VERSION {version_name.identifier()} FROM {stage_path}"
+            ),
+            statement_params=statement_params,
+        ).has_dimensions(expected_rows=1, expected_cols=1).validate()
 
-    def create_from_stage(
+    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,
-        stage_path: str,
         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 {self.fully_qualified_model_name(model_name)} WITH VERSION {version_name.identifier()}"
-                f" FROM {stage_path}"
+                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()
@@ -66,6 +74,8 @@ class ModelVersionSQLClient:
     def add_version_from_stage(
         self,
         *,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
         model_name: sql_identifier.SqlIdentifier,
         version_name: sql_identifier.SqlIdentifier,
         stage_path: str,
@@ -74,8 +84,34 @@ class ModelVersionSQLClient:
         query_result_checker.SqlResultValidator(
             self._session,
             (
-                f"ALTER MODEL {self.fully_qualified_model_name(model_name)} ADD VERSION {version_name.identifier()}"
-                f" FROM {stage_path}"
+                f"ALTER MODEL {self.fully_qualified_object_name(database_name, schema_name, model_name)}"
+                f" ADD VERSION {version_name.identifier()} FROM {stage_path}"
+            ),
+            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()
@@ -83,6 +119,8 @@ class ModelVersionSQLClient:
     def set_default_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,
@@ -90,7 +128,7 @@ class ModelVersionSQLClient:
         query_result_checker.SqlResultValidator(
             self._session,
             (
-                f"ALTER MODEL {self.fully_qualified_model_name(model_name)} "
+                f"ALTER MODEL {self.fully_qualified_object_name(database_name, schema_name, model_name)} "
                 f"SET DEFAULT_VERSION = {version_name.identifier()}"
             ),
             statement_params=statement_params,
@@ -99,6 +137,8 @@ class ModelVersionSQLClient:
     def list_file(
         self,
         *,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
         model_name: sql_identifier.SqlIdentifier,
         version_name: sql_identifier.SqlIdentifier,
         file_path: pathlib.PurePosixPath,
@@ -110,7 +150,10 @@ class ModelVersionSQLClient:
 
         stage_location = (
             pathlib.PurePosixPath(
-                self.fully_qualified_model_name(model_name), "versions", version_name.resolved(), file_path
+                self.fully_qualified_object_name(database_name, schema_name, model_name),
+                "versions",
+                version_name.resolved(),
+                file_path,
             ).as_posix()
             + trailing_slash
         )
@@ -124,13 +167,15 @@ class ModelVersionSQLClient:
                 f"List {_normalize_url_for_sql(stage_location_url)}",
                 statement_params=statement_params,
             )
-            .has_column("name")
+            .has_column("name", allow_empty=True)
             .validate()
         )
 
     def get_file(
         self,
         *,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
         model_name: sql_identifier.SqlIdentifier,
         version_name: sql_identifier.SqlIdentifier,
         file_path: pathlib.PurePosixPath,
@@ -138,7 +183,10 @@ class ModelVersionSQLClient:
         statement_params: Optional[Dict[str, Any]] = None,
     ) -> pathlib.Path:
         stage_location = pathlib.PurePosixPath(
-            self.fully_qualified_model_name(model_name), "versions", version_name.resolved(), file_path
+            self.fully_qualified_object_name(database_name, schema_name, model_name),
+            "versions",
+            version_name.resolved(),
+            file_path,
         ).as_posix()
         stage_location_url = ParseResult(
             scheme="snow", netloc="model", path=stage_location, params="", query="", fragment=""
@@ -149,7 +197,7 @@ class ModelVersionSQLClient:
         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(
@@ -162,6 +210,8 @@ class ModelVersionSQLClient:
     def show_functions(
         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,
@@ -169,7 +219,7 @@ class ModelVersionSQLClient:
         res = query_result_checker.SqlResultValidator(
             self._session,
             (
-                f"SHOW FUNCTIONS IN MODEL {self.fully_qualified_model_name(model_name)}"
+                f"SHOW FUNCTIONS IN MODEL {self.fully_qualified_object_name(database_name, schema_name, model_name)}"
                 f" VERSION {version_name.identifier()}"
             ),
             statement_params=statement_params,
@@ -180,15 +230,17 @@ class ModelVersionSQLClient:
     def set_comment(
         self,
         *,
-        comment: str,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
         model_name: sql_identifier.SqlIdentifier,
         version_name: sql_identifier.SqlIdentifier,
+        comment: str,
         statement_params: Optional[Dict[str, Any]] = None,
     ) -> None:
         query_result_checker.SqlResultValidator(
             self._session,
             (
-                f"ALTER MODEL {self.fully_qualified_model_name(model_name)} "
+                f"ALTER MODEL {self.fully_qualified_object_name(database_name, schema_name, model_name)} "
                 f"MODIFY VERSION {version_name.identifier()} SET COMMENT=$${comment}$$"
             ),
             statement_params=statement_params,
@@ -197,6 +249,8 @@ class ModelVersionSQLClient:
     def invoke_function_method(
         self,
         *,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
         model_name: sql_identifier.SqlIdentifier,
         version_name: sql_identifier.SqlIdentifier,
         method_name: sql_identifier.SqlIdentifier,
@@ -210,10 +264,12 @@ class ModelVersionSQLClient:
             INTERMEDIATE_TABLE_NAME = "SNOWPARK_ML_MODEL_INFERENCE_INPUT"
             with_statements.append(f"{INTERMEDIATE_TABLE_NAME} AS ({input_df.queries['queries'][0]})")
         else:
+            actual_database_name = database_name or self._database_name
+            actual_schema_name = schema_name or self._schema_name
             tmp_table_name = snowpark_utils.random_name_for_temp_object(snowpark_utils.TempObjectType.TABLE)
             INTERMEDIATE_TABLE_NAME = identifier.get_schema_level_object_identifier(
-                self._database_name.identifier(),
-                self._schema_name.identifier(),
+                actual_database_name.identifier(),
+                actual_schema_name.identifier(),
                 tmp_table_name,
             )
             input_df.write.save_as_table(  # type: ignore[call-overload]
@@ -228,7 +284,8 @@ class ModelVersionSQLClient:
         module_version_alias = "MODEL_VERSION_ALIAS"
         with_statements.append(
             f"{module_version_alias} AS "
-            f"MODEL {self.fully_qualified_model_name(model_name)} VERSION {version_name.identifier()}"
+            f"MODEL {self.fully_qualified_object_name(database_name, schema_name, model_name)}"
+            f" VERSION {version_name.identifier()}"
         )
 
         args_sql_list = []
@@ -267,6 +324,8 @@ class ModelVersionSQLClient:
     def invoke_table_function_method(
         self,
         *,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
         model_name: sql_identifier.SqlIdentifier,
         version_name: sql_identifier.SqlIdentifier,
         method_name: sql_identifier.SqlIdentifier,
@@ -281,10 +340,12 @@ class ModelVersionSQLClient:
             INTERMEDIATE_TABLE_NAME = "SNOWPARK_ML_MODEL_INFERENCE_INPUT"
             with_statements.append(f"{INTERMEDIATE_TABLE_NAME} AS ({input_df.queries['queries'][0]})")
         else:
+            actual_database_name = database_name or self._database_name
+            actual_schema_name = schema_name or self._schema_name
             tmp_table_name = snowpark_utils.random_name_for_temp_object(snowpark_utils.TempObjectType.TABLE)
             INTERMEDIATE_TABLE_NAME = identifier.get_schema_level_object_identifier(
-                self._database_name.identifier(),
-                self._schema_name.identifier(),
+                actual_database_name.identifier(),
+                actual_schema_name.identifier(),
                 tmp_table_name,
             )
             input_df.write.save_as_table(  # type: ignore[call-overload]
@@ -297,7 +358,8 @@ class ModelVersionSQLClient:
         module_version_alias = "MODEL_VERSION_ALIAS"
         with_statements.append(
             f"{module_version_alias} AS "
-            f"MODEL {self.fully_qualified_model_name(model_name)} VERSION {version_name.identifier()}"
+            f"MODEL {self.fully_qualified_object_name(database_name, schema_name, model_name)}"
+            f" VERSION {version_name.identifier()}"
         )
 
         partition_by = partition_column.identifier() if partition_column is not None else "1"
@@ -344,6 +406,8 @@ class ModelVersionSQLClient:
         self,
         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,
@@ -352,8 +416,8 @@ class ModelVersionSQLClient:
         query_result_checker.SqlResultValidator(
             self._session,
             (
-                f"ALTER MODEL {self.fully_qualified_model_name(model_name)} MODIFY VERSION {version_name.identifier()}"
-                f" SET METADATA=$${json_metadata}$$"
+                f"ALTER MODEL {self.fully_qualified_object_name(database_name, schema_name, model_name)}"
+                f" MODIFY VERSION {version_name.identifier()} SET METADATA=$${json_metadata}$$"
             ),
             statement_params=statement_params,
         ).has_dimensions(expected_rows=1, expected_cols=1).validate()
@@ -361,12 +425,17 @@ class ModelVersionSQLClient:
     def drop_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:
         query_result_checker.SqlResultValidator(
             self._session,
-            f"ALTER MODEL {self.fully_qualified_model_name(model_name)} DROP VERSION {version_name.identifier()}",
+            (
+                f"ALTER MODEL {self.fully_qualified_object_name(database_name, schema_name, model_name)}"
+                f" DROP VERSION {version_name.identifier()}"
+            ),
             statement_params=statement_params,
         ).has_dimensions(expected_rows=1, expected_cols=1).validate()

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

@@ -1,46 +1,20 @@
 from typing import Any, Dict, Optional
 
-from snowflake.ml._internal.utils import (
-    identifier,
-    query_result_checker,
-    sql_identifier,
-)
-from snowflake.snowpark import session
+from snowflake.ml._internal.utils import query_result_checker, sql_identifier
+from snowflake.ml.model._client.sql import _base
 
 
-class StageSQLClient:
-    def __init__(
-        self,
-        session: session.Session,
-        *,
-        database_name: sql_identifier.SqlIdentifier,
-        schema_name: sql_identifier.SqlIdentifier,
-    ) -> None:
-        self._session = session
-        self._database_name = database_name
-        self._schema_name = schema_name
-
-    def __eq__(self, __value: object) -> bool:
-        if not isinstance(__value, StageSQLClient):
-            return False
-        return self._database_name == __value._database_name and self._schema_name == __value._schema_name
-
-    def fully_qualified_stage_name(
-        self,
-        stage_name: sql_identifier.SqlIdentifier,
-    ) -> str:
-        return identifier.get_schema_level_object_identifier(
-            self._database_name.identifier(), self._schema_name.identifier(), stage_name.identifier()
-        )
-
+class StageSQLClient(_base._BaseSQLClient):
     def create_tmp_stage(
         self,
         *,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
         stage_name: sql_identifier.SqlIdentifier,
         statement_params: Optional[Dict[str, Any]] = None,
     ) -> None:
         query_result_checker.SqlResultValidator(
             self._session,
-            f"CREATE TEMPORARY STAGE {self.fully_qualified_stage_name(stage_name)}",
+            f"CREATE TEMPORARY STAGE {self.fully_qualified_object_name(database_name, schema_name, stage_name)}",
             statement_params=statement_params,
         ).has_dimensions(expected_rows=1, expected_cols=1).validate()

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

@@ -1,52 +1,25 @@
 from typing import Any, Dict, List, Optional
 
-from snowflake.ml._internal.utils import (
-    identifier,
-    query_result_checker,
-    sql_identifier,
-)
-from snowflake.snowpark import row, session
+from snowflake.ml._internal.utils import query_result_checker, sql_identifier
+from snowflake.ml.model._client.sql import _base
+from snowflake.snowpark import row
 
 
-class ModuleTagSQLClient:
-    def __init__(
-        self,
-        session: session.Session,
-        *,
-        database_name: sql_identifier.SqlIdentifier,
-        schema_name: sql_identifier.SqlIdentifier,
-    ) -> None:
-        self._session = session
-        self._database_name = database_name
-        self._schema_name = schema_name
-
-    def __eq__(self, __value: object) -> bool:
-        if not isinstance(__value, ModuleTagSQLClient):
-            return False
-        return self._database_name == __value._database_name and self._schema_name == __value._schema_name
-
-    def fully_qualified_module_name(
-        self,
-        module_name: sql_identifier.SqlIdentifier,
-    ) -> str:
-        return identifier.get_schema_level_object_identifier(
-            self._database_name.identifier(), self._schema_name.identifier(), module_name.identifier()
-        )
-
+class ModuleTagSQLClient(_base._BaseSQLClient):
     def set_tag_on_model(
         self,
-        model_name: sql_identifier.SqlIdentifier,
         *,
-        tag_database_name: sql_identifier.SqlIdentifier,
-        tag_schema_name: sql_identifier.SqlIdentifier,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
+        model_name: sql_identifier.SqlIdentifier,
+        tag_database_name: Optional[sql_identifier.SqlIdentifier],
+        tag_schema_name: Optional[sql_identifier.SqlIdentifier],
         tag_name: sql_identifier.SqlIdentifier,
         tag_value: str,
         statement_params: Optional[Dict[str, Any]] = None,
     ) -> None:
-        fq_model_name = self.fully_qualified_module_name(model_name)
-        fq_tag_name = identifier.get_schema_level_object_identifier(
-            tag_database_name.identifier(), tag_schema_name.identifier(), tag_name.identifier()
-        )
+        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)
         query_result_checker.SqlResultValidator(
             self._session,
             f"ALTER MODEL {fq_model_name} SET TAG {fq_tag_name} = $${tag_value}$$",
@@ -55,17 +28,17 @@ class ModuleTagSQLClient:
 
     def unset_tag_on_model(
         self,
-        model_name: sql_identifier.SqlIdentifier,
         *,
-        tag_database_name: sql_identifier.SqlIdentifier,
-        tag_schema_name: sql_identifier.SqlIdentifier,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
+        model_name: sql_identifier.SqlIdentifier,
+        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,
     ) -> None:
-        fq_model_name = self.fully_qualified_module_name(model_name)
-        fq_tag_name = identifier.get_schema_level_object_identifier(
-            tag_database_name.identifier(), tag_schema_name.identifier(), tag_name.identifier()
-        )
+        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)
         query_result_checker.SqlResultValidator(
             self._session,
             f"ALTER MODEL {fq_model_name} UNSET TAG {fq_tag_name}",
@@ -74,21 +47,21 @@ class ModuleTagSQLClient:
 
     def get_tag_value(
         self,
-        module_name: sql_identifier.SqlIdentifier,
         *,
-        tag_database_name: sql_identifier.SqlIdentifier,
-        tag_schema_name: sql_identifier.SqlIdentifier,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
+        model_name: sql_identifier.SqlIdentifier,
+        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,
     ) -> row.Row:
-        fq_module_name = self.fully_qualified_module_name(module_name)
-        fq_tag_name = identifier.get_schema_level_object_identifier(
-            tag_database_name.identifier(), tag_schema_name.identifier(), tag_name.identifier()
-        )
+        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)
         return (
             query_result_checker.SqlResultValidator(
                 self._session,
-                f"SELECT SYSTEM$GET_TAG($${fq_tag_name}$$, $${fq_module_name}$$, 'MODULE') AS TAG_VALUE",
+                f"SELECT SYSTEM$GET_TAG($${fq_tag_name}$$, $${fq_model_name}$$, 'MODULE') AS TAG_VALUE",
                 statement_params=statement_params,
             )
             .has_dimensions(expected_rows=1, expected_cols=1)
@@ -98,16 +71,19 @@ class ModuleTagSQLClient:
 
     def get_tag_list(
         self,
-        module_name: sql_identifier.SqlIdentifier,
         *,
+        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]:
-        fq_module_name = self.fully_qualified_module_name(module_name)
+        fq_model_name = self.fully_qualified_object_name(database_name, schema_name, model_name)
+        actual_database_name = database_name or self._database_name
         return (
             query_result_checker.SqlResultValidator(
                 self._session,
                 f"""SELECT TAG_DATABASE, TAG_SCHEMA, TAG_NAME, TAG_VALUE
-FROM TABLE({self._database_name.identifier()}.INFORMATION_SCHEMA.TAG_REFERENCES($${fq_module_name}$$, 'MODULE'))""",
+FROM TABLE({actual_database_name.identifier()}.INFORMATION_SCHEMA.TAG_REFERENCES($${fq_model_name}$$, 'MODULE'))""",
                 statement_params=statement_params,
             )
             .has_column("TAG_DATABASE", allow_empty=True)

snowflake/ml/model/_model_composer/model_composer.py

@@ -11,7 +11,7 @@ from packaging import requirements
 from typing_extensions import deprecated
 
 from snowflake.ml._internal import env as snowml_env, env_utils, file_utils
-from snowflake.ml._internal.lineage import data_source
+from snowflake.ml._internal.lineage import data_source, lineage_utils
 from snowflake.ml.model import model_signature, type_hints as model_types
 from snowflake.ml.model._model_composer.model_manifest import model_manifest
 from snowflake.ml.model._packager import model_packager
@@ -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]]:
-        data_sources = getattr(model, "_data_sources", None)
+    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_handlers/mlflow.py

@@ -1,4 +1,5 @@
 import os
+import pathlib
 import tempfile
 from typing import TYPE_CHECKING, Callable, Dict, Optional, Type, cast, final
 
@@ -45,7 +46,7 @@ def _parse_mlflow_env(model_uri: str, env: model_env.ModelEnv) -> model_env.Mode
     if not os.path.exists(conda_env_file_path):
         raise ValueError("Cannot load MLFlow model dependencies.")
 
-    env.load_from_conda_file(conda_env_file_path)
+    env.load_from_conda_file(pathlib.Path(conda_env_file_path))
 
     return env
 

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]