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

snowflake/ml/feature_store/feature_view.py

@@ -126,9 +126,11 @@ class FeatureView(lineage_node.LineageNode):
         name: str,
         entities: List[Entity],
         feature_df: DataFrame,
+        *,
         timestamp_col: Optional[str] = None,
         refresh_freq: Optional[str] = None,
         desc: str = "",
+        warehouse: Optional[str] = None,
         **_kwargs: Any,
     ) -> None:
         """
@@ -149,7 +151,33 @@ class FeatureView(lineage_node.LineageNode):
                 NOTE: If refresh_freq is not provided, then FeatureView will be registered as View on Snowflake backend
                     and there won't be extra storage cost.
             desc: description of the FeatureView.
+            warehouse: warehouse to refresh feature view. Not needed for static feature view (refresh_freq is None).
+                For managed feature view, this warehouse will overwrite the default warehouse of Feature Store if it is
+                specified, otherwise the default warehouse will be used.
             _kwargs: reserved kwargs for system generated args. NOTE: DO NOT USE.
+
+        Example::
+
+            >>> fs = FeatureStore(...)
+            >>> # draft_fv is a local object that hasn't materiaized to Snowflake backend yet.
+            >>> feature_df = session.sql("select f_1, f_2 from source_table")
+            >>> draft_fv = FeatureView(
+            ...     name="my_fv",
+            ...     entities=[e1, e2],
+            ...     feature_df=feature_df,
+            ...     timestamp_col='TS', # optional
+            ...     refresh_freq='1d',  # optional
+            ...     desc='A line about this feature view',  # optional
+            ...     warehouse='WH'      # optional, the warehouse used to refresh (managed) feature view
+            ... )
+            >>> print(draft_fv.status)
+            FeatureViewStatus.DRAFT
+            <BLANKLINE>
+            >>> # registered_fv is a local object that maps to a Snowflake backend object.
+            >>> registered_fv = fs.register_feature_view(draft_fv, "v1")
+            >>> print(registered_fv.status)
+            FeatureViewStatus.ACTIVE
+
         """
 
         self._name: SqlIdentifier = SqlIdentifier(name)
@@ -167,7 +195,7 @@ class FeatureView(lineage_node.LineageNode):
         self._refresh_freq: Optional[str] = refresh_freq
         self._database: Optional[SqlIdentifier] = None
         self._schema: Optional[SqlIdentifier] = None
-        self._warehouse: Optional[SqlIdentifier] = None
+        self._warehouse: Optional[SqlIdentifier] = SqlIdentifier(warehouse) if warehouse is not None else None
         self._refresh_mode: Optional[str] = None
         self._refresh_mode_reason: Optional[str] = None
         self._owner: Optional[str] = None
@@ -185,6 +213,33 @@ class FeatureView(lineage_node.LineageNode):
 
         Raises:
             ValueError: if selected feature names is not found in the FeatureView.
+
+        Example::
+
+            >>> fs = FeatureStore(...)
+            >>> e = fs.get_entity('TRIP_ID')
+            >>> # feature_df contains 3 features and 1 entity
+            >>> feature_df = session.table(source_table).select(
+            ...     'TRIPDURATION',
+            ...     'START_STATION_LATITUDE',
+            ...     'END_STATION_LONGITUDE',
+            ...     'TRIP_ID'
+            ... )
+            >>> darft_fv = FeatureView(name='F_TRIP', entities=[e], feature_df=feature_df)
+            >>> fv = fs.register_feature_view(darft_fv, version='1.0')
+            >>> # shows all 3 features
+            >>> fv.feature_names
+            ['TRIPDURATION', 'START_STATION_LATITUDE', 'END_STATION_LONGITUDE']
+            <BLANKLINE>
+            >>> # slice a subset of features
+            >>> fv_slice = fv.slice(['TRIPDURATION', 'START_STATION_LATITUDE'])
+            >>> fv_slice.names
+            ['TRIPDURATION', 'START_STATION_LATITUDE']
+            <BLANKLINE>
+            >>> # query the full set of features in original feature view
+            >>> fv_slice.feature_view_ref.feature_names
+            ['TRIPDURATION', 'START_STATION_LATITUDE', 'END_STATION_LONGITUDE']
+
         """
 
         res = []
@@ -196,14 +251,30 @@ class FeatureView(lineage_node.LineageNode):
         return FeatureViewSlice(self, res)
 
     def fully_qualified_name(self) -> str:
-        """Returns the fully qualified name (<database_name>.<schema_name>.<feature_view_name>) for the
-            FeatureView in Snowflake.
+        """
+        Returns the fully qualified name (<database_name>.<schema_name>.<feature_view_name>) for the
+        FeatureView in Snowflake.
 
         Returns:
             fully qualified name string.
 
         Raises:
             RuntimeError: if the FeatureView is not registered.
+
+        Example::
+
+            >>> fs = FeatureStore(...)
+            >>> e = fs.get_entity('TRIP_ID')
+            >>> feature_df = session.table(source_table).select(
+            ...     'TRIPDURATION',
+            ...     'START_STATION_LATITUDE',
+            ...     'TRIP_ID'
+            ... )
+            >>> darft_fv = FeatureView(name='F_TRIP', entities=[e], feature_df=feature_df)
+            >>> registered_fv = fs.register_feature_view(darft_fv, version='1.0')
+            >>> registered_fv.fully_qualified_name()
+            'MY_DB.MY_SCHEMA."F_TRIP$1.0"'
+
         """
         if self.status == FeatureViewStatus.DRAFT or self.version is None:
             raise RuntimeError(f"FeatureView {self.name} has not been registered.")
@@ -221,6 +292,22 @@ class FeatureView(lineage_node.LineageNode):
 
         Raises:
             ValueError: if feature name is not found in the FeatureView.
+
+        Example::
+
+            >>> fs = FeatureStore(...)
+            >>> e = fs.get_entity('TRIP_ID')
+            >>> feature_df = session.table(source_table).select('TRIPDURATION', 'START_STATION_LATITUDE', 'TRIP_ID')
+            >>> draft_fv = FeatureView(name='F_TRIP', entities=[e], feature_df=feature_df)
+            >>> draft_fv = draft_fv.attach_feature_desc({
+            ...     "TRIPDURATION": "Duration of a trip.",
+            ...     "START_STATION_LATITUDE": "Latitude of the start station."
+            ... })
+            >>> registered_fv = fs.register_feature_view(draft_fv, version='1.0')
+            >>> registered_fv.feature_descs
+            OrderedDict([('TRIPDURATION', 'Duration of a trip.'),
+                ('START_STATION_LATITUDE', 'Latitude of the start station.')])
+
         """
         for f, d in descs.items():
             f = SqlIdentifier(f)
@@ -254,6 +341,31 @@ class FeatureView(lineage_node.LineageNode):
 
     @desc.setter
     def desc(self, new_value: str) -> None:
+        """Set the description of feature view.
+
+        Args:
+            new_value: new value of description.
+
+        Example::
+
+            >>> fs = FeatureStore(...)
+            >>> e = fs.get_entity('TRIP_ID')
+            >>> darft_fv = FeatureView(
+            ...     name='F_TRIP',
+            ...     entities=[e],
+            ...     feature_df=feature_df,
+            ...     desc='old desc'
+            ... )
+            >>> fv_1 = fs.register_feature_view(darft_fv, version='1.0')
+            >>> print(fv_1.desc)
+            old desc
+            <BLANKLINE>
+            >>> darft_fv.desc = 'NEW DESC'
+            >>> fv_2 = fs.register_feature_view(darft_fv, version='2.0')
+            >>> print(fv_2.desc)
+            NEW DESC
+
+        """
         warnings.warn(
             "You must call register_feature_view() to make it effective. "
             "Or use update_feature_view(desc=<new_value>).",
@@ -288,6 +400,31 @@ class FeatureView(lineage_node.LineageNode):
 
     @refresh_freq.setter
     def refresh_freq(self, new_value: str) -> None:
+        """Set refresh frequency of feature view.
+
+        Args:
+            new_value: The new value of refresh frequency.
+
+        Example::
+
+            >>> fs = FeatureStore(...)
+            >>> e = fs.get_entity('TRIP_ID')
+            >>> darft_fv = FeatureView(
+            ...     name='F_TRIP',
+            ...     entities=[e],
+            ...     feature_df=feature_df,
+            ...     refresh_freq='1d'
+            ... )
+            >>> fv_1 = fs.register_feature_view(darft_fv, version='1.0')
+            >>> print(fv_1.refresh_freq)
+            1 day
+            <BLANKLINE>
+            >>> darft_fv.refresh_freq = '12h'
+            >>> fv_2 = fs.register_feature_view(darft_fv, version='2.0')
+            >>> print(fv_2.refresh_freq)
+            12 hours
+
+        """
         warnings.warn(
             "You must call register_feature_view() to make it effective. "
             "Or use update_feature_view(refresh_freq=<new_value>).",
@@ -310,6 +447,32 @@ class FeatureView(lineage_node.LineageNode):
 
     @warehouse.setter
     def warehouse(self, new_value: str) -> None:
+        """Set warehouse of feature view.
+
+        Args:
+            new_value: The new value of warehouse.
+
+        Example::
+
+            >>> fs = FeatureStore(...)
+            >>> e = fs.get_entity('TRIP_ID')
+            >>> darft_fv = FeatureView(
+            ...     name='F_TRIP',
+            ...     entities=[e],
+            ...     feature_df=feature_df,
+            ...     refresh_freq='1d',
+            ...     warehouse='WH1',
+            ... )
+            >>> fv_1 = fs.register_feature_view(darft_fv, version='1.0')
+            >>> print(fv_1.warehouse)
+            WH1
+            <BLANKLINE>
+            >>> darft_fv.warehouse = 'WH2'
+            >>> fv_2 = fs.register_feature_view(darft_fv, version='2.0')
+            >>> print(fv_2.warehouse)
+            WH2
+
+        """
         warnings.warn(
             "You must call register_feature_view() to make it effective. "
             "Or use update_feature_view(warehouse=<new_value>).",
@@ -456,7 +619,7 @@ Got {len(self._feature_df.queries['queries'])}: {self._feature_df.queries['queri
 
         entities = []
         for e_json in json_dict["_entities"]:
-            e = Entity(e_json["name"], e_json["join_keys"], e_json["desc"])
+            e = Entity(e_json["name"], e_json["join_keys"], desc=e_json["desc"])
             e.owner = e_json["owner"]
             entities.append(e)
 
@@ -504,7 +667,7 @@ Got {len(self._feature_df.queries['queries'])}: {self._feature_df.queries['queri
                 original_exception=ValueError("No active warehouse selected in the current session"),
             )
 
-        fs = feature_store.FeatureStore(session, db_name, feature_store_name, session_warehouse)
+        fs = feature_store.FeatureStore(session, db_name, feature_store_name, default_warehouse=session_warehouse)
         return fs.get_feature_view(feature_view_name, version)  # type: ignore[no-any-return]
 
     @staticmethod

snowflake/ml/fileset/stage_fs.py

@@ -234,21 +234,29 @@ class SFStageFileSystem(fsspec.AbstractFileSystem):
 
         Raises:
             SnowflakeMLException: An error occurred when the given path points to a file that cannot be found.
+            snowpark_exceptions.SnowparkClientException: File access failed with a Snowpark exception
         """
         path = path.lstrip("/")
         if self._USE_FALLBACK_FILE_ACCESS:
             return self._open_with_snowpark(path)
         cached_presigned_url = self._url_cache.get(path, None)
-        if not cached_presigned_url:
-            res = self._fetch_presigned_urls([path])
-            url = res[0][1]
-            expire_at = time.time() + _PRESIGNED_URL_LIFETIME_SEC
-            cached_presigned_url = _PresignedUrl(url, expire_at)
-            self._url_cache[path] = cached_presigned_url
-            logging.debug(f"Retrieved presigned url for {path}.")
-        elif cached_presigned_url.is_expiring():
-            self.optimize_read()
-            cached_presigned_url = self._url_cache[path]
+        try:
+            if not cached_presigned_url:
+                res = self._fetch_presigned_urls([path])
+                url = res[0][1]
+                expire_at = time.time() + _PRESIGNED_URL_LIFETIME_SEC
+                cached_presigned_url = _PresignedUrl(url, expire_at)
+                self._url_cache[path] = cached_presigned_url
+                logging.debug(f"Retrieved presigned url for {path}.")
+            elif cached_presigned_url.is_expiring():
+                self.optimize_read()
+                cached_presigned_url = self._url_cache[path]
+        except snowpark_exceptions.SnowparkClientException as e:
+            if self._USE_FALLBACK_FILE_ACCESS == False:  # noqa: E712 # Fallback disabled
+                raise
+            # This may be an intermittent failure, so don't set _USE_FALLBACK_FILE_ACCESS = True
+            logging.warning(f"Pre-signed URL generation failed with {e.message}, trying fallback file access")
+            return self._open_with_snowpark(path)
         url = cached_presigned_url.url
         try:
             return self._fs._open(url, mode=mode, **kwargs)

snowflake/ml/lineage/lineage_node.py

@@ -118,7 +118,7 @@ class LineageNode:
             )
             domain = lineage_object["domain"].lower()
             if domain_filter is None or domain in domain_filter:
-                if domain in DOMAIN_LINEAGE_REGISTRY:
+                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")

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):

snowflake/ml/model/_model_composer/model_composer.py

@@ -11,7 +11,8 @@ 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
@@ -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=model_metadata,
-            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,
@@ -186,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

snowflake/ml/model/_model_composer/model_manifest/model_manifest.py

@@ -1,11 +1,12 @@
 import collections
 import copy
 import pathlib
+import warnings
 from typing import List, Optional, cast
 
 import yaml
 
-from snowflake.ml._internal.lineage import data_source
+from snowflake.ml.data import data_source
 from snowflake.ml.model import type_hints
 from snowflake.ml.model._model_composer.model_manifest import model_manifest_schema
 from snowflake.ml.model._model_composer.model_method import (
@@ -16,7 +17,6 @@ from snowflake.ml.model._packager.model_meta import (
     model_meta as model_meta_api,
     model_meta_schema,
 )
-from snowflake.snowpark import Session
 
 
 class ModelManifest:
@@ -36,9 +36,8 @@ class ModelManifest:
 
     def save(
         self,
-        session: Session,
         model_meta: model_meta_api.ModelMetadata,
-        model_file_rel_path: pathlib.PurePosixPath,
+        model_rel_path: pathlib.PurePosixPath,
         options: Optional[type_hints.ModelSaveOption] = None,
         data_sources: Optional[List[data_source.DataSource]] = None,
     ) -> None:
@@ -47,10 +46,10 @@ class ModelManifest:
 
         runtime_to_use = copy.deepcopy(model_meta.runtimes["cpu"])
         runtime_to_use.name = self._DEFAULT_RUNTIME_NAME
-        runtime_to_use.imports.append(model_file_rel_path)
+        runtime_to_use.imports.append(str(model_rel_path) + "/")
         runtime_dict = runtime_to_use.save(self.workspace_path)
 
-        self.function_generator = function_generator.FunctionGenerator(model_file_rel_path=model_file_rel_path)
+        self.function_generator = function_generator.FunctionGenerator(model_dir_rel_path=model_rel_path)
         self.methods: List[model_method.ModelMethod] = []
         for target_method in model_meta.signatures.keys():
             method = model_method.ModelMethod(
@@ -75,6 +74,16 @@ class ModelManifest:
                 "In this case, set case_sensitive as True for those methods to distinguish them."
             )
 
+        dependencies = model_manifest_schema.ModelRuntimeDependenciesDict(conda=runtime_dict["dependencies"]["conda"])
+        if options.get("include_pip_dependencies"):
+            warnings.warn(
+                "`include_pip_dependencies` specified as True: pip dependencies will be included and may not"
+                "be warehouse-compabible. The model may need to be run in SPCS.",
+                category=UserWarning,
+                stacklevel=1,
+            )
+            dependencies["pip"] = runtime_dict["dependencies"]["pip"]
+
         manifest_dict = model_manifest_schema.ModelManifestDict(
             manifest_version=model_manifest_schema.MODEL_MANIFEST_VERSION,
             runtimes={
@@ -82,9 +91,7 @@ class ModelManifest:
                     language="PYTHON",
                     version=runtime_to_use.runtime_env.python_version,
                     imports=runtime_dict["imports"],
-                    dependencies=model_manifest_schema.ModelRuntimeDependenciesDict(
-                        conda=runtime_dict["dependencies"]["conda"]
-                    ),
+                    dependencies=dependencies,
                 )
             },
             methods=[
@@ -127,12 +134,13 @@ class ModelManifest:
         result = []
         if data_sources:
             for source in data_sources:
-                result.append(
-                    model_manifest_schema.LineageSourceDict(
-                        # Currently, we only support lineage from Dataset.
-                        type=model_manifest_schema.LineageSourceTypes.DATASET.value,
-                        entity=source.fully_qualified_name,
-                        version=source.version,
+                if isinstance(source, data_source.DatasetInfo):
+                    result.append(
+                        model_manifest_schema.LineageSourceDict(
+                            # Currently, we only support lineage from Dataset.
+                            type=model_manifest_schema.LineageSourceTypes.DATASET.value,
+                            entity=source.fully_qualified_name,
+                            version=source.version,
+                        )
                     )
-                )
         return result

snowflake/ml/model/_model_composer/model_manifest/model_manifest_schema.py

@@ -18,7 +18,8 @@ class ModelMethodFunctionTypes(enum.Enum):
 
 
 class ModelRuntimeDependenciesDict(TypedDict):
-    conda: Required[str]
+    conda: NotRequired[str]
+    pip: NotRequired[str]
 
 
 class ModelRuntimeDict(TypedDict):

snowflake/ml/model/_model_composer/model_method/function_generator.py

@@ -33,9 +33,9 @@ class FunctionGenerator:
 
     def __init__(
         self,
-        model_file_rel_path: pathlib.PurePosixPath,
+        model_dir_rel_path: pathlib.PurePosixPath,
     ) -> None:
-        self.model_file_rel_path = model_file_rel_path
+        self.model_dir_rel_path = model_dir_rel_path
 
     def generate(
         self,
@@ -67,7 +67,7 @@ class FunctionGenerator:
         )
 
         udf_code = function_template.format(
-            model_file_name=self.model_file_rel_path.name,
+            model_dir_name=self.model_dir_rel_path.name,
             target_method=target_method,
             max_batch_size=options.get("max_batch_size", None),
             function_name=FunctionGenerator.FUNCTION_NAME,

snowflake/ml/model/_model_composer/model_method/infer_function.py_template

@@ -1,12 +1,7 @@
-import fcntl
 import functools
 import inspect
 import os
 import sys
-import threading
-import zipfile
-from types import TracebackType
-from typing import Optional, Type
 
 import anyio
 import pandas as pd
@@ -15,42 +10,18 @@ from _snowflake import vectorized
 from snowflake.ml.model._packager import model_packager
 
 
-class FileLock:
-    def __enter__(self) -> None:
-        self._lock = threading.Lock()
-        self._lock.acquire()
-        self._fd = open("/tmp/lockfile.LOCK", "w+")
-        fcntl.lockf(self._fd, fcntl.LOCK_EX)
-
-    def __exit__(
-        self, exc_type: Optional[Type[BaseException]], exc: Optional[BaseException], traceback: Optional[TracebackType]
-    ) -> None:
-        self._fd.close()
-        self._lock.release()
-
-
 # User-defined parameters
-MODEL_FILE_NAME = "{model_file_name}"
+MODEL_DIR_REL_PATH = "{model_dir_name}"
 TARGET_METHOD = "{target_method}"
 MAX_BATCH_SIZE = {max_batch_size}
 
-
 # Retrieve the model
 IMPORT_DIRECTORY_NAME = "snowflake_import_directory"
 import_dir = sys._xoptions[IMPORT_DIRECTORY_NAME]
-
-model_dir_name = os.path.splitext(MODEL_FILE_NAME)[0]
-zip_model_path = os.path.join(import_dir, MODEL_FILE_NAME)
-extracted = "/tmp/models"
-extracted_model_dir_path = os.path.join(extracted, model_dir_name)
-
-with FileLock():
-    if not os.path.isdir(extracted_model_dir_path):
-        with zipfile.ZipFile(zip_model_path, "r") as myzip:
-            myzip.extractall(extracted_model_dir_path)
+model_dir_path = os.path.join(import_dir, MODEL_DIR_REL_PATH)
 
 # Load the model
-pk = model_packager.ModelPackager(extracted_model_dir_path)
+pk = model_packager.ModelPackager(model_dir_path)
 pk.load(as_custom_model=True)
 assert pk.model, "model is not loaded"
 assert pk.meta, "model metadata is not loaded"

snowflake/ml/model/_model_composer/model_method/infer_partitioned.py_template

@@ -15,42 +15,18 @@ from _snowflake import vectorized
 from snowflake.ml.model._packager import model_packager
 
 
-class FileLock:
-    def __enter__(self) -> None:
-        self._lock = threading.Lock()
-        self._lock.acquire()
-        self._fd = open("/tmp/lockfile.LOCK", "w+")
-        fcntl.lockf(self._fd, fcntl.LOCK_EX)
-
-    def __exit__(
-        self, exc_type: Optional[Type[BaseException]], exc: Optional[BaseException], traceback: Optional[TracebackType]
-    ) -> None:
-        self._fd.close()
-        self._lock.release()
-
-
 # User-defined parameters
-MODEL_FILE_NAME = "{model_file_name}"
+MODEL_DIR_REL_PATH = "{model_dir_name}"
 TARGET_METHOD = "{target_method}"
 MAX_BATCH_SIZE = {max_batch_size}
 
-
 # Retrieve the model
 IMPORT_DIRECTORY_NAME = "snowflake_import_directory"
 import_dir = sys._xoptions[IMPORT_DIRECTORY_NAME]
-
-model_dir_name = os.path.splitext(MODEL_FILE_NAME)[0]
-zip_model_path = os.path.join(import_dir, MODEL_FILE_NAME)
-extracted = "/tmp/models"
-extracted_model_dir_path = os.path.join(extracted, model_dir_name)
-
-with FileLock():
-    if not os.path.isdir(extracted_model_dir_path):
-        with zipfile.ZipFile(zip_model_path, "r") as myzip:
-            myzip.extractall(extracted_model_dir_path)
+model_dir_path = os.path.join(import_dir, MODEL_DIR_REL_PATH)
 
 # Load the model
-pk = model_packager.ModelPackager(extracted_model_dir_path)
+pk = model_packager.ModelPackager(model_dir_path)
 pk.load(as_custom_model=True)
 assert pk.model, "model is not loaded"
 assert pk.meta, "model metadata is not loaded"

snowflake/ml/model/_model_composer/model_method/infer_table_function.py_template

@@ -1,12 +1,7 @@
-import fcntl
 import functools
 import inspect
 import os
 import sys
-import threading
-import zipfile
-from types import TracebackType
-from typing import Optional, Type
 
 import anyio
 import pandas as pd
@@ -15,42 +10,18 @@ from _snowflake import vectorized
 from snowflake.ml.model._packager import model_packager
 
 
-class FileLock:
-    def __enter__(self) -> None:
-        self._lock = threading.Lock()
-        self._lock.acquire()
-        self._fd = open("/tmp/lockfile.LOCK", "w+")
-        fcntl.lockf(self._fd, fcntl.LOCK_EX)
-
-    def __exit__(
-        self, exc_type: Optional[Type[BaseException]], exc: Optional[BaseException], traceback: Optional[TracebackType]
-    ) -> None:
-        self._fd.close()
-        self._lock.release()
-
-
 # User-defined parameters
-MODEL_FILE_NAME = "{model_file_name}"
+MODEL_DIR_REL_PATH = "{model_dir_name}"
 TARGET_METHOD = "{target_method}"
 MAX_BATCH_SIZE = {max_batch_size}
 
-
 # Retrieve the model
 IMPORT_DIRECTORY_NAME = "snowflake_import_directory"
 import_dir = sys._xoptions[IMPORT_DIRECTORY_NAME]
-
-model_dir_name = os.path.splitext(MODEL_FILE_NAME)[0]
-zip_model_path = os.path.join(import_dir, MODEL_FILE_NAME)
-extracted = "/tmp/models"
-extracted_model_dir_path = os.path.join(extracted, model_dir_name)
-
-with FileLock():
-    if not os.path.isdir(extracted_model_dir_path):
-        with zipfile.ZipFile(zip_model_path, "r") as myzip:
-            myzip.extractall(extracted_model_dir_path)
+model_dir_path = os.path.join(import_dir, MODEL_DIR_REL_PATH)
 
 # Load the model
-pk = model_packager.ModelPackager(extracted_model_dir_path)
+pk = model_packager.ModelPackager(model_dir_path)
 pk.load(as_custom_model=True)
 assert pk.model, "model is not loaded"
 assert pk.meta, "model metadata is not loaded"

snowflake/ml/model/_model_composer/model_method/model_method.py

@@ -26,11 +26,14 @@ class ModelMethodOptions(TypedDict):
 def get_model_method_options_from_options(
     options: type_hints.ModelSaveOption, target_method: str
 ) -> ModelMethodOptions:
+    default_function_type = model_manifest_schema.ModelMethodFunctionTypes.FUNCTION.value
+    if options.get("enable_explainability", False) and target_method.startswith("explain"):
+        default_function_type = model_manifest_schema.ModelMethodFunctionTypes.TABLE_FUNCTION.value
     method_option = options.get("method_options", {}).get(target_method, {})
-    global_function_type = options.get("function_type", model_manifest_schema.ModelMethodFunctionTypes.FUNCTION.value)
+    global_function_type = options.get("function_type", default_function_type)
     function_type = method_option.get("function_type", global_function_type)
     if function_type not in [function_type.value for function_type in model_manifest_schema.ModelMethodFunctionTypes]:
-        raise NotImplementedError
+        raise NotImplementedError(f"Function type {function_type} is not supported.")
 
     return ModelMethodOptions(
         case_sensitive=method_option.get("case_sensitive", False),