qwak-core 0.1.37__py3-none-any.whl → 0.5.19__py3-none-any.whl

qwak/clients/build_management/client.py

@@ -1,114 +0,0 @@
-import grpc
-from dependency_injector.wiring import Provide
-
-from _qwak_proto.qwak.builds.builds_pb2 import (
-    BuildStatus,
-    ExperimentTrackingValues,
-    GetBuildRequest,
-    GetBuildResponse,
-    ListBuildsRequest,
-    ListBuildsResponse,
-    Metric,
-    Param,
-    RegisterExperimentTrackingValuesRequest,
-    RegisterModelSchemaRequest,
-    UpdateBuildStatusRequest,
-)
-from _qwak_proto.qwak.builds.builds_pb2_grpc import BuildsManagementServiceStub
-from qwak.exceptions import QwakException
-from qwak.inner.di_configuration import QwakContainer
-
-
-class BuildsManagementClient:
-    """
-    Used for interacting with Feature Registry endpoints
-    """
-
-    def __init__(self, grpc_channel=Provide[QwakContainer.core_grpc_channel]):
-        self._builds_management_service = BuildsManagementServiceStub(grpc_channel)
-
-    def update_build_status(self, build_id, build_status, qwak_user_id=""):
-        try:
-            self._builds_management_service.UpdateBuildStatus(
-                UpdateBuildStatusRequest(
-                    build_id=build_id,
-                    build_status=BuildStatus.Value(build_status),
-                    qwak_calling_user_id=qwak_user_id,
-                )
-            )
-
-        except grpc.RpcError as e:
-            raise QwakException(
-                f"Failed to update build status, error is {e.details()}"
-            )
-
-    def get_build(self, build_id) -> GetBuildResponse:
-        try:
-            return self._builds_management_service.GetBuild(
-                GetBuildRequest(build_id=build_id)
-            )
-        except grpc.RpcError as e:
-            raise QwakException(f"Failed to get build, error is {e.details()}")
-
-    def is_build_exists(self, build_id: str) -> bool:
-        try:
-            self._builds_management_service.GetBuild(GetBuildRequest(build_id=build_id))
-            return True
-        except grpc.RpcError as e:
-            if e.code() == grpc.StatusCode.NOT_FOUND:
-                return False
-            raise QwakException(
-                f"Failed to check if build {build_id} is exists, error is {e.details()}"
-            )
-
-    def register_schema(self, build_id, model_schema):
-        try:
-            return self._builds_management_service.RegisterModelSchema(
-                RegisterModelSchemaRequest(build_id=build_id, model_schema=model_schema)
-            )
-
-        except grpc.RpcError as e:
-            raise QwakException(f"Failed to register schema, error is {e.details()}")
-
-    def register_experiment_tracking_values(
-        self, build_id, qwak_params=None, cli_params=None, qwak_metrics=None
-    ):
-        qwak_params = qwak_params if qwak_params else {}
-        cli_params = cli_params if cli_params else {}
-        qwak_metrics = qwak_metrics if qwak_metrics else {}
-
-        try:
-            params = cli_params.copy()
-            params.update(qwak_params)
-
-            experiment_tracking_values = ExperimentTrackingValues(
-                metrics=[
-                    Metric(key=key, value=value) for key, value in qwak_metrics.items()
-                ],
-                params=[
-                    Param(key=key, value=str(value)) for key, value in params.items()
-                ],
-            )
-
-            return self._builds_management_service.RegisterExperimentTrackingValues(
-                RegisterExperimentTrackingValuesRequest(
-                    build_id=build_id,
-                    experiment_tracking_values=experiment_tracking_values,
-                )
-            )
-
-        except grpc.RpcError as e:
-            raise QwakException(
-                f"Failed to register experiment tracking values, error is {e.details()}"
-            )
-
-    def list_builds(self, model_uuid="", **kwargs) -> ListBuildsResponse:
-        try:
-            _model_uuid = model_uuid if model_uuid else kwargs.get("branch_id")
-            if not _model_uuid:
-                raise QwakException("missing argument model uuid or branch id.")
-            return self._builds_management_service.ListBuilds(
-                ListBuildsRequest(model_uuid=_model_uuid)
-            )
-        except grpc.RpcError as e:
-            raise QwakException(f"Failed to list builds, error is {e.details()}")

qwak/feature_store/_common/featureset_asterisk_handler.py

@@ -1,115 +0,0 @@
-import dataclasses
-import itertools
-import re
-from typing import Callable, Dict, List
-
-from qwak.feature_store.offline import OfflineClient
-from qwak.model.schema_entities import BaseFeature, Entity, FeatureStoreInput
-
-
-def is_asterisk(feature: BaseFeature) -> bool:
-    """
-    Checks whether the given Feature is a FeatureStoreInput with an '*'
-    Args:
-        feature: input feature
-    Return
-    """
-    p = "^[^.]+.\\*$"
-    # if this is not a FeatureStoreInput, no need to check
-    if isinstance(feature, FeatureStoreInput):
-        return bool(re.match(p, feature.name))
-    return False
-
-
-def inflate_feature(feature: BaseFeature, offline_fs: OfflineClient) -> List:
-    """
-    Inflates a Feature into a list. if the feature contains an '*', a list
-        a feature for each feature in the associate FeatureSet is returned. else -
-        a list containing the original feature is returned.
-    Args:
-        feature: input feature
-        offline_fs: OfflineFeatureStore to query
-    Return:
-        if an '*' exists, returns a List of FeatureStoreInuput for each feature in the feature set
-        else returns a list containing the feature.
-
-    """
-    if is_asterisk(feature):
-        featureset_name = feature.name.lower().split(".")[0]
-
-        p = f"^{featureset_name}+.*$"
-        feature_names = [
-            feature.lower()
-            for feature in offline_fs.get_columns_by_feature_set(featureset_name)
-            if re.match(p, feature.lower())
-        ]
-        return [dataclasses.replace(feature, name=name) for name in feature_names]
-
-    return [feature]
-
-
-def unpack_asterisk_features(
-    features: List[BaseFeature],
-    feature_store_generator: Callable[[], OfflineClient] = lambda: OfflineClient(),
-) -> List[BaseFeature]:
-    """
-    Handles features with an '*'.
-    If a feature of type FeatureStoreInput that matches '<featureset_name>.*' exists,
-    transforms this feature into the list of features actually present in the FeatureSet.
-    Other features remain unchanged.
-    the original order is maintained, and the internal order (inside a FeatureStoreInput
-     that has an '*') is inferred from the order of columns returned from
-     OfflineFeatureStore.get_columns_by_feature_set(...)
-     Args:
-         features: List of features
-         feature_store_generator: A function that generates an instance of OfflineFeatureStore
-                                  If no feature with an '*' is found, the function is not called.
-
-    Returns:
-        the unpacked features
-    """
-    asterisk_found = any([is_asterisk(f) for f in features])
-    if asterisk_found:
-        # at least 1 feature contains an asterisk, map each feature to
-        # a list of features (features with no '*' are mapped to a list of size 1), then flatten it.
-        offline_feature_store = feature_store_generator()
-        inflated_features = [
-            inflate_feature(feature, offline_feature_store) for feature in features
-        ]
-        return list(itertools.chain.from_iterable(inflated_features))
-
-    return features
-
-
-def unpack_asterisk_features_from_key_mapping(
-    key_to_features: Dict[str, List[str]],
-    feature_store_generator: Callable[[], OfflineClient] = lambda: OfflineClient(),
-) -> dict:
-    """
-    Handles features with an '*'.
-    If a feature name that matches '<featureset_name>.*' exists in one of the lists,
-    transforms this feature into the list of feature names actually present in the FeatureSet.
-    Other features remain unchanged.
-    the original order is maintained, and the internal order (inside a feature
-     that has an '*') is inferred from the order of columns returned from
-     OfflineFeatureStore.get_columns_by_feature_set(...)
-     Args:
-         key_to_features: dictionary of entity keys to requested features (same as in the offline store)
-             >>> key_to_features = {'uuid': ['user_purchases.*',
-             >>>                             'user_purchases.avg_purchase_amount']}
-         feature_store_generator: A function that generates an instance of OfflineFeatureStore
-                                  If no feature with an '*' is found, the function is not called.
-    Returns:
-        the unpacked features
-    """
-    inflated_key_to_features = {}
-
-    for entity, features in key_to_features.items():
-        feature_inputs = [
-            FeatureStoreInput(name=feature, entity=Entity(name=entity, type=str))
-            for feature in features
-        ]
-        unpacked = unpack_asterisk_features(feature_inputs, feature_store_generator)
-        feature_list = [feature.name for feature in unpacked]
-        inflated_key_to_features[entity] = feature_list
-    return inflated_key_to_features

qwak/feature_store/data_sources/batch_sources/__init__.py

@@ -1,6 +0,0 @@
-from qwak.feature_store.data_sources.batch_sources.mongodb import (
-    MongoDbSource as MongoDbSource,
-    MongoDbSource as MongoSource,
-)
-
-__all__ = ["MongoSource", "MongoDbSource"]

qwak/feature_store/data_sources/batch_sources/_batch.py

@@ -1,86 +0,0 @@
-import ast
-from abc import ABC, abstractmethod
-from dataclasses import dataclass
-from typing import TYPE_CHECKING, Optional
-
-from _qwak_proto.qwak.feature_store.sources.batch_pb2 import (
-    BatchSource as ProtoBatchSource,
-)
-from _qwak_proto.qwak.features_operator.v3.features_operator_async_service_pb2 import (
-    DataSourceValidationOptions as ProtoDataSourceValidationOptions,
-)
-from _qwak_proto.qwak.features_operator.v3.features_operator_pb2 import (
-    ValidationSuccessResponse,
-)
-from qwak.clients.feature_store.operator_client import FeaturesOperatorClient
-from qwak.exceptions import QwakException
-from qwak.feature_store._common.functions import print_validation_outputs
-from qwak.feature_store.validation_options import DataSourceValidationOptions
-
-if TYPE_CHECKING:
-    try:
-        import pandas as pd
-    except ImportError:
-        pass
-
-
-@dataclass
-class BaseBatchSource(ABC):
-    name: str
-    date_created_column: str
-    description: str
-
-    @abstractmethod
-    def _to_proto(self) -> ProtoBatchSource:
-        pass
-
-    @classmethod
-    @abstractmethod
-    def _from_proto(cls, proto):
-        pass
-
-    def get_sample(
-        self,
-        number_of_rows: int = 10,
-        validation_options: Optional[DataSourceValidationOptions] = None,
-    ) -> "pd.DataFrame":
-        """
-        Tries to get a sample of length `number_rows` from the data source.
-        Args:
-            number_of_rows: number of rows to get from data source
-            validation_options: validation options
-        Returns:
-            A tuple containing the resulting dataframe and a tuple of the columns names and types.
-            (the types are pyspark dataframe types)
-        """
-
-        try:
-            import pandas as pd
-        except ImportError:
-            raise QwakException("Missing Pandas dependency required for getting sample")
-
-        if number_of_rows > 1000:
-            raise ValueError(f"`number_rows` must be under 1000, got: {number_of_rows}")
-
-        operator_client = FeaturesOperatorClient()
-        ds_spec = self._to_proto()
-        validation_options_proto: Optional[ProtoDataSourceValidationOptions] = (
-            validation_options.to_proto() if validation_options else None
-        )
-
-        result = operator_client.validate_data_source_blocking(
-            data_source_spec=ds_spec,
-            num_samples=number_of_rows,
-            validation_options=validation_options_proto,
-        )
-
-        response = getattr(result, result.WhichOneof("type"))
-        if isinstance(response, ValidationSuccessResponse):
-            print_validation_outputs(response)
-
-            return pd.read_json(
-                path_or_buf=ast.literal_eval(response.sample),
-                dtype=response.spark_column_description,
-            )
-        else:
-            raise QwakException(f"Sampling failed:\n{response}")

qwak/feature_store/feature_sets/transformations.py

@@ -1,48 +0,0 @@
-from abc import ABC, abstractmethod
-
-from _qwak_proto.qwak.feature_store.features.feature_set_types_pb2 import (
-    SqlTransformation as ProtoSqlTransformation,
-    Transformation as ProtoTransformation,
-)
-from qwak.exceptions import QwakException
-
-
-class BatchTransformation(ABC):
-    @abstractmethod
-    def _to_proto(self) -> ProtoTransformation:
-        pass
-
-    @classmethod
-    def _from_proto(cls, proto: "ProtoTransformation"):
-        function_mapping = {
-            "sql_transformation": SparkSqlTransformation,
-        }
-
-        function_type: str = proto.WhichOneof("type")
-        if function_type in function_mapping:
-            function_class = function_mapping.get(function_type)
-            return function_class._from_proto(proto)
-
-        raise QwakException(f"Got unsupported function type: {function_type}")
-
-
-class SparkSqlTransformation(BatchTransformation):
-    """
-    A Spark SQL transformation
-    :param sql: A valid Spark SQL transformation
-    Example transformation:
-    .. code-block:: python
-        SparkSqlTransformation("SELECT user_id, age FROM data_source")
-    """
-
-    def __init__(self, sql):
-        self._sql = sql
-
-    def _to_proto(self) -> ProtoTransformation:
-        return ProtoTransformation(
-            sql_transformation=ProtoSqlTransformation(sql=self._sql, function_names=[]),
-        )
-
-    @classmethod
-    def _from_proto(cls, proto: "ProtoSqlTransformation"):
-        return cls(sql=proto.sql_transformation.sql)

qwak/feature_store/offline/_query_engine.py

@@ -1,32 +0,0 @@
-import uuid
-from abc import ABC, abstractmethod
-
-import pandas as pd
-
-
-class BaseQueryEngine(ABC):
-    @abstractmethod
-    def upload_table(self, df: pd.DataFrame):
-        pass
-
-    @abstractmethod
-    def cleanup(self):
-        pass
-
-    @abstractmethod
-    def run_query(self, query: str):
-        pass
-
-    @abstractmethod
-    def read_pandas_from_query(self, query: str, parse_dates=None):
-        pass
-
-    @staticmethod
-    @abstractmethod
-    def get_quotes():
-        pass
-
-    class JoinTableSpec:
-        def __init__(self, join_tables_db_name: str, quotes: str):
-            self.table_name = str(uuid.uuid4())
-            self.join_table_full_path = f"{quotes}{join_tables_db_name}{quotes}.{quotes}{self.table_name}{quotes}"

qwak/feature_store/offline/athena/athena_query_engine.py

@@ -1,153 +0,0 @@
-import time
-import uuid
-
-import pandas as pd
-from _qwak_proto.qwak.ecosystem.v0.ecosystem_runtime_service_pb2 import (
-    GetCloudCredentialsParameters,
-    GetCloudCredentialsRequest,
-    OfflineFeatureStoreClient,
-    PermissionSet,
-)
-from google.protobuf.duration_pb2 import Duration
-from qwak.clients.administration.eco_system.client import EcosystemClient
-from qwak.exceptions import QwakException
-from qwak.feature_store.offline._query_engine import BaseQueryEngine
-
-RECONNECT_THRESHOLD_SEC = 300
-
-
-class AthenaQueryEngine(BaseQueryEngine):
-    def __init__(self):
-        eco_client = EcosystemClient()
-        self.bucket, environment_id = self._get_env_details(eco_client)
-
-        self.staging_folder_prefix = (
-            f"{environment_id}/tmp/offline_fs/{str(uuid.uuid4())}"  # nosec B108
-        )
-        self.temp_join_table_base_folder = (
-            f"s3://{self.bucket}/{self.staging_folder_prefix}"
-        )
-
-        self.conn, self.expiration_time = self._init_connection()
-        self.cursor = self.conn.cursor()
-
-        self.join_table_specs = []
-
-        self.join_tables_db_name = f"qwak_temp_data_{environment_id.replace('-', '_')}"
-        self.cursor.execute(f"CREATE DATABASE IF NOT EXISTS {self.join_tables_db_name}")
-
-    @staticmethod
-    def _get_env_details(eco_client):
-        environment_configuration = eco_client.get_environment_configuration()
-
-        return (
-            environment_configuration.configuration.object_storage_bucket,
-            environment_configuration.id,
-        )
-
-    def _init_connection(self):
-        try:
-            # obtain credentials through STS
-            eco_client = EcosystemClient()
-            cloud_credentials_response = eco_client.get_cloud_credentials(
-                request=GetCloudCredentialsRequest(
-                    parameters=GetCloudCredentialsParameters(
-                        duration=Duration(seconds=60 * 60, nanos=0),
-                        permission_set=PermissionSet(
-                            offline_feature_store_client=OfflineFeatureStoreClient()
-                        ),
-                    )
-                )
-            )
-
-            aws_credentials = (
-                cloud_credentials_response.cloud_credentials.aws_temporary_credentials
-            )
-
-            try:
-                from pyathena import connect
-                from pyathena.pandas.cursor import PandasCursor
-            except ImportError:
-                raise QwakException(
-                    """
-                    Missing 'pyathena' dependency required for fetching data from the offline store.
-                    Please pip install pyathena
-                """
-                )
-
-            conn = connect(
-                s3_staging_dir=self.temp_join_table_base_folder,
-                aws_access_key_id=aws_credentials.access_key_id,
-                aws_secret_access_key=aws_credentials.secret_access_key,
-                aws_session_token=aws_credentials.session_token,
-                region_name=aws_credentials.region,
-                cursor_class=PandasCursor,
-            )
-
-            return (
-                conn,
-                aws_credentials.expiration_time.seconds,
-            )
-
-        except QwakException as e:
-            raise e
-
-        except Exception as e:
-            raise QwakException(
-                f"Got an error trying to retrieve credentials to query the offline store "
-                f"in the cloud, error is: {e}"
-            )
-
-    def upload_table(self, df: pd.DataFrame):
-        join_table_spec = super().JoinTableSpec(
-            self.join_tables_db_name, AthenaQueryEngine.get_quotes()
-        )
-        self.join_table_specs.append(join_table_spec)
-
-        from pyathena.pandas.util import to_sql
-
-        to_sql(
-            df,
-            join_table_spec.table_name,
-            self.conn,
-            f"{self.temp_join_table_base_folder}/{join_table_spec.table_name}/",
-            schema=self.join_tables_db_name,
-            index=False,
-            if_exists="replace",
-        )
-
-        return join_table_spec.join_table_full_path
-
-    def run_query(self, query: str):
-        self._check_reconnection()
-        return self.cursor.execute(query).fetchall()
-
-    def read_pandas_from_query(self, query: str, parse_dates=None):
-        self._check_reconnection()
-        return pd.read_sql(
-            query,
-            self.conn,
-            parse_dates=parse_dates,
-        )
-
-    def _check_reconnection(self):
-        if self.expiration_time - time.time() < RECONNECT_THRESHOLD_SEC:
-            self.conn, self.expiration_time = self._init_connection()
-            self.cursor = self.conn.cursor()
-
-    def cleanup(self):
-        self._check_reconnection()
-        for join_table_spec in self.join_table_specs:
-            self.cursor.execute(
-                f"""DROP TABLE {join_table_spec.join_table_full_path.replace('"', '`')}"""
-            )
-
-        self.join_table_specs = []
-
-        s3 = self.conn.session.resource("s3")
-        bucket = s3.Bucket(self.bucket)
-        bucket.objects.filter(Prefix=self.staging_folder_prefix).delete()
-
-    @staticmethod
-    def get_quotes():
-        return '"'