snowflake-ml-python 1.5.3__py3-none-any.whl → 1.6.0__py3-none-any.whl

snowflake/ml/lineage/lineage_node.py

@@ -0,0 +1,139 @@
+import json
+from datetime import datetime
+from typing import TYPE_CHECKING, Dict, List, Literal, Optional, Set, Type, Union
+
+from snowflake import snowpark
+from snowflake.ml._internal import telemetry
+from snowflake.ml._internal.utils import identifier
+
+if TYPE_CHECKING:
+    from snowflake.ml import dataset
+    from snowflake.ml.feature_store import feature_view
+    from snowflake.ml.model._client.model import model_version_impl
+
+_PROJECT = "LINEAGE"
+DOMAIN_LINEAGE_REGISTRY: Dict[str, Type["LineageNode"]] = {}
+
+
+class LineageNode:
+    """
+    Represents a node in a lineage graph and serves as the base class for all machine learning objects.
+    """
+
+    def __init__(
+        self,
+        session: snowpark.Session,
+        name: str,
+        domain: Union[Literal["feature_view", "dataset", "model", "table", "view"]],
+        version: Optional[str] = None,
+        status: Optional[Literal["ACTIVE", "DELETED", "MASKED"]] = None,
+        created_on: Optional[datetime] = None,
+    ) -> None:
+        """
+        Initializes a LineageNode instance.
+
+        Args:
+            session : The Snowflake session object.
+            name : Fully qualified name of the lineage node, which is in the format '<db>.<schema>.<object_name>'.
+            domain : The domain of the lineage node.
+            version : The version of the lineage node, if applies.
+            status : The status of the lineage node. Possible values are:
+                      - 'MASKED': The user does not have the privilege to view the node.
+                      - 'DELETED': The node has been deleted.
+                      - 'ACTIVE': The node is currently active.
+            created_on : The creation time of the lineage node.
+
+        Raises:
+            ValueError: If the name is not fully qualified.
+        """
+        if name and not identifier.is_fully_qualified_name(name):
+            raise ValueError("name should be fully qualifed.")
+
+        self._lineage_node_name = name
+        self._lineage_node_domain = domain
+        self._lineage_node_version = version
+        self._lineage_node_status = status
+        self._lineage_node_created_on = created_on
+        self._session = session
+
+    def __repr__(self) -> str:
+        return (
+            f"{self.__class__.__name__}(\n"
+            f"  name='{self._lineage_node_name}',\n"
+            f"  version='{self._lineage_node_version}',\n"
+            f"  domain='{self._lineage_node_domain}',\n"
+            f"  status='{self._lineage_node_status}',\n"
+            f"  created_on='{self._lineage_node_created_on}'\n"
+            f")"
+        )
+
+    @staticmethod
+    def _load_from_lineage_node(session: snowpark.Session, name: str, version: str) -> "LineageNode":
+        """
+        Loads the concrete object.
+
+        Args:
+            session : The Snowflake session object.
+            name : Fully qualified name of the object.
+            version : The version of object.
+
+        Raises:
+            NotImplementedError: If the derived class does not implement this method.
+        """
+        raise NotImplementedError()
+
+    @telemetry.send_api_usage_telemetry(project=_PROJECT)
+    @snowpark._internal.utils.private_preview(version="1.5.3")
+    def lineage(
+        self,
+        direction: Literal["upstream", "downstream"] = "downstream",
+        domain_filter: Optional[Set[Literal["feature_view", "dataset", "model", "table", "view"]]] = None,
+    ) -> List[Union["feature_view.FeatureView", "dataset.Dataset", "model_version_impl.ModelVersion", "LineageNode"]]:
+        """
+        Retrieves the lineage nodes connected to this node.
+
+        Args:
+            direction : The direction to trace lineage. Defaults to "downstream".
+            domain_filter : Set of domains to filter nodes. Defaults to None.
+
+        Returns:
+            List[LineageNode]: A list of connected lineage nodes.
+        """
+        df = self._session.lineage.trace(
+            self._lineage_node_name,
+            self._lineage_node_domain.upper(),
+            object_version=self._lineage_node_version,
+            direction=direction,
+            distance=1,
+        )
+        if domain_filter is not None:
+            domain_filter = {d.lower() for d in domain_filter}  # type: ignore[misc]
+
+        lineage_nodes: List["LineageNode"] = []
+        for row in df.collect():
+            lineage_object = (
+                json.loads(row["TARGET_OBJECT"])
+                if direction.lower() == "downstream"
+                else json.loads(row["SOURCE_OBJECT"])
+            )
+            domain = lineage_object["domain"].lower()
+            if domain_filter is None or domain in domain_filter:
+                if domain in DOMAIN_LINEAGE_REGISTRY and lineage_object["status"] == "ACTIVE":
+                    lineage_nodes.append(
+                        DOMAIN_LINEAGE_REGISTRY[domain]._load_from_lineage_node(
+                            self._session, lineage_object["name"], lineage_object.get("version")
+                        )
+                    )
+                else:
+                    lineage_nodes.append(
+                        LineageNode(
+                            name=lineage_object["name"],
+                            version=lineage_object.get("version"),
+                            domain=domain,
+                            status=lineage_object["status"],
+                            created_on=datetime.strptime(lineage_object["createdOn"], "%Y-%m-%dT%H:%M:%SZ"),
+                            session=self._session,
+                        )
+                    )
+
+        return lineage_nodes

snowflake/ml/model/_client/model/model_impl.py

@@ -9,6 +9,10 @@ from snowflake.ml.model._client.ops import model_ops
 
 _TELEMETRY_PROJECT = "MLOps"
 _TELEMETRY_SUBPROJECT = "ModelManagement"
+SYSTEM_VERSION_ALIAS_DEFAULT = "DEFAULT"
+SYSTEM_VERSION_ALIAS_FIRST = "FIRST"
+SYSTEM_VERSION_ALIAS_LAST = "LAST"
+SYSTEM_VERSION_ALIASES = (SYSTEM_VERSION_ALIAS_DEFAULT, SYSTEM_VERSION_ALIAS_FIRST, SYSTEM_VERSION_ALIAS_LAST)
 
 
 class Model:
@@ -144,12 +148,28 @@ class Model:
         project=_TELEMETRY_PROJECT,
         subproject=_TELEMETRY_SUBPROJECT,
     )
-    def version(self, version_name: str) -> model_version_impl.ModelVersion:
+    def first(self) -> model_version_impl.ModelVersion:
+        """The first version of the model."""
+        return self.version(SYSTEM_VERSION_ALIAS_FIRST)
+
+    @telemetry.send_api_usage_telemetry(
+        project=_TELEMETRY_PROJECT,
+        subproject=_TELEMETRY_SUBPROJECT,
+    )
+    def last(self) -> model_version_impl.ModelVersion:
+        """The latest version of the model."""
+        return self.version(SYSTEM_VERSION_ALIAS_LAST)
+
+    @telemetry.send_api_usage_telemetry(
+        project=_TELEMETRY_PROJECT,
+        subproject=_TELEMETRY_SUBPROJECT,
+    )
+    def version(self, version_or_alias: str) -> model_version_impl.ModelVersion:
         """
-        Get a model version object given a version name in the model.
+        Get a model version object given a version name or version alias in the model.
 
         Args:
-            version_name: The name of the version.
+            version_or_alias: The name of the version or alias to a version.
 
         Raises:
             ValueError: When the requested version does not exist.
@@ -161,23 +181,36 @@ class Model:
             project=_TELEMETRY_PROJECT,
             subproject=_TELEMETRY_SUBPROJECT,
         )
-        version_id = sql_identifier.SqlIdentifier(version_name)
-        if self._model_ops.validate_existence(
+
+        # check with system alias or with user defined alias
+        version_id = self._model_ops.get_version_by_alias(
             database_name=None,
             schema_name=None,
             model_name=self._model_name,
-            version_name=version_id,
+            alias_name=sql_identifier.SqlIdentifier(version_or_alias),
             statement_params=statement_params,
-        ):
-            return model_version_impl.ModelVersion._ref(
-                self._model_ops,
+        )
+
+        # version_id is still None implies version_or_alias is not an alias. So it must be a version name.
+        if version_id is None:
+            version_id = sql_identifier.SqlIdentifier(version_or_alias)
+            if not self._model_ops.validate_existence(
+                database_name=None,
+                schema_name=None,
                 model_name=self._model_name,
                 version_name=version_id,
-            )
-        else:
-            raise ValueError(
-                f"Unable to find version with name {version_id.identifier()} in model {self.fully_qualified_name}"
-            )
+                statement_params=statement_params,
+            ):
+                raise ValueError(
+                    f"Unable to find version or alias with name {version_id.identifier()} "
+                    f"in model {self.fully_qualified_name}"
+                )
+
+        return model_version_impl.ModelVersion._ref(
+            self._model_ops,
+            model_name=self._model_name,
+            version_name=version_id,
+        )
 
     @telemetry.send_api_usage_telemetry(
         project=_TELEMETRY_PROJECT,

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

@@ -8,12 +8,13 @@ import pandas as pd
 
 from snowflake.ml._internal import telemetry
 from snowflake.ml._internal.utils import sql_identifier
+from snowflake.ml.lineage import lineage_node
 from snowflake.ml.model import type_hints as model_types
 from snowflake.ml.model._client.ops import metadata_ops, model_ops
 from snowflake.ml.model._model_composer import model_composer
 from snowflake.ml.model._model_composer.model_manifest import model_manifest_schema
 from snowflake.ml.model._packager.model_handlers import snowmlmodel
-from snowflake.snowpark import dataframe
+from snowflake.snowpark import Session, dataframe
 
 _TELEMETRY_PROJECT = "MLOps"
 _TELEMETRY_SUBPROJECT = "ModelManagement"
@@ -24,7 +25,7 @@ class ExportMode(enum.Enum):
     FULL = "full"
 
 
-class ModelVersion:
+class ModelVersion(lineage_node.LineageNode):
     """Model Version Object representing a specific version of the model that could be run."""
 
     _model_ops: model_ops.ModelOperator
@@ -48,6 +49,15 @@ class ModelVersion:
         self._model_name = model_name
         self._version_name = version_name
         self._functions = self._get_functions()
+        super(cls, cls).__init__(
+            self,
+            session=model_ops._session,
+            name=model_ops._model_client.fully_qualified_object_name(
+                database_name=None, schema_name=None, object_name=model_name
+            ),
+            domain="model",
+            version=version_name,
+        )
         return self
 
     def __eq__(self, __value: object) -> bool:
@@ -59,6 +69,11 @@ class ModelVersion:
             and self._version_name == __value._version_name
         )
 
+    def __repr__(self) -> str:
+        return (
+            f"{self.__class__.__name__}(\n" f"  name='{self.model_name}',\n" f"  version='{self._version_name}',\n" f")"
+        )
+
     @property
     def model_name(self) -> str:
         """Return the name of the model to which the model version belongs, usable as a reference in SQL."""
@@ -198,6 +213,52 @@ class ModelVersion:
             statement_params=statement_params,
         )
 
+    @telemetry.send_api_usage_telemetry(
+        project=_TELEMETRY_PROJECT,
+        subproject=_TELEMETRY_SUBPROJECT,
+    )
+    def set_alias(self, alias_name: str) -> None:
+        """Set alias to a model version.
+
+        Args:
+            alias_name: Alias to the model version.
+        """
+        statement_params = telemetry.get_statement_params(
+            project=_TELEMETRY_PROJECT,
+            subproject=_TELEMETRY_SUBPROJECT,
+        )
+        alias_name = sql_identifier.SqlIdentifier(alias_name)
+        self._model_ops.set_alias(
+            alias_name=alias_name,
+            database_name=None,
+            schema_name=None,
+            model_name=self._model_name,
+            version_name=self._version_name,
+            statement_params=statement_params,
+        )
+
+    @telemetry.send_api_usage_telemetry(
+        project=_TELEMETRY_PROJECT,
+        subproject=_TELEMETRY_SUBPROJECT,
+    )
+    def unset_alias(self, version_or_alias: str) -> None:
+        """unset alias to a model version.
+
+        Args:
+            version_or_alias: The name of the version or alias to a version.
+        """
+        statement_params = telemetry.get_statement_params(
+            project=_TELEMETRY_PROJECT,
+            subproject=_TELEMETRY_SUBPROJECT,
+        )
+        self._model_ops.unset_alias(
+            version_or_alias_name=sql_identifier.SqlIdentifier(version_or_alias),
+            database_name=None,
+            schema_name=None,
+            model_name=self._model_name,
+            statement_params=statement_params,
+        )
+
     @telemetry.send_api_usage_telemetry(
         project=_TELEMETRY_PROJECT,
         subproject=_TELEMETRY_SUBPROJECT,
@@ -451,3 +512,22 @@ class ModelVersion:
             f"model_name={self._model_name}, version_name={self._version_name}, metadata={pk.meta}"
         )
         return pk.model
+
+    @staticmethod
+    def _load_from_lineage_node(session: Session, name: str, version: str) -> "ModelVersion":
+        database_name_id, schema_name_id, model_name_id = sql_identifier.parse_fully_qualified_name(name)
+        if not database_name_id or not schema_name_id:
+            raise ValueError("name should be fully qualifed.")
+
+        return ModelVersion._ref(
+            model_ops.ModelOperator(
+                session,
+                database_name=database_name_id,
+                schema_name=schema_name_id,
+            ),
+            model_name=model_name_id,
+            version_name=sql_identifier.SqlIdentifier(version),
+        )
+
+
+lineage_node.DOMAIN_LINEAGE_REGISTRY["model"] = ModelVersion

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

@@ -1,11 +1,12 @@
 import os
 import pathlib
 import tempfile
+import warnings
 from typing import Any, Dict, List, Literal, Optional, Union, cast
 
 import yaml
 
-from snowflake.ml._internal.utils import identifier, sql_identifier
+from snowflake.ml._internal.utils import formatting, identifier, sql_identifier
 from snowflake.ml.model import model_signature, type_hints
 from snowflake.ml.model._client.ops import metadata_ops
 from snowflake.ml.model._client.sql import (
@@ -311,6 +312,42 @@ class ModelOperator:
                 statement_params=statement_params,
             )
 
+    def set_alias(
+        self,
+        *,
+        alias_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:
+        self._model_version_client.set_alias(
+            alias_name=alias_name,
+            database_name=database_name,
+            schema_name=schema_name,
+            model_name=model_name,
+            version_name=version_name,
+            statement_params=statement_params,
+        )
+
+    def unset_alias(
+        self,
+        *,
+        version_or_alias_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,
+    ) -> None:
+        self._model_version_client.unset_alias(
+            database_name=database_name,
+            schema_name=schema_name,
+            model_name=model_name,
+            version_or_alias_name=version_or_alias_name,
+            statement_params=statement_params,
+        )
+
     def set_default_version(
         self,
         *,
@@ -354,6 +391,28 @@ class ModelOperator:
             res[self._model_client.MODEL_DEFAULT_VERSION_NAME_COL_NAME], case_sensitive=True
         )
 
+    def get_version_by_alias(
+        self,
+        *,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
+        model_name: sql_identifier.SqlIdentifier,
+        alias_name: sql_identifier.SqlIdentifier,
+        statement_params: Optional[Dict[str, Any]] = None,
+    ) -> Optional[sql_identifier.SqlIdentifier]:
+        res = self._model_client.show_versions(
+            database_name=database_name,
+            schema_name=schema_name,
+            model_name=model_name,
+            statement_params=statement_params,
+        )
+        for r in res:
+            if alias_name in r[self._model_client.MODEL_VERSION_ALIASES_COL_NAME]:
+                return sql_identifier.SqlIdentifier(
+                    r[self._model_client.MODEL_VERSION_NAME_COL_NAME], case_sensitive=True
+                )
+        return None
+
     def get_tag_value(
         self,
         *,
@@ -625,10 +684,23 @@ class ModelOperator:
             )
 
         if keep_order:
-            df_res = df_res.sort(
-                "_ID",
-                ascending=True,
-            )
+            # if it's a partitioned table function, _ID will be null and we won't be able to sort.
+            if df_res.select("_ID").limit(1).collect()[0][0] is None:
+                warnings.warn(
+                    formatting.unwrap(
+                        """
+                        When invoking partitioned inference methods, ordering of rows in output dataframe will differ
+                        from that of input dataframe.
+                        """
+                    ),
+                    category=UserWarning,
+                    stacklevel=1,
+                )
+            else:
+                df_res = df_res.sort(
+                    "_ID",
+                    ascending=True,
+                )
 
         if not output_with_input_features:
             cols_to_drop = original_cols

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

@@ -14,6 +14,7 @@ class ModelSQLClient(_base._BaseSQLClient):
     MODEL_VERSION_COMMENT_COL_NAME = "comment"
     MODEL_VERSION_METADATA_COL_NAME = "metadata"
     MODEL_VERSION_MODEL_SPEC_COL_NAME = "model_spec"
+    MODEL_VERSION_ALIASES_COL_NAME = "aliases"
 
     def show_models(
         self,

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

@@ -134,6 +134,43 @@ class ModelVersionSQLClient(_base._BaseSQLClient):
             statement_params=statement_params,
         ).has_dimensions(expected_rows=1, expected_cols=1).validate()
 
+    def set_alias(
+        self,
+        *,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
+        model_name: sql_identifier.SqlIdentifier,
+        version_name: sql_identifier.SqlIdentifier,
+        alias_name: sql_identifier.SqlIdentifier,
+        statement_params: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        query_result_checker.SqlResultValidator(
+            self._session,
+            (
+                f"ALTER MODEL {self.fully_qualified_object_name(database_name, schema_name, model_name)} "
+                f"VERSION {version_name.identifier()} SET ALIAS = {alias_name.identifier()}"
+            ),
+            statement_params=statement_params,
+        ).has_dimensions(expected_rows=1, expected_cols=1).validate()
+
+    def unset_alias(
+        self,
+        *,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
+        model_name: sql_identifier.SqlIdentifier,
+        version_or_alias_name: sql_identifier.SqlIdentifier,
+        statement_params: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        query_result_checker.SqlResultValidator(
+            self._session,
+            (
+                f"ALTER MODEL {self.fully_qualified_object_name(database_name, schema_name, model_name)} "
+                f"VERSION {version_or_alias_name.identifier()} UNSET ALIAS"
+            ),
+            statement_params=statement_params,
+        ).has_dimensions(expected_rows=1, expected_cols=1).validate()
+
     def list_file(
         self,
         *,
@@ -383,9 +420,13 @@ class ModelVersionSQLClient(_base._BaseSQLClient):
         # Prepare the output
         output_cols = []
         output_names = []
+        cols_to_drop = []
 
         for output_name, output_type, output_col_name in returns:
-            output_cols.append(F.col(output_name).astype(output_type))
+            output_identifier = sql_identifier.SqlIdentifier(output_name).identifier()
+            if output_identifier != output_col_name:
+                cols_to_drop.append(output_identifier)
+            output_cols.append(F.col(output_identifier).astype(output_type))
             output_names.append(output_col_name)
 
         if partition_column is not None:
@@ -396,10 +437,12 @@ class ModelVersionSQLClient(_base._BaseSQLClient):
             col_names=output_names,
             values=output_cols,
         )
-
         if statement_params:
             output_df._statement_params = statement_params  # type: ignore[assignment]
 
+        if cols_to_drop:
+            output_df = output_df.drop(cols_to_drop)
+
         return output_df
 
     def set_metadata(

snowflake/ml/model/_deploy_client/image_builds/inference_server/main.py

@@ -85,9 +85,8 @@ def _run_setup() -> None:
 
         TARGET_METHOD = os.getenv("TARGET_METHOD")
 
-        _concurrent_requests_max_env = os.getenv("_CONCURRENT_REQUESTS_MAX", None)
-
-        _CONCURRENT_REQUESTS_MAX = int(_concurrent_requests_max_env) if _concurrent_requests_max_env else None
+        _concurrent_requests_max_env = os.getenv("_CONCURRENT_REQUESTS_MAX", "1")
+        _CONCURRENT_REQUESTS_MAX = int(_concurrent_requests_max_env)
 
         with tempfile.TemporaryDirectory() as tmp_dir:
             if zipfile.is_zipfile(model_zip_stage_path):
@@ -101,7 +100,6 @@ def _run_setup() -> None:
             logger.info(f"Loading model from {extracted_dir} into memory")
 
             sys.path.insert(0, os.path.join(extracted_dir, _MODEL_CODE_DIR))
-            from snowflake.ml.model import type_hints as model_types
 
             # TODO (Server-side Model Rollout):
             # Keep try block only
@@ -114,7 +112,7 @@ def _run_setup() -> None:
                 pk.load(
                     as_custom_model=True,
                     meta_only=False,
-                    options=model_types.ModelLoadOption({"use_gpu": use_gpu}),
+                    options={"use_gpu": use_gpu},
                 )
                 _LOADED_MODEL = pk.model
                 _LOADED_META = pk.meta
@@ -132,7 +130,7 @@ def _run_setup() -> None:
                     _LOADED_MODEL, meta_LOADED_META = model_api._load(
                         local_dir_path=extracted_dir,
                         as_custom_model=True,
-                        options=model_types.ModelLoadOption({"use_gpu": use_gpu}),
+                        options={"use_gpu": use_gpu},
                     )
             _MODEL_LOADING_STATE = _ModelLoadingState.SUCCEEDED
             logger.info("Successfully loaded model into memory")

snowflake/ml/model/_model_composer/model_composer.py

@@ -11,10 +11,12 @@ 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, lineage_utils
+from snowflake.ml._internal.lineage import lineage_utils
+from snowflake.ml.data import data_source
 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
+from snowflake.ml.model._packager.model_meta import model_meta
 from snowflake.snowpark import Session
 from snowflake.snowpark._internal import utils as snowpark_utils
 
@@ -90,7 +92,7 @@ class ModelComposer:
         ext_modules: Optional[List[ModuleType]] = None,
         code_paths: Optional[List[str]] = None,
         options: Optional[model_types.ModelSaveOption] = None,
-    ) -> None:
+    ) -> model_meta.ModelMetadata:
         if not options:
             options = model_types.BaseModelSaveOption()
 
@@ -106,7 +108,7 @@ class ModelComposer:
                 )
                 options["embed_local_ml_library"] = True
 
-        self.packager.save(
+        model_metadata: model_meta.ModelMetadata = self.packager.save(
             name=name,
             model=model,
             signatures=signatures,
@@ -119,7 +121,6 @@ class ModelComposer:
             code_paths=code_paths,
             options=options,
         )
-
         assert self.packager.meta is not None
 
         if not options.get("_legacy_save", False):
@@ -128,16 +129,14 @@ class ModelComposer:
             file_utils.copytree(
                 str(self._packager_workspace_path), str(self.workspace_path / ModelComposer.MODEL_DIR_REL_PATH)
             )
-
-        file_utils.make_archive(self.model_local_path, str(self._packager_workspace_path))
-
-        self.manifest.save(
-            session=self.session,
-            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, sample_input_data),
-        )
+            self.manifest.save(
+                model_meta=self.packager.meta,
+                model_rel_path=pathlib.PurePosixPath(ModelComposer.MODEL_DIR_REL_PATH),
+                options=options,
+                data_sources=self._get_data_sources(model, sample_input_data),
+            )
+        else:
+            file_utils.make_archive(self.model_local_path, str(self._packager_workspace_path))
 
         file_utils.upload_directory_to_stage(
             self.session,
@@ -145,6 +144,7 @@ class ModelComposer:
             stage_path=self.stage_path,
             statement_params=self._statement_params,
         )
+        return model_metadata
 
     @deprecated("Only used by PrPr model registry. Use static method version of load instead.")
     def legacy_load(
@@ -185,6 +185,4 @@ class ModelComposer:
         data_sources = lineage_utils.get_data_sources(model)
         if not data_sources and sample_input_data is not None:
             data_sources = lineage_utils.get_data_sources(sample_input_data)
-        if isinstance(data_sources, list) and all(isinstance(item, data_source.DataSource) for item in data_sources):
-            return data_sources
-        return None
+        return data_sources