hafnia 0.2.4__py3-none-any.whl → 0.3.0__py3-none-any.whl

hafnia/dataset/operations/table_transformations.py

@@ -7,6 +7,7 @@ from tqdm import tqdm
 from hafnia.dataset.dataset_names import (
     FILENAME_ANNOTATIONS_JSONL,
     FILENAME_ANNOTATIONS_PARQUET,
+    ColumnName,
     FieldName,
 )
 from hafnia.dataset.operations import table_transformations
@@ -34,8 +35,12 @@ def create_primitive_table(
 
     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])
+
+        drop_columns_primitives = set(PRIMITIVE_TYPES) - {PrimitiveType, Classification}
+        drop_columns_names = [primitive.column_name() for primitive in drop_columns_primitives]
+        drop_columns_names = [c for c in drop_columns_names if c in remove_no_object_frames.columns]
+
+        remove_no_object_frames = remove_no_object_frames.drop(drop_columns_names)
         # 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"}
@@ -46,6 +51,38 @@ def create_primitive_table(
     return objects_df
 
 
+def merge_samples(samples0: pl.DataFrame, samples1: pl.DataFrame) -> pl.DataFrame:
+    has_same_schema = samples0.schema == samples1.schema
+    if not has_same_schema:
+        shared_columns = []
+        for column_name, column_type in samples0.schema.items():
+            if column_name not in samples1.schema:
+                continue
+
+            if column_type != samples1.schema[column_name]:
+                continue
+            shared_columns.append(column_name)
+
+        dropped_columns0 = [
+            f"{n}[{ctype._string_repr()}]" for n, ctype in samples0.schema.items() if n not in shared_columns
+        ]
+        dropped_columns1 = [
+            f"{n}[{ctype._string_repr()}]" for n, ctype in samples1.schema.items() if n not in shared_columns
+        ]
+        user_logger.warning(
+            "Datasets with different schemas are being merged. "
+            "Only the columns with the same name and type will be kept in the merged dataset.\n"
+            f"Dropped columns in samples0: {dropped_columns0}\n"
+            f"Dropped columns in samples1: {dropped_columns1}\n"
+        )
+
+        samples0 = samples0.select(list(shared_columns))
+        samples1 = samples1.select(list(shared_columns))
+    merged_samples = pl.concat([samples0, samples1], how="vertical")
+    merged_samples = merged_samples.drop(ColumnName.SAMPLE_INDEX).with_row_index(name=ColumnName.SAMPLE_INDEX)
+    return merged_samples
+
+
 def filter_table_for_class_names(
     samples_table: pl.DataFrame, class_names: List[str], PrimitiveType: Type[Primitive]
 ) -> Optional[pl.DataFrame]:

hafnia/dataset/primitives/__init__.py

@@ -14,3 +14,11 @@ from .utils import class_color_by_name  # noqa: F401
 PRIMITIVE_TYPES: List[Type[Primitive]] = [Bbox, Classification, Polygon, Bitmask]
 PRIMITIVE_NAME_TO_TYPE = {cls.__name__: cls for cls in PRIMITIVE_TYPES}
 PRIMITIVE_COLUMN_NAMES: List[str] = [PrimitiveType.column_name() for PrimitiveType in PRIMITIVE_TYPES]
+
+
+def get_primitive_type_from_string(name: str) -> Type[Primitive]:
+    if name not in PRIMITIVE_NAME_TO_TYPE:
+        raise ValueError(
+            f"Primitive '{name}' is not recognized. Available primitives: {list(PRIMITIVE_NAME_TO_TYPE.keys())}"
+        )
+    return PRIMITIVE_NAME_TO_TYPE[name]

hafnia/dataset/primitives/classification.py

@@ -38,7 +38,7 @@ class Classification(Primitive):
             text = class_name
         else:
             text = f"{self.task_name}: {class_name}"
-        image = image_visualizations.append_text_below_frame(image, text=text)
+        image = image_visualizations.append_text_below_frame(image, text=text, text_size_ratio=0.05)
 
         return image
 

hafnia/experiment/hafnia_logger.py

@@ -2,6 +2,7 @@ import json
 import os
 import platform
 import sys
+import textwrap
 from datetime import datetime
 from enum import Enum
 from pathlib import Path
@@ -16,6 +17,16 @@ from hafnia.dataset.hafnia_dataset import HafniaDataset
 from hafnia.log import sys_logger, user_logger
 from hafnia.utils import is_hafnia_cloud_job, now_as_str
 
+try:
+    import mlflow
+    import mlflow.tracking
+    import sagemaker_mlflow  # noqa: F401
+
+    MLFLOW_AVAILABLE = True
+except ImportError:
+    user_logger.warning("MLFlow is not available")
+    MLFLOW_AVAILABLE = False
+
 
 class EntityType(Enum):
     """Types of entities that can be logged."""
@@ -87,11 +98,44 @@ class HafniaLogger:
         for path in create_paths:
             path.mkdir(parents=True, exist_ok=True)
 
+        path_file = self.path_model() / "HOW_TO_STORE_YOUR_MODEL.txt"
+        path_file.write_text(get_instructions_how_to_store_model())
+
         self.dataset_name: Optional[str] = None
         self.log_file = self._path_artifacts() / self.EXPERIMENT_FILE
         self.schema = Entity.create_schema()
+
+        # Initialize MLflow for remote jobs
+        self._mlflow_initialized = False
+        if is_hafnia_cloud_job() and MLFLOW_AVAILABLE:
+            self._init_mlflow()
+
         self.log_environment()
 
+    def _init_mlflow(self):
+        """Initialize MLflow tracking for remote jobs."""
+        try:
+            # Set MLflow tracking URI from environment variable
+            tracking_uri = os.getenv("MLFLOW_TRACKING_URI")
+            if tracking_uri:
+                mlflow.set_tracking_uri(tracking_uri)
+                user_logger.info(f"MLflow tracking URI set to: {tracking_uri}")
+
+            # Set experiment name from environment variable
+            experiment_name = os.getenv("MLFLOW_EXPERIMENT_NAME")
+            if experiment_name:
+                mlflow.set_experiment(experiment_name)
+                user_logger.info(f"MLflow experiment set to: {experiment_name}")
+
+            # Start MLflow run
+            run_name = os.getenv("MLFLOW_RUN_NAME", "undefined")
+            mlflow.start_run(run_name=run_name)
+            self._mlflow_initialized = True
+            user_logger.info("MLflow run started successfully")
+
+        except Exception as e:
+            user_logger.error(f"Failed to initialize MLflow: {e}")
+
     def load_dataset(self, dataset_name: str) -> HafniaDataset:
         """
         Load a dataset from the specified path.
@@ -153,6 +197,14 @@ class HafniaLogger:
         )
         self.write_entity(entity)
 
+        # Also log to MLflow if initialized
+        if not self._mlflow_initialized:
+            return
+        try:
+            mlflow.log_metric(name, value, step=step)
+        except Exception as e:
+            user_logger.error(f"Failed to log metric to MLflow: {e}")
+
     def log_configuration(self, configurations: Dict):
         self.log_hparams(configurations, "configuration.json")
 
@@ -166,6 +218,15 @@ class HafniaLogger:
             existing_params.update(params)
             file_path.write_text(json.dumps(existing_params, indent=2))
             user_logger.info(f"Saved parameters to {file_path}")
+
+            # Also log to MLflow if initialized
+            if not self._mlflow_initialized:
+                return
+            try:
+                mlflow.log_params(params)
+            except Exception as e:
+                user_logger.error(f"Failed to log params to MLflow: {e}")
+
         except Exception as e:
             user_logger.error(f"Failed to save parameters to {file_path}: {e}")
 
@@ -202,3 +263,54 @@ class HafniaLogger:
                 pq.write_table(next_table, self.log_file)
         except Exception as e:
             sys_logger.error(f"Failed to flush logs: {e}")
+
+    def end_run(self) -> None:
+        """End the MLflow run if initialized."""
+        if not self._mlflow_initialized:
+            return
+        try:
+            mlflow.end_run()
+            self._mlflow_initialized = False
+            user_logger.info("MLflow run ended successfully")
+        except Exception as e:
+            user_logger.error(f"Failed to end MLflow run: {e}")
+
+    def __del__(self):
+        """Cleanup when logger is destroyed."""
+        self.end_run()
+
+
+def get_instructions_how_to_store_model() -> str:
+    instructions = textwrap.dedent(
+        """\
+        If you, against your expectations, don't see any models in this folder,
+        we have provided a small guide to help.
+
+        The hafnia TaaS framework expects models to be stored in a folder generated
+        by the hafnia logger. You will need to store models in this folder
+        to  ensure that they are properly stored and accessible after training.
+
+        Please check your recipe script and ensure that the models are being stored
+        as expected by the TaaS framework.
+
+        Below is also a small example to demonstrate:
+
+        ```python
+        from hafnia.experiment import HafniaLogger
+
+        # Initiate Hafnia logger
+        logger = HafniaLogger()
+
+        # Folder path to store models - generated by the hafnia logger. 
+        model_dir = logger.path_model()
+
+        # Example for storing a pytorch based model. Note: the model is stored in 'model_dir'
+        path_pytorch_model = model_dir / "model.pth"
+
+        # Finally save the model to the specified path
+        torch.save(model.state_dict(), path_pytorch_model)
+        ```
+        """
+    )
+
+    return instructions

hafnia/http.py

@@ -1,11 +1,11 @@
 import json
 from pathlib import Path
-from typing import Dict, Optional, Union
+from typing import Dict, List, Optional, Union
 
 import urllib3
 
 
-def fetch(endpoint: str, headers: Dict, params: Optional[Dict] = None) -> Dict:
+def fetch(endpoint: str, headers: Dict, params: Optional[Dict] = None) -> Union[Dict, List]:
     """Fetches data from the API endpoint.
 
     Args:
@@ -81,3 +81,17 @@ def post(endpoint: str, headers: Dict, data: Union[Path, Dict, bytes, str], mult
         return json.loads(response.data.decode("utf-8"))
     finally:
         http.clear()
+
+
+def delete(endpoint: str, headers: Dict) -> Dict:
+    """Sends a DELETE request to the specified endpoint."""
+    http = urllib3.PoolManager(retries=urllib3.Retry(3))
+    try:
+        response = http.request("DELETE", endpoint, headers=headers)
+
+        if response.status not in (200, 204):
+            error_details = response.data.decode("utf-8")
+            raise urllib3.exceptions.HTTPError(f"Request failed with status {response.status}: {error_details}")
+        return json.loads(response.data.decode("utf-8")) if response.data else {}
+    finally:
+        http.clear()

hafnia/platform/__init__.py

@@ -1,3 +1,4 @@
+from hafnia.platform.datasets import get_dataset_id
 from hafnia.platform.download import (
     download_resource,
     download_single_object,
@@ -5,17 +6,22 @@ from hafnia.platform.download import (
 )
 from hafnia.platform.experiment import (
     create_experiment,
-    create_recipe,
-    get_dataset_id,
+    get_environments,
     get_exp_environment_id,
+    pretty_print_training_environments,
 )
+from hafnia.platform.trainer_package import create_trainer_package, get_trainer_package_by_id, get_trainer_packages
 
 __all__ = [
     "get_dataset_id",
-    "create_recipe",
+    "create_trainer_package",
+    "get_trainer_packages",
+    "get_trainer_package_by_id",
     "get_exp_environment_id",
     "create_experiment",
     "download_resource",
     "download_single_object",
     "get_resource_credentials",
+    "pretty_print_training_environments",
+    "get_environments",
 ]

hafnia/platform/builder.py

@@ -14,26 +14,28 @@ from hafnia.log import sys_logger, user_logger
 from hafnia.platform import download_resource
 
 
-def validate_recipe_format(path: Path) -> None:
-    """Validate Hafnia Recipe Format submition"""
+def validate_trainer_package_format(path: Path) -> None:
+    """Validate Hafnia Trainer Package Format submission"""
     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.")
+            user_logger.error(f"Required path {rp} not found in trainer package.")
             errors += 1
     if errors > 0:
-        raise FileNotFoundError("Wrong recipe structure")
+        raise FileNotFoundError("Wrong trainer package structure")
 
 
-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:
+def prepare_trainer_package(
+    trainer_url: str, output_dir: Path, api_key: str, state_file: Optional[Path] = None
+) -> Dict:
+    resource = download_resource(trainer_url, output_dir.as_posix(), api_key)
+    trainer_path = Path(resource["downloaded_files"][0])
+    with zipfile.ZipFile(trainer_path, "r") as zip_ref:
         zip_ref.extractall(output_dir)
 
-    validate_recipe_format(output_dir)
+    validate_trainer_package_format(output_dir)
 
     scripts_dir = output_dir / "scripts"
     if not any(scripts_dir.iterdir()):
@@ -42,7 +44,7 @@ def prepare_recipe(recipe_url: str, output_dir: Path, api_key: str, state_file:
     metadata = {
         "user_data": (output_dir / "src").as_posix(),
         "dockerfile": (output_dir / "Dockerfile").as_posix(),
-        "digest": sha256(recipe_path.read_bytes()).hexdigest()[:8],
+        "digest": sha256(trainer_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:

hafnia/platform/dataset_recipe.py

@@ -0,0 +1,99 @@
+import json
+from pathlib import Path
+from typing import Dict, List, Optional
+
+from flatten_dict import flatten
+
+from hafnia import http
+from hafnia.log import user_logger
+from hafnia.utils import pretty_print_list_as_table, timed
+
+
+@timed("Get or create dataset recipe")
+def get_or_create_dataset_recipe(
+    recipe: dict, endpoint: str, api_key: str, name: Optional[str] = None
+) -> Optional[Dict]:
+    headers = {"Authorization": api_key}
+    data = {"template": {"body": recipe}}
+    if name is not None:
+        data["name"] = name  # type: ignore[assignment]
+    response = http.post(endpoint, headers=headers, data=data)
+    return response
+
+
+def get_or_create_dataset_recipe_by_dataset_name(dataset_name: str, endpoint: str, api_key: str) -> Dict:
+    return get_or_create_dataset_recipe(recipe=dataset_name, endpoint=endpoint, api_key=api_key)
+
+
+def get_dataset_recipes(endpoint: str, api_key: str) -> List[Dict]:
+    headers = {"Authorization": api_key}
+    dataset_recipes: List[Dict] = http.fetch(endpoint, headers=headers)  # type: ignore[assignment]
+    return dataset_recipes
+
+
+def get_dataset_recipe_by_id(dataset_recipe_id: str, endpoint: str, api_key: str) -> Dict:
+    headers = {"Authorization": api_key}
+    full_url = f"{endpoint}/{dataset_recipe_id}"
+    dataset_recipe_info: Dict = http.fetch(full_url, headers=headers)  # type: ignore[assignment]
+    if not dataset_recipe_info:
+        raise ValueError(f"Dataset recipe with ID '{dataset_recipe_id}' was not found.")
+    return dataset_recipe_info
+
+
+def get_or_create_dataset_recipe_from_path(
+    path_recipe_json: Path, endpoint: str, api_key: str, name: Optional[str] = None
+) -> Dict:
+    path_recipe_json = Path(path_recipe_json)
+    if not path_recipe_json.exists():
+        raise FileNotFoundError(f"Dataset recipe file '{path_recipe_json}' does not exist.")
+    json_dict = json.loads(path_recipe_json.read_text())
+    return get_or_create_dataset_recipe(json_dict, endpoint=endpoint, api_key=api_key, name=name)
+
+
+def delete_dataset_recipe_by_id(id: str, endpoint: str, api_key: str) -> Dict:
+    headers = {"Authorization": api_key}
+    full_url = f"{endpoint}/{id}"
+    response = http.delete(endpoint=full_url, headers=headers)
+    return response
+
+
+@timed("Get dataset recipe")
+def get_dataset_recipe_by_name(name: str, endpoint: str, api_key: str) -> Optional[Dict]:
+    headers = {"Authorization": api_key}
+    full_url = f"{endpoint}?name__iexact={name}"
+    dataset_recipes: List[Dict] = http.fetch(full_url, headers=headers)  # type: ignore[assignment]
+    if len(dataset_recipes) == 0:
+        return None
+
+    if len(dataset_recipes) > 1:
+        user_logger.warning(f"Found {len(dataset_recipes)} dataset recipes called '{name}'. Using the first one.")
+
+    dataset_recipe = dataset_recipes[0]
+    return dataset_recipe
+
+
+def delete_dataset_recipe_by_name(name: str, endpoint: str, api_key: str) -> Optional[Dict]:
+    recipe_response = get_dataset_recipe_by_name(name, endpoint=endpoint, api_key=api_key)
+
+    if recipe_response:
+        return delete_dataset_recipe_by_id(recipe_response["id"], endpoint=endpoint, api_key=api_key)
+    return recipe_response
+
+
+def pretty_print_dataset_recipes(recipes: List[Dict]) -> None:
+    recipes = [flatten(recipe, reducer="dot", max_flatten_depth=2) for recipe in recipes]  # noqa: F821
+    for recipe in recipes:
+        recipe["recipe_json"] = json.dumps(recipe["template.body"])[:20]
+
+    RECIPE_FIELDS = {
+        "ID": "id",
+        "Name": "name",
+        "Recipe": "recipe_json",
+        "Created": "created_at",
+        "IsDataset": "template.is_direct_dataset_reference",
+    }
+    pretty_print_list_as_table(
+        table_title="Available Dataset Recipes",
+        dict_items=recipes,
+        column_name_to_key_mapping=RECIPE_FIELDS,
+    )

hafnia/platform/datasets.py

@@ -8,10 +8,11 @@ from pathlib import Path
 from typing import Any, Dict, List, Optional
 
 import rich
+from rich import print as rprint
 from tqdm import tqdm
 
 from cli.config import Config
-from hafnia import utils
+from hafnia import http, utils
 from hafnia.dataset.dataset_names import DATASET_FILENAMES_REQUIRED, ColumnName
 from hafnia.dataset.dataset_recipe.dataset_recipe import (
     DatasetRecipe,
@@ -20,13 +21,12 @@ from hafnia.dataset.dataset_recipe.dataset_recipe import (
 from hafnia.dataset.hafnia_dataset import HafniaDataset
 from hafnia.http import fetch
 from hafnia.log import sys_logger, 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]]:
+def get_datasets(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")
@@ -38,6 +38,19 @@ def dataset_list(cfg: Optional[Config] = None) -> List[Dict[str, str]]:
     return datasets
 
 
+@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_responses: List[Dict] = http.fetch(full_url, headers=headers)  # type: ignore[assignment]
+    if not dataset_responses:
+        raise ValueError(f"Dataset '{dataset_name}' was not found in the dataset library.")
+    try:
+        return dataset_responses[0]["id"]
+    except (IndexError, KeyError) as e:
+        raise ValueError("Dataset information is missing or invalid") from e
+
+
 def download_or_get_dataset_path(
     dataset_name: str,
     cfg: Optional[Config] = None,
@@ -131,6 +144,28 @@ def fast_copy_files_s3(
     return lines
 
 
+def find_s5cmd() -> Optional[str]:
+    """Locate the s5cmd executable across different installation methods.
+
+    Searches for s5cmd in:
+    1. System PATH (via shutil.which)
+    2. Python bin directory (Unix-like systems)
+    3. Python executable directory (direct installs)
+
+    Returns:
+        str: Absolute path to s5cmd executable if found, None otherwise.
+    """
+    result = shutil.which("s5cmd")
+    if result:
+        return result
+    python_dir = Path(sys.executable).parent
+    locations = (python_dir / "Scripts" / "s5cmd.exe", python_dir / "bin" / "s5cmd", python_dir / "s5cmd")
+    for loc in locations:
+        if loc.exists():
+            return str(loc)
+    return None
+
+
 def execute_s5cmd_commands(
     commands: List[str],
     append_envs: Optional[Dict[str, str]] = None,
@@ -142,7 +177,10 @@ def execute_s5cmd_commands(
     with tempfile.TemporaryDirectory() as temp_dir:
         tmp_file_path = Path(temp_dir, f"{uuid.uuid4().hex}.txt")
         tmp_file_path.write_text("\n".join(commands))
-        s5cmd_bin = (Path(sys.executable).parent / "s5cmd").absolute().as_posix()
+
+        s5cmd_bin = find_s5cmd()
+        if s5cmd_bin is None:
+            raise ValueError("Can not find s5cmd executable.")
         run_cmds = [s5cmd_bin, "run", str(tmp_file_path)]
         sys_logger.debug(run_cmds)
         envs = os.environ.copy()
@@ -185,7 +223,7 @@ TABLE_FIELDS = {
 }
 
 
-def create_rich_table_from_dataset(datasets: List[Dict[str, str]]) -> rich.table.Table:
+def pretty_print_datasets(datasets: List[Dict[str, str]]) -> None:
     datasets = extend_dataset_details(datasets)
     datasets = sorted(datasets, key=lambda x: x["name"].lower())
 
@@ -197,7 +235,7 @@ def create_rich_table_from_dataset(datasets: List[Dict[str, str]]) -> rich.table
         row = [str(dataset.get(field, "")) for field in TABLE_FIELDS.values()]
         table.add_row(*row)
 
-    return table
+    rprint(table)
 
 
 def extend_dataset_details(datasets: List[Dict[str, Any]]) -> List[Dict[str, Any]]:

hafnia/platform/download.py

@@ -96,7 +96,8 @@ def get_resource_credentials(endpoint: str, api_key: str) -> ResourceCredentials
         RuntimeError: If the call to fetch the credentials fails for any reason.
     """
     try:
-        credentials_dict = fetch(endpoint, headers={"Authorization": api_key, "accept": "application/json"})
+        headers = {"Authorization": api_key, "accept": "application/json"}
+        credentials_dict: Dict = fetch(endpoint, headers=headers)  # type: ignore[assignment]
         credentials = ResourceCredentials.fix_naming(credentials_dict)
         sys_logger.debug("Successfully retrieved credentials from DIP endpoint.")
         return credentials