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

cli/__main__.py

@@ -1,11 +1,21 @@
 #!/usr/bin/env python
 import click
 
-from cli import consts, dataset_cmds, experiment_cmds, profile_cmds, recipe_cmds, runc_cmds
+import hafnia
+from cli import (
+    consts,
+    dataset_cmds,
+    dataset_recipe_cmds,
+    experiment_cmds,
+    profile_cmds,
+    runc_cmds,
+    trainer_package_cmds,
+)
 from cli.config import Config, ConfigSchema
 
 
 @click.group()
+@click.version_option(version=hafnia.__version__)
 @click.pass_context
 def main(ctx: click.Context) -> None:
     """Hafnia CLI."""
@@ -27,7 +37,9 @@ def configure(cfg: Config) -> None:
 
     platform_url = click.prompt("Hafnia Platform URL", type=str, default=consts.DEFAULT_API_URL)
 
-    cfg_profile = ConfigSchema(api_key=api_key, platform_url=platform_url)
+    use_keychain = click.confirm("Store API key in system keychain?", default=False)
+
+    cfg_profile = ConfigSchema(platform_url=platform_url, api_key=api_key, use_keychain=use_keychain)
     cfg.add_profile(profile_name, cfg_profile, set_active=True)
     cfg.save_config()
     profile_cmds.profile_show(cfg)
@@ -45,7 +57,8 @@ main.add_command(profile_cmds.profile)
 main.add_command(dataset_cmds.dataset)
 main.add_command(runc_cmds.runc)
 main.add_command(experiment_cmds.experiment)
-main.add_command(recipe_cmds.recipe)
+main.add_command(trainer_package_cmds.trainer_package)
+main.add_command(dataset_recipe_cmds.dataset_recipe)
 
 if __name__ == "__main__":
     main(max_content_width=120)

cli/config.py

@@ -6,10 +6,12 @@ from typing import Dict, List, Optional
 from pydantic import BaseModel, field_validator
 
 import cli.consts as consts
+import cli.keychain as keychain
 from hafnia.log import sys_logger, user_logger
 
 PLATFORM_API_MAPPING = {
-    "recipes": "/api/v1/recipes",
+    "trainers": "/api/v1/trainers",
+    "dataset_recipes": "/api/v1/dataset-recipes",
     "experiments": "/api/v1/experiments",
     "experiment_environments": "/api/v1/experiment-environments",
     "experiment_runs": "/api/v1/experiment-runs",
@@ -18,9 +20,18 @@ PLATFORM_API_MAPPING = {
 }
 
 
+class SecretStr(str):
+    def __repr__(self):
+        return "********"
+
+    def __str__(self):
+        return "********"
+
+
 class ConfigSchema(BaseModel):
     platform_url: str = ""
     api_key: Optional[str] = None
+    use_keychain: bool = False
 
     @field_validator("api_key")
     def validate_api_key(cls, value: Optional[str]) -> Optional[str]:
@@ -34,7 +45,7 @@ class ConfigSchema(BaseModel):
             sys_logger.warning("API key is missing the 'ApiKey ' prefix. Prefix is being added automatically.")
             value = f"ApiKey {value}"
 
-        return value
+        return SecretStr(value)  # Keeps the API key masked in logs and repr
 
 
 class ConfigFileSchema(BaseModel):
@@ -69,13 +80,32 @@ class Config:
 
     @property
     def api_key(self) -> str:
+        # Check keychain first if enabled
+        if self.config.use_keychain:
+            keychain_key = keychain.get_api_key(self.active_profile)
+            if keychain_key is not None:
+                return keychain_key
+
+        # Fall back to config file
         if self.config.api_key is not None:
             return self.config.api_key
+
         raise ValueError(consts.ERROR_API_KEY_NOT_SET)
 
     @api_key.setter
     def api_key(self, value: str) -> None:
-        self.config.api_key = value
+        # Store in keychain if enabled
+        if self.config.use_keychain:
+            if keychain.store_api_key(self.active_profile, value):
+                # Successfully stored in keychain, don't store in config
+                self.config.api_key = None
+            else:
+                # Keychain storage failed, fall back to config file
+                sys_logger.warning("Failed to store in keychain, falling back to config file")
+                self.config.api_key = value
+        else:
+            # Not using keychain, store in config file
+            self.config.api_key = value
 
     @property
     def platform_url(self) -> str:
@@ -151,8 +181,19 @@ class Config:
             raise ValueError("Failed to parse configuration file")
 
     def save_config(self) -> None:
+        # Create a copy to avoid modifying the original data
+        config_to_save = self.config_data.model_dump()
+
+        # Store API key in keychain if enabled, and don't write to file
+        for profile_name, profile_data in config_to_save["profiles"].items():
+            if profile_data.get("use_keychain", False):
+                api_key = profile_data.get("api_key")
+                if api_key:
+                    keychain.store_api_key(profile_name, api_key)
+                profile_data["api_key"] = None
+
         with open(self.config_path, "w") as f:
-            json.dump(self.config_data.model_dump(), f, indent=4)
+            json.dump(config_to_save, f, indent=4)
 
     def remove_profile(self, profile_name: str) -> None:
         if profile_name not in self.config_data.profiles:

cli/consts.py

@@ -10,7 +10,7 @@ ERROR_CREATE_PROFILE: str = "Failed to create profile. Profile name must be uniq
 ERROR_GET_RESOURCE: str = "Failed to get the data from platform. Verify url or api key."
 
 ERROR_EXPERIMENT_DIR: str = "Source directory does not exist"
-ERROR_RECIPE_FILE_FORMAT: str = "Recipe filename must be a '.zip' file"
+ERROR_TRAINER_PACKAGE_FILE_FORMAT: str = "Trainer package must be a '.zip' file"
 
 PROFILE_SWITCHED_SUCCESS: str = "Switched to profile:"
 PROFILE_REMOVED_SUCCESS: str = "Removed profile:"

cli/dataset_cmds.py

@@ -2,12 +2,10 @@ from pathlib import Path
 from typing import Optional
 
 import click
-from rich import print as rprint
 
-import cli.consts as consts
+from cli import consts
 from cli.config import Config
 from hafnia import utils
-from hafnia.platform.datasets import create_rich_table_from_dataset
 
 
 @click.group()
@@ -18,18 +16,12 @@ def dataset():
 
 @dataset.command("ls")
 @click.pass_obj
-def dataset_list(cfg: Config) -> None:
+def cmd_list_datasets(cfg: Config) -> None:
     """List available datasets on Hafnia platform"""
+    from hafnia.platform.datasets import get_datasets, pretty_print_datasets
 
-    from hafnia.platform.datasets import dataset_list
-
-    try:
-        datasets = dataset_list(cfg=cfg)
-    except Exception:
-        raise click.ClickException(consts.ERROR_GET_RESOURCE)
-
-    table = create_rich_table_from_dataset(datasets)
-    rprint(table)
+    datasets = get_datasets(cfg=cfg)
+    pretty_print_datasets(datasets)
 
 
 @dataset.command("download")
@@ -43,7 +35,7 @@ def dataset_list(cfg: Config) -> None:
 )
 @click.option("--force", "-f", is_flag=True, default=False, help="Flag to enable force redownload")
 @click.pass_obj
-def data_download(cfg: Config, dataset_name: str, destination: Optional[click.Path], force: bool) -> Path:
+def cmd_dataset_download(cfg: Config, dataset_name: str, destination: Optional[click.Path], force: bool) -> Path:
     """Download dataset from Hafnia platform"""
 
     from hafnia.platform import datasets

cli/dataset_recipe_cmds.py

@@ -0,0 +1,78 @@
+from pathlib import Path
+from typing import Dict, Optional
+
+import click
+from rich import print as rprint
+
+from cli.config import Config
+
+
+@click.group(name="dataset-recipe")
+def dataset_recipe() -> None:
+    """Dataset recipe commands"""
+    pass
+
+
+@dataset_recipe.command(name="create")
+@click.argument("path_json_recipe", required=True)
+@click.option(
+    "-n",
+    "--name",
+    type=str,
+    default=None,
+    show_default=True,
+    help="Name of the dataset recipe.",
+)
+@click.pass_obj
+def cmd_get_or_create_dataset_recipe(cfg: Config, path_json_recipe: Path, name: Optional[str]) -> None:
+    """Create Hafnia dataset recipe from dataset recipe JSON file"""
+    from hafnia.platform.dataset_recipe import get_or_create_dataset_recipe_from_path
+
+    endpoint = cfg.get_platform_endpoint("dataset_recipes")
+    recipe = get_or_create_dataset_recipe_from_path(path_json_recipe, endpoint=endpoint, api_key=cfg.api_key, name=name)
+
+    if recipe is None:
+        raise click.ClickException("Failed to create dataset recipe.")
+
+    rprint(recipe)
+
+
+@dataset_recipe.command(name="ls")
+@click.pass_obj
+@click.option("-l", "--limit", type=int, default=None, help="Limit number of listed dataset recipes.")
+def cmd_list_dataset_recipes(cfg: Config, limit: Optional[int]) -> None:
+    """List available dataset recipes"""
+    from hafnia.platform.dataset_recipe import get_dataset_recipes, pretty_print_dataset_recipes
+
+    endpoint = cfg.get_platform_endpoint("dataset_recipes")
+    recipes = get_dataset_recipes(endpoint=endpoint, api_key=cfg.api_key)
+    # Sort recipes to have the most recent first
+    recipes = sorted(recipes, key=lambda x: x["created_at"], reverse=True)
+    if limit is not None:
+        recipes = recipes[:limit]
+    pretty_print_dataset_recipes(recipes)
+
+
+@dataset_recipe.command(name="rm")
+@click.option("-i", "--id", type=str, help="Dataset recipe ID to delete.")
+@click.option("-n", "--name", type=str, help="Dataset recipe name to delete.")
+@click.pass_obj
+def cmd_delete_dataset_recipe(cfg: Config, id: Optional[str], name: Optional[str]) -> Dict:
+    """Delete a dataset recipe by ID or name"""
+    from hafnia.platform.dataset_recipe import delete_dataset_recipe_by_id, delete_dataset_recipe_by_name
+
+    endpoint = cfg.get_platform_endpoint("dataset_recipes")
+
+    if id is not None:
+        return delete_dataset_recipe_by_id(id=id, endpoint=endpoint, api_key=cfg.api_key)
+    if name is not None:
+        dataset_recipe = delete_dataset_recipe_by_name(name=name, endpoint=endpoint, api_key=cfg.api_key)
+        if dataset_recipe is None:
+            raise click.ClickException(f"Dataset recipe with name '{name}' was not found.")
+
+        return dataset_recipe
+
+    raise click.MissingParameter(
+        "No dataset recipe identifier have been given. Provide either --id or --name. "
+        "Get available recipes with 'hafnia dataset-recipe ls'."
+    )

cli/experiment_cmds.py

@@ -1,10 +1,16 @@
 from pathlib import Path
+from typing import Dict, Optional
 
 import click
-from rich import print as rprint
 
-import cli.consts as consts
 from cli.config import Config
+from hafnia import utils
+from hafnia.platform.dataset_recipe import (
+    get_dataset_recipe_by_id,
+    get_dataset_recipe_by_name,
+    get_or_create_dataset_recipe_by_dataset_name,
+)
+from hafnia.platform.trainer_package import create_trainer_package
 
 
 @click.group(name="experiment")
@@ -13,48 +19,225 @@ def experiment() -> None:
     pass
 
 
+@experiment.command(name="environments")
+@click.pass_obj
+def cmd_view_environments(cfg: Config):
+    """
+    View available experiment training environments.
+    """
+    from hafnia.platform import get_environments, pretty_print_training_environments
+
+    envs = get_environments(cfg.get_platform_endpoint("experiment_environments"), cfg.api_key)
+
+    pretty_print_training_environments(envs)
+
+
+def default_experiment_run_name():
+    return f"run-{utils.now_as_str()}"
+
+
 @experiment.command(name="create")
-@click.argument("name")
-@click.argument("source_dir", type=Path)
-@click.argument("exec_cmd", type=str)
-@click.argument("dataset_name")
-@click.argument("env_name")
+@click.option(
+    "-n",
+    "--name",
+    type=str,
+    default=default_experiment_run_name(),
+    required=False,
+    help=f"Name of the experiment. [default: run-[DATETIME] e.g. {default_experiment_run_name()}] ",
+)
+@click.option(
+    "-c",
+    "--cmd",
+    type=str,
+    default="python scripts/train.py",
+    show_default=True,
+    help="Command to run the experiment.",
+)
+@click.option(
+    "-p",
+    "--trainer-path",
+    type=Path,
+    default=None,
+    help="Path to the trainer package directory. ",
+)
+@click.option(
+    "-i",
+    "--trainer-id",
+    type=str,
+    default=None,
+    help="ID of the trainer package. View available trainers with 'hafnia trainer ls'",
+)
+@click.option(
+    "-d",
+    "--dataset",
+    type=str,
+    default=None,
+    required=False,
+    help="DatasetIdentifier: Name of the dataset. View Available datasets with 'hafnia dataset ls'",
+)
+@click.option(
+    "-r",
+    "--dataset-recipe",
+    type=str,
+    default=None,
+    required=False,
+    help="DatasetIdentifier: Name of the dataset recipe. View available dataset recipes with 'hafnia dataset-recipe ls'",
+)
+@click.option(
+    "--dataset-recipe-id",
+    type=str,
+    default=None,
+    required=False,
+    help="DatasetIdentifier: ID of the dataset recipe. View dataset recipes with 'hafnia dataset-recipe ls'",
+)
+@click.option(
+    "-e",
+    "--environment",
+    type=str,
+    default="Free Tier",
+    show_default=True,
+    help="Experiment environment name. View available environments with 'hafnia experiment environments'",
+)
 @click.pass_obj
-def create(cfg: Config, name: str, source_dir: Path, exec_cmd: str, dataset_name: str, env_name: str) -> None:
-    """Create a new experiment run"""
-    from hafnia.platform import create_experiment, create_recipe, get_dataset_id, get_exp_environment_id
-
-    if not source_dir.exists():
-        raise click.ClickException(consts.ERROR_EXPERIMENT_DIR)
-
-    try:
-        dataset_id = get_dataset_id(dataset_name, cfg.get_platform_endpoint("datasets"), cfg.api_key)
-    except Exception:
-        raise click.ClickException(f"Error retrieving dataset '{dataset_name}'.")
-
-    try:
-        recipe_id = create_recipe(source_dir, cfg.get_platform_endpoint("recipes"), cfg.api_key)
-    except Exception:
-        raise click.ClickException(f"Failed to create recipe from '{source_dir}'")
-
-    try:
-        env_id = get_exp_environment_id(env_name, cfg.get_platform_endpoint("experiment_environments"), cfg.api_key)
-    except Exception:
-        raise click.ClickException(f"Environment '{env_name}' not found")
-
-    try:
-        experiment_id = create_experiment(
-            name, dataset_id, recipe_id, exec_cmd, env_id, cfg.get_platform_endpoint("experiments"), cfg.api_key
+def cmd_create_experiment(
+    cfg: Config,
+    name: str,
+    cmd: str,
+    trainer_path: Path,
+    trainer_id: Optional[str],
+    dataset: Optional[str],
+    dataset_recipe: Optional[str],
+    dataset_recipe_id: Optional[str],
+    environment: str,
+) -> None:
+    """
+    Create and launch a new experiment run
+
+    Requires one dataset recipe and one trainer package:.
+        - One dataset identifier is required either '--dataset', '--dataset-recipe' or '--dataset-recipe-id'.
+        - One trainer identifier is required either '--trainer-path' or '--trainer-id'.
+
+    \b
+    Examples:
+    # Launch an experiment with a dataset and a trainer package from local path
+    hafnia experiment create --dataset mnist --trainer-path ../trainer-classification
+
+    \b
+    # Launch experiment with dataset recipe by name and trainer package by id
+    hafnia experiment create --dataset-recipe mnist-recipe --trainer-id 5e454c0d-fdf1-4d1f-9732-771d7fecd28e
+
+    \b
+    # Show available options:
+    hafnia experiment create --name "My Experiment" -d mnist --cmd "python scripts/train.py" -e "Free Tier" -p ../trainer-classification
+    """
+    from hafnia.platform import create_experiment, get_exp_environment_id
+
+    dataset_recipe_response = get_dataset_recipe_by_dataset_identifies(
+        cfg=cfg,
+        dataset_name=dataset,
+        dataset_recipe_name=dataset_recipe,
+        dataset_recipe_id=dataset_recipe_id,
+    )
+    dataset_recipe_id = dataset_recipe_response["id"]
+
+    trainer_id = get_trainer_package_by_identifies(
+        cfg=cfg,
+        trainer_path=trainer_path,
+        trainer_id=trainer_id,
+    )
+
+    env_id = get_exp_environment_id(environment, cfg.get_platform_endpoint("experiment_environments"), cfg.api_key)
+
+    experiment = create_experiment(
+        experiment_name=name,
+        dataset_recipe_id=dataset_recipe_id,
+        trainer_id=trainer_id,
+        exec_cmd=cmd,
+        environment_id=env_id,
+        endpoint=cfg.get_platform_endpoint("experiments"),
+        api_key=cfg.api_key,
+    )
+
+    experiment_properties = {
+        "ID": experiment.get("id", "N/A"),
+        "Name": experiment.get("name", "N/A"),
+        "State": experiment.get("state", "N/A"),
+        "Trainer Package ID": experiment.get("trainer", "N/A"),
+        "Dataset Recipe ID": experiment.get("dataset_recipe", "N/A"),
+        "Dataset ID": experiment.get("dataset", "N/A"),
+        "Created At": experiment.get("created_at", "N/A"),
+    }
+    print("Successfully created experiment: ")
+    for key, value in experiment_properties.items():
+        print(f"  {key}: {value}")
+
+
+def get_dataset_recipe_by_dataset_identifies(
+    cfg: Config,
+    dataset_name: Optional[str],
+    dataset_recipe_name: Optional[str],
+    dataset_recipe_id: Optional[str],
+) -> Dict:
+    dataset_identifiers = [dataset_name, dataset_recipe_name, dataset_recipe_id]
+    n_dataset_identifies_defined = sum([bool(identifier) for identifier in dataset_identifiers])
+
+    if n_dataset_identifies_defined > 1:
+        raise click.ClickException(
+            "Multiple dataset identifiers have been provided. Define only one dataset identifier."
+        )
+
+    dataset_recipe_endpoint = cfg.get_platform_endpoint("dataset_recipes")
+    if dataset_name:
+        return get_or_create_dataset_recipe_by_dataset_name(dataset_name, dataset_recipe_endpoint, cfg.api_key)
+
+    if dataset_recipe_name:
+        recipe = get_dataset_recipe_by_name(dataset_recipe_name, dataset_recipe_endpoint, cfg.api_key)
+        if recipe is None:
+            raise click.ClickException(f"Dataset recipe '{dataset_recipe_name}' was not found in the dataset library.")
+        return recipe
+
+    if dataset_recipe_id:
+        return get_dataset_recipe_by_id(dataset_recipe_id, dataset_recipe_endpoint, cfg.api_key)
+
+    raise click.MissingParameter(
+        "At least one dataset identifier must be provided. Set one of the following:\n"
+        "  --dataset <name>  -- E.g. '--dataset mnist'\n"
+        "  --dataset-recipe <name>  -- E.g. '--dataset-recipe my-recipe'\n"
+        "  --dataset-recipe-id <id>  -- E.g. '--dataset-recipe-id 5e454c0d-fdf1-4d1f-9732-771d7fecd28e'\n"
+    )
+
+
+def get_trainer_package_by_identifies(
+    cfg: Config,
+    trainer_path: Optional[Path],
+    trainer_id: Optional[str],
+) -> str:
+    from hafnia.platform import get_trainer_package_by_id
+
+    if trainer_path is not None and trainer_id is not None:
+        raise click.ClickException(
+            "Multiple trainer identifiers (--trainer-path, --trainer-id) have been provided. Define only one."
         )
-    except Exception:
-        raise click.ClickException(f"Failed to create experiment '{name}'")
-
-    rprint(
-        {
-            "dataset_id": dataset_id,
-            "recipe_id": recipe_id,
-            "environment_id": env_id,
-            "experiment_id": experiment_id,
-            "status": "CREATED",
-        }
+
+    if trainer_path is not None:
+        trainer_path = Path(trainer_path)
+        if not trainer_path.exists():
+            raise click.ClickException(f"Trainer package path '{trainer_path}' does not exist.")
+        trainer_id = create_trainer_package(
+            trainer_path,
+            cfg.get_platform_endpoint("trainers"),
+            cfg.api_key,
+        )
+        return trainer_id
+
+    if trainer_id:
+        trainer_response = get_trainer_package_by_id(
+            id=trainer_id, endpoint=cfg.get_platform_endpoint("trainers"), api_key=cfg.api_key
+        )
+        return trainer_response["id"]
+
+    raise click.MissingParameter(
+        "At least one trainer identifier must be provided. Set one of the following:\n"
+        "  --trainer-path <path>  -- E.g. '--trainer-path .'\n"
+        "  --trainer-id <id>  -- E.g. '--trainer-id 5e454c0d-fdf1-4d1f-9732-771d7fecd28e'\n"
     )

cli/keychain.py

@@ -0,0 +1,88 @@
+"""Keychain storage for API keys using the system keychain."""
+
+from typing import Optional
+
+from hafnia.log import sys_logger
+
+# Keyring is optional - gracefully degrade if not available
+try:
+    import keyring
+
+    KEYRING_AVAILABLE = True
+except ImportError:
+    KEYRING_AVAILABLE = False
+    sys_logger.debug("keyring library not available, keychain storage disabled")
+
+KEYRING_SERVICE_NAME = "hafnia-cli"
+
+
+def store_api_key(profile_name: str, api_key: str) -> bool:
+    """
+    Store an API key in the system keychain.
+
+    Args:
+        profile_name: The profile name to associate with the key
+        api_key: The API key to store
+
+    Returns:
+        True if successfully stored, False otherwise
+    """
+    if not KEYRING_AVAILABLE:
+        sys_logger.warning("Keyring library not available, cannot store API key in keychain")
+        return False
+
+    try:
+        keyring.set_password(KEYRING_SERVICE_NAME, profile_name, api_key)
+        sys_logger.debug(f"Stored API key for profile '{profile_name}' in keychain")
+        return True
+    except Exception as e:
+        sys_logger.warning(f"Failed to store API key in keychain: {e}")
+        return False
+
+
+def get_api_key(profile_name: str) -> Optional[str]:
+    """
+    Retrieve an API key from the system keychain.
+
+    Args:
+        profile_name: The profile name to retrieve the key for
+
+    Returns:
+        The API key if found, None otherwise
+    """
+    if not KEYRING_AVAILABLE:
+        return None
+
+    try:
+        api_key = keyring.get_password(KEYRING_SERVICE_NAME, profile_name)
+        if api_key:
+            sys_logger.debug(f"Retrieved API key for profile '{profile_name}' from keychain")
+        return api_key
+    except Exception as e:
+        sys_logger.warning(f"Failed to retrieve API key from keychain: {e}")
+        return None
+
+
+def delete_api_key(profile_name: str) -> bool:
+    """
+    Delete an API key from the system keychain.
+
+    Args:
+        profile_name: The profile name to delete the key for
+
+    Returns:
+        True if successfully deleted or didn't exist, False on error
+    """
+    if not KEYRING_AVAILABLE:
+        return False
+
+    try:
+        keyring.delete_password(KEYRING_SERVICE_NAME, profile_name)
+        sys_logger.debug(f"Deleted API key for profile '{profile_name}' from keychain")
+        return True
+    except keyring.errors.PasswordDeleteError:
+        # Key didn't exist, which is fine
+        return True
+    except Exception as e:
+        sys_logger.warning(f"Failed to delete API key from keychain: {e}")
+        return False