snowflake-ml-python 1.8.6__py3-none-any.whl → 1.9.1__py3-none-any.whl

snowflake/ml/_internal/env_utils.py

@@ -337,13 +337,54 @@ def get_package_spec_with_supported_ops_only(req: requirements.Requirement) -> r
     Returns:
         A requirements.Requirement object with supported ops only
     """
+
+    if req.name == "numpy":
+        import numpy as np
+
+        package_specifiers = get_numpy_specifiers(req, version.Version(np.__version__).major)
+    else:
+        package_specifiers = [spec for spec in req.specifier if spec.operator in _SUPPORTED_PACKAGE_SPEC_OPS]
+
     new_req = copy.deepcopy(req)
-    new_req.specifier = specifiers.SpecifierSet(
-        specifiers=",".join([str(spec) for spec in req.specifier if spec.operator in _SUPPORTED_PACKAGE_SPEC_OPS])
-    )
+    new_req.specifier = specifiers.SpecifierSet(specifiers=",".join([str(spec) for spec in package_specifiers]))
     return new_req
 
 
+def get_numpy_specifiers(
+    req: requirements.Requirement,
+    client_numpy_major_version: int,
+) -> list[specifiers.Specifier]:
+    """Get the package spec with supported ops only including ==, >=, <=, > and < based on the client numpy
+    major version.
+
+    Args:
+        req: A requirements.Requirement object showing the requirement.
+        client_numpy_major_version: The major version of numpy to be used.
+
+    Returns:
+        A list of specifiers with supported ops only
+    """
+    req_specifiers = []
+    for org_spec in req.specifier:
+        # check specifier that provides upper bound
+        if org_spec.operator in ["<", "<="]:
+            client_version = version.Version(str(client_numpy_major_version))
+            org_spec_version = version.Version(org_spec.version)
+            # check if the client's numpy major version is less than the specifier's upper bound
+            # if so, pin to max possible client major version
+            if client_version.major < org_spec_version.major:
+                modified_spec = specifiers.Specifier(f"<{client_version.major + 1}")
+                req_specifiers.append(modified_spec)
+            else:
+                # use the original specifier
+                req_specifiers.append(org_spec)
+        else:
+            # use the original specifier
+            req_specifiers.append(org_spec)
+
+    return req_specifiers
+
+
 def _relax_specifier_set(
     specifier_set: specifiers.SpecifierSet, strategy: relax_version_strategy.RelaxVersionStrategy
 ) -> specifiers.SpecifierSet:

snowflake/ml/_internal/platform_capabilities.py

@@ -3,7 +3,9 @@ from contextlib import contextmanager
 from typing import Any, Optional
 
 from absl import logging
+from packaging import version
 
+from snowflake.ml import version as snowml_version
 from snowflake.ml._internal.exceptions import error_codes, exceptions
 from snowflake.ml._internal.utils import query_result_checker
 from snowflake.snowpark import (
@@ -12,7 +14,7 @@ from snowflake.snowpark import (
 )
 
 LIVE_COMMIT_PARAMETER = "ENABLE_LIVE_VERSION_IN_SDK"
-INLINE_DEPLOYMENT_SPEC_PARAMETER = "ENABLE_INLINE_DEPLOYMENT_SPEC"
+INLINE_DEPLOYMENT_SPEC_PARAMETER = "ENABLE_INLINE_DEPLOYMENT_SPEC_FROM_CLIENT_VERSION"
 
 
 class PlatformCapabilities:
@@ -67,7 +69,7 @@ class PlatformCapabilities:
             cls.clear_mock_features()
 
     def is_inlined_deployment_spec_enabled(self) -> bool:
-        return self._get_bool_feature(INLINE_DEPLOYMENT_SPEC_PARAMETER, False)
+        return self._is_version_feature_enabled(INLINE_DEPLOYMENT_SPEC_PARAMETER)
 
     def is_live_commit_enabled(self) -> bool:
         return self._get_bool_feature(LIVE_COMMIT_PARAMETER, False)
@@ -126,3 +128,51 @@ class PlatformCapabilities:
             else:
                 raise ValueError(f"Invalid boolean string: {value} for feature {feature_name}")
         raise ValueError(f"Invalid boolean feature value: {value} for feature {feature_name}")
+
+    def _get_version_feature(self, feature_name: str) -> version.Version:
+        """Get a version feature value, returning a large version number on failure or missing feature.
+
+        Args:
+            feature_name: The name of the feature to retrieve.
+
+        Returns:
+            version.Version: The parsed version, or a large version number (999.999.999) if parsing fails
+            or the feature is missing.
+        """
+        # Large version number to use as fallback
+        large_version = version.Version("999.999.999")
+
+        value = self.features.get(feature_name)
+        if value is None:
+            logging.debug(f"Feature {feature_name} not found, returning large version number")
+            return large_version
+
+        try:
+            # Convert to string if it's not already
+            version_str = str(value)
+            return version.Version(version_str)
+        except (version.InvalidVersion, ValueError, TypeError) as e:
+            logging.debug(
+                f"Failed to parse version from feature {feature_name} with value '{value}': {e}. "
+                f"Returning large version number"
+            )
+            return large_version
+
+    def _is_version_feature_enabled(self, feature_name: str) -> bool:
+        """Check if the current package version is greater than or equal to the version feature.
+
+        Args:
+            feature_name: The name of the version feature to compare against.
+
+        Returns:
+            bool: True if current package version >= feature version, False otherwise.
+        """
+        current_version = version.Version(snowml_version.VERSION)
+        feature_version = self._get_version_feature(feature_name)
+
+        result = current_version >= feature_version
+        logging.debug(
+            f"Version comparison for feature {feature_name}: "
+            f"current={current_version}, feature={feature_version}, enabled={result}"
+        )
+        return result

snowflake/ml/_internal/type_utils.py

@@ -66,4 +66,4 @@ class LazyType(Generic[T]):
             return False
 
 
-LiteralNDArrayType = Union[npt.NDArray[np.int_], npt.NDArray[np.float_], npt.NDArray[np.str_], npt.NDArray[np.bool_]]
+LiteralNDArrayType = Union[npt.NDArray[np.int_], npt.NDArray[np.float64], npt.NDArray[np.str_], npt.NDArray[np.bool_]]

snowflake/ml/_internal/utils/identifier.py

@@ -240,7 +240,7 @@ def get_schema_level_object_identifier(
     """
 
     for identifier in (db, schema, object_name):
-        if identifier is not None and SF_IDENTIFIER_RE.match(identifier) is None:
+        if identifier is not None and SF_IDENTIFIER_RE.fullmatch(identifier) is None:
             raise ValueError(f"Invalid identifier {identifier}")
 
     if others is None:

snowflake/ml/_internal/utils/mixins.py

@@ -0,0 +1,71 @@
+from typing import Any, Optional
+
+from snowflake.ml._internal.utils import identifier
+from snowflake.snowpark import session as snowpark_session
+
+_SESSION_KEY = "_session"
+_SESSION_ACCOUNT_KEY = "session$account"
+_SESSION_ROLE_KEY = "session$role"
+_SESSION_DATABASE_KEY = "session$database"
+_SESSION_SCHEMA_KEY = "session$schema"
+
+
+def _identifiers_match(saved: Optional[str], current: Optional[str]) -> bool:
+    saved_resolved = identifier.resolve_identifier(saved) if saved is not None else saved
+    current_resolved = identifier.resolve_identifier(current) if current is not None else current
+    return saved_resolved == current_resolved
+
+
+class SerializableSessionMixin:
+    """Mixin that provides pickling capabilities for objects with Snowpark sessions."""
+
+    def __getstate__(self) -> dict[str, Any]:
+        """Customize pickling to exclude non-serializable session and related components."""
+        if hasattr(super(), "__getstate__"):
+            state: dict[str, Any] = super().__getstate__()  # type: ignore[misc]
+        else:
+            state = self.__dict__.copy()
+
+        # Save session metadata for validation during unpickling
+        session = state.pop(_SESSION_KEY, None)
+        if session is not None:
+            state[_SESSION_ACCOUNT_KEY] = session.get_current_account()
+            state[_SESSION_ROLE_KEY] = session.get_current_role()
+            state[_SESSION_DATABASE_KEY] = session.get_current_database()
+            state[_SESSION_SCHEMA_KEY] = session.get_current_schema()
+
+        return state
+
+    def __setstate__(self, state: dict[str, Any]) -> None:
+        """Restore session from context during unpickling."""
+        saved_account = state.pop(_SESSION_ACCOUNT_KEY, None)
+        saved_role = state.pop(_SESSION_ROLE_KEY, None)
+        saved_database = state.pop(_SESSION_DATABASE_KEY, None)
+        saved_schema = state.pop(_SESSION_SCHEMA_KEY, None)
+
+        if hasattr(super(), "__setstate__"):
+            super().__setstate__(state)  # type: ignore[misc]
+        else:
+            self.__dict__.update(state)
+
+        if saved_account is not None:
+            active_sessions = snowpark_session._get_active_sessions()
+            if len(active_sessions) == 0:
+                raise RuntimeError("No active Snowpark session available. Please create a session.")
+
+            # Best effort match: Find the session with the most matching identifiers
+            setattr(
+                self,
+                _SESSION_KEY,
+                max(
+                    active_sessions,
+                    key=lambda s: sum(
+                        (
+                            _identifiers_match(saved_account, s.get_current_account()),
+                            _identifiers_match(saved_role, s.get_current_role()),
+                            _identifiers_match(saved_database, s.get_current_database()),
+                            _identifiers_match(saved_schema, s.get_current_schema()),
+                        )
+                    ),
+                ),
+            )

snowflake/ml/_internal/utils/service_logger.py

@@ -10,6 +10,10 @@ class LogColor(enum.Enum):
     YELLOW = "\x1b[33;20m"
     BLUE = "\x1b[34;20m"
     GREEN = "\x1b[32;20m"
+    ORANGE = "\x1b[38;5;214m"
+    BOLD_ORANGE = "\x1b[38;5;214;1m"
+    PURPLE = "\x1b[35;20m"
+    BOLD_PURPLE = "\x1b[35;1m"
 
 
 class CustomFormatter(logging.Formatter):
@@ -55,9 +59,7 @@ class CustomFormatter(logging.Formatter):
 
 def get_logger(logger_name: str, info_color: LogColor) -> logging.Logger:
     logger = logging.getLogger(logger_name)
-    logger.setLevel(logging.INFO)
     handler = logging.StreamHandler(sys.stdout)
-    handler.setLevel(logging.INFO)
     handler.setFormatter(CustomFormatter(info_color))
     logger.addHandler(handler)
     return logger

snowflake/ml/data/_internal/arrow_ingestor.py

@@ -2,7 +2,7 @@ import collections
 import logging
 import os
 import time
-from typing import Any, Deque, Iterator, Optional, Sequence, Union
+from typing import TYPE_CHECKING, Any, Deque, Iterator, Optional, Sequence, Union
 
 import numpy as np
 import numpy.typing as npt
@@ -10,6 +10,9 @@ import pandas as pd
 import pyarrow as pa
 import pyarrow.dataset as pds
 
+if TYPE_CHECKING:
+    import ray
+
 from snowflake import snowpark
 from snowflake.ml.data import data_ingestor, data_source, ingestor_utils
 
@@ -70,6 +73,13 @@ class ArrowIngestor(data_ingestor.DataIngestor):
     def from_sources(cls, session: snowpark.Session, sources: Sequence[data_source.DataSource]) -> "ArrowIngestor":
         return cls(session, sources)
 
+    @classmethod
+    def from_ray_dataset(
+        cls,
+        ray_ds: "ray.data.Dataset",
+    ) -> "ArrowIngestor":
+        raise NotImplementedError
+
     @property
     def data_sources(self) -> list[data_source.DataSource]:
         return self._data_sources

snowflake/ml/data/data_connector.py

@@ -6,6 +6,7 @@ from typing_extensions import deprecated
 
 from snowflake import snowpark
 from snowflake.ml._internal import env, telemetry
+from snowflake.ml._internal.utils import mixins
 from snowflake.ml.data import data_ingestor, data_source
 from snowflake.ml.data._internal.arrow_ingestor import ArrowIngestor
 from snowflake.snowpark import context as sf_context
@@ -21,11 +22,13 @@ if TYPE_CHECKING:
     from snowflake.ml import dataset
 
 _PROJECT = "DataConnector"
+_INGESTOR_KEY = "_ingestor"
+_INGESTOR_SOURCES_KEY = "ingestor$sources"
 
 DataConnectorType = TypeVar("DataConnectorType", bound="DataConnector")
 
 
-class DataConnector:
+class DataConnector(mixins.SerializableSessionMixin):
     """Snowflake data reader which provides application integration connectors"""
 
     DEFAULT_INGESTOR_CLASS: type[data_ingestor.DataIngestor] = ArrowIngestor
@@ -33,8 +36,11 @@ class DataConnector:
     def __init__(
         self,
         ingestor: data_ingestor.DataIngestor,
+        *,
+        session: Optional[snowpark.Session] = None,
         **kwargs: Any,
     ) -> None:
+        self._session = session
         self._ingestor = ingestor
         self._kwargs = kwargs
 
@@ -75,6 +81,17 @@ class DataConnector:
         )
         return cls.from_sources(ds._session, [source], ingestor_class=ingestor_class, **kwargs)
 
+    @classmethod
+    def from_ray_dataset(
+        cls: type[DataConnectorType],
+        ray_ds: "ray.data.Dataset",
+        ingestor_class: Optional[type[data_ingestor.DataIngestor]] = None,
+        **kwargs: Any,
+    ) -> DataConnectorType:
+        ingestor_class = ingestor_class or cls.DEFAULT_INGESTOR_CLASS
+        ray_ingestor = ingestor_class.from_ray_dataset(ray_ds=ray_ds)
+        return cls(ray_ingestor, **kwargs)
+
     @classmethod
     @telemetry.send_api_usage_telemetry(
         project=_PROJECT,
@@ -90,7 +107,31 @@ class DataConnector:
     ) -> DataConnectorType:
         ingestor_class = ingestor_class or cls.DEFAULT_INGESTOR_CLASS
         ingestor = ingestor_class.from_sources(session, sources)
-        return cls(ingestor, **kwargs)
+        return cls(ingestor, **kwargs, session=session)
+
+    def __getstate__(self) -> dict[str, Any]:
+        """Customize pickling to exclude non-serializable session and related components."""
+        if hasattr(super(), "__getstate__"):
+            state = super().__getstate__()
+        else:
+            state = self.__dict__.copy()
+
+        ingestor = state.pop(_INGESTOR_KEY)
+        state[_INGESTOR_SOURCES_KEY] = ingestor.data_sources
+
+        return state
+
+    def __setstate__(self, state: dict[str, Any]) -> None:
+        """Restore session from context during unpickling."""
+        data_sources = state.pop(_INGESTOR_SOURCES_KEY)
+
+        if hasattr(super(), "__setstate__"):
+            super().__setstate__(state)
+        else:
+            self.__dict__.update(state)
+
+        assert self._session is not None
+        self._ingestor = self.DEFAULT_INGESTOR_CLASS.from_sources(self._session, data_sources)
 
     @property
     def data_sources(self) -> list[data_source.DataSource]:

snowflake/ml/data/data_ingestor.py

@@ -7,6 +7,7 @@ from snowflake.ml.data import data_source
 
 if TYPE_CHECKING:
     import pandas as pd
+    import ray
 
 
 DataIngestorType = TypeVar("DataIngestorType", bound="DataIngestor")
@@ -19,6 +20,13 @@ class DataIngestor(Protocol):
     ) -> DataIngestorType:
         raise NotImplementedError
 
+    @classmethod
+    def from_ray_dataset(
+        cls: type[DataIngestorType],
+        ray_ds: "ray.data.Dataset",
+    ) -> DataIngestorType:
+        raise NotImplementedError
+
     @property
     def data_sources(self) -> list[data_source.DataSource]:
         raise NotImplementedError

snowflake/ml/data/torch_utils.py

@@ -95,6 +95,6 @@ def _preprocess_array(
         array_list = arr.tolist()
         # If this is an array of arrays, convert the dtype to match the underlying array.
         # Otherwise, if this is a numpy array of strings, convert the array to a list.
-        arr = np.array(array_list, dtype=arr.flat[0].dtype) if isinstance(arr.flat[0], np.ndarray) else array_list
+        arr = np.array(array_list, dtype=arr.item(0).dtype) if isinstance(arr.item(0), np.ndarray) else array_list
 
     return arr

snowflake/ml/dataset/dataset.py

@@ -14,6 +14,7 @@ from snowflake.ml._internal.exceptions import (
 from snowflake.ml._internal.utils import (
     formatting,
     identifier,
+    mixins,
     query_result_checker,
     snowpark_dataframe_utils,
 )
@@ -27,7 +28,7 @@ _METADATA_MAX_QUERY_LENGTH = 10000
 _DATASET_VERSION_NAME_COL = "version"
 
 
-class DatasetVersion:
+class DatasetVersion(mixins.SerializableSessionMixin):
     """Represents a version of a Snowflake Dataset"""
 
     @telemetry.send_api_usage_telemetry(project=_PROJECT)
@@ -176,7 +177,7 @@ class Dataset(lineage_node.LineageNode):
                 original_exception=RuntimeError("No Dataset version selected."),
             )
         if self._reader is None:
-            self._reader = dataset_reader.DatasetReader.from_dataset(self, snowpark_session=self._session)
+            self._reader = dataset_reader.DatasetReader.from_dataset(self)
         return self._reader
 
     @staticmethod

snowflake/ml/dataset/dataset_reader.py

@@ -1,8 +1,10 @@
 from typing import Any, Optional
+from warnings import warn
 
 from snowflake import snowpark
 from snowflake.ml._internal import telemetry
 from snowflake.ml._internal.lineage import lineage_utils
+from snowflake.ml._internal.utils import mixins
 from snowflake.ml.data import data_connector, data_ingestor, data_source, ingestor_utils
 from snowflake.ml.fileset import snowfs
 from snowflake.snowpark._internal import utils as snowpark_utils
@@ -11,7 +13,7 @@ _PROJECT = "Dataset"
 _SUBPROJECT = "DatasetReader"
 
 
-class DatasetReader(data_connector.DataConnector):
+class DatasetReader(data_connector.DataConnector, mixins.SerializableSessionMixin):
     """Snowflake Dataset abstraction which provides application integration connectors"""
 
     @telemetry.send_api_usage_telemetry(project=_PROJECT, subproject=_SUBPROJECT)
@@ -19,14 +21,26 @@ class DatasetReader(data_connector.DataConnector):
         self,
         ingestor: data_ingestor.DataIngestor,
         *,
-        snowpark_session: snowpark.Session,
+        session: snowpark.Session,
+        snowpark_session: Optional[snowpark.Session] = None,
     ) -> None:
-        super().__init__(ingestor)
-
-        self._session: snowpark.Session = snowpark_session
-        self._fs: snowfs.SnowFileSystem = ingestor_utils.get_dataset_filesystem(self._session)
+        if snowpark_session is not None:
+            warn(
+                "Argument snowpark_session is deprecated and will be removed in a future release. Use session instead."
+            )
+            session = snowpark_session
+        super().__init__(ingestor, session=session)
+
+        self._fs_cached: Optional[snowfs.SnowFileSystem] = None
         self._files: Optional[list[str]] = None
 
+    @property
+    def _fs(self) -> snowfs.SnowFileSystem:
+        if self._fs_cached is None:
+            assert self._session is not None
+            self._fs_cached = ingestor_utils.get_dataset_filesystem(self._session)
+        return self._fs_cached
+
     @classmethod
     def from_dataframe(
         cls, df: snowpark.DataFrame, ingestor_class: Optional[type[data_ingestor.DataIngestor]] = None, **kwargs: Any
@@ -42,6 +56,7 @@ class DatasetReader(data_connector.DataConnector):
         files: list[str] = []
         for source in self.data_sources:
             assert isinstance(source, data_source.DatasetInfo)
+            assert self._session is not None
             files.extend(ingestor_utils.get_dataset_files(self._session, source, filesystem=self._fs))
         files.sort()
 
@@ -95,6 +110,7 @@ class DatasetReader(data_connector.DataConnector):
         dfs: list[snowpark.DataFrame] = []
         for source in self.data_sources:
             assert isinstance(source, data_source.DatasetInfo) and source.url is not None
+            assert self._session is not None
             stage_reader = self._session.read.option("pattern", file_path_pattern)
             if "INFER_SCHEMA_OPTIONS" in snowpark_utils.NON_FORMAT_TYPE_OPTIONS:
                 stage_reader = stage_reader.option("INFER_SCHEMA_OPTIONS", {"MAX_FILE_COUNT": 1})

snowflake/ml/experiment/_client/experiment_tracking_sql_client.py

@@ -0,0 +1,98 @@
+from typing import Optional
+
+from snowflake.ml._internal.utils import query_result_checker, sql_identifier
+from snowflake.ml.model._client.sql import _base
+from snowflake.ml.utils import sql_client
+from snowflake.snowpark import row, session
+
+
+class ExperimentTrackingSQLClient(_base._BaseSQLClient):
+
+    RUN_NAME_COL_NAME = "name"
+    RUN_METADATA_COL_NAME = "metadata"
+
+    def __init__(
+        self,
+        session: session.Session,
+        *,
+        database_name: sql_identifier.SqlIdentifier,
+        schema_name: sql_identifier.SqlIdentifier,
+    ) -> None:
+        """Snowflake SQL Client to manage experiment tracking.
+
+        Args:
+            session: Active snowpark session.
+            database_name: Name of the Database where experiment tracking resources are provisioned.
+            schema_name: Name of the Schema where experiment tracking resources are provisioned.
+        """
+        super().__init__(session, database_name=database_name, schema_name=schema_name)
+
+    def create_experiment(
+        self,
+        experiment_name: sql_identifier.SqlIdentifier,
+        creation_mode: sql_client.CreationMode,
+    ) -> None:
+        experiment_fqn = self.fully_qualified_object_name(self._database_name, self._schema_name, experiment_name)
+        if_not_exists_sql = "IF NOT EXISTS" if creation_mode.if_not_exists else ""
+        query_result_checker.SqlResultValidator(
+            self._session, f"CREATE EXPERIMENT {if_not_exists_sql} {experiment_fqn}"
+        ).has_dimensions(expected_rows=1, expected_cols=1).validate()
+
+    def drop_experiment(self, *, experiment_name: sql_identifier.SqlIdentifier) -> None:
+        experiment_fqn = self.fully_qualified_object_name(self._database_name, self._schema_name, experiment_name)
+        query_result_checker.SqlResultValidator(self._session, f"DROP EXPERIMENT {experiment_fqn}").has_dimensions(
+            expected_rows=1, expected_cols=1
+        ).validate()
+
+    def add_run(
+        self,
+        *,
+        experiment_name: sql_identifier.SqlIdentifier,
+        run_name: sql_identifier.SqlIdentifier,
+        live: bool = True,
+    ) -> None:
+        experiment_fqn = self.fully_qualified_object_name(self._database_name, self._schema_name, experiment_name)
+        query_result_checker.SqlResultValidator(
+            self._session, f"ALTER EXPERIMENT {experiment_fqn} ADD {'LIVE' if live else ''} RUN {run_name}"
+        ).has_dimensions(expected_rows=1, expected_cols=1).validate()
+
+    def commit_run(
+        self,
+        *,
+        experiment_name: sql_identifier.SqlIdentifier,
+        run_name: sql_identifier.SqlIdentifier,
+    ) -> None:
+        experiment_fqn = self.fully_qualified_object_name(self._database_name, self._schema_name, experiment_name)
+        query_result_checker.SqlResultValidator(
+            self._session, f"ALTER EXPERIMENT {experiment_fqn} COMMIT RUN {run_name}"
+        ).has_dimensions(expected_rows=1, expected_cols=1).validate()
+
+    def drop_run(
+        self, *, experiment_name: sql_identifier.SqlIdentifier, run_name: sql_identifier.SqlIdentifier
+    ) -> None:
+        experiment_fqn = self.fully_qualified_object_name(self._database_name, self._schema_name, experiment_name)
+        query_result_checker.SqlResultValidator(
+            self._session, f"ALTER EXPERIMENT {experiment_fqn} DROP RUN {run_name}"
+        ).has_dimensions(expected_rows=1, expected_cols=1).validate()
+
+    def modify_run(
+        self,
+        *,
+        experiment_name: sql_identifier.SqlIdentifier,
+        run_name: sql_identifier.SqlIdentifier,
+        run_metadata: str,
+    ) -> None:
+        experiment_fqn = self.fully_qualified_object_name(self._database_name, self._schema_name, experiment_name)
+        query_result_checker.SqlResultValidator(
+            self._session,
+            f"ALTER EXPERIMENT {experiment_fqn} MODIFY RUN {run_name} SET METADATA=$${run_metadata}$$",
+        ).has_dimensions(expected_rows=1, expected_cols=1).validate()
+
+    def show_runs_in_experiment(
+        self, *, experiment_name: sql_identifier.SqlIdentifier, like: Optional[str] = None
+    ) -> list[row.Row]:
+        experiment_fqn = self.fully_qualified_object_name(self._database_name, self._schema_name, experiment_name)
+        like_clause = f"LIKE '{like}'" if like else ""
+        return query_result_checker.SqlResultValidator(
+            self._session, f"SHOW RUNS {like_clause} IN EXPERIMENT {experiment_fqn}"
+        ).validate()

snowflake/ml/experiment/_entities/__init__.py

@@ -0,0 +1,4 @@
+from snowflake.ml.experiment._entities.experiment import Experiment
+from snowflake.ml.experiment._entities.run import Run
+
+__all__ = ["Experiment", "Run"]

snowflake/ml/experiment/_entities/experiment.py

@@ -0,0 +1,10 @@
+from snowflake.ml._internal.utils import sql_identifier
+
+
+class Experiment:
+    def __init__(
+        self,
+        *,
+        experiment_name: sql_identifier.SqlIdentifier,
+    ) -> None:
+        self.name = experiment_name

snowflake/ml/experiment/_entities/run.py

@@ -0,0 +1,62 @@
+import json
+import types
+from typing import TYPE_CHECKING, Optional
+
+from snowflake.ml._internal.utils import sql_identifier
+from snowflake.ml.experiment import _experiment_info as experiment_info
+from snowflake.ml.experiment._client import experiment_tracking_sql_client
+from snowflake.ml.experiment._entities import run_metadata
+
+if TYPE_CHECKING:
+    from snowflake.ml.experiment import experiment_tracking
+
+
+class Run:
+    def __init__(
+        self,
+        experiment_tracking: "experiment_tracking.ExperimentTracking",
+        *,
+        experiment_name: sql_identifier.SqlIdentifier,
+        run_name: sql_identifier.SqlIdentifier,
+    ) -> None:
+        self._experiment_tracking = experiment_tracking
+        self.experiment_name = experiment_name
+        self.name = run_name
+
+        self._patcher = experiment_info.ExperimentInfoPatcher(
+            experiment_info=self._get_experiment_info(),
+        )
+
+    def __enter__(self) -> "Run":
+        self._patcher.__enter__()
+        return self
+
+    def __exit__(
+        self,
+        exc_type: Optional[type[BaseException]],
+        exc_value: Optional[BaseException],
+        traceback: Optional[types.TracebackType],
+    ) -> None:
+        self._patcher.__exit__(exc_type, exc_value, traceback)
+        if self._experiment_tracking._run is self:
+            self._experiment_tracking.end_run()
+
+    def _get_metadata(
+        self,
+    ) -> run_metadata.RunMetadata:
+        runs = self._experiment_tracking._sql_client.show_runs_in_experiment(
+            experiment_name=self.experiment_name, like=str(self.name)
+        )
+        if not runs:
+            raise RuntimeError(f"Run {self.name} not found in experiment {self.experiment_name}.")
+        return run_metadata.RunMetadata.from_dict(
+            json.loads(runs[0][experiment_tracking_sql_client.ExperimentTrackingSQLClient.RUN_METADATA_COL_NAME])
+        )
+
+    def _get_experiment_info(self) -> experiment_info.ExperimentInfo:
+        return experiment_info.ExperimentInfo(
+            fully_qualified_name=self._experiment_tracking._sql_client.fully_qualified_object_name(
+                self._experiment_tracking._database_name, self._experiment_tracking._schema_name, self.experiment_name
+            ),
+            run_name=self.name.identifier(),
+        )