modelcar-maker 0.4.0__py3-none-any.whl

modelcar_maker/__init__.py

@@ -0,0 +1,154 @@
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Optional
+
+from .download import hf_download
+from .image.types import Backend
+from .util import normalize
+from .util import settings
+
+
+@dataclass
+class ProcessResult:
+    """Represents the state of the process workflow."""
+
+    skipped: bool = False
+    downloaded_to: Optional[Path] = None
+    image: Optional[str] = None
+    image_built: bool = False
+    image_pushed: bool = False
+    image_cleaned_up: bool = False
+    model_cleaned_up: bool = False
+
+
+def cleanup(path: Path, skip: list[str] = ["Containerfile"]) -> bool:
+    """Remove all children of the provided path, except for top-level files provided in skip."""
+    changed = False
+    for subpath in path.iterdir():
+        if subpath.name in skip:
+            continue
+        if subpath.is_dir():
+            _ = cleanup(subpath, skip=[])
+            subpath.rmdir()
+            changed = True
+        else:
+            subpath.unlink()
+            changed = True
+    return changed
+
+
+def process(
+    model: str,
+    image_repo: str = f"{settings.image.registry}/{settings.image.repository}",
+    backend: Backend | str = settings.image.backend,
+    base_image: str = settings.image.base_image,
+    architectures: list[str] = settings.image.architectures,
+    authfile: Path | None = settings.image.get("authfile"),
+    push: bool = settings.image.push,
+    image_cleanup: bool = settings.image.cleanup,
+    model_cleanup: bool = settings.models.cleanup,
+    skip_if_exists: bool = settings.image.skip_if_exists,
+    pull: bool = settings.image.pull,
+) -> ProcessResult:
+    """Run through the entire process of downloading, packaging, and publishing a Model Car image."""
+    result = ProcessResult()
+
+    if isinstance(backend, str):
+        try:
+            backend = Backend(backend)
+        except ValueError:
+            raise NotImplementedError(f"Backend {backend!r} is not supported") from None
+
+    from .image.types import BuildArgs
+    from .image.types import PushArgs
+    from .image.types import RmArgs
+
+    if backend is Backend.PODMAN:
+        from .image.podman import do_build
+        from .image.podman import do_image_rm
+        from .image.podman import do_push
+        from .image.podman import image_exists
+    else:
+        from .image.olot import do_build
+        from .image.olot import do_image_rm
+        from .image.olot import do_push
+        from .image.olot import image_exists
+
+    if skip_if_exists:
+        if image_exists(model, image_repo):
+            # It was requested that we skip the build, and the image exists.
+            # We should still check if a cleanup is called for.
+            result.skipped = True
+            if image_cleanup:
+                oci_layout_dir = Path("tmp").joinpath(normalize(model)) if backend is not Backend.PODMAN else None
+                if do_image_rm(
+                    RmArgs(model=model, repo=image_repo, oci_layout_dir=oci_layout_dir, architectures=architectures)
+                ):
+                    result.image_cleaned_up = True
+            if model_cleanup:
+                download_dir = Path("models").joinpath(normalize(model))
+                result.model_cleaned_up = cleanup(download_dir)
+            return result
+
+    # Ensure that the model is downloaded before the build
+    download_dir, commit = hf_download(model)
+    result.downloaded_to = download_dir
+
+    # Build the image
+    build_result = do_build(
+        BuildArgs(
+            model=model,
+            repo=image_repo,
+            model_dir=download_dir,
+            base_image=base_image,
+            commit=commit,
+            pull=pull,
+            architectures=architectures,
+        )
+    )
+    result.image = build_result.image
+    result.image_built = True
+
+    if push:
+        do_push(
+            PushArgs(
+                model=model,
+                repo=image_repo,
+                authfile=authfile,
+                oci_layout_dir=build_result.oci_layout_dir,
+                architectures=architectures,
+                manifest_list=build_result.manifest_list,
+            )
+        )
+        result.image_pushed = True
+
+    if image_cleanup:
+        if backend is Backend.PODMAN:
+            if image_exists(model, image_repo):
+                if do_image_rm(
+                    RmArgs(
+                        model=model,
+                        repo=image_repo,
+                        oci_layout_dir=build_result.oci_layout_dir,
+                        architectures=architectures,
+                        manifest_list=build_result.manifest_list,
+                    )
+                ):
+                    result.image_cleaned_up = True
+        else:
+            if do_image_rm(
+                RmArgs(
+                    model=model,
+                    repo=image_repo,
+                    oci_layout_dir=build_result.oci_layout_dir,
+                    architectures=architectures,
+                    manifest_list=build_result.manifest_list,
+                )
+            ):
+                result.image_cleaned_up = True
+
+    if model_cleanup and result.image_built:
+        # Only clean up the files if it was skipped (above) or an image definitely got built
+        result.model_cleaned_up = cleanup(download_dir)
+
+    return result

modelcar_maker/__version__.py

@@ -0,0 +1,24 @@
+# file generated by vcs-versioning
+# don't change, don't track in version control
+from __future__ import annotations
+
+__all__ = [
+    "__version__",
+    "__version_tuple__",
+    "version",
+    "version_tuple",
+    "__commit_id__",
+    "commit_id",
+]
+
+version: str
+__version__: str
+__version_tuple__: tuple[int | str, ...]
+version_tuple: tuple[int | str, ...]
+commit_id: str | None
+__commit_id__: str | None
+
+__version__ = version = '0.4.0'
+__version_tuple__ = version_tuple = (0, 4, 0)
+
+__commit_id__ = commit_id = None

modelcar_maker/cli/__init__.py

@@ -0,0 +1,3 @@
+from .cli import cli
+
+__all__ = ["cli"]

modelcar_maker/cli/cli.py

@@ -0,0 +1,177 @@
+from pathlib import Path
+from typing import Optional
+
+import typer
+from rich import print as rprint
+from typing_extensions import Annotated
+
+from .. import process
+from ..image.types import Backend
+from ..util import make_logger
+from ..util import settings
+
+
+def version_callback(value: bool):
+    if value:
+        from ..__version__ import version
+
+        rprint(version)
+        raise typer.Exit()
+
+
+cli = typer.Typer(
+    context_settings={"help_option_names": ["-h", "--help"]},
+    pretty_exceptions_show_locals=False,
+    add_completion=False,
+)
+
+
+@cli.command()
+def build(
+    ctx: typer.Context,
+    model: Annotated[
+        Optional[str],
+        typer.Argument(help="The model to download (otherwise ensure all defaults are downloaded)"),
+    ] = None,
+    registry: Annotated[
+        str,
+        typer.Option(
+            "--registry",
+            help="The registry to include in the image reference",
+        ),
+    ] = settings.image.registry,
+    repository: Annotated[
+        str,
+        typer.Option(
+            "-r",
+            "--repository",
+            help="The repository, within the registry, to include in the image reference",
+        ),
+    ] = settings.image.repository,
+    backend: Annotated[
+        Backend,
+        typer.Option(
+            "--backend",
+            help="Build backend to use: olot (default) or podman",
+        ),
+    ] = Backend(settings.image.backend),
+    base_image: Annotated[
+        str,
+        typer.Option(
+            "--base-image",
+            help="Base image to use for the modelcar (applies to both backends)",
+        ),
+    ] = settings.image.base_image,
+    arch: Annotated[
+        list[str],
+        typer.Option(
+            "--arch",
+            help="Target architecture(s) to build for (repeat for multiple)",
+        ),
+    ] = settings.image.architectures,
+    pull: Annotated[
+        bool,
+        typer.Option(
+            "--pull/--no-pull",
+            help="Pull the base image before building if a newer version is available",
+        ),
+    ] = settings.image.pull,
+    push: Annotated[
+        bool,
+        typer.Option(
+            "--push/--no-push",
+            help="Push the image(s) after building",
+        ),
+    ] = settings.image.push,
+    authfile: Annotated[
+        Optional[Path],
+        typer.Option(
+            "-a",
+            "--authfile",
+            help="The authfile to use for the push",
+        ),
+    ] = settings.image.get("authfile"),
+    image_cleanup: Annotated[
+        bool,
+        typer.Option(
+            "--image-clean-up/--no-image-clean-up",
+            help="Clean up the container image after push to free up space",
+        ),
+    ] = settings.image.cleanup,
+    model_cleanup: Annotated[
+        bool,
+        typer.Option(
+            "--model-clean-up/--no-model-clean-up",
+            help="Clean up the downloaded model images after build to free up space",
+        ),
+    ] = settings.models.cleanup,
+    skip_if_exists: Annotated[
+        bool,
+        typer.Option(
+            "--skip-if-exists/--no-skip-if-exists",
+            help="Skip downloading and publishing the image if it already exists",
+        ),
+    ] = settings.image.skip_if_exists,
+    verbose: Annotated[
+        int,
+        typer.Option(
+            "--verbose",
+            "-v",
+            count=True,
+            help="Increase logging verbosity (repeat for more)",
+            show_default=False,
+            metavar="",
+        ),
+    ] = settings.meta.verbosity,
+    _: Annotated[
+        bool,
+        typer.Option(
+            "--version",
+            "-V",
+            callback=version_callback,
+            help="Print the version and exit",
+            is_eager=True,
+            show_default=False,
+        ),
+    ] = False,
+):
+    """Modelcar Maker model download, build, and push"""
+
+    logger = make_logger(verbose)
+    image_repo = f"{registry}/{repository}"
+    logger.debug(f"Backend: {backend}, Architectures: {arch}, Push {image_repo}: {push}")
+    if model is None:
+        if not isinstance(settings.models.default, list):
+            ctx.fail(f"Default models should be a list, not {type(settings.models.default)}")
+        if len(settings.models.default) < 1:
+            ctx.fail("No model provided, default models list is empty")
+        models = settings.models.default
+    else:
+        models = [model]
+
+    for model in models:
+        result = process(
+            str(model),
+            image_repo,
+            backend=backend,
+            base_image=base_image,
+            architectures=arch,
+            authfile=authfile,
+            push=push,
+            image_cleanup=image_cleanup,
+            model_cleanup=model_cleanup,
+            skip_if_exists=skip_if_exists,
+            pull=pull,
+        )
+        cleanup_str = f"Cleanup: {'✅' if result.image_cleaned_up else '❌'} Image, {'✅' if result.model_cleaned_up else '❌'} Model"
+
+        if result.skipped:
+            rprint(f"{model} was skipped as it already exists at {image_repo} - {cleanup_str}.")
+        if result.image_pushed:
+            rprint(
+                f"{model} was downloaded to {result.downloaded_to}, built as {result.image}, and pushed - {cleanup_str}."
+            )
+        elif result.image_built:
+            rprint(
+                f"{model} was downloaded to {result.downloaded_to} and built as {result.image}, but not pushed - {cleanup_str}."
+            )

modelcar_maker/download/__init__.py

@@ -0,0 +1,3 @@
+from .hf_download import hf_download
+
+__all__ = ["hf_download"]

modelcar_maker/download/hf_download.py

@@ -0,0 +1,35 @@
+import warnings
+from pathlib import Path
+
+from huggingface_hub import snapshot_download
+from rich import print as rprint
+from tqdm import TqdmExperimentalWarning
+from tqdm.rich import tqdm_rich
+
+from ..util import logger
+from ..util import normalize
+
+warnings.filterwarnings("ignore", category=TqdmExperimentalWarning)
+
+
+def hf_download(repo_id: str) -> tuple[Path, str]:
+    """Download the given model repo_id into the models/ directory, returning the path."""
+    from huggingface_hub import HfApi
+
+    normalized = normalize(repo_id)
+    download_dir = Path(f"models/{normalized}")
+    logger.info(f"Downloading {repo_id} to {download_dir}")
+    if not download_dir.is_dir():
+        download_dir.mkdir(parents=True, exist_ok=True)
+
+    rprint(f"Downloading {repo_id} to {download_dir}")
+    snapshot_download(repo_id, local_dir=download_dir, tqdm_class=tqdm_rich)  # type: ignore[arg-type]
+
+    # Get the commit hash from the public API instead of internal cache metadata.
+    sha = HfApi().model_info(repo_id).sha
+    assert sha is not None
+
+    return (
+        download_dir,
+        sha,
+    )

modelcar_maker/image/__init__.py

@@ -0,0 +1,17 @@
+from .common import _image
+from .common import list_model_files
+from .template import render
+from .types import BuildArgs
+from .types import BuildResult
+from .types import PushArgs
+from .types import RmArgs
+
+__all__ = [
+    "_image",
+    "list_model_files",
+    "render",
+    "BuildArgs",
+    "BuildResult",
+    "PushArgs",
+    "RmArgs",
+]

modelcar_maker/image/common.py

@@ -0,0 +1,41 @@
+from pathlib import Path
+
+from ..util import normalize
+
+
+def _image(model: str, repo: str) -> str:
+    """Given a model repo id and container image repository, render the full image name."""
+    tag = normalize(model) + "-modelcar"
+    return f"{repo}:{tag}"
+
+
+def should_include(file: str | Path) -> bool:
+    """Simple filter function to determine if a file should be included in an image definition."""
+    if isinstance(file, Path):
+        file = file.name
+
+    if file.startswith("."):
+        return False
+    if file == "Containerfile":
+        return False
+    if file == "original":
+        return False
+    if file == "consolidated.safetensors":
+        return False
+    return True
+
+
+def list_model_files(model_dir: Path) -> tuple[list[str], str | None]:
+    """List files to include in the modelcar and identify the modelcard source.
+
+    Returns:
+        A tuple of (sorted list of included filenames, modelcard_source or None).
+    """
+    model_files = [file.name for file in model_dir.iterdir() if should_include(file)]
+    model_files.sort()
+    modelcard_source = None
+    for file in model_files:
+        if file.startswith("README"):
+            modelcard_source = file
+            break
+    return model_files, modelcard_source