add index and oci-layout file

Signed-off-by: Jaideep Rao <[email protected]>

olot/oci/oci_image_index.py

@@ -9,7 +9,7 @@
 
 from pydantic import BaseModel, Field
 
-from olot.oci.oci_common import MediaType, Digest, Urls
+from olot.oci.oci_common import MediaTypes, MediaType, Digest, Urls
 from olot.utils.types import  Int64, Base64, Annotations
 
 class Platform(BaseModel):
@@ -199,3 +199,24 @@ def read_ocilayout_root_index(ocilayout: Path) -> OCIImageIndex:
     with open(ocilayout / "index.json", "r") as f:
         ocilayout_root_index = OCIImageIndex.model_validate_json(f.read())
     return ocilayout_root_index
+
+
+def create_oci_image_index(
+    schemaVersion: int = 2,
+    mediaType: Optional[str] = MediaTypes.index,
+    artifactType: Optional[str] = None,
+    subject: Optional[ContentDescriptor] = None,
+    manifests: List[Manifest] = [],
+    annotations: Optional[Annotations] = None
+) -> OCIImageIndex:
+    """
+    Create an OCI image index object.
+    """
+    return OCIImageIndex(
+        schemaVersion=schemaVersion,
+        mediaType=mediaType,
+        artifactType=artifactType,
+        subject=subject,
+        manifests=manifests,
+        annotations=annotations
+    )

olot/oci/oci_image_layout.py

@@ -19,10 +19,16 @@ class OCIImageLayout(BaseModel):
         ..., description='version of the OCI Image Layout (in the oci-layout file)'
     )
 
+    class Config:
+        use_enum_values = True
+
 def verify_ocilayout(ocilayout: Path):
     with open(ocilayout / "oci-layout", "r") as f:
         m = OCIImageLayout.model_validate_json(f.read())
-        if not m.imageLayoutVersion == ImageLayoutVersion.field_1_0_0:
+        if not m.imageLayoutVersion == ImageLayoutVersion.field_1_0_0.value:
             raise ValueError(f"Unexpected ocilayout in {ocilayout}")
         else:
             return True
+
+def create_ocilayout() -> OCIImageLayout:
+    return OCIImageLayout(imageLayoutVersion=ImageLayoutVersion.field_1_0_0)

olot/oci_artifact.py

@@ -4,11 +4,20 @@
 import json
 
 from olot.oci.oci_image_manifest import create_oci_image_manifest, create_manifest_layers
-from olot.oci.oci_common import Keys
+from olot.oci.oci_image_layout import create_ocilayout
+from olot.oci.oci_common import Keys, MediaTypes
+from olot.oci.oci_image_index import Manifest, create_oci_image_index
 from olot.utils.files import MIMETypes, tarball_from_file, targz_from_file
 from olot.utils.types import compute_hash_of_str
 
 def create_oci_artifact_from_model(source_dir: Path, dest_dir: Path):
+    """
+    Create an OCI artifact from a model directory.
+
+    Args:
+        source_dir: The directory containing the model files.
+        dest_dir: The directory to write the OCI artifact to. If None, a directory named 'oci' will be created in the source directory.
+    """
     if not source_dir.exists():
         raise NotADirectoryError(f"Input directory '{source_dir}' does not exist.")
 
@@ -39,13 +48,39 @@ def create_oci_artifact_from_model(source_dir: Path, dest_dir: Path):
         layers=manifest_layers,
         annotations=annotations
     )
-    manifest_json = json.dumps(manifest.dict(), indent=4, sort_keys=True)
+    manifest_json = json.dumps(manifest.model_dump(), indent=4, sort_keys=True)
     manifest_SHA = compute_hash_of_str(manifest_json)
     with open(sha256_path / manifest_SHA, "w") as f:
         f.write(manifest_json)
 
+    # Create the OCI image index
+    index = create_oci_image_index(
+        manifests = [
+            Manifest(
+                mediaType=MediaTypes.manifest,
+                size=os.stat(sha256_path / manifest_SHA).st_size,
+                digest=f"sha256:{manifest_SHA}",
+                urls=None,
+                platform=None,
+                annotations=None
+            )
+        ]
+    )
+    index_json = json.dumps(index.model_dump(), indent=4, sort_keys=True)
+    with open(dest_dir / "index.json", "w") as f:
+        f.write(index_json)
+
+
+    # Create the OCI-layout file
+    oci_layout = create_ocilayout()
+    with open(dest_dir / "oci-layout", "w") as f:
+        f.write(json.dumps(oci_layout.model_dump(), indent=4, sort_keys=True))
+
 
 def create_blobs(source_dir: Path, dest_dir: Path):
+    """
+    Create the blobs directory for an OCI artifact.
+    """
     layers = {} # layer digest : diff_id
     sha256_path = dest_dir / "blobs" / "sha256"
 
@@ -64,5 +99,10 @@ def create_blobs(source_dir: Path, dest_dir: Path):
             layers[checksum] = checksum
     return layers
 
+# create a main function to test the function
+def main():
+    source_dir = Path("/Users/jrao/.cache/instructlab/models/huggingface.co/instructlab/granite-7b-lab/main")
+    create_oci_artifact_from_model(source_dir, None)
 
-
+if __name__ == "__main__":
+    main()