frogml-core 0.0.1__py3-none-any.whl → 0.0.2__py3-none-any.whl

frogml_core/inner/instance_template/verify_template_id.py

@@ -0,0 +1,50 @@
+from typing import Optional
+
+from frogml_proto.qwak.instance_template.instance_template_pb2 import InstanceType
+from frogml_core.clients.instance_template.client import (
+    InstanceTemplateManagementClient,
+)
+from frogml_core.exceptions import FrogmlException
+from frogml_core.inner.provider import Provider
+
+INVALID_TEMPLATE_ID_ERROR_FORMAT = (
+    "Invalid instance: {template_id}. Valid instances are: {cpu_templates}. "
+    "For GPU instances you can use {gpu_templates}."
+)
+
+
+def verify_template_id(
+    template_id: str,
+    instance_template_client: InstanceTemplateManagementClient,
+    provider: Optional[Provider] = None,
+) -> None:
+    all_templates = instance_template_client.list_instance_templates()
+    valid_templates = [template for template in all_templates if template.enabled]
+    if provider == Provider.AWS:
+        valid_templates = [
+            template for template in valid_templates if template.aws_supported
+        ]
+    elif provider == Provider.GCP:
+        valid_templates = [
+            template for template in valid_templates if template.gcp_supported
+        ]
+
+    existing_templates = {template.id: template for template in valid_templates}
+    if template_id not in existing_templates.keys():
+        cpu_template_ids = [
+            template.id
+            for template in sorted(valid_templates, key=lambda template: template.order)
+            if template.instance_type == InstanceType.INSTANCE_TYPE_CPU
+        ]
+        gpu_template_ids = [
+            template.id
+            for template in sorted(valid_templates, key=lambda template: template.order)
+            if template.instance_type == InstanceType.INSTANCE_TYPE_GPU
+        ]
+        raise FrogmlException(
+            INVALID_TEMPLATE_ID_ERROR_FORMAT.format(
+                template_id=template_id,
+                cpu_templates=", ".join(cpu_template_ids),
+                gpu_templates=", ".join(gpu_template_ids),
+            )
+        )

frogml_core/inner/model_loggers_utils.py

@@ -0,0 +1,110 @@
+import os
+import re
+from typing import Optional
+
+import requests
+
+from frogml_proto.qwak.builds.builds_orchestrator_service_pb2 import (
+    AuthenticationDetail,
+)
+from frogml_core.clients.model_management.client import ModelsManagementClient
+from frogml_core.exceptions import FrogmlException
+
+TAG_REGEX = re.compile(r"^[a-zA-Z0-9_-]+$")
+
+
+def validate_tag(tag: str) -> bool:
+    """
+    Check if tag exists.
+
+    Args:
+        tag: tag to check
+
+    Returns: If tag is valid
+    """
+    return re.match(TAG_REGEX, tag) is not None
+
+
+def fetch_model_id() -> Optional[str]:
+    """
+    Get model id from environment.
+
+    Returns: model id if found.
+
+    Notes:
+        1. Checking if called inside a model - then model id saved as environment variable.
+    """
+    # Checking if called inside a model - then model id saved as environment variable
+    return os.getenv("QWAK_MODEL_ID", None)
+
+
+def validate_model(model_id: str) -> str:
+    """
+    Validate a model ID validity and existence
+    """
+    if not model_id:
+        model_id = fetch_model_id()
+        if not model_id:
+            raise FrogmlException("Failed to determined model ID.")
+
+    try:
+        ModelsManagementClient().get_model(model_id=model_id)
+    except Exception:
+        raise FrogmlException("Failed to find model.")
+
+    return model_id
+
+
+def fetch_build_id() -> Optional[str]:
+    """
+    Get Build id from environment
+
+    Returns: Build id if found
+
+    Notes:
+        1. Checking if called inside a model - then build id saved as environment variable.
+    """
+    # Checking if called inside a model - then model id saved as environment variable
+    return os.getenv("QWAK_BUILD_ID", None)
+
+
+def upload_data(
+    upload_url: str,
+    data: bytes,
+    authentication_details: Optional[AuthenticationDetail],
+    content_type: str = "text/plain",
+) -> None:
+    """
+    Upload data
+    Args:
+        upload_url: the url to upload to.
+        data: the data to upload
+        authentication_details: authentication details for upload data
+        content_type: Uploaded content-type
+    """
+    try:
+        auth = None
+        if (
+            authentication_details.WhichOneof("integration_type")
+            == "jfrog_authentication_detail"
+        ):
+            auth = (
+                authentication_details.jfrog_authentication_detail.username,
+                authentication_details.jfrog_authentication_detail.token,
+            )
+
+        http_response = requests.put(  # nosec B113
+            upload_url,
+            data=data,
+            headers={"content-type": content_type},
+            auth=auth,
+        )
+
+        if http_response.status_code not in [200, 201]:
+            raise FrogmlException(
+                f"Failed to upload data. "
+                f"Status: [{http_response.status_code}], "
+                f"reason: [{http_response.reason}]"
+            )
+    except Exception as e:
+        raise FrogmlException(f"Failed to upload data. Error is {e}")

frogml_core/inner/runtime_di/__init__.py

@@ -0,0 +1,16 @@
+from frogml_core.inner.runtime_di.containers import (
+    ContainerLock,
+    FrogmlRuntimeContainer,
+)
+
+
+def wire_runtime():
+    container = FrogmlRuntimeContainer()
+    from frogml_core.model import decorators
+
+    container.wire(
+        packages=[
+            decorators,
+        ]
+    )
+    return container

frogml_core/inner/runtime_di/containers.py

@@ -0,0 +1,19 @@
+from dependency_injector import containers, providers
+
+from frogml_core.model.decorators.impl.api_implementation import (
+    create_decorator_function,
+)
+from frogml_core.model.decorators.impl.timer_implementation import create_frogml_timer
+
+
+class FrogmlRuntimeContainer(containers.DeclarativeContainer):
+    """
+    Frogml Core's Runtime Dependency Injection Container
+    """
+
+    api_decorator_function_creator = providers.Object(create_decorator_function)
+    timer_function_creator = providers.Object(create_frogml_timer)
+
+
+class ContainerLock:
+    locked: bool = False

frogml_core/inner/tool/__init__.py

@@ -0,0 +1 @@
+from .auth import FrogMLAuthClient

frogml_core/inner/tool/auth.py

@@ -0,0 +1,70 @@
+from typing import Dict, Optional, Union, cast
+
+import requests
+from frogml_storage._utils import BearerAuth
+from frogml_storage.authentication._authentication_utils import get_credentials
+from frogml_storage.authentication.models._auth_config import AuthConfig
+from requests.auth import AuthBase, HTTPBasicAuth
+
+from frogml_core.exceptions import FrogmlLoginException
+
+
+class FrogMLAuthClient:
+    def __init__(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) -> str:
+        if not self._token:
+            self.login()
+
+        return cast(str, self._token)
+
+    def get_tenant_id(self) -> str:
+        if not self._tenant_id:
+            self.login()
+
+        return cast(str, self._tenant_id)
+
+    def login(self):
+        artifactory_url, auth = get_credentials(self.auth_config)
+
+        if isinstance(auth, BearerAuth):  # For BearerAuth
+            self._token = auth.token
+
+        self.__get_tenant_id(artifactory_url, auth)
+
+    def __get_tenant_id(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"
+        kwargs: Dict[str, Union[HTTPBasicAuth, str]] = {}
+
+        if isinstance(auth, HTTPBasicAuth):
+            kwargs["auth"] = auth
+        elif isinstance(auth, BearerAuth):
+            kwargs["headers"] = {"Authorization": f"Bearer {auth.token}"}
+        else:
+            raise login_exception
+
+        try:
+            response = requests.get(url, timeout=60, **kwargs)
+            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

frogml_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)

frogml_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)

frogml_core/inner/tool/grpc/grpc_try_wrapping.py

@@ -0,0 +1,16 @@
+import grpc
+
+from frogml_core.exceptions import FrogmlException
+
+
+def grpc_try_catch_wrapper(exception_message):
+    def decorator(function):
+        def _inner_wrapper(*args, **kwargs):
+            try:
+                return function(*args, **kwargs)
+            except grpc.RpcError as e:
+                raise FrogmlException(exception_message + f". Error is: {e.details()}.")
+
+        return _inner_wrapper
+
+    return decorator

frogml_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