snowflake-ml-python 1.8.1__py3-none-any.whl → 1.8.3__py3-none-any.whl

snowflake/ml/_internal/env_utils.py

@@ -6,12 +6,13 @@ import textwrap
 import warnings
 from enum import Enum
 from importlib import metadata as importlib_metadata
-from typing import Any, DefaultDict, Dict, List, Optional, Tuple
+from typing import Any, DefaultDict, Optional
 
 import yaml
 from packaging import requirements, specifiers, version
 
 import snowflake.connector
+from snowflake.ml import version as snowml_version
 from snowflake.ml._internal import env as snowml_env, relax_version_strategy
 from snowflake.ml._internal.utils import query_result_checker
 from snowflake.snowpark import context, exceptions, session
@@ -27,8 +28,8 @@ class CONDA_OS(Enum):
 
 
 _NODEFAULTS = "nodefaults"
-_SNOWFLAKE_INFO_SCHEMA_PACKAGE_CACHE: Dict[str, List[version.Version]] = {}
-_SNOWFLAKE_CONDA_PACKAGE_CACHE: Dict[str, List[version.Version]] = {}
+_SNOWFLAKE_INFO_SCHEMA_PACKAGE_CACHE: dict[str, list[version.Version]] = {}
+_SNOWFLAKE_CONDA_PACKAGE_CACHE: dict[str, list[version.Version]] = {}
 _SUPPORTED_PACKAGE_SPEC_OPS = ["==", ">=", "<=", ">", "<"]
 
 DEFAULT_CHANNEL_NAME = ""
@@ -64,7 +65,7 @@ def _validate_pip_requirement_string(req_str: str) -> requirements.Requirement:
     return r
 
 
-def _validate_conda_dependency_string(dep_str: str) -> Tuple[str, requirements.Requirement]:
+def _validate_conda_dependency_string(dep_str: str) -> tuple[str, requirements.Requirement]:
     """Validate conda dependency string like `pytorch == 1.12.1` or `conda-forge::transformer` and split the channel
         name before the double colon and requirement specification after that.
 
@@ -115,7 +116,7 @@ class DuplicateDependencyInMultipleChannelsError(Exception):
     ...
 
 
-def append_requirement_list(req_list: List[requirements.Requirement], p_req: requirements.Requirement) -> None:
+def append_requirement_list(req_list: list[requirements.Requirement], p_req: requirements.Requirement) -> None:
     """Append a requirement to an existing requirement list. If need and able to merge, merge it, otherwise, append it.
 
     Args:
@@ -134,7 +135,7 @@ def append_requirement_list(req_list: List[requirements.Requirement], p_req: req
 
 
 def append_conda_dependency(
-    conda_chan_deps: DefaultDict[str, List[requirements.Requirement]], p_chan_dep: Tuple[str, requirements.Requirement]
+    conda_chan_deps: DefaultDict[str, list[requirements.Requirement]], p_chan_dep: tuple[str, requirements.Requirement]
 ) -> None:
     """Append a conda dependency to an existing conda dependencies dict, if not existed in any channel.
         To avoid making unnecessary modification to dict, we check the existence first, then try to merge, then append,
@@ -164,45 +165,73 @@ def append_conda_dependency(
     conda_chan_deps[p_channel].append(p_req)
 
 
-def validate_pip_requirement_string_list(req_str_list: List[str]) -> List[requirements.Requirement]:
-    """Validate the a list of pip requirement string according to PEP 508.
+def validate_pip_requirement_string_list(
+    req_str_list: list[str], add_local_version_specifier: bool = False
+) -> list[requirements.Requirement]:
+    """Validate the list of pip requirement strings according to PEP 508.
 
     Args:
-        req_str_list: The list of string contains the pip requirement specification.
+        req_str_list: The list of strings containing the pip requirement specification.
+        add_local_version_specifier: if True, add the version specifier of the locally installed package version to
+            requirements without version specifiers.
 
     Returns:
         A requirements.Requirement list containing the requirement information.
     """
-    seen_pip_requirement_list: List[requirements.Requirement] = []
+    seen_pip_requirement_list: list[requirements.Requirement] = []
     for req_str in req_str_list:
         append_requirement_list(seen_pip_requirement_list, _validate_pip_requirement_string(req_str=req_str))
 
+    if add_local_version_specifier:
+        # For any requirement string that does not contain a specifier, add the specifier of a locally installed version
+        # if it exists.
+        seen_pip_requirement_list = list(
+            map(
+                lambda req: req if req.specifier else get_local_installed_version_of_pip_package(req),
+                seen_pip_requirement_list,
+            )
+        )
+
     return seen_pip_requirement_list
 
 
-def validate_conda_dependency_string_list(dep_str_list: List[str]) -> DefaultDict[str, List[requirements.Requirement]]:
+def validate_conda_dependency_string_list(
+    dep_str_list: list[str], add_local_version_specifier: bool = False
+) -> DefaultDict[str, list[requirements.Requirement]]:
     """Validate a list of conda dependency string, find any duplicate package across different channel and create a dict
         to represent the whole dependencies.
 
     Args:
         dep_str_list: The list of string contains the conda dependency specification.
+        add_local_version_specifier: if True, add the version specifier of the locally installed package version to
+            requirements without version specifiers.
 
     Returns:
         A dict mapping from the channel name to the list of requirements from that channel.
     """
     validated_conda_dependency_list = list(map(_validate_conda_dependency_string, dep_str_list))
-    ret_conda_dependency_dict: DefaultDict[str, List[requirements.Requirement]] = collections.defaultdict(list)
+    ret_conda_dependency_dict: DefaultDict[str, list[requirements.Requirement]] = collections.defaultdict(list)
     for p_channel, p_req in validated_conda_dependency_list:
         append_conda_dependency(ret_conda_dependency_dict, (p_channel, p_req))
 
+    if add_local_version_specifier:
+        # For any conda dependency string that does not contain a specifier, add the specifier of a locally installed
+        # version if it exists. This is best-effort: if the conda package does not have the same name as the pip
+        # package, it won't be found in the local environment.
+        for channel_str, reqs in ret_conda_dependency_dict.items():
+            reqs = list(
+                map(lambda req: req if req.specifier else get_local_installed_version_of_pip_package(req), reqs)
+            )
+            ret_conda_dependency_dict[channel_str] = reqs
+
     return ret_conda_dependency_dict
 
 
 def get_local_installed_version_of_pip_package(pip_req: requirements.Requirement) -> requirements.Requirement:
     """Get the local installed version of a given pip package requirement.
-        If the package is locally installed, and the local version meet the specifier of the requirements, return a new
+        If the package is locally installed, and the local version meets the specifier of the requirements, return a new
         requirement specifier that pins the version.
-        If the local version does not meet the specifier of the requirements, a warn will be omitted and returns
+        If the local version does not meet the specifier of the requirements, a warning will be emitted and returns
         the original package requirement.
         If the package is not locally installed or not found, the original package requirement is returned.
 
@@ -217,7 +246,7 @@ def get_local_installed_version_of_pip_package(pip_req: requirements.Requirement
         local_dist_version = local_dist.version
     except importlib_metadata.PackageNotFoundError:
         if pip_req.name == SNOWPARK_ML_PKG_NAME:
-            local_dist_version = snowml_env.VERSION
+            local_dist_version = snowml_version.VERSION
         else:
             return pip_req
     new_pip_req = copy.deepcopy(pip_req)
@@ -372,8 +401,8 @@ def relax_requirement_version(req: requirements.Requirement) -> requirements.Req
 
 
 def get_matched_package_versions_in_information_schema_with_active_session(
-    reqs: List[requirements.Requirement], python_version: str
-) -> Dict[str, List[version.Version]]:
+    reqs: list[requirements.Requirement], python_version: str
+) -> dict[str, list[version.Version]]:
     try:
         session = context.get_active_session()
     except exceptions.SnowparkSessionException:
@@ -383,10 +412,10 @@ def get_matched_package_versions_in_information_schema_with_active_session(
 
 def get_matched_package_versions_in_information_schema(
     session: session.Session,
-    reqs: List[requirements.Requirement],
+    reqs: list[requirements.Requirement],
     python_version: str,
-    statement_params: Optional[Dict[str, Any]] = None,
-) -> Dict[str, List[version.Version]]:
+    statement_params: Optional[dict[str, Any]] = None,
+) -> dict[str, list[version.Version]]:
     """Look up the information_schema table to check if a package with the specified specifier exists in the Snowflake
     Conda channel. Note that this is not the source of truth due to the potential delay caused by a package that might
     exist in the information_schema table but has not yet become available in the Snowflake Conda channel.
@@ -400,8 +429,8 @@ def get_matched_package_versions_in_information_schema(
     Returns:
         A Dict, whose key is the package name, and value is a list of versions match the requirements.
     """
-    ret_dict: Dict[str, List[version.Version]] = {}
-    reqs_to_request: List[requirements.Requirement] = []
+    ret_dict: dict[str, list[version.Version]] = {}
+    reqs_to_request: list[requirements.Requirement] = []
     for req in reqs:
         if req.name in _SNOWFLAKE_INFO_SCHEMA_PACKAGE_CACHE:
             available_versions = list(
@@ -457,7 +486,7 @@ def get_matched_package_versions_in_information_schema(
 
 def save_conda_env_file(
     path: pathlib.Path,
-    conda_chan_deps: DefaultDict[str, List[requirements.Requirement]],
+    conda_chan_deps: DefaultDict[str, list[requirements.Requirement]],
     python_version: str,
     cuda_version: Optional[str] = None,
     default_channel_override: str = SNOWFLAKE_CONDA_CHANNEL_URL,
@@ -478,7 +507,7 @@ def save_conda_env_file(
     """
     assert path.suffix in [".yml", ".yaml"], "Conda environment file should have extension of yml or yaml."
     path.parent.mkdir(parents=True, exist_ok=True)
-    env: Dict[str, Any] = dict()
+    env: dict[str, Any] = dict()
     env["name"] = "snow-env"
     # Get all channels in the dependencies, ordered by the number of the packages which belongs to and put into
     # channels section.
@@ -505,7 +534,7 @@ def save_conda_env_file(
         yaml.safe_dump(env, stream=f, default_flow_style=False)
 
 
-def save_requirements_file(path: pathlib.Path, pip_deps: List[requirements.Requirement]) -> None:
+def save_requirements_file(path: pathlib.Path, pip_deps: list[requirements.Requirement]) -> None:
     """Generate Python requirements.txt file in the given directory path.
 
     Args:
@@ -521,9 +550,9 @@ def save_requirements_file(path: pathlib.Path, pip_deps: List[requirements.Requi
 
 def load_conda_env_file(
     path: pathlib.Path,
-) -> Tuple[
-    DefaultDict[str, List[requirements.Requirement]],
-    Optional[List[requirements.Requirement]],
+) -> tuple[
+    DefaultDict[str, list[requirements.Requirement]],
+    Optional[list[requirements.Requirement]],
     Optional[str],
     Optional[str],
 ]:
@@ -601,7 +630,7 @@ def load_conda_env_file(
     return conda_dep_dict, pip_deps_list if pip_deps_list else None, python_version, cuda_version
 
 
-def load_requirements_file(path: pathlib.Path) -> List[requirements.Requirement]:
+def load_requirements_file(path: pathlib.Path) -> list[requirements.Requirement]:
     """Load Python requirements.txt file from the given directory path.
 
     Args:
@@ -641,8 +670,8 @@ def parse_python_version_string(dep: str) -> Optional[str]:
 
 
 def _find_conda_dep_spec(
-    conda_chan_deps: DefaultDict[str, List[requirements.Requirement]], pkg_name: str
-) -> Optional[Tuple[str, requirements.Requirement]]:
+    conda_chan_deps: DefaultDict[str, list[requirements.Requirement]], pkg_name: str
+) -> Optional[tuple[str, requirements.Requirement]]:
     for channel in conda_chan_deps:
         spec = next(filter(lambda req: req.name == pkg_name, conda_chan_deps[channel]), None)
         if spec:
@@ -650,14 +679,14 @@ def _find_conda_dep_spec(
     return None
 
 
-def _find_pip_req_spec(pip_reqs: List[requirements.Requirement], pkg_name: str) -> Optional[requirements.Requirement]:
+def _find_pip_req_spec(pip_reqs: list[requirements.Requirement], pkg_name: str) -> Optional[requirements.Requirement]:
     spec = next(filter(lambda req: req.name == pkg_name, pip_reqs), None)
     return spec
 
 
 def find_dep_spec(
-    conda_chan_deps: DefaultDict[str, List[requirements.Requirement]],
-    pip_reqs: List[requirements.Requirement],
+    conda_chan_deps: DefaultDict[str, list[requirements.Requirement]],
+    pip_reqs: list[requirements.Requirement],
     conda_pkg_name: str,
     pip_pkg_name: Optional[str] = None,
     remove_spec: bool = False,

snowflake/ml/_internal/file_utils.py

@@ -11,18 +11,7 @@ import sys
 import tarfile
 import tempfile
 import zipfile
-from typing import (
-    Any,
-    Callable,
-    Dict,
-    Generator,
-    List,
-    Literal,
-    Optional,
-    Set,
-    Tuple,
-    Union,
-)
+from typing import Any, Callable, Generator, Literal, Optional, Union
 from urllib import parse
 
 import cloudpickle
@@ -37,7 +26,7 @@ GENERATED_PY_FILE_EXT = (".pyc", ".pyo", ".pyd", ".pyi")
 def copytree(
     src: "Union[str, os.PathLike[str]]",
     dst: "Union[str, os.PathLike[str]]",
-    ignore: Optional[Callable[..., Set[str]]] = None,
+    ignore: Optional[Callable[..., set[str]]] = None,
     dirs_exist_ok: bool = False,
 ) -> "Union[str, os.PathLike[str]]":
     """This is a forked version of shutil.copytree that remove all copystat, to make sure it works in Sproc.
@@ -170,7 +159,7 @@ def zip_python_package(zipfile_path: str, package_name: str, ignore_generated_py
 
 
 def hash_directory(
-    directory: Union[str, pathlib.Path], *, ignore_hidden: bool = False, excluded_files: Optional[List[str]] = None
+    directory: Union[str, pathlib.Path], *, ignore_hidden: bool = False, excluded_files: Optional[list[str]] = None
 ) -> str:
     """Hash the **content** of a folder recursively using SHA-1.
 
@@ -186,7 +175,7 @@ def hash_directory(
         excluded_files = []
 
     def _update_hash_from_dir(
-        directory: Union[str, pathlib.Path], hash: "hashlib._Hash", *, ignore_hidden: bool, excluded_files: List[str]
+        directory: Union[str, pathlib.Path], hash: "hashlib._Hash", *, ignore_hidden: bool, excluded_files: list[str]
     ) -> "hashlib._Hash":
         assert pathlib.Path(directory).is_dir(), "Provided path is not a directory."
         for path in sorted(pathlib.Path(directory).iterdir(), key=lambda p: str(p).lower()):
@@ -208,7 +197,7 @@ def hash_directory(
     ).hexdigest()
 
 
-def get_all_modules(dirname: str, prefix: str = "") -> List[str]:
+def get_all_modules(dirname: str, prefix: str = "") -> list[str]:
     modules = [mod.name for mod in pkgutil.iter_modules([dirname], prefix=prefix)]
     subdirs = [f.path for f in os.scandir(dirname) if f.is_dir()]
     for sub_dirname in subdirs:
@@ -248,7 +237,7 @@ def _create_tar_gz_stream(source_dir: str, arcname: Optional[str] = None) -> Gen
         yield output_stream
 
 
-def get_package_path(package_name: str, strategy: Literal["first", "last"] = "first") -> Tuple[str, str]:
+def get_package_path(package_name: str, strategy: Literal["first", "last"] = "first") -> tuple[str, str]:
     """[Obsolete]Return the path to where a package is defined and its start location.
     Example 1: snowflake.ml -> path/to/site-packages/snowflake/ml, path/to/site-packages
     Example 2: zip_imported_module -> path/to/some/zipfile.zip/zip_imported_module, path/to/some/zipfile.zip
@@ -267,7 +256,7 @@ def get_package_path(package_name: str, strategy: Literal["first", "last"] = "fi
     return pkg_path, pkg_start_path
 
 
-def stage_object(session: snowpark.Session, object: object, stage_location: str) -> List[snowpark.PutResult]:
+def stage_object(session: snowpark.Session, object: object, stage_location: str) -> list[snowpark.PutResult]:
     temp_file = tempfile.NamedTemporaryFile(delete=False)
     temp_file_path = temp_file.name
     temp_file.close()
@@ -279,7 +268,7 @@ def stage_object(session: snowpark.Session, object: object, stage_location: str)
 
 
 def stage_file_exists(
-    session: snowpark.Session, stage_location: str, file_name: str, statement_params: Dict[str, Any]
+    session: snowpark.Session, stage_location: str, file_name: str, statement_params: dict[str, Any]
 ) -> bool:
     try:
         res = session.sql(f"list {stage_location}/{file_name}").collect(statement_params=statement_params)
@@ -297,7 +286,7 @@ def upload_directory_to_stage(
     local_path: pathlib.Path,
     stage_path: Union[pathlib.PurePosixPath, parse.ParseResult],
     *,
-    statement_params: Optional[Dict[str, Any]] = None,
+    statement_params: Optional[dict[str, Any]] = None,
 ) -> None:
     """Upload a local folder recursively to a stage and keep the structure.
 
@@ -350,7 +339,7 @@ def download_directory_from_stage(
     stage_path: pathlib.PurePosixPath,
     local_path: pathlib.Path,
     *,
-    statement_params: Optional[Dict[str, Any]] = None,
+    statement_params: Optional[dict[str, Any]] = None,
 ) -> None:
     """Upload a folder in stage recursively to a folder in local and keep the structure.
 

snowflake/ml/_internal/human_readable_id/hrid_generator_base.py

@@ -15,7 +15,6 @@ In this module you will find:
 
 import math
 from abc import ABC, abstractmethod
-from typing import Dict, List, Tuple
 
 
 class HRIDBase(ABC):
@@ -28,12 +27,11 @@ class HRIDBase(ABC):
     @abstractmethod
     def __id_generator__(self) -> int:
         """The generator to use to generate new IDs. The implementer needs to provide this."""
-        pass
 
-    __hrid_structure__: Tuple[str, ...]
+    __hrid_structure__: tuple[str, ...]
     """The HRID structure to be generated. The implementer needs to provide this."""
 
-    __hrid_words__: Dict[str, Tuple[str, ...]]
+    __hrid_words__: dict[str, tuple[str, ...]]
     """The mapping between the HRID parts and the words to use. The implementer needs to provide this."""
 
     __separator__ = "_"
@@ -82,7 +80,7 @@ class HRIDBase(ABC):
             hrid.append(str(values[idxs[i]]))
         return self.__separator__.join(hrid)
 
-    def generate(self) -> Tuple[int, str]:
+    def generate(self) -> tuple[int, str]:
         """Generate an ID and the corresponding HRID.
 
         Returns:
@@ -92,7 +90,7 @@ class HRIDBase(ABC):
         hrid = self.id_to_hrid(id)
         return (id, hrid)
 
-    def _id_to_idxs(self, id: int) -> List[int]:
+    def _id_to_idxs(self, id: int) -> list[int]:
         """Take the ID and convert it to indices into the HRID words.
 
         Args:
@@ -109,7 +107,7 @@ class HRIDBase(ABC):
             idxs.append((id & mask) >> shift)
         return idxs
 
-    def _hrid_to_idxs(self, hrid: str) -> List[int]:
+    def _hrid_to_idxs(self, hrid: str) -> list[int]:
         """Take the HRID and convert it to indices into the HRID words.
 
         Args:

snowflake/ml/_internal/init_utils.py

@@ -2,10 +2,9 @@ import importlib
 import inspect
 import pkgutil
 from types import FunctionType
-from typing import Dict
 
 
-def fetch_classes_from_modules_in_pkg_dir(pkg_dir: str, pkg_name: str) -> Dict[str, type]:
+def fetch_classes_from_modules_in_pkg_dir(pkg_dir: str, pkg_name: str) -> dict[str, type]:
     """Finds classes defined all the python modules in the given package directory.
 
     Args:
@@ -36,7 +35,7 @@ def fetch_classes_from_modules_in_pkg_dir(pkg_dir: str, pkg_name: str) -> Dict[s
     return exportable_classes
 
 
-def fetch_functions_from_modules_in_pkg_dir(pkg_dir: str, pkg_name: str) -> Dict[str, FunctionType]:
+def fetch_functions_from_modules_in_pkg_dir(pkg_dir: str, pkg_name: str) -> dict[str, FunctionType]:
     """Finds functions defined all the python modules in the given package directory.
 
     Args:

snowflake/ml/_internal/lineage/lineage_utils.py

@@ -1,6 +1,6 @@
 import copy
 import functools
-from typing import Any, Callable, List, Optional, get_args
+from typing import Any, Callable, Optional, get_args
 
 from snowflake import snowpark
 from snowflake.ml.data import data_source
@@ -9,7 +9,7 @@ _DATA_SOURCES_ATTR = "_data_sources"
 
 
 def _wrap_func(
-    fn: Callable[..., snowpark.DataFrame], data_sources: List[data_source.DataSource]
+    fn: Callable[..., snowpark.DataFrame], data_sources: list[data_source.DataSource]
 ) -> Callable[..., snowpark.DataFrame]:
     """Wrap a DataFrame transform function to propagate data_sources to derived DataFrames."""
 
@@ -34,9 +34,9 @@ def _wrap_class_func(fn: Callable[..., snowpark.DataFrame]) -> Callable[..., sno
     return wrapped
 
 
-def get_data_sources(*args: Any) -> Optional[List[data_source.DataSource]]:
+def get_data_sources(*args: Any) -> Optional[list[data_source.DataSource]]:
     """Helper method for extracting data sources attribute from DataFrames in an argument list"""
-    result: Optional[List[data_source.DataSource]] = None
+    result: Optional[list[data_source.DataSource]] = None
     for arg in args:
         srcs = getattr(arg, _DATA_SOURCES_ATTR, None)
         if isinstance(srcs, list) and all(isinstance(s, get_args(data_source.DataSource)) for s in srcs):
@@ -46,7 +46,7 @@ def get_data_sources(*args: Any) -> Optional[List[data_source.DataSource]]:
     return result
 
 
-def set_data_sources(obj: Any, data_sources: Optional[List[data_source.DataSource]]) -> None:
+def set_data_sources(obj: Any, data_sources: Optional[list[data_source.DataSource]]) -> None:
     """Helper method for attaching data sources to an object"""
     if data_sources:
         assert all(isinstance(ds, get_args(data_source.DataSource)) for ds in data_sources)
@@ -54,7 +54,7 @@ def set_data_sources(obj: Any, data_sources: Optional[List[data_source.DataSourc
 
 
 def patch_dataframe(
-    df: snowpark.DataFrame, data_sources: List[data_source.DataSource], inplace: bool = False
+    df: snowpark.DataFrame, data_sources: list[data_source.DataSource], inplace: bool = False
 ) -> snowpark.DataFrame:
     """
     Monkey patch a DataFrame to add attach the provided data_sources as an attribute of the DataFrame.

snowflake/ml/_internal/platform_capabilities.py

@@ -1,5 +1,6 @@
 import json
-from typing import Any, Dict, Optional
+from contextlib import contextmanager
+from typing import Any, Optional
 
 from absl import logging
 
@@ -27,21 +28,50 @@ class PlatformCapabilities:
     """
 
     _instance: Optional["PlatformCapabilities"] = None
+    # Used for unittesting only. This is to avoid the need to mock the session object or reaching out to Snowflake
+    _mock_features: Optional[dict[str, Any]] = None
 
     @classmethod
     def get_instance(cls, session: Optional[snowpark_session.Session] = None) -> "PlatformCapabilities":
+        # Used for unittesting only. In this situation, _instance is not initialized.
+        if cls._mock_features is not None:
+            return cls(features=cls._mock_features)
         if not cls._instance:
-            cls._instance = cls(session)
+            cls._instance = cls(session=session)
         return cls._instance
 
+    @classmethod
+    def set_mock_features(cls, features: Optional[dict[str, Any]] = None) -> None:
+        cls._mock_features = features
+
+    @classmethod
+    def clear_mock_features(cls) -> None:
+        cls._mock_features = None
+
+    # For contextmanager, we need to have return type Iterator[Never]. However, Never type is introduced only in
+    # Python 3.11. So, we are ignoring the type for this method.
+    @classmethod  # type: ignore[arg-type]
+    @contextmanager
+    def mock_features(cls, features: dict[str, Any]) -> None:  # type: ignore[misc]
+        logging.debug(f"Setting mock features: {features}")
+        cls.set_mock_features(features)
+        try:
+            yield
+        finally:
+            logging.debug(f"Clearing mock features: {features}")
+            cls.clear_mock_features()
+
     def is_nested_function_enabled(self) -> bool:
         return self._get_bool_feature("SPCS_MODEL_ENABLE_EMBEDDED_SERVICE_FUNCTIONS", False)
 
+    def is_inlined_deployment_spec_enabled(self) -> bool:
+        return self._get_bool_feature("ENABLE_INLINE_DEPLOYMENT_SPEC", False)
+
     def is_live_commit_enabled(self) -> bool:
         return self._get_bool_feature("ENABLE_BUNDLE_MODULE_CHECKOUT", False)
 
     @staticmethod
-    def _get_features(session: snowpark_session.Session) -> Dict[str, Any]:
+    def _get_features(session: snowpark_session.Session) -> dict[str, Any]:
         try:
             result = (
                 query_result_checker.SqlResultValidator(
@@ -68,11 +98,17 @@ class PlatformCapabilities:
             # This can happen is server side is older than 9.2. That is fine.
         return {}
 
-    def __init__(self, session: Optional[snowpark_session.Session] = None) -> None:
+    def __init__(
+        self, *, session: Optional[snowpark_session.Session] = None, features: Optional[dict[str, Any]] = None
+    ) -> None:
+        # This is for testing purposes only.
+        if features:
+            self.features = features
+            return
         if not session:
             session = next(iter(snowpark_session._get_active_sessions()))
             assert session, "Missing active session object"
-        self.features: Dict[str, Any] = PlatformCapabilities._get_features(session)
+        self.features = PlatformCapabilities._get_features(session)
 
     def _get_bool_feature(self, feature_name: str, default_value: bool) -> bool:
         value = self.features.get(feature_name, default_value)