snowflake-ml-python 1.2.0__py3-none-any.whl → 1.2.2__py3-none-any.whl

snowflake/ml/feature_store/feature_view.py

@@ -0,0 +1,459 @@
+from __future__ import annotations
+
+import json
+import re
+from collections import OrderedDict
+from dataclasses import dataclass
+from enum import Enum
+from typing import Dict, List, Optional
+
+from snowflake.ml._internal.exceptions import (
+    error_codes,
+    exceptions as snowml_exceptions,
+)
+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.entity import Entity
+from snowflake.snowpark import DataFrame, Session
+from snowflake.snowpark.types import (
+    DateType,
+    StructType,
+    TimestampType,
+    TimeType,
+    _NumericType,
+)
+
+_FEATURE_VIEW_NAME_DELIMITER = "$"
+_TIMESTAMP_COL_PLACEHOLDER = "FS_TIMESTAMP_COL_PLACEHOLDER_VAL"
+_FEATURE_OBJ_TYPE = "FEATURE_OBJ_TYPE"
+_FEATURE_VIEW_VERSION_RE = re.compile("^([A-Za-z0-9_]*)$")
+
+
+class FeatureViewVersion(str):
+    def __new__(cls, version: str) -> FeatureViewVersion:
+        if not _FEATURE_VIEW_VERSION_RE.match(version):
+            raise snowml_exceptions.SnowflakeMLException(
+                error_code=error_codes.INVALID_ARGUMENT,
+                original_exception=ValueError(
+                    f"`{version}` is not a valid feature view version. Only letter, number and underscore is allowed."
+                ),
+            )
+        return super().__new__(cls, version.upper())
+
+    def __init__(self, version: str) -> None:
+        return super().__init__()
+
+
+class FeatureViewStatus(Enum):
+    DRAFT = "DRAFT"
+    STATIC = "STATIC"
+    RUNNING = "RUNNING"
+    SUSPENDED = "SUSPENDED"
+
+
+@dataclass(frozen=True)
+class FeatureViewSlice:
+    feature_view_ref: FeatureView
+    names: List[SqlIdentifier]
+
+    def __repr__(self) -> str:
+        states = (f"{k}={v}" for k, v in vars(self).items())
+        return f"{type(self).__name__}({', '.join(states)})"
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, FeatureViewSlice):
+            return False
+
+        return self.names == other.names and self.feature_view_ref == other.feature_view_ref
+
+    def to_json(self) -> str:
+        fvs_dict = {
+            "feature_view_ref": self.feature_view_ref.to_json(),
+            "names": self.names,
+            _FEATURE_OBJ_TYPE: self.__class__.__name__,
+        }
+        return json.dumps(fvs_dict)
+
+    @classmethod
+    def from_json(cls, json_str: str, session: Session) -> FeatureViewSlice:
+        json_dict = json.loads(json_str)
+        if _FEATURE_OBJ_TYPE not in json_dict or json_dict[_FEATURE_OBJ_TYPE] != cls.__name__:
+            raise ValueError(f"Invalid json str for {cls.__name__}: {json_str}")
+        del json_dict[_FEATURE_OBJ_TYPE]
+        json_dict["feature_view_ref"] = FeatureView.from_json(json_dict["feature_view_ref"], session)
+        return cls(**json_dict)
+
+
+class FeatureView:
+    """
+    A FeatureView instance encapsulates a logical group of features.
+    """
+
+    def __init__(
+        self,
+        name: str,
+        entities: List[Entity],
+        feature_df: DataFrame,
+        timestamp_col: Optional[str] = None,
+        refresh_freq: Optional[str] = None,
+        desc: str = "",
+    ) -> None:
+        """
+        Create a FeatureView instance.
+
+        Args:
+            name: name of the FeatureView. NOTE: FeatureView name will be capitalized.
+            entities: entities that the FeatureView is associated with.
+            feature_df: Snowpark DataFrame containing data source and all feature feature_df logics.
+                Final projection of the DataFrame should contain feature names, join keys and timestamp(if applicable).
+            timestamp_col: name of the timestamp column for point-in-time lookup when consuming the
+                feature values.
+            refresh_freq: Time unit defining how often the new feature data should be generated.
+                Valid args are { <num> { seconds | minutes | hours | days } | DOWNSTREAM | <cron expr> <time zone>}.
+                NOTE: Currently minimum refresh frequency is 1 minute.
+                NOTE: If refresh_freq is in cron expression format, there must be a valid time zone as well.
+                    E.g. * * * * * UTC
+                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.
+        """
+
+        self._name: SqlIdentifier = SqlIdentifier(name)
+        self._entities: List[Entity] = entities
+        self._feature_df: DataFrame = feature_df
+        self._timestamp_col: Optional[SqlIdentifier] = (
+            SqlIdentifier(timestamp_col) if timestamp_col is not None else None
+        )
+        self._desc: str = desc
+        self._query: str = self._get_query()
+        self._version: Optional[FeatureViewVersion] = None
+        self._status: FeatureViewStatus = FeatureViewStatus.DRAFT
+        self._feature_desc: OrderedDict[SqlIdentifier, str] = OrderedDict((f, "") for f in self._get_feature_names())
+        self._refresh_freq: Optional[str] = refresh_freq
+        self._database: Optional[SqlIdentifier] = None
+        self._schema: Optional[SqlIdentifier] = None
+        self._warehouse: Optional[SqlIdentifier] = None
+        self._refresh_mode: Optional[str] = None
+        self._refresh_mode_reason: Optional[str] = None
+        self._validate()
+
+    def slice(self, names: List[str]) -> FeatureViewSlice:
+        """
+        Select a subset of features within the FeatureView.
+
+        Args:
+            names: feature names to select.
+
+        Returns:
+            FeatureViewSlice instance containing selected features.
+
+        Raises:
+            ValueError: if selected feature names is not found in the FeatureView.
+        """
+
+        res = []
+        for name in names:
+            name = SqlIdentifier(name)
+            if name not in self.feature_names:
+                raise ValueError(f"Feature name {name} not found in FeatureView {self.name}.")
+            res.append(name)
+        return FeatureViewSlice(self, res)
+
+    def physical_name(self) -> SqlIdentifier:
+        """Returns the physical name for this feature in Snowflake.
+
+        Returns:
+            Physical name string.
+
+        Raises:
+            RuntimeError: if the FeatureView is not materialized.
+        """
+        if self.status == FeatureViewStatus.DRAFT or self.version is None:
+            raise RuntimeError(f"FeatureView {self.name} has not been materialized.")
+        return FeatureView._get_physical_name(self.name, self.version)
+
+    def fully_qualified_name(self) -> str:
+        """Returns the fully qualified name (<database_name>.<schema_name>.<feature_view_name>) for the
+            FeatureView in Snowflake.
+
+        Returns:
+            fully qualified name string.
+        """
+        return f"{self._database}.{self._schema}.{self.physical_name()}"
+
+    def attach_feature_desc(self, descs: Dict[str, str]) -> FeatureView:
+        """
+        Associate feature level descriptions to the FeatureView.
+
+        Args:
+            descs: Dictionary contains feature name and corresponding descriptions.
+
+        Returns:
+            FeatureView with feature level desc attached.
+
+        Raises:
+            ValueError: if feature name is not found in the FeatureView.
+        """
+        for f, d in descs.items():
+            f = SqlIdentifier(f)
+            if f not in self._feature_desc:
+                raise ValueError(
+                    f"Feature name {f} is not found in FeatureView {self.name}, "
+                    f"valid feature names are: {self.feature_names}"
+                )
+            self._feature_desc[f] = d
+        return self
+
+    @property
+    def name(self) -> SqlIdentifier:
+        return self._name
+
+    @property
+    def entities(self) -> List[Entity]:
+        return self._entities
+
+    @property
+    def feature_df(self) -> DataFrame:
+        return self._feature_df
+
+    @property
+    def timestamp_col(self) -> Optional[SqlIdentifier]:
+        return self._timestamp_col
+
+    @property
+    def desc(self) -> str:
+        return self._desc
+
+    @property
+    def query(self) -> str:
+        return self._query
+
+    @property
+    def version(self) -> Optional[FeatureViewVersion]:
+        return self._version
+
+    @property
+    def status(self) -> FeatureViewStatus:
+        return self._status
+
+    @property
+    def feature_names(self) -> List[SqlIdentifier]:
+        return list(self._feature_desc.keys())
+
+    @property
+    def feature_descs(self) -> Dict[SqlIdentifier, str]:
+        return self._feature_desc
+
+    @property
+    def refresh_freq(self) -> Optional[str]:
+        return self._refresh_freq
+
+    @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."
+            )
+        self._refresh_freq = new_value
+
+    @property
+    def database(self) -> Optional[SqlIdentifier]:
+        return self._database
+
+    @property
+    def schema(self) -> Optional[SqlIdentifier]:
+        return self._schema
+
+    @property
+    def warehouse(self) -> Optional[SqlIdentifier]:
+        return self._warehouse
+
+    @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."
+            )
+        self._warehouse = SqlIdentifier(new_value)
+
+    @property
+    def output_schema(self) -> StructType:
+        return self._feature_df.schema
+
+    @property
+    def refresh_mode(self) -> Optional[str]:
+        return self._refresh_mode
+
+    @property
+    def refresh_mode_reason(self) -> Optional[str]:
+        return self._refresh_mode_reason
+
+    def _get_query(self) -> str:
+        if len(self._feature_df.queries["queries"]) != 1:
+            raise ValueError(
+                f"""feature_df dataframe must contain only 1 query.
+Got {len(self._feature_df.queries['queries'])}: {self._feature_df.queries['queries']}
+"""
+            )
+        return str(self._feature_df.queries["queries"][0])
+
+    def _validate(self) -> None:
+        if _FEATURE_VIEW_NAME_DELIMITER in self._name:
+            raise ValueError(
+                f"FeatureView name `{self._name}` contains invalid character `{_FEATURE_VIEW_NAME_DELIMITER}`."
+            )
+
+        unescaped_df_cols = to_sql_identifiers(self._feature_df.columns)
+        for e in self._entities:
+            for k in e.join_keys:
+                if k not in unescaped_df_cols:
+                    raise ValueError(
+                        f"join_key {k} in Entity {e.name} is not found in input dataframe: {unescaped_df_cols}"
+                    )
+
+        if self._timestamp_col is not None:
+            ts_col = self._timestamp_col
+            if ts_col == SqlIdentifier(_TIMESTAMP_COL_PLACEHOLDER):
+                raise ValueError(f"Invalid timestamp_col name, cannot be {_TIMESTAMP_COL_PLACEHOLDER}.")
+            if ts_col not in to_sql_identifiers(self._feature_df.columns):
+                raise ValueError(f"timestamp_col {ts_col} is not found in input dataframe.")
+
+            col_type = self._feature_df.schema[ts_col].datatype
+            if not isinstance(col_type, (DateType, TimeType, TimestampType, _NumericType)):
+                raise ValueError(f"Invalid data type for timestamp_col {ts_col}: {col_type}.")
+
+    def _get_feature_names(self) -> List[SqlIdentifier]:
+        join_keys = [k for e in self._entities for k in e.join_keys]
+        ts_col = [self._timestamp_col] if self._timestamp_col is not None else []
+        feature_names = to_sql_identifiers(self._feature_df.columns, case_sensitive=True)
+        return [c for c in feature_names if c not in join_keys + ts_col]
+
+    def __repr__(self) -> str:
+        states = (f"{k}={v}" for k, v in vars(self).items())
+        return f"{type(self).__name__}({', '.join(states)})"
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, FeatureView):
+            return False
+
+        return (
+            self.name == other.name
+            and self.version == other.version
+            and self.timestamp_col == other.timestamp_col
+            and self.entities == other.entities
+            and self.desc == other.desc
+            and self.feature_descs == other.feature_descs
+            and self.feature_names == other.feature_names
+            and self.query == other.query
+            and self.refresh_freq == other.refresh_freq
+            and str(self.status) == str(other.status)
+            and self.database == other.database
+            and self.warehouse == other.warehouse
+            and self.refresh_mode == other.refresh_mode
+            and self.refresh_mode_reason == other.refresh_mode_reason
+        )
+
+    def _to_dict(self) -> Dict[str, str]:
+        fv_dict = self.__dict__.copy()
+        if "_feature_df" in fv_dict:
+            fv_dict.pop("_feature_df")
+        fv_dict["_entities"] = [e._to_dict() for e in self._entities]
+        fv_dict["_status"] = str(self._status)
+        fv_dict["_name"] = str(self._name) if self._name is not None else None
+        fv_dict["_version"] = str(self._version) if self._version is not None else None
+        fv_dict["_database"] = str(self._database) if self._database is not None else None
+        fv_dict["_schema"] = str(self._schema) if self._schema is not None else None
+        fv_dict["_warehouse"] = str(self._warehouse) if self._warehouse is not None else None
+        fv_dict["_timestamp_col"] = str(self._timestamp_col) if self._timestamp_col is not None else None
+
+        feature_desc_dict = {}
+        for k, v in self._feature_desc.items():
+            feature_desc_dict[k.identifier()] = v
+        fv_dict["_feature_desc"] = feature_desc_dict
+
+        return fv_dict
+
+    def to_df(self, session: Session) -> DataFrame:
+        values = list(self._to_dict().values())
+        schema = [x.lstrip("_") for x in list(self._to_dict().keys())]
+        values.append(str(self.physical_name()))
+        schema.append("physical_name")
+        return session.create_dataframe([values], schema=schema)
+
+    def to_json(self) -> str:
+        state_dict = self._to_dict()
+        state_dict[_FEATURE_OBJ_TYPE] = self.__class__.__name__
+        return json.dumps(state_dict)
+
+    @classmethod
+    def from_json(cls, json_str: str, session: Session) -> FeatureView:
+        json_dict = json.loads(json_str)
+        if _FEATURE_OBJ_TYPE not in json_dict or json_dict[_FEATURE_OBJ_TYPE] != cls.__name__:
+            raise ValueError(f"Invalid json str for {cls.__name__}: {json_str}")
+
+        return FeatureView._construct_feature_view(
+            name=json_dict["_name"],
+            entities=[Entity(**e) for e in json_dict["_entities"]],
+            feature_df=session.sql(json_dict["_query"]),
+            timestamp_col=json_dict["_timestamp_col"],
+            desc=json_dict["_desc"],
+            version=json_dict["_version"],
+            status=json_dict["_status"],
+            feature_descs=json_dict["_feature_desc"],
+            refresh_freq=json_dict["_refresh_freq"],
+            database=json_dict["_database"],
+            schema=json_dict["_schema"],
+            warehouse=json_dict["_warehouse"],
+            refresh_mode=json_dict["_refresh_mode"],
+            refresh_mode_reason=json_dict["_refresh_mode_reason"],
+        )
+
+    @staticmethod
+    def _get_physical_name(fv_name: SqlIdentifier, fv_version: FeatureViewVersion) -> SqlIdentifier:
+        return SqlIdentifier(
+            concat_names(
+                [
+                    str(fv_name),
+                    _FEATURE_VIEW_NAME_DELIMITER,
+                    str(fv_version),
+                ]
+            )
+        )
+
+    @staticmethod
+    def _construct_feature_view(
+        name: str,
+        entities: List[Entity],
+        feature_df: DataFrame,
+        timestamp_col: Optional[str],
+        desc: str,
+        version: str,
+        status: FeatureViewStatus,
+        feature_descs: Dict[str, str],
+        refresh_freq: Optional[str],
+        database: Optional[str],
+        schema: Optional[str],
+        warehouse: Optional[str],
+        refresh_mode: Optional[str],
+        refresh_mode_reason: Optional[str],
+    ) -> FeatureView:
+        fv = FeatureView(
+            name=name,
+            entities=entities,
+            feature_df=feature_df,
+            timestamp_col=timestamp_col,
+            desc=desc,
+        )
+        fv._version = FeatureViewVersion(version) if version is not None else None
+        fv._status = status
+        fv._refresh_freq = refresh_freq
+        fv._database = SqlIdentifier(database) if database is not None else None
+        fv._schema = SqlIdentifier(schema) if schema is not None else None
+        fv._warehouse = SqlIdentifier(warehouse) if warehouse is not None else None
+        fv._refresh_mode = refresh_mode
+        fv._refresh_mode_reason = refresh_mode_reason
+        fv.attach_feature_desc(feature_descs)
+        return fv

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

@@ -4,9 +4,8 @@ import tempfile
 from typing import Any, Dict, List, Optional, Union, cast
 
 import yaml
-from packaging import version
 
-from snowflake.ml._internal.utils import identifier, snowflake_env, sql_identifier
+from snowflake.ml._internal.utils import identifier, sql_identifier
 from snowflake.ml.model import model_signature, type_hints
 from snowflake.ml.model._client.ops import metadata_ops
 from snowflake.ml.model._client.sql import (
@@ -25,8 +24,6 @@ from snowflake.ml.model._signatures import snowpark_handler
 from snowflake.snowpark import dataframe, row, session
 from snowflake.snowpark._internal import utils as snowpark_utils
 
-_TAG_ON_MODEL_AVAILABLE_VERSION = version.parse("8.2.0")
-
 
 class ModelOperator:
     def __init__(
@@ -296,21 +293,14 @@ class ModelOperator:
         tag_value: str,
         statement_params: Optional[Dict[str, Any]] = None,
     ) -> None:
-        sf_version = snowflake_env.get_current_snowflake_version(self._session, statement_params=statement_params)
-        if sf_version >= _TAG_ON_MODEL_AVAILABLE_VERSION:
-            self._tag_client.set_tag_on_model(
-                model_name=model_name,
-                tag_database_name=tag_database_name,
-                tag_schema_name=tag_schema_name,
-                tag_name=tag_name,
-                tag_value=tag_value,
-                statement_params=statement_params,
-            )
-        else:
-            raise NotImplementedError(
-                f"`set_tag` won't work before Snowflake version {_TAG_ON_MODEL_AVAILABLE_VERSION},"
-                f" currently is {sf_version}"
-            )
+        self._tag_client.set_tag_on_model(
+            model_name=model_name,
+            tag_database_name=tag_database_name,
+            tag_schema_name=tag_schema_name,
+            tag_name=tag_name,
+            tag_value=tag_value,
+            statement_params=statement_params,
+        )
 
     def unset_tag(
         self,
@@ -321,20 +311,13 @@ class ModelOperator:
         tag_name: sql_identifier.SqlIdentifier,
         statement_params: Optional[Dict[str, Any]] = None,
     ) -> None:
-        sf_version = snowflake_env.get_current_snowflake_version(self._session, statement_params=statement_params)
-        if sf_version >= _TAG_ON_MODEL_AVAILABLE_VERSION:
-            self._tag_client.unset_tag_on_model(
-                model_name=model_name,
-                tag_database_name=tag_database_name,
-                tag_schema_name=tag_schema_name,
-                tag_name=tag_name,
-                statement_params=statement_params,
-            )
-        else:
-            raise NotImplementedError(
-                f"`unset_tag` won't work before Snowflake version {_TAG_ON_MODEL_AVAILABLE_VERSION},"
-                f" currently is {sf_version}"
-            )
+        self._tag_client.unset_tag_on_model(
+            model_name=model_name,
+            tag_database_name=tag_database_name,
+            tag_schema_name=tag_schema_name,
+            tag_name=tag_name,
+            statement_params=statement_params,
+        )
 
     def get_model_version_manifest(
         self,
@@ -382,11 +365,6 @@ class ModelOperator:
         version_name: sql_identifier.SqlIdentifier,
         statement_params: Optional[Dict[str, Any]] = None,
     ) -> model_manifest_schema.SnowparkMLDataDict:
-        if (
-            snowflake_env.get_current_snowflake_version(self._session)
-            < model_manifest_schema.MANIFEST_USER_DATA_ENABLE_VERSION
-        ):
-            raise NotImplementedError("User_data has not been supported yet.")
         raw_user_data_json_string = self._model_client.show_versions(
             model_name=model_name,
             version_name=version_name,

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

@@ -3,10 +3,8 @@ from typing import Any, Dict, List, Optional
 from snowflake.ml._internal.utils import (
     identifier,
     query_result_checker,
-    snowflake_env,
     sql_identifier,
 )
-from snowflake.ml.model._model_composer.model_manifest import model_manifest_schema
 from snowflake.snowpark import row, session
 
 
@@ -89,12 +87,8 @@ class ModelSQLClient:
             .has_column(ModelSQLClient.MODEL_VERSION_NAME_COL_NAME, allow_empty=True)
             .has_column(ModelSQLClient.MODEL_VERSION_COMMENT_COL_NAME, allow_empty=True)
             .has_column(ModelSQLClient.MODEL_VERSION_METADATA_COL_NAME, allow_empty=True)
+            .has_column(ModelSQLClient.MODEL_VERSION_USER_DATA_COL_NAME, allow_empty=True)
         )
-        if (
-            snowflake_env.get_current_snowflake_version(self._session)
-            >= model_manifest_schema.MANIFEST_USER_DATA_ENABLE_VERSION
-        ):
-            res = res.has_column(ModelSQLClient.MODEL_VERSION_USER_DATA_COL_NAME, allow_empty=True)
         if validate_result and version_name:
             res = res.has_dimensions(expected_rows=1)
 

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

@@ -146,24 +146,29 @@ class ModelVersionSQLClient:
         returns: List[Tuple[str, spt.DataType, sql_identifier.SqlIdentifier]],
         statement_params: Optional[Dict[str, Any]] = None,
     ) -> dataframe.DataFrame:
-        tmp_table_name = snowpark_utils.random_name_for_temp_object(snowpark_utils.TempObjectType.TABLE)
-        INTERMEDIATE_TABLE_NAME = identifier.get_schema_level_object_identifier(
-            self._database_name.identifier(),
-            self._schema_name.identifier(),
-            tmp_table_name,
-        )
-        input_df.write.save_as_table(  # type: ignore[call-overload]
-            table_name=INTERMEDIATE_TABLE_NAME,
-            mode="errorifexists",
-            table_type="temporary",
-            statement_params=statement_params,
-        )
+        with_statements = []
+        if len(input_df.queries["queries"]) == 1 and len(input_df.queries["post_actions"]) == 0:
+            INTERMEDIATE_TABLE_NAME = "SNOWPARK_ML_MODEL_INFERENCE_INPUT"
+            with_statements.append(f"{INTERMEDIATE_TABLE_NAME} AS ({input_df.queries['queries'][0]})")
+        else:
+            tmp_table_name = snowpark_utils.random_name_for_temp_object(snowpark_utils.TempObjectType.TABLE)
+            INTERMEDIATE_TABLE_NAME = identifier.get_schema_level_object_identifier(
+                self._database_name.identifier(),
+                self._schema_name.identifier(),
+                tmp_table_name,
+            )
+            input_df.write.save_as_table(  # type: ignore[call-overload]
+                table_name=INTERMEDIATE_TABLE_NAME,
+                mode="errorifexists",
+                table_type="temporary",
+                statement_params=statement_params,
+            )
 
         INTERMEDIATE_OBJ_NAME = "TMP_RESULT"
 
         module_version_alias = "MODEL_VERSION_ALIAS"
-        model_version_alias_sql = (
-            f"WITH {module_version_alias} AS "
+        with_statements.append(
+            f"{module_version_alias} AS "
             f"MODEL {self.fully_qualified_model_name(model_name)} VERSION {version_name.identifier()}"
         )
 
@@ -174,7 +179,7 @@ class ModelVersionSQLClient:
         args_sql = ", ".join(args_sql_list)
 
         sql = textwrap.dedent(
-            f"""{model_version_alias_sql}
+            f"""WITH {','.join(with_statements)}
                 SELECT *,
                     {module_version_alias}!{method_name.identifier()}({args_sql}) AS {INTERMEDIATE_OBJ_NAME}
                 FROM {INTERMEDIATE_TABLE_NAME}"""

snowflake/ml/model/_deploy_client/image_builds/server_image_builder.py

@@ -2,6 +2,7 @@ import logging
 import os
 import posixpath
 from string import Template
+from typing import List
 
 import importlib_resources
 
@@ -36,6 +37,7 @@ class ServerImageBuilder(base_image_builder.ImageBuilder):
         session: snowpark.Session,
         artifact_stage_location: str,
         compute_pool: str,
+        external_access_integrations: List[str],
     ) -> None:
         """Initialization
 
@@ -47,6 +49,7 @@ class ServerImageBuilder(base_image_builder.ImageBuilder):
             artifact_stage_location: Spec file and future deployment related artifacts will be stored under
                 {stage}/models/{model_id}
             compute_pool: The compute pool used to run docker image build workload.
+            external_access_integrations: EAIs for network connection.
         """
         self.context_dir = context_dir
         self.image_repo = image_repo
@@ -54,6 +57,7 @@ class ServerImageBuilder(base_image_builder.ImageBuilder):
         self.session = session
         self.artifact_stage_location = artifact_stage_location
         self.compute_pool = compute_pool
+        self.external_access_integrations = external_access_integrations
         self.client = snowservice_client.SnowServiceClient(session)
 
         assert artifact_stage_location.startswith(
@@ -202,4 +206,8 @@ class ServerImageBuilder(base_image_builder.ImageBuilder):
 
     def _launch_kaniko_job(self, spec_stage_location: str) -> None:
         logger.debug("Submitting job for building docker image with kaniko")
-        self.client.create_job(compute_pool=self.compute_pool, spec_stage_location=spec_stage_location)
+        self.client.create_job(
+            compute_pool=self.compute_pool,
+            spec_stage_location=spec_stage_location,
+            external_access_integrations=self.external_access_integrations,
+        )

snowflake/ml/model/_deploy_client/snowservice/deploy.py

@@ -465,6 +465,7 @@ class SnowServiceDeployment:
                 session=self.session,
                 artifact_stage_location=self._model_artifact_stage_location,
                 compute_pool=self.options.compute_pool,
+                external_access_integrations=self.options.external_access_integrations,
             )
         else:
             image_builder = client_image_builder.ClientImageBuilder(
@@ -587,6 +588,7 @@ class SnowServiceDeployment:
             spec_stage_location=spec_stage_location,
             min_instances=self.options.min_instances,
             max_instances=self.options.max_instances,
+            external_access_integrations=self.options.external_access_integrations,
         )
         logger.info(f"Wait for service {self._service_name} to become ready...")
         client.block_until_resource_is_ready(