qwak-core 0.4.265__py3-none-any.whl → 0.4.266__py3-none-any.whl

frogml_storage/_artifactory_api.py

@@ -0,0 +1,315 @@
+import json
+import os
+from typing import Optional
+
+from requests import Response
+from tqdm.auto import tqdm
+from tqdm.utils import CallbackIOWrapper
+from urllib3 import Retry
+
+from frogml_storage._log_config import logger
+from frogml_storage._utils import join_url
+from frogml_storage.constants import CHECKSUM_SHA2_HEADER
+from frogml_storage.entity_manifest import Checksums, EntityManifest
+from frogml_storage.frogml_entity_type_info import FrogMLEntityTypeInfo
+from frogml_storage.http.http_client import HTTPClient
+from frogml_storage.models._download_context import DownloadContext
+
+
+class StartTransactionResponse:
+    files_upload_path: str
+    lead_upload_path: str
+    dependencies_upload_path: str
+    code_upload_path: str
+    transaction_id: str
+
+    def __init__(
+        self,
+        files_upload_path,
+        lead_upload_path,
+        dependencies_upload_path,
+        code_upload_path,
+        transaction_id,
+    ):
+        self.files_upload_path = files_upload_path
+        self.lead_upload_path = lead_upload_path
+        self.dependencies_upload_path = dependencies_upload_path
+        self.code_upload_path = code_upload_path
+        self.transaction_id = transaction_id
+
+
+class ArtifactoryApi:
+    def __init__(self, uri, auth=None, http_client=None):
+        self.uri = uri
+        if http_client is not None:
+            self.http_client = http_client
+        else:
+            self.auth = auth
+            self.http_client = HTTPClient(auth=auth)
+
+    def start_transaction(
+        self,
+        entity_type_info: FrogMLEntityTypeInfo,
+        repository: str,
+        entity_name: str,
+        version: Optional[str],
+    ) -> StartTransactionResponse:
+        """
+        Initializes an upload. Returns transaction ID and upload path
+        """
+        if version is None:
+            start_transaction_url = (
+                f"{self.uri}/api/machinelearning/{repository}/"
+                f"{entity_type_info.entity_type}/{entity_name}/start-transaction"
+            )
+        else:
+            start_transaction_url = (
+                f"{self.uri}/api/machinelearning/{repository}/{entity_type_info.entity_type}"
+                f"/{entity_name}/start-transaction/{version}"
+            )
+        try:
+            response = self.http_client.post(start_transaction_url)
+            response.raise_for_status()
+            files_upload_path = response.json()["filesUploadPath"]
+            lead_upload_path = response.json()["leadUploadPath"]
+            dependencies_upload_path = response.json()["dependenciesUploadPath"]
+            code_upload_path = response.json()["codeUploadPath"]
+            transaction_id = response.json()["transactionId"]
+        except Exception as exception:
+            err = (
+                f"Error occurred while trying to start an upload transaction for "
+                f"{entity_type_info.entity_type}: '{entity_name}'"
+                f" Error: '{exception}'"
+            )
+            logger.error(err, exc_info=False)
+            raise exception
+        return StartTransactionResponse(
+            files_upload_path=files_upload_path,
+            lead_upload_path=lead_upload_path,
+            dependencies_upload_path=dependencies_upload_path,
+            code_upload_path=code_upload_path,
+            transaction_id=transaction_id,
+        )
+
+    def end_transaction(
+        self,
+        entity_type: FrogMLEntityTypeInfo,
+        repository: str,
+        entity_name: str,
+        entity_manifest: EntityManifest,
+        transaction_id: str,
+        version: str,
+        properties: Optional[dict[str, str]],
+    ) -> None:
+        """
+        Upload model-manifest.json | dataset-manifest.json file, makes the model | dataset available in the repository
+        """
+        filename = entity_type.metadata_file_name
+
+        url = join_url(
+            self.uri,
+            "api",
+            "machinelearning",
+            repository,
+            entity_type.entity_type,
+            "entity-manifest",
+            entity_name,
+            version,
+            transaction_id,
+        )
+
+        json_entity_manifest = entity_manifest.to_json()
+        self.upload_entity_manifest(
+            entity_type=entity_type,
+            filename=filename,
+            payload=json_entity_manifest,
+            url=url,
+            properties=properties,
+        )
+
+    def download_file(self, args: DownloadContext) -> None:
+        filename = os.path.basename(args.target_path)
+        try:
+            url = f"{self.uri}/{args.repo_key}/{args.source_url}"
+            with self.http_client.get(url=url, stream=True) as response:
+                response.raise_for_status()
+                total_size = int(response.headers.get("content-length", 0))
+                with open(args.target_path, "wb") as file:
+                    with self.__initialize_progress_bar(total_size, filename) as pbar:
+                        for chunk in response.iter_content(chunk_size=8192):
+                            if chunk:
+                                file.write(chunk)
+                                pbar.update(len(chunk))
+
+        except Exception as exception:
+            self.__handle_download_exception(exception, args.target_path, filename)
+
+    def get_entity_manifest(
+        self,
+        entity_type_info: FrogMLEntityTypeInfo,
+        repository: str,
+        entity_name: str,
+        namespace: Optional[str],
+        version: Optional[str],
+    ) -> dict:
+        url = join_url(
+            self.uri,
+            "api",
+            "machinelearning",
+            repository,
+            entity_type_info.entity_type,
+            "entity-manifest",
+            namespace,
+            entity_name,
+            version,
+        )
+        try:
+            with self.http_client.get(url=url) as r:
+                r.raise_for_status()
+                return r.json()
+        except Exception as exception:
+            err = f"Error occurred while trying to get {entity_type_info.entity_type} info file. Error: '{exception}'"
+            logger.error(err, exc_info=False)
+            raise exception
+
+    @staticmethod
+    def __handle_download_exception(
+        exception: Exception, target_path: str, filename: str
+    ) -> None:
+        if os.path.exists(target_path):
+            os.remove(target_path)
+        err = f"Error occurred while trying to download file: '{filename}' Error: '{exception}'"
+        logger.error(err, exc_info=False)
+        raise exception
+
+    def get_artifact_checksum(self, download_context: DownloadContext) -> str:
+        url = f"{self.uri}/{download_context.repo_key}/{download_context.source_url}"
+        try:
+            with self.http_client.head(url=url, stream=True) as response:
+                response.raise_for_status()
+                return response.headers.get(CHECKSUM_SHA2_HEADER)
+
+        except Exception as exception:
+            logger.error(exception.__cause__, exc_info=False)
+            raise exception
+
+    def upload_entity_manifest(
+        self,
+        entity_type: FrogMLEntityTypeInfo,
+        filename: str,
+        payload: str,
+        url: str,
+        properties: Optional[dict[str, str]],
+        stream: bool = False,
+    ) -> None:
+        body_part_name = f"{entity_type.body_part_stream}"
+
+        try:
+            files = {
+                f"{body_part_name}": (
+                    f"{body_part_name}",
+                    payload,
+                    "application/octet-stream",
+                ),  # Include the InputStream
+                "additionalData": (
+                    "additionalData",
+                    json.dumps(properties),
+                    "application/octet-stream",
+                ),  # Include the object
+            }
+            with self.http_client.put(url=url, files=files, stream=stream) as response:
+                response.raise_for_status()
+        except Exception as exception:
+            err = f"Error occurred while trying to upload file: '{filename}' Error: '{exception}'"
+            logger.error(err, exc_info=False)
+            raise exception
+
+    def upload_file(self, file_path: str, url: str) -> None:
+        wrapped_file = None
+        try:
+            file_size = os.stat(file_path).st_size
+            with (
+                self.__initialize_progress_bar(file_size, file_path) as pbar,
+                open(file_path, "rb") as file,
+            ):
+                wrapped_file = CallbackIOWrapper(pbar.update, file, "read")
+                with self.http_client.put(url=url, payload=wrapped_file) as response:
+                    response.raise_for_status()
+        except Exception as exception:
+            err = f"Error occurred while trying to upload file: '{file_path}' Error: '{exception}'"
+            logger.error(err, exc_info=False)
+            raise type(exception)(f"{err} File: {file_path}") from exception
+        finally:
+            if wrapped_file is not None:
+                wrapped_file.close()
+
+    def checksum_deployment(
+        self,
+        checksum: Checksums,
+        url: str,
+        full_path: str,
+        stream: bool = False,
+    ) -> bool:
+        response = self.http_client.put(
+            url=url,
+            headers={"X-Checksum-Sha256": checksum.sha2, "X-Checksum-Deploy": "true"},
+            stream=stream,
+        )
+        if response.status_code != 200 and response.status_code != 201:
+            return False
+        else:
+            file_size = os.path.getsize(full_path)
+            pbar = self.__initialize_progress_bar(file_size, full_path)
+            pbar.update(file_size)
+            pbar.close()
+            return True
+
+    @staticmethod
+    def __initialize_progress_bar(total_size: int, filename: str) -> tqdm:
+        return tqdm(
+            total=total_size, unit="B", unit_scale=True, desc=filename, initial=0
+        )
+
+    def encrypt_password(self) -> Response:
+        """
+        returns encrypted password as text
+        """
+        return self.http_client.get(
+            url=join_url(self.uri, "/api/security/encryptedPassword")
+        )
+
+    def ping(self) -> Response:
+        """
+        Sends a ping to Artifactory to validate login status
+        """
+        url = join_url(self.uri, "api/system/ping")
+        return self.http_client.get(url=url)
+
+    def get_artifactory_version(self) -> Response:
+        return self.http_client.get(url=join_url(self.uri, "/api/system/version"))
+
+    def create_machinelearning_local_repo(self, repo_name: str) -> Response:
+        data = {
+            "rclass": "local",
+            "packageType": "machinelearning",
+        }
+        return self.http_client.put(
+            url=join_url(self.uri, "/api/repositories/" + repo_name), json=data
+        )
+
+    def delete_frogml_local_repo(self, repo_name: str) -> Response:
+        return self.http_client.delete(
+            url=join_url(self.uri, "/api/repositories/" + repo_name)
+        )
+
+
+class RetryWithLog(Retry):
+    """
+    Adding extra logs before making a retry request
+    """
+
+    def __init__(self, *args, **kwargs):
+        history = kwargs.get("history")
+        if history is not None:
+            logger.debug(f"Error: ${history[-1].error}\nretrying...")
+        super().__init__(*args, **kwargs)

frogml_storage/_environment.py

@@ -0,0 +1,22 @@
+import importlib.metadata
+import platform
+
+
+class Recorder:
+    @staticmethod
+    def get_environment_dependencies():
+        distributions = importlib.metadata.distributions()
+        return sorted(
+            [f"{dist.metadata['Name']}=={dist.version}" for dist in distributions]
+        )
+
+    @staticmethod
+    def get_environment_details():
+        return [
+            f"arch={platform.architecture()[0]}",
+            f"cpu={platform.processor()}",
+            f"platform={platform.platform()}",
+            f"python_version={platform.python_version()}",
+            f"python_implementation={platform.python_implementation()}",
+            f"python_compiler={platform.python_compiler()}",
+        ]

frogml_storage/_log_config.py

@@ -0,0 +1,45 @@
+import logging.config
+import os
+import sys
+
+log_level = (
+    "DEBUG"
+    if os.getenv("JFML_DEBUG", "false").casefold() == "true".casefold()
+    else "INFO"
+)
+log_file = f'{os.path.expanduser("~")}/.frogml/frogml-log-history.log'
+os.makedirs(os.path.dirname(log_file), exist_ok=True)
+
+DEFAULT_LOGGING = {
+    "version": 1,
+    "formatters": {
+        "standard": {
+            "format": "%(asctime)s - %(levelname)s - %(name)s.%(module)s.%(funcName)s:%(lineno)d - %(message)s"
+        },
+    },
+    "handlers": {
+        "console": {
+            "class": "logging.StreamHandler",
+            "formatter": "standard",
+            "stream": sys.stdout,
+        },
+        "file": {
+            "class": "logging.FileHandler",
+            "formatter": "standard",
+            "filename": log_file,
+        },
+    },
+    "loggers": {
+        __name__: {
+            "level": log_level,
+            "handlers": ["console", "file"],
+            "propagate": False,
+        },
+    },
+}
+
+if os.getenv("IS_LOGGER_SHADED") is not None:
+    logger = logging.getLogger(__name__)
+else:
+    logging.config.dictConfig(DEFAULT_LOGGING)
+    logger = logging.getLogger(__name__)

frogml_storage/_storage_utils.py

@@ -0,0 +1,15 @@
+import hashlib
+
+
+def calculate_sha2(path: str) -> str:
+    sha256 = hashlib.sha256()
+    with open(path, "rb") as f:
+        while chunk := f.read(8192):
+            sha256.update(chunk)
+    return sha256.hexdigest()
+
+
+def calc_content_sha2(content: str) -> str:
+    sha256 = hashlib.sha256()
+    sha256.update(content.encode("utf-8"))
+    return sha256.hexdigest()

frogml_storage/_utils.py

@@ -0,0 +1,69 @@
+from typing import Optional
+from urllib.parse import urlparse
+
+from requests.auth import AuthBase
+
+from frogml_storage.frogml_entity_type_info import FrogMLEntityTypeInfo
+
+
+class BearerAuth(AuthBase):
+    def __init__(self, token):
+        self.token = token
+
+    def __call__(self, r):
+        r.headers["Authorization"] = f"Bearer {self.token}"
+        return r
+
+    def __eq__(self, other):
+        if not isinstance(other, BearerAuth):
+            return False
+        return self.token == other.token
+
+
+class EmptyAuth(AuthBase):
+    def __call__(self, r):
+        return r
+
+
+def join_url(base_uri: str, *parts: Optional[str]) -> str:
+    if base_uri.endswith("/"):
+        base_uri = base_uri[:-1]
+
+    cleaned_parts = [
+        part.strip("/") for part in parts if part is not None and part.strip("/")
+    ]
+    return f"{base_uri}/{'/'.join(cleaned_parts)}"
+
+
+def assemble_artifact_url(uri: Optional[str]) -> str:
+    if uri is None:
+        raise Exception("Artifactory URI is required")
+
+    parsed_url = urlparse(uri)
+    if parsed_url.scheme not in ["http", "https"]:
+        raise Exception(
+            f"Not a valid Artifactory URI: {uri}. "
+            f"Artifactory URI example: `https://frogger.jfrog.io/artifactory/ml-local`"
+        )
+
+    return f"{parsed_url.scheme}://{parsed_url.netloc}/artifactory"
+
+
+# The following method affect e2e tests.
+def build_download_success_log(
+    entity_type_info: FrogMLEntityTypeInfo, entity_name: str, version: str
+) -> str:
+    return (
+        f'{entity_type_info.entity_type.capitalize()}: "{entity_name}", version: "{version}"'
+        f" has been downloaded successfully"
+    )
+
+
+# The following method affect e2e tests.
+def build_upload_success_log(
+    entity_type_info: FrogMLEntityTypeInfo, entity_name: str, version: str
+) -> str:
+    return (
+        f'{entity_type_info.entity_type.capitalize()}: "{entity_name}", version: "{version}"'
+        f" has been uploaded successfully"
+    )