hafnia 0.1.8__py3-none-any.whl

cli/__main__.py

@@ -0,0 +1,63 @@
+#!/usr/bin/env python
+import click
+
+from cli import consts, data_cmds, experiment_cmds, profile_cmds, runc_cmds
+from cli.config import Config, ConfigSchema
+
+
+@click.group()
+@click.pass_context
+def main(ctx: click.Context) -> None:
+    """MDI CLI."""
+    ctx.obj = Config()
+
+
+@main.command("configure")
+@click.pass_obj
+def configure(cfg: Config) -> None:
+    """Configure MDI CLI settings."""
+
+    from hafnia.platform.api import get_organization_id
+
+    profile_name = click.prompt("Profile Name", type=str, default="default")
+    profile_name = profile_name.strip()
+    try:
+        cfg.add_profile(profile_name, ConfigSchema(), set_active=True)
+    except ValueError:
+        raise click.ClickException(consts.ERROR_CREATE_PROFILE)
+
+    api_key = click.prompt("MDI API Key", type=str, hide_input=True)
+    try:
+        cfg.api_key = api_key.strip()
+    except ValueError as e:
+        click.echo(f"Error: {str(e)}", err=True)
+        return
+    platform_url = click.prompt(
+        "MDI Platform URL", type=str, default="https://api.mdi.milestonesys.com"
+    )
+    cfg.platform_url = platform_url.strip()
+    try:
+        cfg.organization_id = get_organization_id(
+            cfg.get_platform_endpoint("organizations"), cfg.api_key
+        )
+    except Exception:
+        raise click.ClickException(consts.ERROR_ORG_ID)
+    cfg.save_config()
+    profile_cmds.profile_show(cfg)
+
+
+@main.command("clear")
+@click.pass_obj
+def clear(cfg: Config) -> None:
+    """Remove stored configuration."""
+    cfg.clear()
+    click.echo("Successfully cleared MDI configuration.")
+
+
+main.add_command(profile_cmds.profile)
+main.add_command(data_cmds.data)
+main.add_command(runc_cmds.runc)
+main.add_command(experiment_cmds.experiment)
+
+if __name__ == "__main__":
+    main()

cli/config.py

@@ -0,0 +1,148 @@
+import json
+import os
+from pathlib import Path
+from typing import Dict, List, Optional
+
+from pydantic import BaseModel, field_validator
+
+import cli.consts as consts
+from hafnia.log import logger
+
+
+class ConfigSchema(BaseModel):
+    organization_id: str = ""
+    platform_url: str = ""
+    api_key: Optional[str] = None
+    api_mapping: Optional[Dict[str, str]] = None
+
+    @field_validator("api_key")
+    def validate_api_key(cls, value: str) -> str:
+        if value is not None and len(value) < 10:
+            raise ValueError("API key is too short.")
+        return value
+
+
+class ConfigFileSchema(BaseModel):
+    active_profile: Optional[str] = None
+    profiles: Dict[str, ConfigSchema] = {}
+
+
+class Config:
+    @property
+    def available_profiles(self) -> List[str]:
+        return list(self.config_data.profiles.keys())
+
+    @property
+    def active_profile(self) -> str:
+        if self.config_data.active_profile is None:
+            raise ValueError(consts.ERROR_PROFILE_NOT_EXIST)
+        return self.config_data.active_profile
+
+    @active_profile.setter
+    def active_profile(self, value: str) -> None:
+        profile_name = value.strip()
+        if profile_name not in self.config_data.profiles:
+            raise ValueError(f"Profile '{profile_name}' does not exist.")
+        self.config_data.active_profile = profile_name
+
+    @property
+    def config(self) -> ConfigSchema:
+        if not self.config_data.active_profile:
+            raise ValueError(consts.ERROR_PROFILE_NOT_EXIST)
+        return self.config_data.profiles[self.config_data.active_profile]
+
+    @property
+    def api_key(self) -> str:
+        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
+
+    @property
+    def organization_id(self) -> str:
+        return self.config.organization_id
+
+    @organization_id.setter
+    def organization_id(self, value: str) -> None:
+        self.config.organization_id = value
+
+    @property
+    def platform_url(self) -> str:
+        return self.config.platform_url
+
+    @platform_url.setter
+    def platform_url(self, value: str) -> None:
+        base_url = value.rstrip("/")
+        self.config.platform_url = base_url
+        self.config.api_mapping = self.get_api_mapping(base_url)
+
+    def __init__(self, config_path: Optional[Path] = None) -> None:
+        self.config_path = self.resolve_config_path(config_path)
+        self.config_path.parent.mkdir(parents=True, exist_ok=True)
+        self.config_data = self.load_config()
+
+    def resolve_config_path(self, path: Optional[Path] = None) -> Path:
+        if path:
+            return Path(path).expanduser()
+
+        config_env_path = os.getenv("MDI_CONFIG_PATH")
+        if config_env_path:
+            return Path(config_env_path).expanduser()
+
+        return Path.home() / ".mdi" / "config.json"
+
+    def add_profile(
+        self, profile_name: str, profile: ConfigSchema, set_active: bool = False
+    ) -> None:
+        profile_name = profile_name.strip()
+        self.config_data.profiles[profile_name] = profile
+        if set_active:
+            self.config_data.active_profile = profile_name
+        self.save_config()
+
+    def get_api_mapping(self, base_url: str) -> Dict:
+        return {
+            "organizations": f"{base_url}/api/v1/organizations",
+            "recipes": f"{base_url}/api/v1/recipes",
+            "experiments": f"{base_url}/api/v1/experiments",
+            "experiment_environments": f"{base_url}/api/v1/experiment-environments",
+            "experiment_runs": f"{base_url}/api/v1/experiment-runs",
+            "runs": f"{base_url}/api/v1/experiments-runs",
+            "datasets": f"{base_url}/api/v1/datasets",
+        }
+
+    def get_platform_endpoint(self, method: str) -> str:
+        """Get specific API endpoint"""
+        if not self.config.api_mapping or method not in self.config.api_mapping:
+            raise ValueError(f"{method} is not supported.")
+        return self.config.api_mapping[method]
+
+    def load_config(self) -> ConfigFileSchema:
+        """Load configuration from file."""
+        if not self.config_path.exists():
+            return ConfigFileSchema()
+        try:
+            with open(self.config_path.as_posix(), "r") as f:
+                data = json.load(f)
+            return ConfigFileSchema(**data)
+        except json.JSONDecodeError:
+            logger.error("Error decoding JSON file.")
+            raise ValueError("Failed to parse configuration file")
+
+    def save_config(self) -> None:
+        with open(self.config_path, "w") as f:
+            json.dump(self.config_data.model_dump(), f, indent=4)
+
+    def remove_profile(self, profile_name: str) -> None:
+        if profile_name not in self.config_data.profiles:
+            raise ValueError(f"Profile '{profile_name}' does not exist.")
+        del self.config_data.profiles[profile_name]
+        self.save_config()
+
+    def clear(self) -> None:
+        self.config_data = ConfigFileSchema(active_profile=None, profiles={})
+        if self.config_path.exists():
+            self.config_path.unlink()

cli/consts.py

@@ -0,0 +1,18 @@
+ERROR_CONFIGURE: str = "Please configure the CLI with `mdi configure`"
+ERROR_PROFILE_NOT_EXIST: str = (
+    "No active profile configured. Please configure the CLI with `mdi configure`"
+)
+ERROR_PROFILE_REMOVE_ACTIVE: str = (
+    "Cannot remove active profile. Please switch to another profile first."
+)
+ERROR_API_KEY_NOT_SET: str = "API key not set. Please configure the CLI with `mdi configure`."
+ERROR_ORG_ID: str = "Failed to fetch organization ID. Verify platform URL and API key."
+ERROR_CREATE_PROFILE: str = "Failed to create profile. Profile name must be unique and not empty."
+
+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"
+
+PROFILE_SWITCHED_SUCCESS: str = "Switched to profile:"
+PROFILE_REMOVED_SUCCESS: str = "Removed profile:"
+PROFILE_TABLE_HEADER: str = "MDI Platform Profile:"

cli/data_cmds.py

@@ -0,0 +1,58 @@
+from typing import Optional
+
+import click
+from rich import print as rprint
+
+import cli.consts as consts
+from cli.config import Config
+
+
+@click.group()
+def data():
+    """Manage data interaction"""
+    pass
+
+
+@data.command("get")
+@click.argument("url")
+@click.argument("destination")
+@click.pass_obj
+def data_get(cfg: Config, url: str, destination: click.Path) -> None:
+    """Download resource from MDI platform"""
+
+    from hafnia.platform import download_resource
+
+    try:
+        result = download_resource(
+            resource_url=url, destination=str(destination), api_key=cfg.api_key
+        )
+    except Exception:
+        raise click.ClickException(consts.ERROR_GET_RESOURCE)
+
+    rprint(result)
+
+
+@data.command("download")
+@click.argument("dataset_name")
+@click.argument("destination", default=None, required=False)
+@click.option("--force", is_flag=True, default=False, help="Force download")
+@click.pass_obj
+def data_download(
+    cfg: Config, dataset_name: str, destination: Optional[click.Path], force: bool
+) -> None:
+    """Download dataset from MDI platform"""
+
+    from hafnia.data.factory import download_or_get_dataset_path
+
+    try:
+        endpoint_dataset = cfg.get_platform_endpoint("datasets")
+        api_key = cfg.api_key
+        download_or_get_dataset_path(
+            dataset_name=dataset_name,
+            endpoint=endpoint_dataset,
+            api_key=api_key,
+            output_dir=destination,
+            force=force,
+        )
+    except Exception:
+        raise click.ClickException(consts.ERROR_GET_RESOURCE)

cli/experiment_cmds.py

@@ -0,0 +1,82 @@
+from pathlib import Path
+
+import click
+from rich import print as rprint
+
+import cli.consts as consts
+from cli.config import Config
+
+
+@click.group(name="experiment")
+def experiment() -> None:
+    """Experiment management commands"""
+    pass
+
+
+@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.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 (IndexError, KeyError):
+        raise click.ClickException(f"Dataset '{dataset_name}' not found.")
+    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, cfg.organization_id
+        )
+    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,
+            cfg.organization_id,
+        )
+    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",
+        }
+    )

cli/profile_cmds.py

@@ -0,0 +1,86 @@
+import click
+from rich.console import Console
+from rich.table import Table
+
+import cli.consts as consts
+from cli.config import Config
+
+
+@click.group()
+def profile():
+    """Manage profile."""
+    pass
+
+
+@profile.command("ls")
+@click.pass_obj
+def profile_ls(cfg: Config) -> None:
+    """List all available profiles."""
+    profiles = cfg.available_profiles
+    if not profiles:
+        raise click.ClickException(consts.ERROR_CONFIGURE)
+    active = cfg.active_profile
+
+    for profile in profiles:
+        status = "* " if profile == active else "  "
+        print(f"{status}{profile}")
+
+    print(f"\nActive profile: {active}")
+
+
+@profile.command("use")
+@click.argument("profile_name", required=True)
+@click.pass_obj
+def profile_use(cfg: Config, profile_name: str) -> None:
+    """Switch to a different profile."""
+    if len(cfg.available_profiles) == 0:
+        raise click.ClickException(consts.ERROR_CONFIGURE)
+    try:
+        cfg.active_profile = profile_name
+        cfg.save_config()
+    except ValueError:
+        raise click.ClickException(consts.ERROR_PROFILE_NOT_EXIST)
+    click.echo(f"{consts.PROFILE_SWITCHED_SUCCESS} {profile_name}")
+
+
+@profile.command("rm")
+@click.argument("profile_name", required=True)
+@click.pass_obj
+def profile_rm(cfg: Config, profile_name: str) -> None:
+    """Remove a profile."""
+    if len(cfg.available_profiles) == 0:
+        raise click.ClickException(consts.ERROR_CONFIGURE)
+
+    if profile_name == cfg.active_profile:
+        raise click.ClickException(consts.ERROR_PROFILE_REMOVE_ACTIVE)
+
+    try:
+        cfg.remove_profile(profile_name)
+        cfg.save_config()
+    except ValueError:
+        raise click.ClickException(consts.ERROR_PROFILE_NOT_EXIST)
+    click.echo(f"{consts.PROFILE_REMOVED_SUCCESS} {profile_name}")
+
+
+@profile.command("active")
+@click.pass_obj
+def profile_active(cfg: Config) -> None:
+    try:
+        profile_show(cfg)
+    except Exception as e:
+        raise click.ClickException(str(e))
+
+
+def profile_show(cfg: Config) -> None:
+    masked_key = f"{cfg.api_key[:4]}...{cfg.api_key[-4:]}" if len(cfg.api_key) > 8 else "****"
+    console = Console()
+
+    table = Table(title=f"{consts.PROFILE_TABLE_HEADER} {cfg.active_profile}", show_header=False)
+    table.add_column("Property", style="cyan")
+    table.add_column("Value")
+
+    table.add_row("API Key", masked_key)
+    table.add_row("Organization", cfg.organization_id)
+    table.add_row("Platform URL", cfg.platform_url)
+    table.add_row("Config File", cfg.config_path.as_posix())
+    console.print(table)

cli/runc_cmds.py

@@ -0,0 +1,70 @@
+from hashlib import sha256
+from pathlib import Path
+from tempfile import TemporaryDirectory
+
+import click
+
+from cli.config import Config
+
+
+@click.group(name="runc")
+def runc():
+    """Experiment management commands"""
+    pass
+
+
+@runc.command(name="launch")
+@click.argument("task", required=True)
+def launch(task: str) -> None:
+    """Launch a job within the image."""
+    from hafnia.platform.executor import handle_launch
+
+    handle_launch(task)
+
+
+@runc.command(name="build")
+@click.argument("recipe_url")
+@click.argument("state_file", default="state.json")
+@click.argument("ecr_repository", default="localhost")
+@click.argument("image_name", default="recipe")
+@click.pass_obj
+def build(
+    cfg: Config, recipe_url: str, state_file: str, ecr_repository: str, image_name: str
+) -> None:
+    """Build docker image with a given recipe."""
+    from hafnia.platform.builder import build_image, prepare_recipe
+
+    with TemporaryDirectory() as temp_dir:
+        image_info = prepare_recipe(recipe_url, Path(temp_dir), cfg.api_key)
+        image_info["name"] = image_name
+        build_image(image_info, ecr_repository, state_file)
+
+
+@runc.command(name="build-local")
+@click.argument("recipe")
+@click.argument("state_file", default="state.json")
+@click.argument("image_name", default="recipe")
+def build_local(recipe: str, state_file: str, image_name: str) -> None:
+    """Build recipe from local path as image with prefix - localhost"""
+
+    from hafnia.platform.builder import build_image, validate_recipe
+    from hafnia.utils import archive_dir
+
+    recipe_zip = Path(recipe)
+    recipe_created = False
+    if not recipe_zip.suffix == ".zip" and recipe_zip.is_dir():
+        recipe_zip = archive_dir(recipe_zip)
+        recipe_created = True
+
+    validate_recipe(recipe_zip)
+    click.echo("Recipe successfully validated")
+    image_info = {
+        "name": image_name,
+        "dockerfile": f"{recipe_zip.parent}/Dockerfile",
+        "docker_context": f"{recipe_zip.parent}",
+        "hash": sha256(recipe_zip.read_bytes()).hexdigest()[:8],
+    }
+    click.echo("Start building image")
+    build_image(image_info, "localhost", state_file=state_file)
+    if recipe_created:
+        recipe_zip.unlink()

hafnia/__init__.py

@@ -0,0 +1,4 @@
+from importlib.metadata import version
+
+__package_name__ = "hafnia"
+__version__ = version(__package_name__)

hafnia/data/__init__.py

@@ -0,0 +1,3 @@
+from hafnia.data.factory import load_dataset
+
+__all__ = ["load_dataset"]

hafnia/data/factory.py

@@ -0,0 +1,88 @@
+import os
+import shutil
+from pathlib import Path
+from typing import Optional, Union
+
+from datasets import Dataset, DatasetDict, load_from_disk
+
+from cli.config import Config
+from hafnia import utils
+from hafnia.log import logger
+from hafnia.platform import download_resource, get_dataset_id
+
+
+def load_local(dataset_path: Path) -> Union[Dataset, DatasetDict]:
+    """Load a Hugging Face dataset from a local directory path."""
+    if not dataset_path.exists():
+        raise ValueError(f"Can not load dataset, directory does not exist -- {dataset_path}")
+    logger.info(f"Loading data from {dataset_path.as_posix()}")
+    return load_from_disk(dataset_path.as_posix())
+
+
+def download_or_get_dataset_path(
+    dataset_name: str,
+    endpoint: str,
+    api_key: str,
+    output_dir: Optional[str] = None,
+    force_redownload: bool = False,
+) -> Path:
+    """Download or get the path of the dataset."""
+    output_dir = output_dir or str(utils.PATH_DATASET)
+    dataset_path_base = Path(output_dir).absolute() / dataset_name
+    dataset_path_base.mkdir(exist_ok=True, parents=True)
+    dataset_path_sample = dataset_path_base / "sample"
+
+    if dataset_path_sample.exists() and not force_redownload:
+        logger.info(
+            "Dataset found locally. Set 'force=True' or add `--force` flag with cli to re-download"
+        )
+        return dataset_path_sample
+
+    dataset_id = get_dataset_id(dataset_name, endpoint, api_key)
+    dataset_access_info_url = f"{endpoint}/{dataset_id}/temporary-credentials"
+
+    if force_redownload and dataset_path_sample.exists():
+        # Remove old files to avoid old files conflicting with new files
+        shutil.rmtree(dataset_path_sample, ignore_errors=True)
+    status = download_resource(dataset_access_info_url, dataset_path_base, api_key)
+    if status:
+        return dataset_path_sample
+    raise RuntimeError("Failed to download dataset")
+
+
+def load_from_platform(
+    dataset_name: str,
+    endpoint: str,
+    api_key: str,
+    output_dir: Optional[str] = None,
+    force_redownload: bool = False,
+) -> Union[Dataset, DatasetDict]:
+    path_dataset = download_or_get_dataset_path(
+        dataset_name=dataset_name,
+        endpoint=endpoint,
+        api_key=api_key,
+        output_dir=output_dir,
+        force_redownload=force_redownload,
+    )
+    return load_local(path_dataset)
+
+
+def load_dataset(dataset_name: str, force_redownload: bool = False) -> Union[Dataset, DatasetDict]:
+    """Load a dataset either from a local path or from the MDI platform."""
+
+    if utils.is_remote_job():
+        path_dataset = Path(os.getenv("MDI_DATASET_DIR", "/opt/ml/input/data/training"))
+        return load_local(path_dataset)
+
+    cfg = Config()
+    endpoint_dataset = cfg.get_platform_endpoint("datasets")
+    api_key = cfg.api_key
+    dataset = load_from_platform(
+        dataset_name=dataset_name,
+        endpoint=endpoint_dataset,
+        api_key=api_key,
+        output_dir=None,
+        force_redownload=force_redownload,
+    )
+
+    return dataset

hafnia/experiment/__init__.py

@@ -0,0 +1,3 @@
+from hafnia.experiment.mdi_logger import MDILogger
+
+__all__ = ["MDILogger"]