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

snowflake/ml/feature_store/feature_view.py

@@ -2,17 +2,25 @@ from __future__ import annotations
 
 import json
 import re
+import warnings
 from collections import OrderedDict
 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 import identifier
 from snowflake.ml._internal.utils.identifier import concat_names
 from snowflake.ml._internal.utils.sql_identifier import (
     SqlIdentifier,
     to_sql_identifiers,
 )
+from snowflake.ml.feature_store import feature_store
 from snowflake.ml.feature_store.entity import Entity
+from snowflake.ml.lineage import lineage_node
 from snowflake.snowpark import DataFrame, Session
 from snowflake.snowpark.types import (
     DateType,
@@ -67,6 +75,7 @@ class FeatureViewVersion(str):
 
 
 class FeatureViewStatus(Enum):
+    MASKED = "MASKED"  # for shared feature views where scheduling state is not available
     DRAFT = "DRAFT"
     STATIC = "STATIC"
     RUNNING = "RUNNING"  # This can be deprecated after BCR 2024_02 gets fully deployed
@@ -107,7 +116,7 @@ class FeatureViewSlice:
         return cls(**json_dict)
 
 
-class FeatureView:
+class FeatureView(lineage_node.LineageNode):
     """
     A FeatureView instance encapsulates a logical group of features.
     """
@@ -243,6 +252,16 @@ class FeatureView:
     def desc(self) -> str:
         return self._desc
 
+    @desc.setter
+    def desc(self, new_value: str) -> None:
+        warnings.warn(
+            "You must call register_feature_view() to make it effective. "
+            "Or use update_feature_view(desc=<new_value>).",
+            stacklevel=2,
+            category=UserWarning,
+        )
+        self._desc = new_value
+
     @property
     def query(self) -> str:
         return self._query
@@ -269,10 +288,12 @@ class FeatureView:
 
     @refresh_freq.setter
     def refresh_freq(self, new_value: str) -> None:
-        if self.status == FeatureViewStatus.DRAFT or self.status == FeatureViewStatus.STATIC:
-            raise RuntimeError(
-                f"Feature view {self.name}/{self.version} must be registered and non-static to update refresh_freq."
-            )
+        warnings.warn(
+            "You must call register_feature_view() to make it effective. "
+            "Or use update_feature_view(refresh_freq=<new_value>).",
+            stacklevel=2,
+            category=UserWarning,
+        )
         self._refresh_freq = new_value
 
     @property
@@ -289,10 +310,12 @@ class FeatureView:
 
     @warehouse.setter
     def warehouse(self, new_value: str) -> None:
-        if self.status == FeatureViewStatus.DRAFT or self.status == FeatureViewStatus.STATIC:
-            raise RuntimeError(
-                f"Feature view {self.name}/{self.version} must be registered and non-static to update warehouse."
-            )
+        warnings.warn(
+            "You must call register_feature_view() to make it effective. "
+            "Or use update_feature_view(warehouse=<new_value>).",
+            stacklevel=2,
+            category=UserWarning,
+        )
         self._warehouse = SqlIdentifier(new_value)
 
     @property
@@ -406,6 +429,11 @@ Got {len(self._feature_df.queries['queries'])}: {self._feature_df.queries['queri
             feature_desc_dict[k.identifier()] = v
         fv_dict["_feature_desc"] = feature_desc_dict
 
+        lineage_node_keys = [key for key in fv_dict if key.startswith("_node") or key == "_session"]
+
+        for key in lineage_node_keys:
+            fv_dict.pop(key)
+
         return fv_dict
 
     def to_df(self, session: Session) -> DataFrame:
@@ -449,6 +477,7 @@ Got {len(self._feature_df.queries['queries'])}: {self._feature_df.queries['queri
             refresh_mode_reason=json_dict["_refresh_mode_reason"],
             owner=json_dict["_owner"],
             infer_schema_df=session.sql(json_dict.get("_infer_schema_query", None)),
+            session=session,
         )
 
     @staticmethod
@@ -463,6 +492,21 @@ Got {len(self._feature_df.queries['queries'])}: {self._feature_df.queries['queri
             )
         )
 
+    @staticmethod
+    def _load_from_lineage_node(session: Session, name: str, version: str) -> FeatureView:
+        db_name, feature_store_name, feature_view_name, _ = identifier.parse_schema_level_object_identifier(name)
+
+        session_warehouse = session.get_current_warehouse()
+
+        if not session_warehouse:
+            raise snowml_exceptions.SnowflakeMLException(
+                error_code=error_codes.NOT_FOUND,
+                original_exception=ValueError("No active warehouse selected in the current session"),
+            )
+
+        fs = feature_store.FeatureStore(session, db_name, feature_store_name, session_warehouse)
+        return fs.get_feature_view(feature_view_name, version)  # type: ignore[no-any-return]
+
     @staticmethod
     def _construct_feature_view(
         name: str,
@@ -481,6 +525,7 @@ Got {len(self._feature_df.queries['queries'])}: {self._feature_df.queries['queri
         refresh_mode_reason: Optional[str],
         owner: Optional[str],
         infer_schema_df: Optional[DataFrame],
+        session: Session,
     ) -> FeatureView:
         fv = FeatureView(
             name=name,
@@ -500,4 +545,11 @@ Got {len(self._feature_df.queries['queries'])}: {self._feature_df.queries['queri
         fv._refresh_mode_reason = refresh_mode_reason
         fv._owner = owner
         fv.attach_feature_desc(feature_descs)
+
+        lineage_node.LineageNode.__init__(
+            fv, session=session, name=f"{fv.database}.{fv._schema}.{name}", domain="feature_view", version=version
+        )
         return fv
+
+
+lineage_node.DOMAIN_LINEAGE_REGISTRY["feature_view"] = FeatureView

snowflake/ml/fileset/embedded_stage_fs.py

@@ -11,11 +11,17 @@ from snowflake.ml._internal.exceptions import (
     fileset_errors,
 )
 from snowflake.ml._internal.utils import identifier
+from snowflake.ml.fileset import stage_fs
 from snowflake.snowpark import exceptions as snowpark_exceptions
 
-from . import stage_fs
-
-_SNOWURL_PATH_RE = re.compile(r"versions/(?P<version>[^/]+)(?:/+(?P<filepath>.*))?")
+PROTOCOL_NAME = "snow"
+_SNOWURL_ENTITY_PATTERN = (
+    f"(?:{PROTOCOL_NAME}://)?"
+    r"(?<!@)(?P<domain>\w+)/"
+    rf"(?P<name>(?:{identifier._SF_IDENTIFIER}\.){{,2}}{identifier._SF_IDENTIFIER})/"
+)
+_SNOWURL_VERSION_PATTERN = r"(?P<path>versions/(?:(?P<version>[^/]+)(?:/+(?P<relpath>.*))?)?)"
+_SNOWURL_PATH_RE = re.compile(f"(?:{_SNOWURL_ENTITY_PATTERN})?" + _SNOWURL_VERSION_PATTERN)
 
 
 class SFEmbeddedStageFileSystem(stage_fs.SFStageFileSystem):
@@ -76,8 +82,8 @@ class SFEmbeddedStageFileSystem(stage_fs.SFStageFileSystem):
         versions_dict = defaultdict(list)
         for file in files:
             match = _SNOWURL_PATH_RE.fullmatch(file)
-            assert match is not None and match.group("filepath") is not None
-            versions_dict[match.group("version")].append(match.group("filepath"))
+            assert match is not None and match.group("relpath") is not None
+            versions_dict[match.group("version")].append(match.group("relpath"))
         try:
             async_jobs: List[snowpark.AsyncJob] = []
             for version, version_files in versions_dict.items():
@@ -98,10 +104,8 @@ class SFEmbeddedStageFileSystem(stage_fs.SFStageFileSystem):
                 (r["NAME"], r["URL"]) for job in async_jobs for r in stage_fs._resolve_async_job(job)
             ]
             return presigned_urls
-        except snowpark_exceptions.SnowparkClientException as e:
-            if e.message.startswith(fileset_errors.ERRNO_DOMAIN_NOT_EXIST) or e.message.startswith(
-                fileset_errors.ERRNO_STAGE_NOT_EXIST
-            ):
+        except snowpark_exceptions.SnowparkSQLException as e:
+            if e.sql_error_code in {fileset_errors.ERRNO_DOMAIN_NOT_EXIST, fileset_errors.ERRNO_STAGE_NOT_EXIST}:
                 raise snowml_exceptions.SnowflakeMLException(
                     error_code=error_codes.SNOWML_NOT_FOUND,
                     original_exception=fileset_errors.StageNotFoundError(
@@ -118,7 +122,7 @@ class SFEmbeddedStageFileSystem(stage_fs.SFStageFileSystem):
     def _parent(cls, path: str) -> str:
         """Get parent of specified path up to minimally valid root path.
 
-        For SnowURL, the minimum valid path is snow://<domain>/<entity>/versions/<version>
+        For SnowURL, the minimum valid relative path is versions/<version>
 
         Args:
             path: File or directory path
@@ -128,22 +132,22 @@ class SFEmbeddedStageFileSystem(stage_fs.SFStageFileSystem):
 
         Examples:
         ----
-        >>> fs._parent("snow://dataset/my_ds/versions/my_version/file.ext")
-        "snow://dataset/my_ds/versions/my_version/"
-        >>> fs._parent("snow://dataset/my_ds/versions/my_version/subdir/file.ext")
-        "snow://dataset/my_ds/versions/my_version/subdir/"
-        >>> fs._parent("snow://dataset/my_ds/versions/my_version/")
-        "snow://dataset/my_ds/versions/my_version/"
-        >>> fs._parent("snow://dataset/my_ds/versions/my_version")
-        "snow://dataset/my_ds/versions/my_version"
+        >>> fs._parent("versions/my_version/file.ext")
+        "versions/my_version"
+        >>> fs._parent("versions/my_version/subdir/file.ext")
+        "versions/my_version/subdir"
+        >>> fs._parent("versions/my_version/")
+        "versions/my_version"
+        >>> fs._parent("versions/my_version")
+        "versions/my_version"
         """
         path_match = _SNOWURL_PATH_RE.fullmatch(path)
         if not path_match:
             return super()._parent(path)  # type: ignore[no-any-return]
-        filepath: str = path_match.group("filepath") or ""
-        root: str = path[: path_match.start("filepath")] if filepath else path
+        filepath: str = path_match.group("relpath") or ""
+        root: str = path[: path_match.start("relpath")] if filepath else path
         if "/" in filepath:
             parent = filepath.rsplit("/", 1)[0]
             return root + parent
         else:
-            return root
+            return root.rstrip("/")

snowflake/ml/fileset/fileset.py

@@ -256,9 +256,9 @@ class FileSet:
                     api_calls=[snowpark.DataFrameWriter.copy_into_location],
                 ),
             )
-        except snowpark_exceptions.SnowparkClientException as e:
+        except snowpark_exceptions.SnowparkSQLException as e:
             # Snowpark wraps the Python Connector error code in the head of the error message.
-            if e.message.startswith(fileset_errors.ERRNO_FILE_EXIST_IN_STAGE):
+            if e.sql_error_code == fileset_errors.ERRNO_FILE_EXIST_IN_STAGE:
                 raise fileset_errors.FileSetExistError(fileset_error_messages.FILESET_ALREADY_EXISTS.format(name))
             else:
                 raise fileset_errors.FileSetError(str(e))

snowflake/ml/fileset/snowfs.py

@@ -14,18 +14,10 @@ from snowflake.ml._internal.exceptions import (
 from snowflake.ml._internal.utils import identifier
 from snowflake.ml.fileset import embedded_stage_fs, sfcfs
 
-PROTOCOL_NAME = "snow"
-
 _SFFileEntityPath = collections.namedtuple(
     "_SFFileEntityPath", ["domain", "name", "filepath", "version", "relative_path"]
 )
-_PROJECT = "FileSet"
-_SNOWURL_PATTERN = re.compile(
-    f"({PROTOCOL_NAME}://)?"
-    r"(?<!@)(?P<domain>\w+)/"
-    rf"(?P<name>(?:{identifier._SF_IDENTIFIER}\.){{,2}}{identifier._SF_IDENTIFIER})/"
-    r"(?P<path>versions/(?:(?P<version>[^/]+)(?:/(?P<relpath>.*))?)?)"
-)
+_SNOWURL_PATTERN = re.compile(embedded_stage_fs._SNOWURL_ENTITY_PATTERN + embedded_stage_fs._SNOWURL_VERSION_PATTERN)
 
 
 class SnowFileSystem(sfcfs.SFFileSystem):
@@ -38,7 +30,7 @@ class SnowFileSystem(sfcfs.SFFileSystem):
     See `sfcfs.SFFileSystem` documentation for example usage patterns.
     """
 
-    protocol = PROTOCOL_NAME
+    protocol = embedded_stage_fs.PROTOCOL_NAME
     _IS_BUGGED_VERSION = None
 
     def __init__(
@@ -75,10 +67,7 @@ class SnowFileSystem(sfcfs.SFFileSystem):
         """Convert the relative path in a stage to an absolute path starts with the location of the stage."""
         # Strip protocol from absolute path, since backend needs snow:// prefix to resolve correctly
         # but fsspec logic strips protocol when doing any searching and globbing
-        stage_name = stage_fs.stage_name
-        protocol = f"{PROTOCOL_NAME}://"
-        if stage_name.startswith(protocol):
-            stage_name = stage_name[len(protocol) :]
+        stage_name: str = self._strip_protocol(stage_fs.stage_name)
         abs_path = stage_name + "/" + path
         return abs_path
 
@@ -128,4 +117,4 @@ class SnowFileSystem(sfcfs.SFFileSystem):
             )
 
 
-fsspec.register_implementation(PROTOCOL_NAME, SnowFileSystem)
+fsspec.register_implementation(SnowFileSystem.protocol, SnowFileSystem)

snowflake/ml/fileset/stage_fs.py

@@ -170,8 +170,8 @@ class SFStageFileSystem(fsspec.AbstractFileSystem):
             path = path.lstrip("/")
             async_job: snowpark.AsyncJob = self._session.sql(f"LIST '{loc}/{path}'").collect(block=False)
             objects: List[snowpark.Row] = _resolve_async_job(async_job)
-        except snowpark_exceptions.SnowparkClientException as e:
-            if e.message.startswith(fileset_errors.ERRNO_DOMAIN_NOT_EXIST):
+        except snowpark_exceptions.SnowparkSQLException as e:
+            if e.sql_error_code == fileset_errors.ERRNO_DOMAIN_NOT_EXIST:
                 raise snowml_exceptions.SnowflakeMLException(
                     error_code=error_codes.SNOWML_NOT_FOUND,
                     original_exception=fileset_errors.StageNotFoundError(
@@ -387,10 +387,8 @@ class SFStageFileSystem(fsspec.AbstractFileSystem):
                     api_calls=[snowpark.DataFrame.collect],
                 ),
             )
-        except snowpark_exceptions.SnowparkClientException as e:
-            if e.message.startswith(fileset_errors.ERRNO_DOMAIN_NOT_EXIST) or e.message.startswith(
-                fileset_errors.ERRNO_STAGE_NOT_EXIST
-            ):
+        except snowpark_exceptions.SnowparkSQLException as e:
+            if e.sql_error_code in {fileset_errors.ERRNO_DOMAIN_NOT_EXIST, fileset_errors.ERRNO_STAGE_NOT_EXIST}:
                 raise snowml_exceptions.SnowflakeMLException(
                     error_code=error_codes.SNOWML_NOT_FOUND,
                     original_exception=fileset_errors.StageNotFoundError(
@@ -406,9 +404,9 @@ class SFStageFileSystem(fsspec.AbstractFileSystem):
 
 
 def _match_error_code(ex: snowpark_exceptions.SnowparkSQLException, error_code: int) -> bool:
-    # Snowpark writes error code to message instead of populating e.error_code
+    # Snowpark writes error code to message instead of populating e.sql_error_code
     error_code_str = str(error_code)
-    return ex.error_code == error_code_str or error_code_str in ex.message
+    return ex.sql_error_code == error_code_str or error_code_str in ex.message
 
 
 @snowflake_plan.SnowflakePlan.Decorator.wrap_exception  # type: ignore[misc]

snowflake/ml/lineage/__init__.py

@@ -0,0 +1,3 @@
+from .lineage_node import LineageNode
+
+__all__ = ["LineageNode"]

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:
+                    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,