qwak-sdk 0.1.0__py3-none-any.whl → 0.2.20rc0__py3-none-any.whl

qwak_sdk/commands/feature_store/register/_logic.py

@@ -0,0 +1,289 @@
+from typing import List, Optional
+
+from _qwak_proto.qwak.feature_store.features.feature_set_pb2 import Feature
+from _qwak_proto.qwak.features_operator.v3.features_operator_pb2 import (
+    SparkColumnDescription,
+    ValidationSuccessResponse,
+)
+from qwak.clients.feature_store import FeatureRegistryClient
+from qwak.clients.feature_store.operator_client import FeaturesOperatorClient
+from qwak.exceptions import QwakException
+from qwak.feature_store.data_sources.batch_sources._batch import BaseBatchSource
+from qwak.feature_store.entities.entity import Entity
+from qwak.feature_store.feature_sets.batch import BatchFeatureSetV1
+from tabulate import tabulate
+
+from qwak_sdk.inner.file_registry import extract_class_objects
+from qwak_sdk.inner.tools.cli_tools import ask_yesno
+from qwak_sdk.tools.utils import qwak_spinner
+
+DELIMITER = "----------------------------------------"
+
+
+def _register_entities(
+    qwak_python_files: List[str], registry: FeatureRegistryClient, force: bool
+):
+    """
+    Register Feature Store Entity Objects
+
+    Args:
+        qwak_python_files: a list of python files containing qwak package imports
+        registry: FeatureRegistryClient
+        force: boolean determining if to force register all encountered Entity objects
+    """
+    with qwak_spinner(begin_text="Finding Entities to register", print_callback=print):
+        qwak_entities: List[Entity] = extract_class_objects(qwak_python_files, Entity)
+
+    print(f"👀 Found {len(qwak_entities)} Entities")
+    for entity, source_file_path in qwak_entities:
+        existing_entity = registry.get_entity_by_name(entity.name)
+        if existing_entity:
+            if ask_yesno(
+                f"Update existing Entity '{entity.name}' from source file '{source_file_path}'?",
+                force,
+            ):
+                registry.update_entity(
+                    existing_entity.entity.entity_definition.entity_id,
+                    entity.to_proto(),
+                )
+        else:
+            if ask_yesno(
+                f"Create new Entity '{entity.name}' from source file '{source_file_path}'?",
+                force,
+            ):
+                registry.create_entity(entity.to_proto())
+    print(DELIMITER)
+
+
+def _register_data_sources(
+    qwak_python_files: List[str],
+    registry: FeatureRegistryClient,
+    operator_client: FeaturesOperatorClient,
+    force: bool,
+    no_validation: bool,
+    ignore_validation_errors: bool,
+):
+    """
+    Register Feature Store Data Source Objects
+
+    Args:
+        qwak_python_files: a list of python files containing qwak package imports
+        registry: FeatureRegistryClient
+        operator_client: features operator grpc client
+        force: boolean determining if to force register all encountered Data Source objects
+        no_validation: whether to validate entities
+        ignore_validation_errors: whether to ignore and continue registering objects after encountering validation errors
+    """
+    with qwak_spinner(
+        begin_text="Finding Data Sources to register", print_callback=print
+    ):
+        qwak_sources = extract_class_objects(qwak_python_files, BaseBatchSource)
+
+    print(f"👀 Found {len(qwak_sources)} Data Sources")
+    for data_source, source_file_path in qwak_sources:
+        validation_failed = False
+
+        if no_validation:
+            print("Skipping data source validation")
+        else:
+            try:
+                _handle_data_source_validation(operator_client, data_source)
+            except Exception as e:
+                print(str(e))
+                validation_failed = True
+
+        if validation_failed and not ignore_validation_errors:
+            print("Not continuing to registration due to failure in validation")
+            exit(1)
+
+        existing_source = registry.get_data_source_by_name(data_source.name)
+        if existing_source:
+            if ask_yesno(
+                f"Update existing Data Source '{data_source.name}' from source file '{source_file_path}'?",
+                force,
+            ):
+                registry.update_data_source(
+                    existing_source.data_source.data_source_definition.data_source_id,
+                    data_source._to_proto(),
+                )
+        else:
+            if ask_yesno(
+                f"Create Data Source '{data_source.name}' from source file '{source_file_path}'?",
+                force,
+            ):
+                registry.create_data_source(data_source._to_proto())
+    print(DELIMITER)
+
+
+def _handle_data_source_validation(
+    operator_client: FeaturesOperatorClient,
+    data_source: BaseBatchSource,
+):
+    print(f"Validating '{data_source.name}' data source")
+    with qwak_spinner(begin_text="", print_callback=print):
+        result = operator_client.validate_data_source_blocking(
+            data_source_spec=data_source._to_proto(), num_samples=10
+        )
+
+    response = getattr(result, result.WhichOneof("type"))
+    if isinstance(response, ValidationSuccessResponse):
+        print("✅ Validation completed successfully, got data source columns:")
+
+        table = [
+            (x.column_name, x.spark_type) for x in response.spark_column_description
+        ]
+        print(tabulate(table, headers=["column name", "type"]))
+    else:
+        raise QwakException(f"🧨 Validation failed:\n{response}")
+
+
+def _register_features_sets(
+    qwak_python_files: List[str],
+    registry: FeatureRegistryClient,
+    operator_client: FeaturesOperatorClient,
+    force: bool,
+    git_commit: str,
+    no_validation: bool,
+    ignore_validation_errors: bool,
+):
+    """
+    Register Feature Store Feature Set Objects
+
+    Args:
+        qwak_python_files: a list of python files containing qwak package imports
+        registry: FeatureRegistryClient
+        operator_client: features operator grpc client
+        force: boolean determining if to force register all encountered Feature Set objects
+        git_commit: the git commit of the parent folder
+        no_validation: whether to validate entities
+        ignore_validation_errors: whether to ignore and continue registering objects after encountering validation errors
+    """
+    with qwak_spinner(
+        begin_text="Finding Feature Sets to register", print_callback=print
+    ):
+        qwak_feature_sets = extract_class_objects(qwak_python_files, BatchFeatureSetV1)
+
+    print(f"👀 Found {len(qwak_feature_sets)} Feature Set(s)")
+
+    for feature_set, source_file_path in qwak_feature_sets:
+        existing_feature_set = registry.get_feature_set_by_name(feature_set.name)
+
+        registration: bool = False
+        if existing_feature_set:
+            registration = ask_yesno(
+                f"Update existing feature set '{feature_set.name}' from source file '{source_file_path}'?",  # nosec B608
+                force,
+            )
+        else:
+            registration = ask_yesno(
+                f"Create new feature set '{feature_set.name}' from source file '{source_file_path}'?",
+                force,
+            )
+
+        if registration:
+            spark_columns_features = _batch_v1_feature_set_validation(
+                feature_set=feature_set,
+                no_validation=no_validation,
+                operator=operator_client,
+                registry=registry,
+            )
+
+            proto_feature_set = feature_set._to_proto(
+                git_commit, spark_columns_features, registry
+            )
+
+            if existing_feature_set:
+                registry.update_feature_set(
+                    existing_feature_set.feature_set.feature_set_definition.feature_set_id,
+                    proto_feature_set,
+                )
+            else:
+                registry.create_feature_set(proto_feature_set)
+
+
+def _handle_batch_v1_feature_set_validation(
+    operator_client: FeaturesOperatorClient,
+    registry_client: FeatureRegistryClient,
+    feature_set: BatchFeatureSetV1,
+) -> List[SparkColumnDescription]:
+    print(f"Validating '{feature_set.name}' feature set")
+    with qwak_spinner(begin_text="", print_callback=print):
+        feature_set_spec = feature_set._to_proto(
+            feature_registry=registry_client, features=None, git_commit=None
+        )
+        result = operator_client.validate_featureset_blocking(
+            feature_set_spec, resource_path=None, num_samples=10
+        )
+
+    response = getattr(result, result.WhichOneof("type"))
+    if isinstance(response, ValidationSuccessResponse):
+        print("✅ Validation completed successfully, got data source columns:")
+        table = [
+            (x.column_name, x.spark_type) for x in response.spark_column_description
+        ]
+        print(tabulate(table, headers=["column name", "type"]))
+        return response.spark_column_description
+
+    else:
+        raise QwakException(f"🧨 Validation failed:\n{response}")
+
+
+def _get_features_from_spark_columns_description(
+    spark_columns_description: List[SparkColumnDescription],
+) -> Optional[List[Feature]]:
+    if spark_columns_description:
+        spark_features = []
+        for spark_column in spark_columns_description:
+            feature_col_name_list = str(spark_column.column_name).split(".")
+            feature_name = (
+                feature_col_name_list[1]
+                if feature_col_name_list.__len__() == 2
+                else spark_column.column_name
+            )
+            spark_features.append(
+                Feature(
+                    feature_name=feature_name,
+                    feature_type=spark_column.spark_type,
+                )
+            )
+        return spark_features
+    else:
+        return None
+
+
+def _batch_v1_feature_set_validation(
+    feature_set: BatchFeatureSetV1,
+    no_validation: bool,
+    operator: FeaturesOperatorClient,
+    registry: FeatureRegistryClient,
+) -> Optional[List[Feature]]:
+    """
+    Validates featureset transformation
+    Args:
+        artifact_path: uploaded artifact zip path
+        feature_set: BatchFeatureSetV1 featureset
+        no_validation: skip validation
+        operator: Operator client
+        registry: Registry client
+    Returns:
+        Optional list of features returned from validation
+    """
+    spark_columns_features: Optional[List[Feature]] = []
+
+    if not no_validation:
+        try:
+            spark_columns_description = _handle_batch_v1_feature_set_validation(
+                operator_client=operator,
+                registry_client=registry,
+                feature_set=feature_set,
+            )
+            spark_columns_features = _get_features_from_spark_columns_description(
+                spark_columns_description
+            )
+        except Exception as e:
+            print(str(e))
+            print("Not continuing to registration due to failure in validation")
+            exit(1)
+    else:
+        print(f"Skipping validation for '{feature_set.name}' feature set")
+    return spark_columns_features

qwak_sdk/commands/feature_store/register/ui.py

@@ -0,0 +1,105 @@
+import os
+from pathlib import Path
+
+import click
+from qwak.clients.feature_store import FeatureRegistryClient
+from qwak.clients.feature_store.operator_client import FeaturesOperatorClient
+
+from qwak_sdk.commands.feature_store.register._logic import (
+    _register_data_sources,
+    _register_entities,
+    _register_features_sets,
+)
+from qwak_sdk.inner.file_registry import list_qwak_python_files
+from qwak_sdk.inner.tools.cli_tools import QwakCommand
+from qwak_sdk.tools.utils import qwak_spinner
+
+
+@click.command(
+    "register",
+    cls=QwakCommand,
+    help="Register and deploy all feature store object under the given path. Registered "
+    "features will be visible on the Qwak management platform after registration",
+)
+@click.option(
+    "--path",
+    "-p",
+    type=click.Path(exists=True),
+    metavar="PATH",
+    help="Directory / module where qwak feature store objects are stored",
+)
+@click.option(
+    "--force",
+    "-f",
+    default=False,
+    is_flag=True,
+    metavar="FLAG",
+    help="Force register all found qwak Feature Store objects",
+)
+@click.option(
+    "--no-validation",
+    "-nv",
+    default=False,
+    is_flag=True,
+    metavar="FLAG",
+    help="Skip validation for all found qwak Feature Store objects",
+)
+@click.option(
+    "--ignore-validation-errors",
+    "-ive",
+    default=False,
+    is_flag=True,
+    metavar="FLAG",
+    help="Ignore validation errors. Continue registering even if an error occurs",
+)
+def register_fs_objects(
+    path: Path,
+    force: bool,
+    no_validation: bool,
+    ignore_validation_errors: bool,
+    **kwargs,
+):
+    if not path:
+        path = Path.cwd()
+    else:
+        path = Path(path)
+
+    if path.is_file():
+        qwak_python_files = [(str(path), os.path.abspath(path))]
+    elif Path.is_dir(path):
+        with qwak_spinner(
+            begin_text="Recursively looking for python files in input dir",
+            print_callback=print,
+        ) as sp:
+            qwak_python_files = list_qwak_python_files(path, sp)
+
+    try:
+        import git
+
+        git_commit = git.Repo(path, search_parent_directories=True).head.commit.hexsha
+    except Exception:
+        # be super defensive on Git errors. Failing to fetch anything git related should not fail the registration
+        git_commit = None
+
+    registry_client = FeatureRegistryClient()
+    operator_client = FeaturesOperatorClient()
+    _register_entities(qwak_python_files, registry_client, force)
+
+    _register_data_sources(
+        qwak_python_files,
+        registry_client,
+        operator_client,
+        force,
+        no_validation,
+        ignore_validation_errors,
+    )
+
+    _register_features_sets(
+        qwak_python_files,
+        registry_client,
+        operator_client,
+        force,
+        git_commit,
+        no_validation,
+        ignore_validation_errors,
+    )

qwak_sdk/commands/feature_store/resume/ui.py

@@ -0,0 +1,18 @@
+import click
+from qwak.clients.feature_store import FeatureRegistryClient
+from qwak.exceptions import QwakException
+
+from qwak_sdk.inner.tools.cli_tools import QwakCommand
+from qwak_sdk.tools.colors import Color
+
+
+@click.command("resume", cls=QwakCommand, help="Resume a paused feature set")
+@click.argument("name")
+def resume_feature_set(name, **kwargs):
+    try:
+        FeatureRegistryClient().resume_feature_set(feature_set_name=name)
+    except Exception as e:
+        print(f"{Color.RED} Failed to resume feature set {name} {Color.END}")
+        raise QwakException(f"Failed to resume feature set {name}") from e
+
+    print(f"{Color.GREEN}Successfully resume feature set {Color.YELLOW}{name}")

qwak_sdk/commands/feature_store/trigger/ui.py

@@ -0,0 +1,32 @@
+import click
+from qwak.clients.feature_store import FeatureRegistryClient
+from qwak.exceptions import QwakException
+
+from qwak_sdk.inner.tools.cli_tools import QwakCommand
+from qwak_sdk.tools.colors import Color
+
+
+@click.command(
+    "trigger", cls=QwakCommand, help="Trigger a batch feature set job ingestion job"
+)
+@click.argument("name")
+def trigger_feature_set(name, **kwargs):
+    """
+    Trigger a batch feature set ingestion job
+
+    Args:
+        name: feature set name that will be triggered
+    """
+    try:
+        FeatureRegistryClient().run_feature_set(feature_set_name=name)
+    except Exception as e:
+        print(
+            f"{Color.RED} Failed to trigger a batch feature set ingestion for feature set {name} {Color.END}"
+        )
+        raise QwakException(
+            f"Failed to trigger a batch feature set ingestion for feature set {name}"
+        ) from e
+
+    print(
+        f"{Color.GREEN}Successfully triggered a batch feature set ingestion for feature set {Color.YELLOW}{name}"
+    )

qwak_sdk/commands/models/_logic/variations.py

@@ -0,0 +1,55 @@
+from __future__ import annotations
+
+from typing import List
+
+from _qwak_proto.qwak.deployment.deployment_pb2 import TrafficSpec, Variation
+from qwak.exceptions import QwakException
+
+from qwak_sdk.commands.models.deployments.deploy._logic.deploy_config import (
+    DeployConfig,
+)
+
+
+def validate_percentages(variations: List[Variation]):
+    for variation in variations:
+        if not (0 <= variation.traffic.percentage <= 100):
+            raise QwakException(
+                f"The variation '{variation.name}' contains invalid value '{variation.traffic.percentage}'. "
+                f"Value must be between 0 and 100."
+            )
+
+    non_shadow_variations = list(filter(lambda v: not v.traffic.is_shadow, variations))
+    non_shadow_variations_percentage_sum = sum(
+        map(lambda v: v.traffic.percentage, non_shadow_variations)
+    )
+    if non_shadow_variations and non_shadow_variations_percentage_sum != 100:
+        raise QwakException(
+            "The variations do not sum to 100 percent. Please go over the variations config file."
+        )
+
+
+def validate_shadow_variation(variations: List[Variation]):
+    shadow_variations = list(filter(lambda v: v.traffic.is_shadow, variations))
+    if len(shadow_variations) > 1:
+        shadow_variations_names = list(map(lambda v: v.name, shadow_variations))
+        raise QwakException(
+            f"The variations contain more than one shadow variation {shadow_variations_names}. "
+            f"Please go over the variations config file."
+        )
+
+
+def create_variation_from_variation_config(
+    variation_def: DeployConfig.Realtime.VariationConfig,
+) -> Variation:
+    try:
+        return Variation(
+            name=variation_def.name,
+            traffic=TrafficSpec(
+                percentage=variation_def.traffic.percentage,
+                is_shadow=variation_def.traffic.shadow,
+            ),
+        )
+    except Exception:
+        raise QwakException(
+            f"Could not parse variation {variation_def}. Please check you variation configuration file."
+        )

qwak_sdk/commands/models/build/_logic/build_steps.py

@@ -0,0 +1,68 @@
+from typing import List
+
+from .config.config_v1 import ConfigV1
+from .constant.step_description import BuildPhase, PhaseDetails
+from .context import Context
+from .interface.step_inteface import Step
+from .phase.a_fetch_model_code import get_fetch_model_code_steps
+from .phase.b_remote_register_qwak_build import get_remote_register_qwak_build_steps
+from .phase.c_deploy import get_deploy_stepss
+
+
+class StepsPipeline:
+    def __init__(
+        self,
+        config: ConfigV1,
+        context: Context = None,
+        build_phase: BuildPhase = None,
+    ) -> None:
+        self._phases: List[StepsPipeline] = []
+        self._steps: List[Step] = []
+        if not context:
+            self.context = Context()
+        self._config = config
+        self.build_phase = build_phase
+
+    @property
+    def phases(self) -> List["StepsPipeline"]:
+        return self._phases
+
+    @property
+    def steps(self) -> List[Step]:
+        return self._steps
+
+    def add_phase(self, steps: List[Step], build_phase: BuildPhase):
+        phase = StepsPipeline(
+            config=self._config, context=self.context, build_phase=build_phase
+        )
+
+        for step in steps:
+            step.context = self.context
+            step.config = self._config
+            step.build_phase = build_phase
+            phase._steps.append(step)
+
+        self._phases.append(phase)
+
+    def get_phase_details(self) -> PhaseDetails:
+        return PhaseDetails(self.build_phase)
+
+
+def remote_build_steps(config: ConfigV1) -> StepsPipeline:
+    steps_root = StepsPipeline(config=config)
+    steps_root.add_phase(
+        steps=get_fetch_model_code_steps(),
+        build_phase=BuildPhase.FETCHING_MODEL_CODE,
+    )
+    steps_root.add_phase(
+        steps=get_remote_register_qwak_build_steps(),
+        build_phase=BuildPhase.REGISTERING_QWAK_BUILD,
+    )
+
+    if config.deploy:
+        steps_root.add_phase(
+            steps=get_deploy_stepss(),
+            build_phase=BuildPhase.DEPLOY_PHASE,
+        )
+
+    return steps_root