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

hafnia/dataset/table_transformations.py

@@ -0,0 +1,183 @@
+from pathlib import Path
+from typing import List, Optional, Type
+
+import polars as pl
+from tqdm import tqdm
+
+from hafnia.dataset import table_transformations
+from hafnia.dataset.dataset_names import (
+    FILENAME_ANNOTATIONS_JSONL,
+    FILENAME_ANNOTATIONS_PARQUET,
+    FieldName,
+)
+from hafnia.dataset.primitives import PRIMITIVE_TYPES
+from hafnia.dataset.primitives.classification import Classification
+from hafnia.dataset.primitives.primitive import Primitive
+from hafnia.log import user_logger
+
+
+def create_primitive_table(
+    samples_table: pl.DataFrame, PrimitiveType: Type[Primitive], keep_sample_data: bool = False
+) -> Optional[pl.DataFrame]:
+    """
+    Returns a DataFrame with objects of the specified primitive type.
+    """
+    column_name = PrimitiveType.column_name()
+    has_primitive_column = (column_name in samples_table.columns) and (
+        samples_table[column_name].dtype == pl.List(pl.Struct)
+    )
+    if not has_primitive_column:
+        return None
+
+    # Remove frames without objects
+    remove_no_object_frames = samples_table.filter(pl.col(column_name).list.len() > 0)
+
+    if keep_sample_data:
+        # Drop other primitive columns to avoid conflicts
+        drop_columns = set(PRIMITIVE_TYPES) - {PrimitiveType, Classification}
+        remove_no_object_frames = remove_no_object_frames.drop(*[primitive.column_name() for primitive in drop_columns])
+        # Rename columns "height", "width" and "meta" for sample to avoid conflicts with object fields names
+        remove_no_object_frames = remove_no_object_frames.rename(
+            {"height": "image.height", "width": "image.width", "meta": "image.meta"}
+        )
+        objects_df = remove_no_object_frames.explode(column_name).unnest(column_name)
+    else:
+        objects_df = remove_no_object_frames.select(pl.col(column_name).explode().struct.unnest())
+    return objects_df
+
+
+def filter_table_for_class_names(
+    samples_table: pl.DataFrame, class_names: List[str], PrimitiveType: Type[Primitive]
+) -> Optional[pl.DataFrame]:
+    table_with_selected_class_names = samples_table.filter(
+        pl.col(PrimitiveType.column_name())
+        .list.eval(pl.element().struct.field(FieldName.CLASS_NAME).is_in(class_names))
+        .list.any()
+    )
+
+    return table_with_selected_class_names
+
+
+def split_primitive_columns_by_task_name(
+    samples_table: pl.DataFrame,
+    coordinate_types: Optional[List[Type[Primitive]]] = None,
+) -> pl.DataFrame:
+    """
+    Convert Primitive columns such as "objects" (Bbox) into a column for each task name.
+    For example, if the "objects" column (containing Bbox objects) has tasks "task1" and "task2".
+
+
+    This:
+    ─┬────────────┬─
+     ┆ objects    ┆
+     ┆ ---        ┆
+     ┆ list[struc ┆
+     ┆ t[11]]     ┆
+    ═╪════════════╪═
+    becomes this:
+    ─┬────────────┬────────────┬─
+     ┆ objects.   ┆ objects.   ┆
+     ┆ task1      ┆ task2      ┆
+     ┆ ---        ┆ ---        ┆
+     ┆ list[struc ┆ list[struc ┆
+     ┆ t[11]]     ┆ t[13]]     ┆
+    ═╪════════════╪════════════╪═
+
+    """
+    coordinate_types = coordinate_types or PRIMITIVE_TYPES
+    for PrimitiveType in coordinate_types:
+        col_name = PrimitiveType.column_name()
+
+        if col_name not in samples_table.columns:
+            continue
+
+        if samples_table[col_name].dtype != pl.List(pl.Struct):
+            continue
+
+        task_names = samples_table[col_name].explode().struct.field(FieldName.TASK_NAME).unique().to_list()
+        samples_table = samples_table.with_columns(
+            [
+                pl.col(col_name)
+                .list.filter(pl.element().struct.field(FieldName.TASK_NAME).eq(task_name))
+                .alias(f"{col_name}.{task_name}")
+                for task_name in task_names
+            ]
+        )
+        samples_table = samples_table.drop(col_name)
+    return samples_table
+
+
+def read_table_from_path(path: Path) -> pl.DataFrame:
+    path_annotations = path / FILENAME_ANNOTATIONS_PARQUET
+    if path_annotations.exists():
+        user_logger.info(f"Reading dataset annotations from Parquet file: {path_annotations}")
+        return pl.read_parquet(path_annotations)
+
+    path_annotations_jsonl = path / FILENAME_ANNOTATIONS_JSONL
+    if path_annotations_jsonl.exists():
+        user_logger.info(f"Reading dataset annotations from JSONL file: {path_annotations_jsonl}")
+        return pl.read_ndjson(path_annotations_jsonl)
+
+    raise FileNotFoundError(
+        f"Unable to read annotations. No json file '{path_annotations.name}' or Parquet file '{{path_annotations.name}} in in '{path}'."
+    )
+
+
+def check_image_paths(table: pl.DataFrame) -> bool:
+    missing_files = []
+    for org_path in tqdm(table["file_name"].to_list(), desc="Check image paths"):
+        org_path = Path(org_path)
+        if not org_path.exists():
+            missing_files.append(org_path)
+
+    if len(missing_files) > 0:
+        user_logger.warning(f"Missing files: {len(missing_files)}. Show first 5:")
+        for missing_file in missing_files[:5]:
+            user_logger.warning(f" - {missing_file}")
+        raise FileNotFoundError(f"Some files are missing in the dataset: {len(missing_files)} files not found.")
+
+    return True
+
+
+def unnest_classification_tasks(table: pl.DataFrame, strict: bool = True) -> pl.DataFrame:
+    """
+    Unnest classification tasks in table.
+    Classificiations tasks are all stored in the same column in the HafniaDataset table.
+    This function splits them into separate columns for each task name.
+
+    Type is converted from a list of structs (pl.List[pl.Struct]) to a struct (pl.Struct) column.
+
+    Converts classification column from this:
+       ─┬─────────────────┬─
+        ┆ classifications ┆
+        ┆ ---             ┆
+        ┆ list[struct[6]] ┆
+       ═╪═════════════════╪═
+
+    For example, if the classification column has tasks "task1" and "task2",
+       ─┬──────────────────┬──────────────────┬─
+        ┆ classifications. ┆ classifications. ┆
+        ┆ task1            ┆ task2            ┆
+        ┆ ---              ┆ ---              ┆
+        ┆ struct[6]        ┆ struct[6]        ┆
+       ═╪══════════════════╪══════════════════╪═
+
+    """
+    coordinate_types = [Classification]
+    table_out = table_transformations.split_primitive_columns_by_task_name(table, coordinate_types=coordinate_types)
+
+    classification_columns = [c for c in table_out.columns if c.startswith(Classification.column_name() + ".")]
+    for classification_column in classification_columns:
+        has_multiple_items_per_sample = all(table_out[classification_column].list.len() > 1)
+        if has_multiple_items_per_sample:
+            if strict:
+                raise ValueError(
+                    f"Column {classification_column} has multiple items per sample, but expected only one item."
+                )
+            else:
+                user_logger.warning(
+                    f"Warning: Unnesting of column '{classification_column}' is skipped because it has multiple items per sample."
+                )
+
+    table_out = table_out.with_columns([pl.col(c).list.first() for c in classification_columns])
+    return table_out

hafnia/experiment/hafnia_logger.py

@@ -9,10 +9,10 @@ from typing import Dict, Optional, Union
 
 import pyarrow as pa
 import pyarrow.parquet as pq
-from datasets import DatasetDict
 from pydantic import BaseModel, field_validator
 
 from hafnia.data.factory import load_dataset
+from hafnia.dataset.hafnia_dataset import HafniaDataset
 from hafnia.log import sys_logger, user_logger
 from hafnia.utils import is_remote_job, now_as_str
 
@@ -92,7 +92,7 @@ class HafniaLogger:
         self.schema = Entity.create_schema()
         self.log_environment()
 
-    def load_dataset(self, dataset_name: str) -> DatasetDict:
+    def load_dataset(self, dataset_name: str) -> HafniaDataset:
         """
         Load a dataset from the specified path.
         """

hafnia/helper_testing.py

@@ -0,0 +1,63 @@
+from pathlib import Path
+
+from hafnia import utils
+from hafnia.dataset.dataset_names import FILENAME_ANNOTATIONS_JSONL, DatasetVariant
+from hafnia.dataset.hafnia_dataset import HafniaDataset, Sample
+
+MICRO_DATASETS = {
+    "tiny-dataset": utils.PATH_DATASETS / "tiny-dataset",
+    "coco-2017": utils.PATH_DATASETS / "coco-2017",
+}
+
+
+def get_path_workspace() -> Path:
+    return Path(__file__).parents[2]
+
+
+def get_path_expected_images() -> Path:
+    return get_path_workspace() / "tests" / "data" / "expected_images"
+
+
+def get_path_test_data() -> Path:
+    return get_path_workspace() / "tests" / "data"
+
+
+def get_path_micro_hafnia_dataset_no_check() -> Path:
+    return get_path_test_data() / "micro_test_datasets"
+
+
+def get_path_micro_hafnia_dataset(dataset_name: str, force_update=False) -> Path:
+    import pytest
+
+    if dataset_name not in MICRO_DATASETS:
+        raise ValueError(f"Dataset name '{dataset_name}' is not recognized. Available options: {list(MICRO_DATASETS)}")
+    path_dataset = MICRO_DATASETS[dataset_name]
+
+    path_test_dataset = get_path_micro_hafnia_dataset_no_check() / dataset_name
+    path_test_dataset_annotations = path_test_dataset / FILENAME_ANNOTATIONS_JSONL
+    if path_test_dataset_annotations.exists() and not force_update:
+        return path_test_dataset
+
+    hafnia_dataset = HafniaDataset.read_from_path(path_dataset / DatasetVariant.SAMPLE.value)
+    hafnia_dataset = hafnia_dataset.sample(n_samples=3, seed=42)
+    hafnia_dataset.write(path_test_dataset)
+
+    if force_update:
+        pytest.fail(
+            "Sample image and metadata have been updated using 'force_update=True'. Set 'force_update=False' and rerun the test."
+        )
+    pytest.fail("Missing test sample image. Please rerun the test.")
+    return path_test_dataset
+
+
+def get_sample_micro_hafnia_dataset(dataset_name: str, force_update=False) -> Sample:
+    micro_dataset = get_micro_hafnia_dataset(dataset_name=dataset_name, force_update=force_update)
+    sample_dict = micro_dataset[0]
+    sample = Sample(**sample_dict)
+    return sample
+
+
+def get_micro_hafnia_dataset(dataset_name: str, force_update: bool = False) -> HafniaDataset:
+    path_dataset = get_path_micro_hafnia_dataset(dataset_name=dataset_name, force_update=force_update)
+    hafnia_dataset = HafniaDataset.read_from_path(path_dataset)
+    return hafnia_dataset

hafnia/http.py

@@ -31,7 +31,7 @@ def fetch(endpoint: str, headers: Dict, params: Optional[Dict] = None) -> Dict:
         http.clear()
 
 
-def post(endpoint: str, headers: Dict, data: Union[Path, Dict, bytes], multipart: bool = False) -> Dict:
+def post(endpoint: str, headers: Dict, data: Union[Path, Dict, bytes, str], multipart: bool = False) -> Dict:
     """Posts data to backend endpoint.
 
     Args:
@@ -64,9 +64,11 @@ def post(endpoint: str, headers: Dict, data: Union[Path, Dict, bytes], multipart
                 with open(data, "rb") as f:
                     body = f.read()
                 response = http.request("POST", endpoint, body=body, headers=headers)
-            elif isinstance(data, dict):
+            elif isinstance(data, (str, dict)):
+                if isinstance(data, dict):
+                    data = json.dumps(data)
                 headers["Content-Type"] = "application/json"
-                response = http.request("POST", endpoint, body=json.dumps(data), headers=headers)
+                response = http.request("POST", endpoint, body=data, headers=headers)
             elif isinstance(data, bytes):
                 response = http.request("POST", endpoint, body=data, headers=headers)
             else:

hafnia/platform/__init__.py

@@ -1,7 +1,7 @@
 from hafnia.platform.download import (
     download_resource,
     download_single_object,
-    get_resource_creds,
+    get_resource_credentials,
 )
 from hafnia.platform.experiment import (
     create_experiment,
@@ -17,5 +17,5 @@ __all__ = [
     "create_experiment",
     "download_resource",
     "download_single_object",
-    "get_resource_creds",
+    "get_resource_credentials",
 ]

hafnia/platform/builder.py

@@ -1,7 +1,7 @@
 import json
 import os
+import re
 import subprocess
-import tempfile
 import zipfile
 from hashlib import sha256
 from pathlib import Path
@@ -58,7 +58,7 @@ def buildx_available() -> bool:
         return False
 
 
-def build_dockerfile(dockerfile: str, docker_context: str, docker_tag: str, meta_file: str) -> None:
+def build_dockerfile(dockerfile: str, docker_context: str, docker_tag: str) -> None:
     """
     Build a Docker image using the provided Dockerfile.
 
@@ -73,12 +73,12 @@ def build_dockerfile(dockerfile: str, docker_context: str, docker_tag: str, meta
 
     cmd = ["docker", "build", "--platform", "linux/amd64", "-t", docker_tag, "-f", dockerfile]
 
-    remote_cache = os.getenv("REMOTE_CACHE_REPO")
+    remote_cache = os.getenv("EXPERIMENT_CACHE_ECR")
     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]
+        cmd += ["--build-arg", "BUILDKIT_INLINE_CACHE=1"]
         if cloud_mode:
             cmd += ["--push"]
         if remote_cache:
@@ -91,11 +91,27 @@ def build_dockerfile(dockerfile: str, docker_context: str, docker_tag: str, meta
     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'}")
+    result = None
+    output = ""
+    errors = []
     try:
-        subprocess.run(cmd, check=True)
+        result = subprocess.run(cmd, check=True, capture_output=True, text=True)
+        output = (result.stdout or "") + (result.stderr or "")
     except subprocess.CalledProcessError as e:
-        sys_logger.error(f"Docker build failed: {e}")
-        raise RuntimeError(f"Docker build failed: {e}")
+        output = (e.stdout or "") + (e.stderr or "")
+        error_pattern = r"ERROR: (.+?)(?:\n|$)"
+        errors = re.findall(error_pattern, output)
+        if not errors:
+            raise RuntimeError(f"Docker build failed: {output}")
+        if re.search(r"image tag '([^']+)' already exists", errors[-1]):
+            sys_logger.warning("Image {} already exists in the registry.".format(docker_tag.rsplit("/")[-1]))
+            return
+        raise RuntimeError(f"Docker build failed: {output}")
+    finally:
+        stage_pattern = r"^.*\[\d+/\d+\][^\n]*"
+        stages = re.findall(stage_pattern, output, re.MULTILINE)
+        user_logger.info("\n".join(stages))
+        sys_logger.debug(output)
 
 
 def check_registry(docker_image: str) -> Optional[str]:
@@ -127,18 +143,8 @@ def build_image(metadata: Dict, registry_repo: str, state_file: str = "state.jso
     docker_image = f"{registry_repo}:{metadata['digest']}"
     image_exists = check_registry(docker_image) is not None
     if image_exists:
-        sys_logger.info(f"Tag already in ECR – skipping build of {docker_image}.")
+        sys_logger.info("Image {} already exists in the registry.".format(docker_image.rsplit("/")[-1]))
     else:
-        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"] = ""
+        build_dockerfile(metadata["dockerfile"], Path(metadata["dockerfile"]).parent.as_posix(), docker_image)
     metadata.update({"image_tag": docker_image, "image_exists": image_exists})
     Path(state_file).write_text(json.dumps(metadata, indent=2))

hafnia/platform/datasets.py

@@ -0,0 +1,184 @@
+import os
+import shutil
+import subprocess
+import tempfile
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+
+import rich
+from tqdm import tqdm
+
+from cli.config import Config
+from hafnia import utils
+from hafnia.dataset import dataset_names
+from hafnia.dataset.hafnia_dataset import HafniaDataset
+from hafnia.http import fetch
+from hafnia.log import user_logger
+from hafnia.platform import get_dataset_id
+from hafnia.platform.download import get_resource_credentials
+from hafnia.utils import timed
+
+
+@timed("Fetching dataset list.")
+def dataset_list(cfg: Optional[Config] = None) -> List[Dict[str, str]]:
+    """List available datasets on the Hafnia platform."""
+    cfg = cfg or Config()
+    endpoint_dataset = cfg.get_platform_endpoint("datasets")
+    header = {"Authorization": cfg.api_key}
+    datasets: List[Dict[str, str]] = fetch(endpoint_dataset, headers=header)  # type: ignore
+    if not datasets:
+        raise ValueError("No datasets found on the Hafnia platform.")
+
+    return datasets
+
+
+def download_or_get_dataset_path(
+    dataset_name: str,
+    cfg: Optional[Config] = None,
+    path_datasets_folder: Optional[str] = None,
+    force_redownload: bool = False,
+) -> Path:
+    """Download or get the path of the dataset."""
+    if utils.is_remote_job():
+        return Path(os.getenv("MDI_DATASET_DIR", "/opt/ml/input/data/training"))
+
+    path_datasets_folder = path_datasets_folder or str(utils.PATH_DATASETS)
+    path_dataset = Path(path_datasets_folder).absolute() / dataset_name
+
+    is_dataset_valid = HafniaDataset.check_dataset_path(path_dataset, raise_error=False)
+    if is_dataset_valid and not force_redownload:
+        user_logger.info("Dataset found locally. Set 'force=True' or add `--force` flag with cli to re-download")
+        return path_dataset
+
+    cfg = cfg or Config()
+    api_key = cfg.api_key
+
+    shutil.rmtree(path_dataset, ignore_errors=True)
+
+    endpoint_dataset = cfg.get_platform_endpoint("datasets")
+    dataset_id = get_dataset_id(dataset_name=dataset_name, endpoint=endpoint_dataset, api_key=api_key)
+    access_dataset_endpoint = f"{endpoint_dataset}/{dataset_id}/temporary-credentials"
+
+    download_dataset_from_access_endpoint(
+        endpoint=access_dataset_endpoint,
+        api_key=api_key,
+        path_dataset=path_dataset,
+    )
+    return path_dataset
+
+
+def download_dataset_from_access_endpoint(endpoint: str, api_key: str, path_dataset: Path) -> None:
+    resource_credentials = get_resource_credentials(endpoint, api_key)
+
+    local_dataset_paths = [str(path_dataset / filename) for filename in dataset_names.DATASET_FILENAMES]
+    s3_uri = resource_credentials.s3_uri()
+    s3_dataset_files = [f"{s3_uri}/{filename}" for filename in dataset_names.DATASET_FILENAMES]
+
+    envs = resource_credentials.aws_credentials()
+    fast_copy_files_s3(
+        src_paths=s3_dataset_files,
+        dst_paths=local_dataset_paths,
+        append_envs=envs,
+        description="Downloading annotations",
+    )
+
+    dataset = HafniaDataset.read_from_path(path_dataset, check_for_images=False)
+    fast_copy_files_s3(
+        src_paths=dataset.samples[dataset_names.ColumnName.REMOTE_PATH].to_list(),
+        dst_paths=dataset.samples[dataset_names.ColumnName.FILE_NAME].to_list(),
+        append_envs=envs,
+        description="Downloading images",
+    )
+
+
+def fast_copy_files_s3(
+    src_paths: List[str],
+    dst_paths: List[str],
+    append_envs: Optional[Dict[str, str]] = None,
+    description: str = "Copying files",
+) -> List[str]:
+    if len(src_paths) != len(dst_paths):
+        raise ValueError("Source and destination paths must have the same length.")
+
+    cmds = [f"cp {src} {dst}" for src, dst in zip(src_paths, dst_paths)]
+    lines = execute_s5cmd_commands(cmds, append_envs=append_envs, description=description)
+    return lines
+
+
+def execute_s5cmd_commands(
+    commands: List[str],
+    append_envs: Optional[Dict[str, str]] = None,
+    description: str = "Executing s5cmd commands",
+) -> List[str]:
+    append_envs = append_envs or {}
+    with tempfile.NamedTemporaryFile(suffix=".txt") as tmp_file:
+        tmp_file_path = Path(tmp_file.name)
+        tmp_file_path.write_text("\n".join(commands))
+        run_cmds = [
+            "s5cmd",
+            "run",
+            str(tmp_file_path),
+        ]
+        envs = os.environ.copy()
+        envs.update(append_envs)
+
+        process = subprocess.Popen(
+            run_cmds,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.STDOUT,
+            universal_newlines=True,
+            env=envs,
+        )
+
+        error_lines = []
+        lines = []
+        for line in tqdm(process.stdout, total=len(commands), desc=description):
+            if "ERROR" in line or "error" in line:
+                error_lines.append(line.strip())
+            lines.append(line.strip())
+
+        if len(error_lines) > 0:
+            show_n_lines = min(5, len(error_lines))
+            str_error_lines = "\n".join(error_lines[:show_n_lines])
+            user_logger.error(
+                f"Detected {len(error_lines)} errors occurred while executing a total of {len(commands)} "
+                f" commands with s5cmd. The first {show_n_lines} is printed below:\n{str_error_lines}"
+            )
+            raise RuntimeError("Errors occurred during s5cmd execution.")
+    return lines
+
+
+TABLE_FIELDS = {
+    "ID": "id",
+    "Hidden\nSamples": "hidden.samples",
+    "Hidden\nSize": "hidden.size",
+    "Sample\nSamples": "sample.samples",
+    "Sample\nSize": "sample.size",
+    "Name": "name",
+    "Title": "title",
+}
+
+
+def create_rich_table_from_dataset(datasets: List[Dict[str, str]]) -> rich.table.Table:
+    datasets = extend_dataset_details(datasets)
+    datasets = sorted(datasets, key=lambda x: x["name"].lower())
+
+    table = rich.table.Table(title="Available Datasets")
+    for i_dataset, dataset in enumerate(datasets):
+        if i_dataset == 0:
+            for column_name, _ in TABLE_FIELDS.items():
+                table.add_column(column_name, justify="left", style="cyan", no_wrap=True)
+        row = [str(dataset.get(field, "")) for field in TABLE_FIELDS.values()]
+        table.add_row(*row)
+
+    return table
+
+
+def extend_dataset_details(datasets: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+    """Extends dataset details with number of samples and size"""
+    for dataset in datasets:
+        for variant in dataset["dataset_variants"]:
+            variant_type = variant["variant_type"]
+            dataset[f"{variant_type}.samples"] = variant["number_of_data_items"]
+            dataset[f"{variant_type}.size"] = utils.size_human_readable(variant["size_bytes"])
+    return datasets