← back to invincible-jha__aumai-modeloci

Function bodies 10 total

All specs Real LLM only Function bodies
pack_command function · python · L48-L81 (34 LOC)
src/aumai_modeloci/cli.py
def pack_command(
    model_dir: str,
    name: str,
    model_version: str,
    framework: str,
    architecture: str,
    metadata_json: str,
) -> None:
    """Package a model directory into an OCI-compliant tar archive."""
    try:
        metadata = json.loads(metadata_json)
    except json.JSONDecodeError as exc:
        click.echo(f"Error: invalid JSON for --metadata: {exc}", err=True)
        sys.exit(1)

    config = OCIConfig(
        model_name=name,
        version=model_version,
        framework=framework,
        architecture=architecture,
        metadata=metadata,
    )
    packager = ModelPackager()
    try:
        archive_path = packager.package(model_dir, config)
    except (NotADirectoryError, FileNotFoundError) as exc:
        click.echo(f"Error: {exc}", err=True)
        sys.exit(1)

    click.echo(f"Packaged model: {archive_path}")
    click.echo(f"  Name        : {name}")
    click.echo(f"  Version     : {model_version}")
    click.echo(f"  Framework   : {frame
unpack_command function · python · L99-L110 (12 LOC)
src/aumai_modeloci/cli.py
def unpack_command(archive_path: str, output_dir: str) -> None:
    """Unpack an OCI model archive."""
    unpacker = ModelUnpackager()
    try:
        config = unpacker.unpack(archive_path, output_dir)
    except (FileNotFoundError, Exception) as exc:
        click.echo(f"Error: {exc}", err=True)
        sys.exit(1)

    click.echo(f"Unpacked to: {output_dir}")
    click.echo(f"  Model   : {config.model_name} v{config.version}")
    click.echo(f"  Framework: {config.framework}")
inspect_command function · python · L121-L178 (58 LOC)
src/aumai_modeloci/cli.py
def inspect_command(archive_path: str) -> None:
    """Inspect an OCI model archive without extracting it."""
    import tarfile

    try:
        with tarfile.open(archive_path, "r") as tar:
            members = {m.name: m for m in tar.getmembers()}

            # Read config
            cfg_member = members.get("config.json")
            if cfg_member:
                fh = tar.extractfile(cfg_member)
                if fh:
                    config = OCIConfig.model_validate(json.loads(fh.read()))
                    click.echo(f"Model    : {config.model_name}")
                    click.echo(f"Version  : {config.version}")
                    click.echo(f"Framework: {config.framework}")
                    click.echo(f"Arch     : {config.architecture}")
                    if config.metadata:
                        click.echo(f"Metadata : {json.dumps(config.metadata)}")

            # Read manifest
            man_member = members.get("manifest.json")
            if man_member:
 
_sha256_file function · python · L26-L32 (7 LOC)
src/aumai_modeloci/core.py
def _sha256_file(path: str) -> str:
    """Return 'sha256:<hex>' digest for the file at *path*."""
    h = hashlib.sha256()
    with open(path, "rb") as fh:
        for chunk in iter(lambda: fh.read(65536), b""):
            h.update(chunk)
    return f"sha256:{h.hexdigest()}"
ModelPackager.package method · python · L51-L97 (47 LOC)
src/aumai_modeloci/core.py
    def package(self, model_dir: str, config: OCIConfig) -> str:
        """
        Create an OCI-compliant tar archive from *model_dir*.

        Returns the path to the created archive.
        """
        model_path = Path(model_dir)
        if not model_path.is_dir():
            raise NotADirectoryError(f"{model_dir!r} is not a directory.")

        output_archive = str(
            model_path.parent / f"{config.model_name}-{config.version}.tar"
        )

        layers: list[ModelLayer] = []
        with tempfile.TemporaryDirectory() as tmp_dir:
            tmp = Path(tmp_dir)
            blobs_dir = tmp / _LAYERS_DIR
            blobs_dir.mkdir(parents=True)

            # Walk model_dir and create one layer per file
            for file_path in sorted(model_path.rglob("*")):
                if file_path.is_file():
                    layer = self._create_layer_blob(
                        str(file_path), str(blobs_dir), str(model_path)
                    )
                 
ModelPackager.create_manifest method · python · L99-L123 (25 LOC)
src/aumai_modeloci/core.py
    def create_manifest(
        self, config: OCIConfig, layers: list[ModelLayer]
    ) -> OCIManifest:
        """Build an OCIManifest from config and layer list."""
        config_bytes = config.model_dump_json(indent=2).encode("utf-8")
        config_digest = _sha256_bytes(config_bytes)

        config_descriptor: dict[str, Any] = {
            "mediaType": "application/vnd.oci.image.config.v1+json",
            "digest": config_digest,
            "size": len(config_bytes),
        }
        layer_descriptors: list[dict[str, Any]] = [
            {
                "mediaType": layer.media_type,
                "digest": layer.digest,
                "size": layer.size,
                "annotations": layer.annotations,
            }
            for layer in layers
        ]
        return OCIManifest(
            config=config_descriptor,
            layers=layer_descriptors,
        )
ModelPackager.add_layer method · python · L125-L152 (28 LOC)
src/aumai_modeloci/core.py
    def add_layer(
        self, archive_path: str, file_path: str
    ) -> ModelLayer:
        """
        Add a file as a new layer to an existing archive.

        Returns a ``ModelLayer`` descriptor for the added file.
        """
        if not Path(file_path).is_file():
            raise FileNotFoundError(f"File not found: {file_path!r}")

        with tempfile.TemporaryDirectory() as tmp_dir:
            tmp = Path(tmp_dir)
            blobs_dir = tmp / _LAYERS_DIR
            blobs_dir.mkdir(parents=True)

            # Create a layer blob for the new file
            layer = self._create_layer_blob(file_path, str(blobs_dir), "")

            # Append blob to the archive
            with tarfile.open(archive_path, "a") as tar:
                blob_name = layer.digest.split(":")[1]
                tar.add(
                    str(blobs_dir / blob_name),
                    arcname=f"{_LAYERS_DIR}/{blob_name}",
                )

        return layer
Repobility · MCP-ready · https://repobility.com
ModelPackager._create_layer_blob method · python · L158-L182 (25 LOC)
src/aumai_modeloci/core.py
    def _create_layer_blob(
        self, file_path: str, blobs_dir: str, base_dir: str
    ) -> ModelLayer:
        """
        Create a gzip-compressed tar layer from a single file.

        The layer blob is written to *blobs_dir* with its SHA-256 as filename.
        """
        buf = io.BytesIO()
        with tarfile.open(fileobj=buf, mode="w:gz") as layer_tar:
            arcname = (
                os.path.relpath(file_path, base_dir) if base_dir else Path(file_path).name
            )
            layer_tar.add(file_path, arcname=arcname)
        blob_bytes = buf.getvalue()
        digest = _sha256_bytes(blob_bytes)
        hex_digest = digest.split(":")[1]
        blob_path = Path(blobs_dir) / hex_digest
        blob_path.write_bytes(blob_bytes)
        rel_path = os.path.relpath(file_path, base_dir) if base_dir else Path(file_path).name
        return ModelLayer(
            digest=digest,
            size=len(blob_bytes),
            annotations={"org.opencontainers.image.tit
ModelUnpackager.unpack method · python · L190-L222 (33 LOC)
src/aumai_modeloci/core.py
    def unpack(self, archive_path: str, output_dir: str) -> OCIConfig:
        """
        Extract *archive_path* into *output_dir*.

        Reads ``config.json`` from the archive root and returns the
        parsed ``OCIConfig``.
        """
        output = Path(output_dir)
        output.mkdir(parents=True, exist_ok=True)

        with tarfile.open(archive_path, "r") as tar:
            try:
                tar.extractall(path=str(output), filter="data")
            except TypeError:
                # Python <3.12 does not support the filter argument.
                # Manually validate each member to prevent path traversal.
                resolved_output = output.resolve()
                for member in tar.getmembers():
                    member_path = (resolved_output / member.name).resolve()
                    if not str(member_path).startswith(str(resolved_output)):
                        raise ValueError(
                            f"Attempted path traversal in archive me
ModelUnpackager.verify_layers method · python · L224-L266 (43 LOC)
src/aumai_modeloci/core.py
    def verify_layers(
        self, archive_path: str
    ) -> list[tuple[str, bool]]:
        """
        Verify the SHA-256 digest of every layer blob in the archive.

        Returns a list of (digest, is_valid) tuples.  A layer is valid
        when its stored filename matches the SHA-256 of its content.
        """
        results: list[tuple[str, bool]] = []

        with tarfile.open(archive_path, "r") as tar:
            members = {m.name: m for m in tar.getmembers()}

            # Read manifest to know which digests to expect
            manifest_member = members.get(_MANIFEST_FILENAME)
            if manifest_member is None:
                raise FileNotFoundError(
                    f"manifest.json not found in {archive_path!r}."
                )
            manifest_fh = tar.extractfile(manifest_member)
            if manifest_fh is None:
                raise RuntimeError("Could not read manifest.json from archive.")
            manifest = OCIManifest.model_validate(