snowflake-ml-python 1.2.1__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/_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(

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

@@ -1,6 +1,6 @@
 import inspect
 import logging
-from typing import Any, Dict, Optional
+from typing import Any, Dict, List, Optional
 
 from snowflake.ml._internal.exceptions import (
     error_codes,
@@ -16,6 +16,7 @@ class SnowServiceDeployOptions:
         self,
         compute_pool: str,
         *,
+        external_access_integrations: List[str],
         image_repo: Optional[str] = None,
         min_instances: Optional[int] = 1,
         max_instances: Optional[int] = 1,
@@ -34,7 +35,15 @@ class SnowServiceDeployOptions:
 
         Args:
             compute_pool: SnowService compute pool name. Please refer to official doc for how to create a
-                compute pool: https://docs.snowflake.com/LIMITEDACCESS/snowpark-containers/reference/compute-pool
+                compute pool:
+                https://docs.snowflake.com/en/developer-guide/snowpark-container-services/working-with-compute-pool
+            external_access_integrations: External Access Integrations name used to build image and deploy the model.
+                Please refer to the doc for how to create an External Access Integrations: https://docs.snowflake.com/
+                developer-guide/snowpark-container-services/additional-considerations-services-jobs
+                #configuring-network-capabilities .
+                To make sure your image could be built, access to the following endpoint must be allowed.
+                docker.com:80, docker.com:443, anaconda.com:80, anaconda.com:443, anaconda.org:80, anaconda.org:443,
+                pypi.org:80, pypi.org:443
             image_repo: SnowService image repo path. e.g. "<image_registry>/<db>/<schema>/<repo>". Default to auto
                 inferred based on session information.
             min_instances: Minimum number of service replicas. Default to 1.
@@ -70,6 +79,7 @@ class SnowServiceDeployOptions:
         self.model_in_image = model_in_image
         self.debug_mode = debug_mode
         self.enable_ingress = enable_ingress
+        self.external_access_integrations = external_access_integrations
 
         if self.num_workers is None and self.use_gpu:
             logger.info("num_workers has been defaulted to 1 when using GPU.")

snowflake/ml/model/_deploy_client/utils/snowservice_client.py

@@ -2,7 +2,7 @@ import json
 import logging
 import textwrap
 import time
-from typing import Optional
+from typing import List, Optional
 
 from snowflake.ml._internal.exceptions import (
     error_codes,
@@ -36,6 +36,7 @@ class SnowServiceClient:
         service_name: str,
         compute_pool: str,
         spec_stage_location: str,
+        external_access_integrations: List[str],
         *,
         min_instances: Optional[int] = 1,
         max_instances: Optional[int] = 1,
@@ -48,6 +49,7 @@ class SnowServiceClient:
             service_name: Name of the service.
             min_instances: Minimum number of service replicas.
             max_instances: Maximum number of service replicas.
+            external_access_integrations: EAIs for network connection.
             compute_pool: Name of the compute pool.
             spec_stage_location: Stage path for the service spec.
         """
@@ -61,13 +63,14 @@ class SnowServiceClient:
                 SPEC = '{path}'
                 MIN_INSTANCES={min_instances}
                 MAX_INSTANCES={max_instances}
+                EXTERNAL_ACCESS_INTEGRATIONS = ({', '.join(external_access_integrations)})
             """
         )
         logger.info(f"Creating service {service_name}")
         logger.debug(f"Create service with SQL: \n {sql}")
         self.session.sql(sql).collect()
 
-    def create_job(self, compute_pool: str, spec_stage_location: str) -> None:
+    def create_job(self, compute_pool: str, spec_stage_location: str, external_access_integrations: List[str]) -> None:
         """Execute the job creation SQL command. Note that the job creation is synchronous, hence we execute it in a
         async way so that we can query the log in the meantime.
 
@@ -76,7 +79,7 @@ class SnowServiceClient:
         Args:
             compute_pool: name of the compute pool
             spec_stage_location: path to the stage location where the spec is located at.
-
+            external_access_integrations: EAIs for network connection.
         """
         stage, path = uri.get_stage_and_path(spec_stage_location)
         sql = textwrap.dedent(
@@ -85,6 +88,7 @@ class SnowServiceClient:
             IN COMPUTE POOL {compute_pool}
             FROM {stage}
             SPEC = '{path}'
+            EXTERNAL_ACCESS_INTEGRATIONS = ({', '.join(external_access_integrations)})
             """
         )
         logger.debug(f"Create job with SQL: \n {sql}")