Compression frames: consecutive new sub-chunks compressed together (B2→B8 28.3 → 3.6 MB, store 17.0 → 16.7 GB) (#52)

README.md

@@ -67,7 +67,8 @@ export ONGGI_STORE='s3://models/onggi?endpoint=s3.lab.example:8333&insecure=1' A
 On the lab node, the 264 GB ATOM Llama-3.1-8B build set is stored in **15.9 GB** (push 3 min, pull 5 min,
 byte-identical). Two steps get it there. Two-level chunking re-splits new ATOM chunks into ~4 KiB sub-chunks, so prefill
 variants that hold the same weights in a different order dedup (24.9 GB new instead of 46.8 GB). Per-chunk compression
-then takes 1.6× off. `--compress none` and `--sub 0` turn the two off.
+then takes 1.6× off. `--compress none` and `--sub 0` turn the two off. New chunks of sub-chunked files are compressed together in frames
+of up to 255 KiB where that is smaller (#52; `--frame 0` turns it off).
 
 With `--cache DIR` a pull reads chunks the node already holds from a local cache store and adds what it fetches. On
 the lab node, after one ATOM decode variant (10.2 GB read), a second one read 28 MB from the store (0.3 %), a prefill

cmd/onggi/store.go

@@ -25,7 +25,7 @@ import (
 )
 
 type storeFlags struct {
-	url, durability, min, avg, max, sub, subKinds, packSize *string
+	url, durability, min, avg, max, sub, subKinds, packSize, frame *string
 }
 
 func addStoreFlags(fs *flag.FlagSet) storeFlags {
@@ -38,6 +38,7 @@ func addStoreFlags(fs *flag.FlagSet) storeFlags {
 		sub:        fs.String("sub", "4KiB", "average sub-chunk size for two-level chunking, 0 for one level (only when creating a store)"),
 		subKinds:   fs.String("sub-kinds", "atom-prefill,atom-decode,atom-other", "file kinds that get sub-chunks, or all (only when creating a store)"),
 		packSize:   fs.String("pack-size", "64MiB", "pack object size (only when creating a store)"),
+		frame:      fs.String("frame", "256KiB", "compress new chunks of sub-chunked files together up to this size, 0 for one by one (only when creating a store)"),
 	}
 }
 
@@ -87,6 +88,13 @@ func (f storeFlags) open(ctx context.Context) (*store.Store, backend.Refs, error
 			return nil, nil, err
 		}
 		c.Sub = &chunker.Params{Min: avg / 4, Avg: avg, Max: avg * 4}
+		if *f.frame != "0" {
+			fr, err := size(*f.frame)
+			if err != nil {
+				return nil, nil, err
+			}
+			c.Frame = int64(fr)
+		}
 		if *f.subKinds != "all" {
 			for _, k := range strings.Split(*f.subKinds, ",") {
 				c.SubKinds = append(c.SubKinds, layout.Kind(strings.TrimSpace(k)))

docs/evaluation.md

@@ -15,7 +15,8 @@ sequence buckets, each embedding the weights).
 | fixed-block dedup, 64 KiB | 113.4 GB | 2.33× | 27 s |
 | FastCDC 64 KiB dedup | 46.8 GB | 5.64× | — |
 | Xet-style: FastCDC 64 KiB + per-chunk compression (`onggi --sub 0`) | 29.5 GB | 8.95× | 1.0 min |
-| **onggi**: two-level chunking + per-chunk codecs | **17.0 GB** | **15.6×** | 1.3 min |
+| onggi before #52: two-level chunking + per-chunk codecs | 17.0 GB | 15.6× | 1.3 min |
+| **onggi**: + compression frames (#52) | **16.7 GB** | **15.8×** | 1.5 min |
 
 Whole-file methods find nothing to share: every file differs. A long-range window of 2 GiB cannot reach the copy of
 the weights in the previous 15 GB file. Fixed blocks lose the weights as soon as a header shifts them. CDC finds the
@@ -32,12 +33,15 @@ Bytes that cross the network when a node that holds B2 needs B8 (15.2 GB):
 | method | transferred |
 |---|---|
 | OCI layer (zstd -3) | 11.46 GB |
-| Xet-style CDC + per-chunk compression | **6.1 MB** |
-| onggi | 28.3 MB |
+| Xet-style CDC + per-chunk compression | 6.1 MB |
+| onggi before #52 (each ~4 KiB sub-chunk compressed alone) | 28.3 MB |
+| **onggi** with compression frames (#52) | **3.6 MB** |
 
-The 218 MB that differ compress well. As 64 KiB chunks they shrink to 6.1 MB; onggi splits new chunks of ATOM files
-into ~4 KiB sub-chunks, which compress worse, so it sends 28.3 MB. #52 tracks compressing a chunk's new sub-chunks
-together.
+The 218 MB that differ compress well. onggi splits new chunks of ATOM files into ~4 KiB sub-chunks for dedup; compressed
+one by one they sent 28.3 MB. Since #52 push compresses consecutive new chunks together in frames of up to 255 KiB,
+keeping a frame only where it beats its chunks compressed alone; every chunk stays addressable and hash-checked on its
+own.
 
 Applying the switch on the node (README, deploy/csi): a plain pull through the node cache takes 73.8 s; delta apply
-(#41) takes 5.7 s without and 8.9 s with the node cache, and a CSI pod is Running 12 s after creation.
+(#41) takes 5.7 s without and 8.9 s with the node cache (4.8 s, 3.6 MB read, with frames), and a CSI pod is Running
+12 s after creation.

pkg/store/diff.go

@@ -266,7 +266,7 @@ func (s *Store) Diff(ctx context.Context, base, m digest.Digest, o DiffOptions)
 	}
 	for i := range nodes {
 		if nodes[i].stored && !inBase(nodes[i].chunk) {
-			st.NewStored += int64(nodes[i].loc.Stored)
+			st.NewStored += nodes[i].loc.share()
 		}
 	}
 	return st, nil

pkg/store/format.go

@@ -19,7 +19,10 @@ import (
 //	                       (chunk digest [32], offset uint32, stored length uint32,
 //	                        raw length uint32, codec uint8, 3 zero bytes)
 //	                       Version 1 ("ONGGIIX1", entries of digest, offset, length; every chunk
-//	                       raw) is still read.
+//	                       raw) is still read. A pack with chunks in compression frames (#52) has
+//	                       version 3 ("ONGGIIX3"): each entry then adds offset in the frame
+//	                       uint32 and the frame's raw length uint32 (0: not in a frame), and
+//	                       offset / stored length / codec are the frame's.
 //	list   (one per file): magic "ONGGICL1", count uint32, then count × (chunk digest [32], length uint32)
 //	composites (index kind, one per push that made any): magic "ONGGICX1", count uint32, then count ×
 //	                       (chunk digest [32], sub-chunk count uint32, then that many
@@ -32,6 +35,7 @@ import (
 var (
 	indexMagic   = [8]byte{'O', 'N', 'G', 'G', 'I', 'I', 'X', '2'}
 	indexMagicV1 = [8]byte{'O', 'N', 'G', 'G', 'I', 'I', 'X', '1'}
+	indexMagicV3 = [8]byte{'O', 'N', 'G', 'G', 'I', 'I', 'X', '3'}
 	listMagic    = [8]byte{'O', 'N', 'G', 'G', 'I', 'C', 'L', '1'}
 	compMagic    = [8]byte{'O', 'N', 'G', 'G', 'I', 'C', 'X', '1'}
 )
@@ -44,6 +48,10 @@ type Loc struct {
 	Stored uint32 // bytes in the pack
 	Length uint32 // raw bytes
 	Codec  codec.ID
+	// Frame, when not 0, is the raw length of the compression frame the chunk lies in (#52):
+	// Offset, Stored and Codec are then the frame's, and Inner is the chunk's offset in it.
+	Inner uint32
+	Frame uint32
 }
 
 type indexEntry struct {
@@ -52,8 +60,16 @@ type indexEntry struct {
 }
 
 func encodeIndex(pack digest.Digest, entries []indexEntry) []byte {
+	v3 := false
+	for _, e := range entries {
+		v3 = v3 || e.Loc.Frame != 0
+	}
 	var b bytes.Buffer
-	b.Write(indexMagic[:])
+	if v3 {
+		b.Write(indexMagicV3[:])
+	} else {
+		b.Write(indexMagic[:])
+	}
 	b.Write(pack[:])
 	binary.Write(&b, binary.LittleEndian, uint32(len(entries)))
 	for _, e := range entries {
@@ -62,6 +78,10 @@ func encodeIndex(pack digest.Digest, entries []indexEntry) []byte {
 		binary.Write(&b, binary.LittleEndian, e.Loc.Stored)
 		binary.Write(&b, binary.LittleEndian, e.Loc.Length)
 		b.Write([]byte{byte(e.Loc.Codec), 0, 0, 0})
+		if v3 {
+			binary.Write(&b, binary.LittleEndian, e.Loc.Inner)
+			binary.Write(&b, binary.LittleEndian, e.Loc.Frame)
+		}
 	}
 	return b.Bytes()
 }
@@ -72,12 +92,16 @@ func decodeIndex(data []byte) (digest.Digest, []indexEntry, error) {
 		return pack, nil, errors.New("not an onggi index")
 	}
 	v1 := bytes.Equal(data[:8], indexMagicV1[:])
-	if !v1 && !bytes.Equal(data[:8], indexMagic[:]) {
+	v3 := bytes.Equal(data[:8], indexMagicV3[:])
+	if !v1 && !v3 && !bytes.Equal(data[:8], indexMagic[:]) {
 		return pack, nil, errors.New("not an onggi index")
 	}
 	size := uint64(48)
-	if v1 {
+	switch {
+	case v1:
 		size = 40
+	case v3:
+		size = 56
 	}
 	copy(pack[:], data[8:40])
 	n := binary.LittleEndian.Uint32(data[40:44])
@@ -95,6 +119,12 @@ func decodeIndex(data []byte) (digest.Digest, []indexEntry, error) {
 			l.Length = binary.LittleEndian.Uint32(p[40:44])
 			l.Codec = codec.ID(p[44])
 		}
+		if v3 {
+			l.Inner, l.Frame = binary.LittleEndian.Uint32(p[48:52]), binary.LittleEndian.Uint32(p[52:56])
+			if l.Frame != 0 && uint64(l.Inner)+uint64(l.Length) > uint64(l.Frame) {
+				return pack, nil, fmt.Errorf("index: chunk %d lies outside its frame", i)
+			}
+		}
 		out[i].Loc = l
 		p = p[size:]
 	}
@@ -203,7 +233,10 @@ type Config struct {
 	// SubKinds limits sub-chunking to these file kinds (by name, pkg/layout); empty means all.
 	// Weights that equal their HF source dedup whole at 64 KiB, and finer chunks only add metadata.
 	SubKinds []layout.Kind `json:"sub_kinds,omitempty"`
-	Created  time.Time     `json:"created"`
+	// Frame, when set, makes push compress consecutive new chunks of a sub-chunked file
+	// together, up to this many raw bytes (#52). Stores created before it have none.
+	Frame   int64     `json:"frame,omitempty"`
+	Created time.Time `json:"created"`
 }
 
 // File is one file of an artifact.

pkg/store/frame.go

@@ -0,0 +1,147 @@
+package store
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/ai-computing/onggi/pkg/codec"
+	"github.com/ai-computing/onggi/pkg/digest"
+)
+
+// Compression frames (#52). Two-level chunking splits new chunks into ~4 KiB sub-chunks, and a
+// sub-chunk compressed alone compresses badly: switching a node between two decode variants sent
+// 28.3 MB where 64 KiB chunks sent 6.1 MB. Push therefore compresses consecutive new chunks of a
+// sub-chunked file together, up to Config.Frame raw bytes, as one frame in the pack. Each chunk
+// keeps its own index entry, so dedup and the hash check per chunk are unchanged: its Loc names
+// the frame (Offset, Stored, Codec), the frame's raw length (Frame) and the chunk's place in it
+// (Inner). On the decode B2 → B8 switch, 256 KiB frames send 3.3 MB.
+
+// DefaultFrame is the largest raw size of a compression frame in new two-level stores. It stays
+// below 256 KiB so that the byte planes of a frame fit one Huffman block and weights keep the
+// Shuffle2Huffman codec (pkg/codec).
+const DefaultFrame = 255 << 10
+
+// share is the part of the stored bytes that is the chunk's: for a chunk in a frame, the frame's
+// stored bytes split by raw length. GC sums it per pack.
+func (l Loc) share() int64 {
+	if l.Frame == 0 {
+		return int64(l.Stored)
+	}
+	return int64(l.Stored) * int64(l.Length) / int64(l.Frame)
+}
+
+// frameReader decodes each frame once while chunks are read in pack order.
+type frameReader struct {
+	pack, off uint32
+	raw       []byte
+	ok        bool
+}
+
+// chunk returns the raw bytes of the chunk at l from stored, the chunk's stored bytes or its
+// frame's. The result of a framed chunk aliases the decoded frame.
+func (fr *frameReader) chunk(stored []byte, l Loc) ([]byte, error) {
+	if l.Frame == 0 {
+		return codec.Decode(stored, l.Codec, int(l.Length))
+	}
+	if !fr.ok || fr.pack != l.Pack || fr.off != l.Offset {
+		fr.ok = false
+		raw, err := codec.Decode(stored, l.Codec, int(l.Frame))
+		if err != nil {
+			return nil, err
+		}
+		fr.raw, fr.pack, fr.off, fr.ok = raw, l.Pack, l.Offset, true
+	}
+	if uint64(l.Inner)+uint64(l.Length) > uint64(len(fr.raw)) {
+		return nil, fmt.Errorf("chunk at %d+%d lies outside its frame of %d bytes", l.Inner, l.Length, len(fr.raw))
+	}
+	return fr.raw[l.Inner : l.Inner+l.Length], nil
+}
+
+// alone returns the stored form of a chunk on its own: its stored bytes, or for a chunk in a
+// frame its raw bytes encoded again. Node caches and peers hold and send chunks one by one.
+func alone(raw, stored []byte, l Loc) ([]byte, codec.ID) {
+	if l.Frame == 0 {
+		return stored, l.Codec
+	}
+	data, id := codec.Encode(raw, codec.ModeAuto)
+	if id == codec.None {
+		data = append([]byte(nil), raw...)
+	}
+	return data, id
+}
+
+// member is one chunk of a frame being built.
+type member struct {
+	chunk  digest.Digest
+	inner  uint32
+	length uint32
+}
+
+// frameBuilder collects the new chunks of one file and sends them as frames.
+type frameBuilder struct {
+	limit   int
+	mode    codec.Mode
+	raw     []byte
+	members []member
+}
+
+// add appends a new chunk, sending the frame first if the chunk would take it past limit. It
+// returns the raw bytes added.
+func (fb *frameBuilder) add(ctx context.Context, raw []byte, d digest.Digest, out chan<- pending) (int64, error) {
+	if len(fb.raw)+len(raw) > fb.limit {
+		if err := fb.flush(ctx, out); err != nil {
+			return 0, err
+		}
+	}
+	fb.members = append(fb.members, member{chunk: d, inner: uint32(len(fb.raw)), length: uint32(len(raw))})
+	fb.raw = append(fb.raw, raw...)
+	return int64(len(raw)), nil
+}
+
+// flush sends what the builder holds: one chunk as itself, several as a frame.
+func (fb *frameBuilder) flush(ctx context.Context, out chan<- pending) error {
+	defer func() { fb.raw, fb.members = nil, nil }()
+	switch len(fb.members) {
+	case 0:
+		return nil
+	case 1:
+		_, err := send(ctx, fb.raw, fb.members[0].chunk, fb.mode, out)
+		return err
+	}
+	data, id := codec.Encode(fb.raw, fb.mode)
+	// weights compress about as well chunk by chunk (byte planes, Huffman) and a frame mixing them
+	// with other bytes can do worse: keep the frame only if it beats its chunks stored alone
+	type one struct {
+		data []byte
+		id   codec.ID
+	}
+	alone := make([]one, len(fb.members))
+	var sum int
+	for i, m := range fb.members {
+		d, cid := codec.Encode(fb.raw[m.inner:m.inner+m.length], fb.mode)
+		alone[i], sum = one{d, cid}, sum+len(d)
+	}
+	if sum <= len(data) {
+		for i, m := range fb.members {
+			d := alone[i].data
+			if alone[i].id == codec.None {
+				d = append([]byte(nil), d...)
+			}
+			select {
+			case out <- pending{data: d, raw: m.length, codec: alone[i].id, chunk: m.chunk}:
+			case <-ctx.Done():
+				return context.Cause(ctx)
+			}
+		}
+		return nil
+	}
+	if id == codec.None {
+		data = fb.raw // owned: the builder starts a new buffer
+	}
+	select {
+	case out <- pending{data: data, raw: uint32(len(fb.raw)), codec: id, members: fb.members}:
+		return nil
+	case <-ctx.Done():
+		return context.Cause(ctx)
+	}
+}