hafnia 0.1.24__py3-none-any.whl → 0.1.26__py3-none-any.whl

hafnia/http.py

@@ -20,7 +20,7 @@ def fetch(endpoint: str, headers: Dict, params: Optional[Dict] = None) -> Dict:
         json.JSONDecodeError: On invalid JSON response
     """
     params = {} if params is None else params
-    http = urllib3.PoolManager(timeout=5.0, retries=urllib3.Retry(3))
+    http = urllib3.PoolManager(retries=urllib3.Retry(3))
     try:
         response = http.request("GET", endpoint, fields=params, headers=headers)
         if response.status != 200:
@@ -48,7 +48,7 @@ def post(endpoint: str, headers: Dict, data: Union[Path, Dict, bytes], multipart
         json.JSONDecodeError: If response isn't valid JSON
         ValueError: If data type is unsupported
     """
-    http = urllib3.PoolManager(timeout=5.0, retries=urllib3.Retry(3))
+    http = urllib3.PoolManager(retries=urllib3.Retry(3))
     try:
         if multipart:
             # Remove content-type header if present as urllib3 will set it

hafnia/log.py

@@ -1,32 +1,23 @@
 import logging
+import os
 
-from hafnia import __package_name__
-
-
-class CustomFormatter(logging.Formatter):
-    log_format = "%(asctime)s - %(name)s:%(filename)s @ %(lineno)d - %(levelname)s - %(message)s"
+from rich.logging import RichHandler
 
-    def format(self, record):
-        formatter = logging.Formatter(self.log_format)
-        return formatter.format(record)
-
-
-def create_logger() -> logging.Logger:
-    root_logger = logging.getLogger(__package_name__)
-    if root_logger.hasHandlers():
-        return root_logger
+from hafnia import __package_name__
 
-    ch = logging.StreamHandler()
-    ch.setLevel(logging.INFO)
-    ch.setFormatter(CustomFormatter())
+system_handler = RichHandler(rich_tracebacks=True, show_path=True, show_level=True)
+user_handler = RichHandler(rich_tracebacks=False, show_path=False, show_level=False, log_time_format="[%X]")
 
-    root_logger.propagate = False
-    for handler in root_logger.handlers:
-        root_logger.removeHandler(handler)
 
-    root_logger.addHandler(ch)
-    root_logger.setLevel(logging.INFO)
-    return root_logger
+def create_logger(handler: RichHandler, name: str, log_level: str) -> logging.Logger:
+    logger = logging.getLogger(name)
+    if logger.hasHandlers():
+        logger.handlers.clear()
+    logger.addHandler(handler)
+    logger.setLevel(log_level)
+    logger.propagate = False
+    return logger
 
 
-logger = create_logger()
+sys_logger = create_logger(system_handler, f"{__package_name__}.system", os.getenv("HAFNIA_LOG", "INFO").upper())
+user_logger = create_logger(user_handler, f"{__package_name__}.user", "DEBUG")

hafnia/platform/__init__.py

@@ -1,4 +1,3 @@
-from hafnia.platform.api import get_organization_id
 from hafnia.platform.download import (
     download_resource,
     download_single_object,
@@ -12,7 +11,6 @@ from hafnia.platform.experiment import (
 )
 
 __all__ = [
-    "get_organization_id",
     "get_dataset_id",
     "create_recipe",
     "get_exp_environment_id",

hafnia/platform/builder.py

@@ -1,184 +1,144 @@
 import json
 import os
+import subprocess
+import tempfile
+import zipfile
 from hashlib import sha256
 from pathlib import Path
-from shutil import rmtree
-from typing import Dict, List, Optional
-from zipfile import ZipFile
+from typing import Dict, Optional
 
 import boto3
 from botocore.exceptions import ClientError
 
-from hafnia.log import logger
+from hafnia.log import sys_logger, user_logger
 from hafnia.platform import download_resource
 
 
-def validate_recipe(zip_path: Path, required_paths: Optional[set] = None) -> None:
-    """
-    Validates the structure of a zip archive.
-    Ensures the presence of specific files and directories.
-
-    Args:
-        zip_path (Path): Path to the zip archive.
-        required_paths (set): A set of required paths relative to the archive root.
-
-    Raises:
-        FileNotFoundError: If any required file or directory is missing.
-    """
-    required_paths = {"src/lib/", "src/scripts/", "Dockerfile"} if required_paths is None else required_paths
-    with ZipFile(zip_path, "r") as archive:
-        archive_contents = {Path(file).as_posix() for file in archive.namelist()}
-        missing_paths = {
-            path for path in required_paths if not any(entry.startswith(path) for entry in archive_contents)
-        }
+def validate_recipe_format(path: Path) -> None:
+    """Validate Hafnia Recipe Format submition"""
+    hrf = zipfile.Path(path) if path.suffix == ".zip" else path
+    required = {"src", "scripts", "Dockerfile"}
+    errors = 0
+    for rp in required:
+        if not (hrf / rp).exists():
+            user_logger.error(f"Required path {rp} not found in recipe.")
+            errors += 1
+    if errors > 0:
+        raise FileNotFoundError("Wrong recipe structure")
 
-        if missing_paths:
-            raise FileNotFoundError(f"The following required paths are missing in the zip archive: {missing_paths}")
-
-        script_files = [f for f in archive_contents if f.startswith("src/scripts/") and f.endswith(".py")]
-
-        if not script_files:
-            raise ValueError("No Python script files found in the 'src/scripts/' directory.")
-
-
-def clean_up(files: List[Path], dirs: List[Path], prefix: str = "__") -> None:
-    """
-    Clean up a list of files first, and then remove all folders starting with a specific prefix.
-
-    Args:
-        paths (list[Path]): List of file and directory paths to clean up.
-        prefix (str, optional): Prefix to match for folder removal. Defaults to "__".
-    """
-    for path in files:
-        if path.exists() and path.is_file():
-            path.unlink()
-
-    for path in dirs:
-        if path.exists() and path.is_dir():
-            for sub_dir in path.glob(f"**/{prefix}*"):
-                if sub_dir.is_dir():
-                    rmtree(sub_dir)
-
-
-def get_recipe_content(recipe_url: str, output_dir: Path, state_file: str, api_key: str) -> Dict:
-    """
-    Retrieves and validates the recipe content from an S3 location and extracts it.
-
-    Args:
-        recipe_uuid (str): The unique identifier of the recipe.
-        output_dir (str): Directory to extract the recipe content.
-        state_file (str): File to save the state information.
 
-    Returns:
-        Dict: Metadata about the recipe for further processing.
-    """
-    result = download_resource(recipe_url, output_dir, api_key)
-    recipe_path = Path(result["downloaded_files"][0])
-
-    validate_recipe(recipe_path)
-
-    with ZipFile(recipe_path, "r") as zip_ref:
+def prepare_recipe(recipe_url: str, output_dir: Path, api_key: str, state_file: Optional[Path] = None) -> Dict:
+    resource = download_resource(recipe_url, output_dir.as_posix(), api_key)
+    recipe_path = Path(resource["downloaded_files"][0])
+    with zipfile.ZipFile(recipe_path, "r") as zip_ref:
         zip_ref.extractall(output_dir)
 
-    tag = sha256(recipe_path.read_bytes()).hexdigest()[:8]
+    validate_recipe_format(output_dir)
 
-    scripts_dir = output_dir / "src/scripts"
-    valid_commands = [str(f.name)[:-3] for f in scripts_dir.iterdir() if f.is_file() and f.suffix.lower() == ".py"]
+    scripts_dir = output_dir / "scripts"
+    if not any(scripts_dir.iterdir()):
+        user_logger.warning("Scripts folder is empty")
 
-    if not valid_commands:
-        raise ValueError("No valid Python script commands found in the 'src/scripts' directory.")
-
-    state = {
+    metadata = {
         "user_data": (output_dir / "src").as_posix(),
-        "docker_context": output_dir.as_posix(),
         "dockerfile": (output_dir / "Dockerfile").as_posix(),
-        "docker_tag": f"runtime:{tag}",
-        "hash": tag,
-        "valid_commands": valid_commands,
+        "digest": sha256(recipe_path.read_bytes()).hexdigest()[:8],
     }
+    state_file = state_file if state_file else output_dir / "state.json"
+    with open(state_file, "w", encoding="utf-8") as f:
+        json.dump(metadata, f)
+    return metadata
 
-    try:
-        with open(state_file, "w", encoding="utf-8") as f:
-            json.dump(state, f)
-    except Exception as e:
-        raise RuntimeError(f"Failed to write state file: {e}")
 
-    clean_up([recipe_path], [output_dir])
-
-    return state
+def buildx_available() -> bool:
+    try:
+        result = subprocess.run(["docker", "buildx", "version"], capture_output=True, text=True, check=True)
+        return "buildx" in result.stdout.lower()
+    except (subprocess.CalledProcessError, FileNotFoundError):
+        return False
 
 
-def build_dockerfile(dockerfile: str, docker_context: str, docker_tag: str, secrets: Optional[Dict] = None) -> None:
+def build_dockerfile(dockerfile: str, docker_context: str, docker_tag: str, meta_file: str) -> None:
     """
     Build a Docker image using the provided Dockerfile.
 
     Args:
-        dockerfile (Path): Path to the Dockerfile.
+        dockerfile (str): Path to the Dockerfile.
         docker_context (str): Path to the build context.
         docker_tag (str): Tag for the Docker image.
-        secrets (dict, optional): Dictionary of secrets to pass to docker build.
-            Each key-value pair will be passed as --secret id=key,env=value
+        meta_file (Optional[str]): File to store build metadata.
     """
-
-    import subprocess
-
     if not Path(dockerfile).exists():
         raise FileNotFoundError("Dockerfile not found.")
-    build_cmd = [
-        "docker",
-        "build",
-        "--platform=linux/amd64",
-        "-t",
-        docker_tag,
-        "-f",
-        dockerfile,
-    ]
-    build_cmd.append(docker_context)
-    logger.info(f"Building Docker image: {' '.join(build_cmd)}")
-    subprocess.run(build_cmd, check=True)
-
-
-def check_ecr(repository_name: str, image_tag: str) -> bool:
-    aws_region = os.getenv("AWS_REGION", None)
-    if aws_region is None:
-        logger.warning("ECR registry region is not provided can not look up in the registry.")
-        return False
-    session = boto3.Session(region_name=aws_region)
-    ecr_client = session.client("ecr")
-    try:
-        response = ecr_client.describe_images(repositoryName=repository_name, imageIds=[{"imageTag": image_tag}])
-        if response["imageDetails"]:
-            logger.info(f"Image {image_tag} already exists in ECR.")
-            return True
-        else:
-            return False
-    except ClientError as e:
-        if e.response["Error"]["Code"] == "ImageNotFoundException":
-            logger.info(f"Image {image_tag} does not exist in ECR.")
-            return False
-        else:
-            raise e
 
+    cmd = ["docker", "build", "--platform", "linux/amd64", "-t", docker_tag, "-f", dockerfile]
+
+    remote_cache = os.getenv("REMOTE_CACHE_REPO")
+    cloud_mode = os.getenv("HAFNIA_CLOUD", "false").lower() in ["true", "1", "yes"]
+
+    if buildx_available():
+        cmd.insert(1, "buildx")
+        cmd += ["--build-arg", "BUILDKIT_INLINE_CACHE=1", "--metadata-file", meta_file]
+        if cloud_mode:
+            cmd += ["--push"]
+        if remote_cache:
+            cmd += [
+                "--cache-from",
+                f"type=registry,ref={remote_cache}:buildcache",
+                "--cache-to",
+                f"type=registry,ref={remote_cache}:buildcache,mode=max",
+            ]
+    cmd.append(docker_context)
+    sys_logger.debug("Build cmd: `{}`".format(" ".join(cmd)))
+    sys_logger.info(f"Building and pushing Docker image with BuildKit (buildx); cache repo: {remote_cache or 'none'}")
+    try:
+        subprocess.run(cmd, check=True)
+    except subprocess.CalledProcessError as e:
+        sys_logger.error(f"Docker build failed: {e}")
+        raise RuntimeError(f"Docker build failed: {e}")
 
-def prepare_recipe(recipe_url: str, output_dir: Path, api_key: str) -> Dict:
-    state_file = output_dir / "state.json"
-    get_recipe_content(recipe_url, output_dir, state_file.as_posix(), api_key)
-    with open(state_file.as_posix(), "r") as f:
-        return json.loads(f.read())
 
+def check_registry(docker_image: str) -> Optional[str]:
+    """
+    Returns the remote digest for TAG if it exists, otherwise None.
+    """
+    if "localhost" in docker_image:
+        return None
+
+    region = os.getenv("AWS_REGION")
+    if not region:
+        sys_logger.warning("AWS_REGION environment variable not set. Skip image exist check.")
+        return None
+
+    repo_name, image_tag = docker_image.rsplit(":")
+    if "/" in repo_name:
+        repo_name = repo_name.rsplit("/", 1)[-1]
+    ecr = boto3.client("ecr", region_name=region)
+    try:
+        out = ecr.describe_images(repositoryName=repo_name, imageIds=[{"imageTag": image_tag}])
+        return out["imageDetails"][0]["imageDigest"]
+    except ClientError as e:
+        error_code = e.response["Error"]["Code"]
+        sys_logger.error(f"ECR client error: {error_code}")
+        return None
 
-def build_image(image_info: Dict, ecr_prefix: str, state_file: str = "state.json") -> None:
-    hafnia_tag = f"{ecr_prefix}/{image_info['name']}:{image_info['hash']}"
-    image_exists = False
-    if "localhost" not in ecr_prefix:
-        image_exists = check_ecr(image_info["name"], image_info["hash"])
 
-    image_info.update({"mdi_tag": hafnia_tag, "image_exists": image_exists})
-    state_path = Path(state_file)
-    state_path.parent.mkdir(parents=True, exist_ok=True)
+def build_image(metadata: Dict, registry_repo: str, state_file: str = "state.json") -> None:
+    docker_image = f"{registry_repo}:{metadata['digest']}"
+    image_exists = check_registry(docker_image) is not None
     if image_exists:
-        logger.info(f"Image {hafnia_tag} already exists in ECR. Skipping build.")
+        sys_logger.info(f"Tag already in ECR – skipping build of {docker_image}.")
     else:
-        build_dockerfile(image_info["dockerfile"], image_info["docker_context"], hafnia_tag)
-    with open(state_path.as_posix(), "w") as f:
-        json.dump(image_info, f, indent=4)
+        with tempfile.NamedTemporaryFile() as meta_tmp:
+            meta_file = meta_tmp.name
+            build_dockerfile(
+                metadata["dockerfile"], Path(metadata["dockerfile"]).parent.as_posix(), docker_image, meta_file
+            )
+            with open(meta_file) as m:
+                try:
+                    build_meta = json.load(m)
+                    metadata["local_digest"] = build_meta["containerimage.digest"]
+                except Exception:
+                    metadata["local_digest"] = ""
+    metadata.update({"image_tag": docker_image, "image_exists": image_exists})
+    Path(state_file).write_text(json.dumps(metadata, indent=2))

hafnia/platform/download.py

@@ -6,7 +6,7 @@ from botocore.exceptions import ClientError
 from tqdm import tqdm
 
 from hafnia.http import fetch
-from hafnia.log import logger
+from hafnia.log import sys_logger, user_logger
 
 
 def get_resource_creds(endpoint: str, api_key: str) -> Dict[str, Any]:
@@ -30,11 +30,11 @@ def get_resource_creds(endpoint: str, api_key: str) -> Dict[str, Any]:
         RuntimeError: If the call to fetch the credentials fails for any reason.
     """
     try:
-        creds = fetch(endpoint, headers={"X-APIKEY": api_key, "accept": "application/json"})
-        logger.debug("Successfully retrieved credentials from DIP endpoint.")
+        creds = fetch(endpoint, headers={"Authorization": api_key, "accept": "application/json"})
+        sys_logger.debug("Successfully retrieved credentials from DIP endpoint.")
         return creds
     except Exception as e:
-        logger.error(f"Failed to fetch credentials from endpoint: {e}")
+        sys_logger.error(f"Failed to fetch credentials from endpoint: {e}")
         raise RuntimeError(f"Failed to retrieve credentials: {e}") from e
 
 
@@ -99,12 +99,12 @@ def download_resource(resource_url: str, destination: str, api_key: str) -> Dict
         s3_client.head_object(Bucket=bucket_name, Key=key)
         local_file = download_single_object(s3_client, bucket_name, key, output_path)
         downloaded_files.append(str(local_file))
-        logger.info(f"Downloaded single file: {local_file}")
+        user_logger.info(f"Downloaded single file: {local_file}")
 
     except ClientError as e:
         error_code = e.response.get("Error", {}).get("Code")
         if error_code == "404":
-            logger.debug(f"Object '{key}' not found; trying as a prefix.")
+            sys_logger.debug(f"Object '{key}' not found; trying as a prefix.")
             response = s3_client.list_objects_v2(Bucket=bucket_name, Prefix=key)
             contents = response.get("Contents", [])
 
@@ -118,9 +118,9 @@ def download_resource(resource_url: str, destination: str, api_key: str) -> Dict
                 local_file = download_single_object(s3_client, bucket_name, sub_key, output_path)
                 downloaded_files.append(local_file.as_posix())
 
-            logger.info(f"Downloaded folder/prefix '{key}' with {len(downloaded_files)} object(s).")
+            user_logger.info(f"Downloaded folder/prefix '{key}' with {len(downloaded_files)} object(s).")
         else:
-            logger.error(f"Error checking object or prefix: {e}")
+            user_logger.error(f"Error checking object or prefix: {e}")
             raise RuntimeError(f"Failed to check or download S3 resource: {e}") from e
 
     return {"status": "success", "downloaded_files": downloaded_files}

hafnia/platform/experiment.py

@@ -2,45 +2,53 @@ from pathlib import Path
 from typing import Optional
 
 from hafnia.http import fetch, post
-from hafnia.platform.builder import validate_recipe
-from hafnia.utils import archive_dir, get_recipe_path
+from hafnia.log import user_logger
+from hafnia.utils import archive_dir, get_recipe_path, timed
 
 
-def get_dataset_id(dataset_name: str, endpoint: str, api_key: str) -> Optional[str]:
-    headers = {"X-APIKEY": api_key}
+@timed("Fetching dataset info.")
+def get_dataset_id(dataset_name: str, endpoint: str, api_key: str) -> str:
+    headers = {"Authorization": api_key}
     full_url = f"{endpoint}?name__iexact={dataset_name}"
     dataset_info = fetch(full_url, headers=headers)
     if not dataset_info:
         raise ValueError(f"Dataset '{dataset_name}' was not found in the dataset library.")
-    return dataset_info[0]["id"]
+    try:
+        return dataset_info[0]["id"]
+    except (IndexError, KeyError) as e:
+        raise ValueError("Dataset information is missing or invalid") from e
 
 
-def create_recipe(source_dir: Path, endpoint: str, api_key: str, organization_id: str) -> Optional[str]:
-    headers = {"X-APIKEY": api_key, "accept": "application/json"}
+@timed("Uploading recipe.")
+def create_recipe(source_dir: Path, endpoint: str, api_key: str) -> str:
     source_dir = source_dir.resolve()  # Ensure the path is absolute to handle '.' paths are given an appropriate name.
     path_recipe = get_recipe_path(recipe_name=source_dir.name)
     zip_path = archive_dir(source_dir, output_path=path_recipe)
+    user_logger.info(f"Recipe created and stored in '{path_recipe}'")
 
-    print(f"Recipe created and stored in '{path_recipe}'")
-
-    validate_recipe(zip_path)
-    with open(zip_path, "rb") as zip_file:
-        fields = {
-            "name": path_recipe.name,
-            "description": "Recipe created by Hafnia CLI",
-            "organization": organization_id,
-            "file": (zip_path.name, zip_file.read()),
-        }
-        response = post(endpoint, headers=headers, data=fields, multipart=True)
-        return response["id"]
+    headers = {"Authorization": api_key, "accept": "application/json"}
+    data = {
+        "name": path_recipe.name,
+        "description": "Recipe created by Hafnia CLI",
+        "file": (zip_path.name, Path(zip_path).read_bytes()),
+    }
+    response = post(endpoint, headers=headers, data=data, multipart=True)
+    return response["id"]
 
 
-def get_exp_environment_id(name: str, endpoint: str, api_key: str) -> Optional[str]:
-    headers = {"X-APIKEY": api_key}
+@timed("Fetching environment info.")
+def get_exp_environment_id(name: str, endpoint: str, api_key: str) -> str:
+    headers = {"Authorization": api_key}
     env_info = fetch(endpoint, headers=headers)
-    return next((env["id"] for env in env_info if env["name"] == name), None)
+
+    for env in env_info:
+        if env["name"] == name:
+            return env["id"]
+
+    raise ValueError(f"Environment '{name}' not found")
 
 
+@timed("Creating exepriment.")
 def create_experiment(
     exp_name: str,
     dataset_id: str,
@@ -49,14 +57,12 @@ def create_experiment(
     environment_id: str,
     endpoint: str,
     api_key: str,
-    organization_id: str,
 ) -> Optional[str]:
-    headers = {"X-APIKEY": api_key}
+    headers = {"Authorization": api_key}
     response = post(
         endpoint,
         headers=headers,
         data={
-            "organization": organization_id,
             "name": exp_name,
             "recipe": recipe_id,
             "dataset": dataset_id,