frogml-core 0.0.115__py3-none-any.whl → 0.0.117__py3-none-any.whl

core/inner/tool/auth.py

@@ -0,0 +1,75 @@
+from typing import Optional, Union, cast
+
+import requests
+from frogml.storage.authentication.models import AuthConfig, BearerAuth
+from frogml.storage.authentication.utils import get_credentials
+from requests.auth import AuthBase
+from typing_extensions import Self
+
+from frogml.core.exceptions import FrogmlLoginException
+
+
+class FrogMLAuthClient:
+    __MIN_TOKEN_LENGTH: int = 64
+
+    def __init__(self: Self, auth_config: Optional[AuthConfig] = None):
+        self.auth_config: Optional[AuthConfig] = auth_config
+        self._token: Optional[str] = None
+        self._tenant_id: Optional[str] = None
+
+    def get_token(self: Self) -> str:
+        if not self._token:
+            self.login()
+
+        return cast(str, self._token)
+
+    def get_tenant_id(self: Self) -> str:
+        if not self._tenant_id:
+            self.login()
+
+        return cast(str, self._tenant_id)
+
+    def login(self: Self):
+        artifactory_url, auth = get_credentials(self.auth_config)
+
+        if isinstance(auth, BearerAuth):  # For BearerAuth
+            self.__validate_token(auth.token)
+            self._token = auth.token
+
+        self.__get_tenant_id(artifactory_url, auth)
+
+    def get_auth(self: Self) -> Union[AuthBase]:
+        return get_credentials(self.auth_config)[1]
+
+    def __get_tenant_id(self: Self, artifactory_url: str, auth: AuthBase):
+        login_exception = FrogmlLoginException(
+            "Failed to authenticate with JFrog. Please check your credentials"
+        )
+
+        # Remove '/artifactory/' from the URL
+        if "/artifactory" in artifactory_url:
+            base_url = artifactory_url.replace("/artifactory", "/ui")
+        else:
+            # Remove trailing slash if it exists and append /ui
+            base_url = artifactory_url.rstrip("/") + "/ui"
+
+        url = f"{base_url}/api/v1/system/auth/screen/footer"
+
+        try:
+            response = requests.get(url, timeout=15, auth=auth)
+            response.raise_for_status()  # Raises an HTTPError for bad responses
+            response_data = response.json()
+
+            if "serverId" not in response_data:
+                raise login_exception
+            self._tenant_id = response_data["serverId"]
+
+        except (requests.exceptions.RequestException, ValueError) as exc:
+            raise login_exception from exc
+
+    def __validate_token(self: Self, token: Optional[str]):
+        if token is None or len(token) <= self.__MIN_TOKEN_LENGTH:
+            raise FrogmlLoginException(
+                "Authentication with JFrog failed: Only JWT Access Tokens are supported. "
+                "Please ensure you are using a valid JWT Access Token."
+            )

core/inner/tool/grpc/grpc_auth.py

@@ -0,0 +1,42 @@
+from typing import TYPE_CHECKING, Callable, Optional, Tuple
+
+import grpc
+
+from frogml.core.inner.const import FrogMLConstants
+
+if TYPE_CHECKING:
+    from frogml.core.inner.tool.auth import FrogMLAuthClient
+
+_SIGNATURE_HEADER_KEY = "authorization"
+
+
+class FrogMLAuthMetadataPlugin(grpc.AuthMetadataPlugin):
+    def __init__(self):
+        self._auth_client: Optional["FrogMLAuthClient"] = None
+
+    def __call__(
+        self,
+        context: grpc.AuthMetadataContext,
+        callback: Callable[[Tuple[Tuple[str, str], Tuple[str, str]], None], None],
+    ):
+        """Implements authentication by passing metadata to a callback.
+
+        Args:
+            context: An AuthMetadataContext providing information on the RPC that
+                the plugin is being called to authenticate.
+            callback: A callback that accepts a tuple of metadata key/value pairs and a None
+                parameter.
+        """
+        # Get token from FrogML client
+        if not self._auth_client:
+            from frogml.core.inner.tool.auth import FrogMLAuthClient
+
+            self._auth_client = FrogMLAuthClient()
+
+        token = self._auth_client.get_token()
+        jfrog_tenant_id = self._auth_client.get_tenant_id()
+        metadata = (
+            (_SIGNATURE_HEADER_KEY, f"Bearer {token}"),
+            (FrogMLConstants.JFROG_TENANT_HEADER_KEY.lower(), jfrog_tenant_id),
+        )
+        callback(metadata, None)

core/inner/tool/grpc/grpc_tools.py

@@ -0,0 +1,209 @@
+import logging
+import time
+from abc import ABC, abstractmethod
+from random import randint
+from typing import Callable, Optional, Tuple
+
+import grpc
+
+from frogml.core.exceptions import FrogmlException
+from frogml.core.inner.tool.grpc.grpc_auth import FrogMLAuthMetadataPlugin
+
+logger = logging.getLogger()
+
+
+def create_grpc_channel(
+    url: str,
+    enable_ssl: bool = True,
+    enable_auth: bool = True,
+    auth_metadata_plugin: grpc.AuthMetadataPlugin = None,
+    timeout: int = 100,
+    options=None,
+    backoff_options=None,
+    max_attempts=4,
+    status_for_retry=(grpc.StatusCode.UNAVAILABLE,),
+    attempt=0,
+) -> grpc.Channel:
+    """
+    Create a gRPC channel
+    Args:
+        url: gRPC URL to connect to
+        enable_ssl: Enable TLS/SSL, optionally provide a server side certificate
+        enable_auth: Enable user auth
+        auth_metadata_plugin: Metadata plugin to use to sign requests, only used
+            with "enable auth" when SSL/TLS is enabled
+        timeout: Connection timeout to server
+        options:  An optional list of key-value pairs (channel_arguments in gRPC Core runtime) to configure the channel.
+        backoff_options: dictionary - init_backoff_ms: default 50, max_backoff_ms: default 500, multiplier: default 2
+        max_attempts: max number of retry attempts
+        attempt: current retry attempts
+        status_for_retry: grpc statuses to retry upon
+    Returns: Returns a grpc.Channel
+    """
+    if backoff_options is None:
+        backoff_options = {}
+
+    if not url:
+        raise FrogmlException(
+            "Unable to create gRPC channel. URL has not been defined."
+        )
+
+    if enable_ssl or url.endswith(":443"):
+        credentials: grpc.ChannelCredentials = grpc.ssl_channel_credentials()
+
+        if enable_auth:
+            if auth_metadata_plugin is None:
+                auth_metadata_plugin = FrogMLAuthMetadataPlugin()
+
+            credentials = grpc.composite_channel_credentials(
+                credentials, grpc.metadata_call_credentials(auth_metadata_plugin)
+            )
+
+        channel: grpc.Channel = grpc.secure_channel(
+            url, credentials=credentials, options=options
+        )
+    else:
+        channel = grpc.insecure_channel(url, options=options)
+    try:
+        interceptors: Tuple[RetryOnRpcErrorClientInterceptor] = (
+            RetryOnRpcErrorClientInterceptor(
+                max_attempts=max_attempts,
+                sleeping_policy=ExponentialBackoff(**backoff_options),
+                status_for_retry=status_for_retry,
+            ),
+        )
+
+        intercept_channel: grpc.Channel = grpc.intercept_channel(channel, *interceptors)
+        try:
+            grpc.channel_ready_future(intercept_channel).result(timeout=timeout)
+        except (grpc.FutureTimeoutError, grpc.RpcError, Exception) as e:
+            logger.debug(
+                f"Received error: {repr(e)} attempt #{attempt + 1} of {max_attempts}"
+            )
+            if attempt < max_attempts:
+                return create_grpc_channel(
+                    url,
+                    enable_ssl,
+                    enable_auth,
+                    auth_metadata_plugin,
+                    timeout,
+                    options,
+                    backoff_options,
+                    max_attempts,
+                    status_for_retry,
+                    attempt + 1,
+                )
+            else:
+                raise e
+
+        return intercept_channel
+    except grpc.FutureTimeoutError as e:
+        raise FrogmlException(
+            f"Connection timed out while attempting to connect to {url}, with: {repr(e)}"
+        ) from e
+
+
+def create_grpc_channel_or_none(
+    url: str,
+    enable_ssl: bool = True,
+    enable_auth: bool = True,
+    auth_metadata_plugin: grpc.AuthMetadataPlugin = None,
+    timeout: int = 30,
+    options=None,
+    backoff_options=None,
+    max_attempts=2,
+    status_for_retry=(grpc.StatusCode.UNAVAILABLE,),
+    attempt=0,
+) -> Callable[[Optional[str], Optional[bool]], Optional[grpc.Channel]]:
+    final_backoff_options = backoff_options if backoff_options else {}
+
+    def deferred_channel(
+        url_overwrite: Optional[str] = None, ssl_overwrite: Optional[bool] = None
+    ):
+        try:
+            return create_grpc_channel(
+                url_overwrite if url_overwrite else url,
+                ssl_overwrite if ssl_overwrite else enable_ssl,
+                enable_auth,
+                auth_metadata_plugin,
+                timeout,
+                options,
+                final_backoff_options,
+                max_attempts,
+                status_for_retry,
+                attempt,
+            )
+        except Exception as e:
+            logger.debug(f"create_grpc_channel error: {repr(e)}")
+            return None
+
+    return deferred_channel
+
+
+class SleepingPolicy(ABC):
+    @abstractmethod
+    def sleep(self, try_i: int):
+        """
+        How long to sleep in milliseconds.
+        :param try_i: the number of retry (starting from zero)
+        """
+        if try_i < 0:
+            raise ValueError("Number of retries must be non-negative.")
+
+
+class ExponentialBackoff(SleepingPolicy):
+    def __init__(
+        self,
+        *,
+        init_backoff_ms: int = 50,
+        max_backoff_ms: int = 5000,
+        multiplier: int = 2,
+    ):
+        self.init_backoff = init_backoff_ms
+        self.max_backoff = max_backoff_ms
+        self.multiplier = multiplier
+
+    def sleep(self, try_i: int):
+        sleep_time = min(self.init_backoff * self.multiplier**try_i, self.max_backoff)
+        sleep_ms = sleep_time + randint(0, self.init_backoff)  # nosec B311
+        logger.debug("ExponentialBackoff - Sleeping between retries")
+        time.sleep(sleep_ms / 1000)
+
+
+class RetryOnRpcErrorClientInterceptor(grpc.UnaryUnaryClientInterceptor):
+    def __init__(
+        self,
+        *,
+        max_attempts: int,
+        sleeping_policy: SleepingPolicy,
+        status_for_retry: Optional[Tuple[grpc.StatusCode]] = None,
+    ):
+        self.max_attempts = max_attempts
+        self.sleeping_policy = sleeping_policy
+        self.status_for_retry = status_for_retry
+
+    def _intercept_call(self, continuation, client_call_details, request_or_iterator):
+        for try_i in range(self.max_attempts):
+            response = continuation(client_call_details, request_or_iterator)
+
+            if isinstance(response, grpc.RpcError):
+                # Return if it was last attempt
+                if try_i == (self.max_attempts - 1):
+                    return response
+
+                # If status code is not in retryable status codes
+                if (
+                    self.status_for_retry
+                    and response.code() not in self.status_for_retry
+                ):
+                    return response
+                logger.debug(
+                    f"Retry GRPC call attempt #{try_i} after status {response.code()}"
+                )
+                logger.debug(f"Client call details: {client_call_details}")
+                self.sleeping_policy.sleep(try_i)
+            else:
+                return response
+
+    def intercept_unary_unary(self, continuation, client_call_details, request):
+        return self._intercept_call(continuation, client_call_details, request)

core/inner/tool/grpc/grpc_try_wrapping.py

@@ -0,0 +1,23 @@
+from typing import Callable
+
+import grpc
+from frogml.core.exceptions import FrogmlException
+
+
+def grpc_try_catch_wrapper(exception_message: str):
+    def decorator(function: Callable):
+        def _inner_wrapper(*args, **kwargs):
+            try:
+                return function(*args, **kwargs)
+            except Exception as e:
+                if isinstance(e, grpc.RpcError):
+                    # noinspection PyUnresolvedReferences
+                    raise FrogmlException(
+                        exception_message + f". Error is: {e.details()}."
+                    ) from e
+
+                raise e
+
+        return _inner_wrapper
+
+    return decorator

core/inner/tool/protobuf_factory.py

@@ -0,0 +1,45 @@
+from frogml.core.exceptions import FrogmlException
+
+
+def protobuf_factory(
+    protobuf_class, exclude_fields=None, include_only=None, mapping=None
+):
+    def get_final_field_names(cls):
+        fields = dict(cls.__dataclass_fields__)
+        all_field_names = set(fields.keys())
+        excluded_fields = set(exclude_fields) if exclude_fields else set()
+        included_fields = set(include_only) if include_only else all_field_names
+        unrecognized_fields = included_fields.difference(all_field_names)
+        if unrecognized_fields:
+            raise FrogmlException(
+                f"Unknown fields were included: {list(unrecognized_fields)}"
+            )
+        return included_fields.difference(excluded_fields)
+
+    def protofy_decorator(cls):
+        if not hasattr(cls, "__dataclass_fields__"):
+            raise FrogmlException("Protofy must receive a dataclass")
+
+        fields_mapping = mapping if mapping else {}
+        final_field_names = get_final_field_names(cls)
+
+        def to_proto(self) -> protobuf_class:
+            final_fields = {}
+            for field_name in final_field_names:
+                field_val = self.__getattribute__(field_name)
+                field_name = (
+                    fields_mapping[field_name]
+                    if field_name in fields_mapping
+                    else field_name
+                )
+                final_fields[field_name] = (
+                    field_val.to_proto()
+                    if hasattr(field_val, "to_proto")
+                    else field_val
+                )
+            return protobuf_class(**final_fields)
+
+        setattr(cls, "to_proto", to_proto)
+        return cls
+
+    return protofy_decorator

core/inner/tool/run_config/base.py

@@ -0,0 +1,127 @@
+from abc import ABC, abstractmethod
+from logging import Logger
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Tuple
+
+from marshmallow import RAISE, Schema
+from marshmallow_dataclass import class_schema
+from yaml import SafeLoader, dump, load
+
+from frogml.core.exceptions import FrogmlException, LoadConfigurationException
+
+from .utils import ConfigCliMap, rgetattr, rsetattr
+
+ERROR_OPENING_YAML_FILE_MSG = "Error opening the config file"
+ERROR_PARSING_YAML_FILE_MSG = "Error parsing the config file's yaml"
+ERROR_IN_CONFIG_MSG = "Error loading the config"
+
+
+class FrogmlConfigBase(ABC):
+    """Base frogml config object."""
+
+    @property
+    @abstractmethod
+    def _config_mapping(self) -> List[ConfigCliMap]:
+        """Config mapping, Return a list of ConfigCliMap object in order to create the mapping.
+
+        Returns:
+            list: List of ConfigCliMap to apply.
+        """
+        pass
+
+    def merge_cli_argument(
+        self, sections: Tuple[str, ...] = (), **kwargs: Dict[str, Any]
+    ):
+        """Merge and validate cli arguments by supplied mapping.
+
+        Args:
+            sections: Sections to validate.
+            **kwargs: argument from cli.
+
+        Raises:
+            FrogmlException: In case that the argument is not valid.
+        """
+        for prop_map in self._config_mapping:
+            value = kwargs.get(prop_map.key)
+            if value is not None:
+                if isinstance(value, (list, tuple)):
+                    new_value = list(rgetattr(self, prop_map.prop))
+                    new_value.extend(value)
+                    value = new_value
+                rsetattr(self, prop_map.prop, value)
+            if (
+                not sections
+                or any(
+                    list(
+                        map(
+                            lambda section, _prop_map=prop_map: _prop_map.prop.startswith(
+                                section
+                            ),
+                            sections,
+                        )
+                    )
+                )
+                or "." not in prop_map.prop
+            ):
+                config_value = rgetattr(self, prop_map.prop)
+                if not prop_map.validation_func(config_value, prop_map.is_required):
+                    raise FrogmlException(
+                        f"{prop_map.key} argument contain invalid argument: "
+                        f"{value or config_value}"
+                    )
+        self._post_merge_cli()
+
+    @abstractmethod
+    def _post_merge_cli(self):
+        """Actions to perform after merging cli argument in to properties"""
+        pass
+
+
+class YamlConfigMixin(object):
+    @classmethod
+    def from_yaml(
+        cls, yaml_path: str, unknown: str = RAISE, logger: Optional[Logger] = None
+    ) -> Any:
+        """Create instance of class from yaml and class scheme.
+
+        Args:
+            yaml_path (str): Yaml path.
+
+        Returns:
+            object: Instance of created class.
+        """
+        if not yaml_path:
+            return cls()
+        Schema.Meta.unknown = unknown
+        schema = class_schema(cls)
+        try:
+            yaml_content = Path(yaml_path).read_text()
+        except Exception as e:
+            if logger:
+                logger.error(f"Error opening the config file {yaml_path}")
+            raise LoadConfigurationException(ERROR_OPENING_YAML_FILE_MSG) from e
+
+        try:
+            yaml_parsed = load(stream=yaml_content, Loader=SafeLoader)
+        except Exception as e:
+            if logger:
+                logger.exception(ERROR_PARSING_YAML_FILE_MSG, e)
+            raise LoadConfigurationException(ERROR_PARSING_YAML_FILE_MSG) from e
+
+        try:
+            return schema(unknown=unknown).load(yaml_parsed)
+        except Exception as e:
+            if logger:
+                logger.exception(ERROR_IN_CONFIG_MSG, e)
+            raise LoadConfigurationException(ERROR_IN_CONFIG_MSG) from e
+
+    def to_yaml(self) -> str:
+        """Convert class by scheme to yaml.
+
+        Returns:
+            str: Class as yaml string by scheme.
+        """
+        loaded_type = type(self)
+        schema = class_schema(loaded_type)
+
+        return dump(schema().dump(self))

core/model_loggers/artifact_logger.py

@@ -0,0 +1,128 @@
+import logging
+from datetime import timedelta
+from typing import Optional
+import requests
+from requests import Response
+
+from frogml._proto.qwak.builds.build_url_pb2 import BuildVersioningTagsType
+from frogml.core.clients.build_orchestrator.client import BuildOrchestratorClient
+from frogml.core.clients.file_versioning.client import FileVersioningManagementClient
+from frogml.core.exceptions import FrogmlException
+from frogml.core.inner.model_loggers_utils import (
+    fetch_build_id,
+    upload_data,
+    validate_model,
+    validate_tag,
+)
+from frogml._proto.qwak.builds.builds_orchestrator_service_pb2 import (
+    GetBuildVersioningUploadURLResponse,
+    GetBuildVersioningDownloadURLResponse,
+)
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+MAX_CHUNK_SIZE = 8 * 1_024
+
+
+def log_file(
+    from_path: str,
+    tag: str = "artifact",
+    model_id: Optional[str] = None,
+    build_id: Optional[str] = None,
+) -> None:
+    """
+    Log a file by a given tag
+
+    Args:
+        from_path: file path to log
+        tag: tag to save the file with
+        model_id: optional model id to save data with - if not given found from environment
+        build_id: optional build id - if not given found from environment.
+    """
+    if not validate_tag(tag):
+        raise FrogmlException(
+            "Tag should contain only letters, numbers, underscore or hyphen"
+        )
+
+    model_id = validate_model(model_id)
+    if not build_id:
+        # Checking if called inside a model - then build id saved as environment variable or stays
+        build_id = fetch_build_id()
+
+    upload_url_response: (
+        GetBuildVersioningUploadURLResponse
+    ) = BuildOrchestratorClient().get_build_versioning_upload_url(
+        build_id=build_id,
+        model_id=model_id,
+        tag=tag,
+        tag_type=BuildVersioningTagsType.FILE_TAG_TYPE,
+    )
+
+    FileVersioningManagementClient().register_file_tag(
+        model_id, tag, from_path, build_id
+    )
+
+    with open(from_path, mode="rb") as f:
+        upload_data(
+            upload_url_response.upload_url,
+            f.read(),
+            upload_url_response.headers,
+        )
+
+
+def load_file(
+    to_path: str,
+    tag: str = "artifact",
+    model_id: Optional[str] = None,
+    build_id: Optional[str] = None,
+) -> str:
+    """
+    Load a file by a given tag
+
+    Args:
+        to_path: the local path the downloaded file will be written to
+        tag: load the artifact with this given tag
+        model_id: optional model id to save data with - if not given found from environment
+        build_id: optional build id - if not given found from environment.
+
+    Returns:
+        the path to the newly created data file
+    """
+    logger.info(f"Loading file to {to_path}")
+    if not validate_tag(tag):
+        raise FrogmlException(
+            "Tag should contain only letters, numbers, underscore or hyphen"
+        )
+
+    model_id = validate_model(model_id)
+    download_url_response: (
+        GetBuildVersioningDownloadURLResponse
+    ) = BuildOrchestratorClient().get_build_versioning_download_url(
+        build_id=build_id,
+        model_id=model_id,
+        tag=tag,
+        tag_type=BuildVersioningTagsType.FILE_TAG_TYPE,
+    )
+
+    try:
+        response: Response = requests.get(
+            download_url_response.download_url,
+            headers=download_url_response.headers,
+            stream=True,
+            timeout=(
+                timedelta(seconds=10).total_seconds(),  # timeout to connect
+                timedelta(minutes=20).total_seconds(),  # timeout to read
+            ),
+        )
+        logger.info(f"Downloading file finished with status {response.status_code}")
+        response.raise_for_status()
+
+        with open(to_path, "wb") as f:
+            for chunk in response.iter_content(chunk_size=MAX_CHUNK_SIZE):
+                if chunk:
+                    f.write(chunk)
+
+        return to_path
+    except Exception as error:
+        raise FrogmlException(f"Unable to load save artifact locally: {str(error)}")