frogml-cli 0.0.1__py3-none-any.whl

frogml_cli/commands/feature_store/list/ui.py

@@ -0,0 +1,140 @@
+from datetime import datetime
+from typing import List, Tuple
+
+import click
+from frogml._proto.qwak.feature_store.features.feature_set_pb2 import (
+    FeaturesetSchedulingState,
+    FeatureStatus,
+)
+from frogml._proto.qwak.feature_store.features.feature_set_service_pb2 import (
+    GetFeaturesetSchedulingStateResponse,
+    ListFeatureSetsResponse,
+)
+from frogml.core.clients.feature_store import FeatureRegistryClient
+from tabulate import tabulate
+
+from frogml_cli.inner.tools.cli_tools import FrogMLCommand
+
+
+def _pts_to_datetime(pts) -> datetime:
+    return datetime.fromtimestamp(pts.seconds + pts.nanos / 1e9)
+
+
+def _datetime_to_str(dt: datetime) -> str:
+    return dt.strftime("%Y-%b-%d %H:%M:%S")
+
+
+def get_featureset_scheduling_state_as_str(
+    registry_client: FeatureRegistryClient,
+    featureset_name: str,
+) -> str:
+    try:
+        resp: GetFeaturesetSchedulingStateResponse = (
+            registry_client.get_feature_set_scheduling_state(
+                featureset_name=featureset_name
+            )
+        )
+        if resp.state == FeaturesetSchedulingState.SCHEDULING_STATE_PAUSED:
+            return "Paused"
+        elif resp.state == FeaturesetSchedulingState.SCHEDULING_STATE_ENABLED:
+            return "Enabled"
+        else:
+            return "Unknown"
+    except Exception:
+        return "Unknown"
+
+
+def _list_feature_sets_verbose(
+    resp: ListFeatureSetsResponse, registry_client: FeatureRegistryClient
+) -> Tuple[List[str], List[str]]:
+    columns = [
+        "Name",
+        "Type",
+        "Status",
+        "Scheduling state",
+        "Owner",
+        "Description",
+        "Creation date",
+        "Last updated",
+        "Featureset id",
+    ]
+    data = []
+    for feature_set in resp.feature_families:
+        definition = feature_set.feature_set_definition
+        spec = definition.feature_set_spec
+        metadata = feature_set.metadata
+        if not (spec.feature_set_type.WhichOneof("set_type")):
+            feature_set_type = ""
+        else:
+            feature_set_type = (
+                spec.feature_set_type.WhichOneof("set_type")
+                .replace("_", " ")
+                .capitalize()
+            )
+
+        sched_state: str = "Enabled"
+        if spec.feature_set_type.WhichOneof("set_type") in {
+            "batch_feature_set",
+            "batch_feature_set_v1",
+        }:
+            sched_state = get_featureset_scheduling_state_as_str(
+                registry_client, spec.name
+            )
+
+        data.append(
+            [
+                spec.name,
+                feature_set_type,
+                FeatureStatus.Name(definition.status),
+                sched_state,
+                spec.metadata.owner,
+                spec.metadata.description,
+                _datetime_to_str(_pts_to_datetime(metadata.created_at)),
+                _datetime_to_str(_pts_to_datetime(metadata.last_modified_at)),
+                definition.feature_set_id,
+            ]
+        )
+
+    return (data, columns)
+
+
+def _list_feature_sets(
+    resp: ListFeatureSetsResponse, registry_client: FeatureRegistryClient
+) -> Tuple[List[str], List[str]]:
+    desired_columns = [
+        "Name",
+        "Type",
+        "Status",
+        "Scheduling state",
+        "Last updated",
+    ]
+    verbose_data, verbose_columns = _list_feature_sets_verbose(resp, registry_client)
+    col_indexes = [verbose_columns.index(c) for c in desired_columns]
+
+    def project_lst(lst: List[str]) -> List[str]:
+        return [lst[idx] for idx in col_indexes]
+
+    data = [project_lst(d) for d in verbose_data]
+
+    return (data, desired_columns)
+
+
+@click.command("list", cls=FrogMLCommand, help="List registered feature sets")
+@click.option(
+    "--verbose",
+    "-v",
+    default=False,
+    is_flag=True,
+    metavar="FLAG",
+    help="Verbose output",
+)
+def list_feature_sets(verbose: bool, **kwargs):
+    registry_client = FeatureRegistryClient()
+    response = registry_client.list_feature_sets()
+
+    if verbose:
+        data, columns = _list_feature_sets_verbose(response, registry_client)
+    else:
+        data, columns = _list_feature_sets(response, registry_client)
+
+    print(tabulate(data, headers=columns))

frogml_cli/commands/feature_store/pause/ui.py

@@ -0,0 +1,18 @@
+import click
+from frogml.core.clients.feature_store import FeatureRegistryClient
+from frogml.core.exceptions import FrogmlException
+
+from frogml_cli.inner.tools.cli_tools import FrogMLCommand
+from frogml_cli.tools.colors import Color
+
+
+@click.command("pause", cls=FrogMLCommand, help="Pause a running feature set")
+@click.argument("name")
+def pause_feature_set(name, **kwargs):
+    try:
+        FeatureRegistryClient().pause_feature_set(feature_set_name=name)
+    except Exception as e:
+        print(f"{Color.RED} Failed to pause feature set {name} {Color.END}")
+        raise FrogmlException(f"Failed to pause feature set {name}") from e
+
+    print(f"{Color.GREEN}Successfully paused feature set {Color.YELLOW}{name}")

frogml_cli/commands/feature_store/register/_logic.py

@@ -0,0 +1,369 @@
+from pathlib import Path
+from typing import List, Optional, Tuple, cast
+
+from frogml._proto.qwak.feature_store.entities.entity_pb2 import (
+    EntityDefinition as ProtoEntityDefinition,
+)
+from frogml._proto.qwak.feature_store.features.feature_set_pb2 import (
+    Feature as ProtoFeature,
+)
+from frogml._proto.qwak.feature_store.features.feature_set_service_pb2 import (
+    GetFeatureSetByNameResponse as ProtoGetFeatureSetByNameResponse,
+)
+from frogml.core.clients.feature_store import FeatureRegistryClient
+from frogml.core.exceptions import FrogmlException
+from frogml.core.feature_store.entities.entity import Entity
+from frogml.core.feature_store.validations.validation_options import (
+    FeatureSetValidationOptions,
+)
+from frogml.core.feature_store.validations.validation_response import (
+    SuccessValidationResponse,
+    ValidationResponse,
+)
+from frogml.feature_store.data_sources.base import BaseSource
+from frogml.feature_store.feature_sets.base_feature_set import BaseFeatureSet
+from tabulate import tabulate
+
+from frogml_cli.inner.file_registry import extract_class_objects
+from frogml_cli.inner.tools.cli_tools import ask_yesno
+from frogml_cli.tools.utils import frogml_spinner
+
+DELIMITER = "----------------------------------------"
+
+
+def _register_entities(
+    qwak_python_files: List[Tuple[str, 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 frogml_spinner(
+        begin_text="Finding Entities to register", print_callback=print
+    ):
+        qwak_entities: List[Tuple[Entity, str]] = 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[Tuple[str, str]],
+    registry: FeatureRegistryClient,
+    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
+        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 frogml_spinner(
+        begin_text="Finding Data Sources to register", print_callback=print
+    ):
+        qwak_sources: List[Tuple[BaseSource, str]] = extract_class_objects(
+            qwak_python_files, BaseSource
+        )
+
+    print(f"👀 Found {len(qwak_sources)} Data Sources")
+    for data_source, source_file_path in qwak_sources:
+        validation_failed = False
+        artifact_url: Optional[str] = None
+
+        if no_validation:
+            print("Skipping data source validation")
+        else:
+            try:
+                artifact_url = _handle_data_source_validation(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,
+            ):
+                data_source_proto, _ = data_source._prepare_and_get(
+                    artifact_url=artifact_url,
+                    source_definition_path=Path(source_file_path),
+                )
+                registry.update_data_source(
+                    existing_source.data_source.data_source_definition.data_source_id,
+                    data_source_proto,
+                )
+        else:
+            if ask_yesno(
+                f"Create Data Source '{data_source.name}' from source file '{source_file_path}'?",
+                force,
+            ):
+                data_source_proto, _ = data_source._prepare_and_get(
+                    artifact_url=artifact_url,
+                    source_definition_path=Path(source_file_path),
+                )
+                registry.create_data_source(data_source_proto)
+    print(DELIMITER)
+
+
+def _handle_data_source_validation(
+    data_source: BaseSource,
+) -> Optional[str]:
+    print(f"Validating '{data_source.name}' data source")
+    with frogml_spinner(begin_text="", print_callback=print):
+        from frogml.feature_store.validations.validator import FeaturesOperatorValidator
+
+        v = FeaturesOperatorValidator()
+        response: ValidationResponse
+        artifact_url: Optional[str]
+
+        response, artifact_url = v.validate_data_source(data_source=data_source)
+    if isinstance(response, SuccessValidationResponse):
+        print_validation_outputs(response.stdout, response.stderr)
+        print("✅ Validation completed successfully, got data source columns:")
+
+        table = [(x.feature_name, x.feature_type) for x in response.features]
+        print(tabulate(table, headers=["column name", "type"]))
+
+        return artifact_url
+    else:
+        raise FrogmlException(f"🧨 Validation failed: \n{response}")
+
+
+def _register_features_sets(
+    qwak_python_files: List[Tuple[str, str]],
+    registry: FeatureRegistryClient,
+    force: bool,
+    git_commit: str,
+    no_validation: bool,
+    ignore_validation_errors: bool,
+    data_source_limit: Optional[int] = None,
+):
+    """
+    Register Feature Store Feature Set 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 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 frogml_spinner(
+        begin_text="Finding Feature Sets to register", print_callback=print
+    ):
+        qwak_feature_sets = extract_class_objects(qwak_python_files, BaseFeatureSet)
+
+    print(f"👀 Found {len(qwak_feature_sets)} Feature Set(s)")
+
+    for featureset, source_file_path in qwak_feature_sets:
+        featureset = cast(BaseFeatureSet, featureset)
+        existing_feature_set: ProtoGetFeatureSetByNameResponse = (
+            registry.get_feature_set_by_name(featureset.name)
+        )
+
+        registration: bool = False
+        if existing_feature_set:
+            # Provide entity information of registered feature set before any other operation
+            if featureset.key:
+                featureset.entity = (
+                    existing_feature_set.feature_set.feature_set_definition.feature_set_spec.entity.entity_spec.name
+                )
+
+            registration = ask_yesno(
+                f"Update existing feature set '{featureset.name}' from source file '{source_file_path}'?",  # nosec B608
+                force,
+            )
+        else:
+            registration = ask_yesno(
+                f"Create new feature set '{featureset.name}' from source file '{source_file_path}'?",
+                force,
+            )
+
+        if registration:
+            features: Optional[List[ProtoFeature]] = []
+            artifact_url: Optional[str] = None
+
+            features, artifact_url = _validate_featureset(
+                featureset=featureset,
+                no_validation=no_validation,
+                ignore_validation_errors=ignore_validation_errors,
+                data_source_limit=data_source_limit,
+            )
+
+            proto_feature_set, _ = featureset._to_proto(
+                git_commit=git_commit,
+                features=features,
+                feature_registry=registry,
+                artifact_url=artifact_url,
+            )
+
+            if existing_feature_set:
+                registry.update_feature_set(
+                    existing_feature_set.feature_set.feature_set_definition.feature_set_id,
+                    proto_feature_set,
+                )
+            else:
+                _register_new_feature_set(
+                    new_feature_set=featureset,
+                    entity_key=proto_feature_set.entity,
+                    registry=registry,
+                    git_commit=git_commit,
+                    features=features,
+                    artifact_url=artifact_url,
+                )
+
+
+def _register_new_feature_set(
+    new_feature_set: BaseFeatureSet,
+    entity_key: ProtoEntityDefinition,
+    registry: FeatureRegistryClient,
+    git_commit: str,
+    features: Optional[List[ProtoFeature]],
+    artifact_url: Optional[str] = None,
+):
+    # Create entity for the defined key and set it as the fs's entity
+    if new_feature_set.key:
+        new_entity: Entity = Entity._from_proto(proto=entity_key)
+        try:
+            new_entity.register()
+        except FrogmlException:
+            raise FrogmlException(
+                f"Failed to create key for {new_feature_set.name}, "
+                f"aborting feature set creation"
+            )
+        new_feature_set.entity = new_entity.name
+
+    # this is done to retrieve the entities metadata
+    proto_feature_set, _ = new_feature_set._to_proto(
+        git_commit, features, registry, artifact_url=artifact_url
+    )
+
+    try:
+        registry.create_feature_set(proto_feature_set)
+    # rollback entity creation in case fs registration failed
+    except FrogmlException as e1:
+        try:
+            if new_feature_set.key:
+                registry.delete_entity(entity_id=proto_feature_set.entity.entity_id)
+        except FrogmlException:
+            raise e1
+        raise e1
+
+
+def _handle_featureset_validation(
+    featureset: BaseFeatureSet,
+    data_source_limit: Optional[int] = None,
+) -> Tuple[List[ProtoFeature], Optional[str]]:
+    print(f"Validating '{featureset.name}' feature set")
+    with frogml_spinner(begin_text="", print_callback=print):
+        from frogml.feature_store.validations.validator import FeaturesOperatorValidator
+
+        v = FeaturesOperatorValidator()
+        response: ValidationResponse
+        artifact_url: Optional[str]
+        validation_options = FeatureSetValidationOptions(
+            data_source_limit=data_source_limit
+        )
+
+        response, artifact_url = v.validate_featureset(
+            featureset=featureset, validation_options=validation_options
+        )
+    if isinstance(response, SuccessValidationResponse):
+        print_validation_outputs(response.stdout, response.stderr)
+        print("✅ Validation completed successfully, got data source columns:")
+        table = [(x.feature_name, x.feature_type) for x in response.features]
+        print(tabulate(table, headers=["column name", "type"]))
+        return response.features, artifact_url
+    else:
+        raise FrogmlException(f"🧨 Validation failed: \n{response}")
+
+
+def _validate_featureset(
+    featureset: BaseFeatureSet,
+    no_validation: bool,
+    ignore_validation_errors: bool,
+    data_source_limit: Optional[int] = None,
+) -> Tuple[List[ProtoFeature], Optional[str]]:
+    """
+    Validates featureset transformation
+    Args:
+        featureset: BaseFeatureSet featureset
+        no_validation: skip validation
+        operator: Operator client
+        registry: Registry client
+    Returns:
+        Optional list of features returned from validation
+    """
+    features: List[ProtoFeature] = []
+    artifact_url: Optional[str] = None
+    if not no_validation:
+        try:
+            features, artifact_url = _handle_featureset_validation(
+                featureset=featureset, data_source_limit=data_source_limit
+            )
+        except Exception as e:
+            print(str(e))
+
+            if not ignore_validation_errors:
+                print("Not continuing to registration due to failure in validation")
+                exit(1)
+            else:
+                print("Ignoring validation errors")
+    else:
+        print(f"Skipping validation for '{featureset.name}' feature set")
+    return features, artifact_url
+
+
+def print_validation_outputs(stdout: str, stderr: str) -> bool:
+    did_print = False
+
+    if stdout or stderr:
+        message = "Validation outputs: "
+
+        if stdout:
+            message += f"stdout: {stdout}\n "
+
+        if stderr:
+            message += f"stderr: {stderr}"
+
+        print(message)
+        did_print = True
+
+    return did_print

frogml_cli/commands/feature_store/register/ui.py

@@ -0,0 +1,111 @@
+import os
+from pathlib import Path
+from typing import List, Optional, Tuple
+
+import click
+from frogml.core.clients.feature_store import FeatureRegistryClient
+
+from frogml_cli.commands.feature_store.register._logic import (
+    _register_data_sources,
+    _register_entities,
+    _register_features_sets,
+)
+from frogml_cli.inner.file_registry import list_qwak_python_files
+from frogml_cli.inner.tools.cli_tools import FrogMLCommand
+from frogml_cli.tools.utils import frogml_spinner
+
+
+@click.command(
+    "register",
+    cls=FrogMLCommand,
+    help="Register and deploy all feature store object under the given path. Registered "
+    "features will be visible on the Frogml management platform after registration",
+)
+@click.option(
+    "--path",
+    "-p",
+    type=click.Path(exists=True),
+    metavar="PATH",
+    help="Directory / module where frogml feature store objects are stored",
+)
+@click.option(
+    "--force",
+    "-f",
+    default=False,
+    is_flag=True,
+    metavar="FLAG",
+    help="Force register all found frogml Feature Store objects",
+)
+@click.option(
+    "--no-validation",
+    "-nv",
+    default=False,
+    is_flag=True,
+    metavar="FLAG",
+    help="Skip validation for all found frogml 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",
+)
+@click.option(
+    "--fs-validation-ds-limit",
+    type=click.IntRange(1, 10000),
+    help="Limit the number of records to be fetched from each datasource during featureset validation",
+)
+def register_fs_objects(
+    path: Path,
+    force: bool,
+    no_validation: bool,
+    ignore_validation_errors: bool,
+    fs_validation_ds_limit: Optional[int],
+    **kwargs,
+):
+    qwak_python_files: List[Tuple[str, str]]
+
+    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 frogml_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()
+    _register_entities(qwak_python_files, registry_client, force)
+
+    _register_data_sources(
+        qwak_python_files,
+        registry_client,
+        force,
+        no_validation,
+        ignore_validation_errors,
+    )
+
+    _register_features_sets(
+        qwak_python_files,
+        registry_client,
+        force,
+        git_commit,
+        no_validation,
+        ignore_validation_errors,
+        fs_validation_ds_limit,
+    )

frogml_cli/commands/feature_store/resume/ui.py

@@ -0,0 +1,18 @@
+import click
+from frogml.core.clients.feature_store import FeatureRegistryClient
+from frogml.core.exceptions import FrogmlException
+
+from frogml_cli.inner.tools.cli_tools import FrogMLCommand
+from frogml_cli.tools.colors import Color
+
+
+@click.command("resume", cls=FrogMLCommand, 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 FrogmlException(f"Failed to resume feature set {name}") from e
+
+    print(f"{Color.GREEN}Successfully resume feature set {Color.YELLOW}{name}")