diff --git a/README.md b/README.md index 266f550..abbfb92 100644 --- a/README.md +++ b/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 diff --git a/cmd/onggi/store.go b/cmd/onggi/store.go index 734742c..fb1003c 100644 --- a/cmd/onggi/store.go +++ b/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))) diff --git a/docs/evaluation.md b/docs/evaluation.md index 1ffdcea..cb64d8c 100644 --- a/docs/evaluation.md +++ b/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. diff --git a/pkg/store/diff.go b/pkg/store/diff.go index 4ccb270..809ef30 100644 --- a/pkg/store/diff.go +++ b/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 diff --git a/pkg/store/format.go b/pkg/store/format.go index 348ae25..5a8bc91 100644 --- a/pkg/store/format.go +++ b/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. diff --git a/pkg/store/frame.go b/pkg/store/frame.go new file mode 100644 index 0000000..ef0691b --- /dev/null +++ b/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) + } +} diff --git a/pkg/store/frame_test.go b/pkg/store/frame_test.go new file mode 100644 index 0000000..0ac8462 --- /dev/null +++ b/pkg/store/frame_test.go @@ -0,0 +1,212 @@ +package store + +import ( + "bytes" + "context" + "fmt" + "path/filepath" + "testing" + "time" + + "github.com/ai-computing/onggi/pkg/backend" + bfs "github.com/ai-computing/onggi/pkg/backend/fs" + "github.com/ai-computing/onggi/pkg/chunker" + "github.com/ai-computing/onggi/pkg/codec" + "github.com/ai-computing/onggi/pkg/digest" +) + +func openFramed(t *testing.T, root string, frame int64) *Store { + t.Helper() + b, err := bfs.Open(root, backend.DurabilityNone) + if err != nil { + t.Fatal(err) + } + sub := DefaultSub + s, err := OpenWith(context.Background(), b, Create{Chunker: chunker.Default, Sub: &sub, PackSize: 1 << 20, Frame: frame}) + if err != nil { + t.Fatal(err) + } + return s +} + +// compressible returns data that does not dedup but compresses across chunks: a random 16 KiB +// block repeated with a counter written every 512 bytes, so every sub-chunk differs while a frame +// finds the repeats. +func compressible(seed, n int) []byte { + block := random(int64(seed), 16<<10) + var b bytes.Buffer + for i := 0; b.Len() < n; i++ { + blk := append([]byte(nil), block...) + for o := 0; o < len(blk); o += 512 { + copy(blk[o:], fmt.Sprintf("%08x", i*64+o)) + } + b.Write(blk) + } + return b.Bytes()[:n] +} + +func framedEntries(t *testing.T, s *Store) (framed, all int) { + t.Helper() + if err := s.LoadIndex(context.Background()); err != nil { + t.Fatal(err) + } + for _, e := range s.idx.ent { + if e.state == stateStored { + all++ + if e.loc.Frame != 0 { + framed++ + } + } + } + return framed, all +} + +// Frames: smaller packs than chunk-by-chunk compression, identical pulls, dedup kept (#52). +func TestFramesRoundTripAndCompressBetter(t *testing.T) { + ctx := context.Background() + base := append(random(81, 2<<20), compressible(1, 2<<20)...) + next := append([]byte(nil), base...) + copy(next[3<<20:], compressible(2, 256<<10)) // a compressible region changes + src, src2 := t.TempDir(), t.TempDir() + write(t, filepath.Join(src, "decode_B2.rbln"), base) + write(t, filepath.Join(src2, "decode_B8.rbln"), next) + + var packBytes [2]int64 + for k, frame := range []int64{0, DefaultFrame} { + root := t.TempDir() + s := openFramed(t, root, frame) + m1, st1, err := s.Push(ctx, src, PushOptions{Compress: codec.ModeAuto}) + if err != nil { + t.Fatal(err) + } + m2, st2, err := s.Push(ctx, src2, PushOptions{Compress: codec.ModeAuto}) + if err != nil { + t.Fatal(err) + } + packBytes[k] = st2.PackBytes + t.Logf("frame %d: first push %d B packed, second %d B packed", frame, st1.PackBytes, st2.PackBytes) + for _, c := range []struct { + m digest.Digest + src string + }{{m1, src}, {m2, src2}} { + out := t.TempDir() + g := openFramed(t, root, 0) // a fresh reader of the same store + if _, err := g.Pull(ctx, c.m, out, PullOptions{}); err != nil { + t.Fatal(err) + } + sameTree(t, c.src, out) + } + framed, all := framedEntries(t, openFramed(t, root, 0)) + if (frame == 0) != (framed == 0) || all == 0 { + t.Fatalf("frame %d: %d of %d entries framed", frame, framed, all) + } + if st2.NewBytes > 1<<20 { + t.Fatalf("frame %d: the variant stored %d new bytes; dedup lost", frame, st2.NewBytes) + } + } + if packBytes[1]*2 > packBytes[0] { + t.Fatalf("frames did not help the variant: %d B framed vs %d B chunk by chunk", packBytes[1], packBytes[0]) + } +} + +// A framed store feeds a node cache, peers and repack, each of which takes chunks one by one. +func TestFramesThroughCachePeersAndRepack(t *testing.T) { + ctx := context.Background() + root := t.TempDir() + s := openFramed(t, root, DefaultFrame) + src := t.TempDir() + write(t, filepath.Join(src, "decode_B2.rbln"), append(compressible(3, 1<<20), random(82, 1<<20)...)) + m, _, err := s.Push(ctx, src, PushOptions{Compress: codec.ModeAuto}) + if err != nil { + t.Fatal(err) + } + // node cache: filled from frames, then read alone + cacheDir := t.TempDir() + if _, err := s.Pull(ctx, m, t.TempDir(), PullOptions{Cache: openCache(t, cacheDir, s)}); err != nil { + t.Fatal(err) + } + c := openCache(t, cacheDir, s) + if framed, all := framedEntries(t, c); framed != 0 || all == 0 { + t.Fatalf("cache: %d of %d entries framed", framed, all) + } + out := t.TempDir() + st, err := s.Pull(ctx, m, out, PullOptions{Cache: openCache(t, cacheDir, s)}) + if err != nil { + t.Fatal(err) + } + sameTree(t, src, out) + if st.ReadBytes != 0 { + t.Fatalf("second pull read %d B from the store", st.ReadBytes) + } + // the framed store as a peer + out = t.TempDir() + st, err = s.Pull(ctx, m, out, PullOptions{Peers: []ChunkSource{LocalSource("peer", s)}}) + if err != nil { + t.Fatal(err) + } + sameTree(t, src, out) + if st.PeerChunks == 0 { + t.Fatalf("no chunks from the framed peer: %+v", st) + } + // repack after half the frames' chunks die + refs := s.B.(*bfs.Store) + if err := refs.SetRef(ctx, "keep", "", m.String()); err != nil { + t.Fatal(err) + } + src2 := t.TempDir() + write(t, filepath.Join(src2, "other.rbln"), compressible(4, 1<<20)) + m2, _, err := s.Push(ctx, src2, PushOptions{Compress: codec.ModeAuto}) + if err != nil { + t.Fatal(err) + } + if err := refs.SetRef(ctx, "drop", "", m2.String()); err != nil { + t.Fatal(err) + } + if err := refs.SetRef(ctx, "drop", m2.String(), ""); err != nil { + t.Fatal(err) + } + t0 := time.Now() + for _, after := range []time.Duration{0, 2 * time.Hour, 4 * time.Hour} { + g := openFramed(t, root, 0) + if _, err := g.GC(ctx, refs, GCOptions{Grace: time.Hour, RepackDead: 0.25, Now: clock(t0, after)}); err != nil { + t.Fatal(err) + } + } + out = t.TempDir() + if _, err := openFramed(t, root, 0).Pull(ctx, m, out, PullOptions{}); err != nil { + t.Fatal(err) + } + sameTree(t, src, out) + assertIndexConsistent(t, openFramed(t, root, 0)) +} + +func TestIndexV3RoundTrip(t *testing.T) { + var pack digest.Digest + pack[0] = 7 + in := []indexEntry{ + {Chunk: digest.Digest{1}, Loc: Loc{Offset: 0, Stored: 100, Length: 4096, Codec: codec.Zstd, Inner: 0, Frame: 8192}}, + {Chunk: digest.Digest{2}, Loc: Loc{Offset: 0, Stored: 100, Length: 4096, Codec: codec.Zstd, Inner: 4096, Frame: 8192}}, + {Chunk: digest.Digest{3}, Loc: Loc{Offset: 100, Stored: 50, Length: 60, Codec: codec.None}}, + } + data := encodeIndex(pack, in) + if !bytes.HasPrefix(data, indexMagicV3[:]) { + t.Fatal("framed entries not written as version 3") + } + p, out, err := decodeIndex(data) + if err != nil || p != pack || len(out) != 3 { + t.Fatalf("%v %v %d", err, p, len(out)) + } + for i := range in { + if out[i] != in[i] { + t.Fatalf("entry %d: %+v, want %+v", i, out[i], in[i]) + } + } + if !bytes.HasPrefix(encodeIndex(pack, in[2:]), indexMagic[:]) { + t.Fatal("an index without frames is not version 2") + } + bad := append([]indexEntry(nil), in...) + bad[1].Loc.Inner = 8000 // runs past the frame + if _, _, err := decodeIndex(encodeIndex(pack, bad)); err == nil { + t.Fatal("accepted a chunk outside its frame") + } +} diff --git a/pkg/store/gc.go b/pkg/store/gc.go index 37da8e7..5af8119 100644 --- a/pkg/store/gc.go +++ b/pkg/store/gc.go @@ -14,7 +14,6 @@ import ( "time" "github.com/ai-computing/onggi/pkg/backend" - "github.com/ai-computing/onggi/pkg/codec" "github.com/ai-computing/onggi/pkg/digest" ) @@ -163,7 +162,7 @@ func (s *Store) GC(ctx context.Context, refs backend.Refs, o GCOptions) (GCStats } live[id] = true if e.state == stateStored { - livePack[e.loc.Pack] += int64(e.loc.Stored) + livePack[e.loc.Pack] += e.loc.share() return nil } for _, sc := range s.subRefs(e) { @@ -489,7 +488,13 @@ func (s *Store) repack(ctx context.Context, o GCOptions, live []bool, livePack m defer w.abort() for _, j := range jobs { ids := byPack[j.id] - sort.Slice(ids, func(a, b int) bool { return s.idx.ent[ids[a]].loc.Offset < s.idx.ent[ids[b]].loc.Offset }) + sort.Slice(ids, func(a, b int) bool { + x, y := s.idx.ent[ids[a]].loc, s.idx.ent[ids[b]].loc + if x.Offset != y.Offset { + return x.Offset < y.Offset + } + return x.Inner < y.Inner + }) rc, err := s.B.Get(ctx, backend.Key{Kind: backend.Pack, Name: j.name}, 0, -1) if err != nil { return err @@ -502,19 +507,22 @@ func (s *Store) repack(ctx context.Context, o GCOptions, live []bool, livePack m if digest.Of(data).Hex() != j.name { return fmt.Errorf("pack %s: content does not match its name (repack stopped)", j.name) } + var fr frameReader for _, id := range ids { s.mu.RLock() e, d := s.idx.ent[id], s.idx.dig[id] s.mu.RUnlock() stored := data[e.loc.Offset : e.loc.Offset+e.loc.Stored] - raw, err := codec.Decode(stored, e.loc.Codec, int(e.loc.Length)) + raw, err := fr.chunk(stored, e.loc) if err != nil || digest.Of(raw) != d { return fmt.Errorf("pack %s: chunk %s fails its check (repack stopped)", j.name, d) } - if err := w.add(ctx, d, stored, e.loc.Codec, e.loc.Length); err != nil { + // a live chunk of a frame is copied on its own: the rest of the frame is dead (#52) + out, id := alone(raw, stored, e.loc) + if err := w.add(ctx, d, out, id, e.loc.Length); err != nil { return err } - st.RepackedBytes += int64(e.loc.Stored) + st.RepackedBytes += int64(len(out)) } } if err := w.close(ctx); err != nil { @@ -616,7 +624,7 @@ func (s *Store) markFromLocator(ctx context.Context, root *Store, md digest.Dige continue } live[id] = true - livePack[e.loc.Pack] += int64(e.loc.Stored) + livePack[e.loc.Pack] += e.loc.share() } return nil } diff --git a/pkg/store/locator.go b/pkg/store/locator.go index 1b269b0..18c1179 100644 --- a/pkg/store/locator.go +++ b/pkg/store/locator.go @@ -25,6 +25,8 @@ import ( // nodes uint32, then per node: chunk digest [32], raw length uint32, kind uint8, then // kind 0 (stored): pack number uint32, offset uint32, stored length uint32, codec uint8 // kind 1 (composite): sub-chunk count uint32, then that many earlier node numbers uint32 +// kind 2 (in a compression frame, #52): as kind 0 for the frame, then offset in the frame +// uint32, frame raw length uint32 // packs uint32, then that many pack digests [32] // files uint32, then per file: chunk count uint32, then that many node numbers uint32 @@ -81,11 +83,19 @@ func (s *Store) writeLocator(lists [][]ChunkRef, f *os.File) (int64, error) { } raw(c.Digest[:]) w(c.Length) - one(0) + if loc.Frame != 0 { + one(2) + } else { + one(0) + } w(p) w(loc.Offset) w(loc.Stored) one(byte(loc.Codec)) + if loc.Frame != 0 { + w(loc.Inner) + w(loc.Frame) + } case comp: kids := make([]uint32, len(subs)) for i, sc := range subs { @@ -199,11 +209,17 @@ func readLocator(data []byte, m *Manifest) ([]node, []digest.Digest, error) { copy(nd.chunk[:], r.take(32)) length := r.u32() lengths[i] = length - switch r.u8() { - case 0: + switch kind := r.u8(); kind { + case 0, 2: nd.stored = true nd.loc = Loc{Pack: r.u32(), Offset: r.u32(), Stored: r.u32(), Length: length} nd.loc.Codec = codec.ID(r.u8()) + if kind == 2 { + nd.loc.Inner, nd.loc.Frame = r.u32(), r.u32() + if r.err == nil && (nd.loc.Frame == 0 || uint64(nd.loc.Inner)+uint64(length) > uint64(nd.loc.Frame)) { + return nil, nil, fmt.Errorf("locator: chunk %s lies outside its frame", nd.chunk) + } + } case 1: var off int64 for j, n := 0, r.count(4); j < n; j++ { diff --git a/pkg/store/pull.go b/pkg/store/pull.go index 7df96f7..6f0a3f9 100644 --- a/pkg/store/pull.go +++ b/pkg/store/pull.go @@ -14,7 +14,6 @@ import ( "time" "github.com/ai-computing/onggi/pkg/backend" - "github.com/ai-computing/onggi/pkg/codec" "github.com/ai-computing/onggi/pkg/digest" ) @@ -202,7 +201,10 @@ func (s *Store) Pull(ctx context.Context, manifest digest.Digest, dest string, o if a.Pack != b.Pack { return a.Pack < b.Pack } - return a.Offset < b.Offset + if a.Offset != b.Offset { + return a.Offset < b.Offset + } + return a.Inner < b.Inner }) type span struct{ a, b int } // order[a:b] var ranges []span @@ -211,7 +213,8 @@ func (s *Store) Pull(ctx context.Context, manifest digest.Digest, dest string, o if n := len(ranges); n > 0 { r := &ranges[n-1] first, prev := nodes[order[r.a]].loc, nodes[order[r.b-1]].loc - if nodes[order[r.b-1]].fromCache == nodes[i].fromCache && prev.Pack == l.Pack && prev.Offset+prev.Stored == l.Offset && + // the next chunk follows, or shares the frame of the previous one (#52) + if nodes[order[r.b-1]].fromCache == nodes[i].fromCache && prev.Pack == l.Pack && (prev.Offset+prev.Stored == l.Offset || prev.Frame != 0 && prev.Offset == l.Offset) && int64(l.Offset+l.Stored-first.Offset) <= o.MaxRange { r.b = k + 1 continue @@ -498,9 +501,11 @@ func (s *Store) readRange(ctx context.Context, nodes []node, packs []digest.Dige if err != nil { return 0, fmt.Errorf("pack %s: %w", pack, err) } + var fr frameReader for _, i := range ids { nd := &nodes[i] - raw, err := codec.Decode(b[int64(nd.loc.Offset)-int64(first.Offset):][:nd.loc.Stored], nd.loc.Codec, int(nd.loc.Length)) + stored := b[int64(nd.loc.Offset)-int64(first.Offset):][:nd.loc.Stored] + raw, err := fr.chunk(stored, nd.loc) if err != nil { return 0, fmt.Errorf("pack %s offset %d: chunk %s: %w", pack, nd.loc.Offset, nd.chunk, err) } @@ -511,8 +516,8 @@ func (s *Store) readRange(ctx context.Context, nodes []node, packs []digest.Dige return 0, err } if fill != nil { - stored := b[int64(nd.loc.Offset)-int64(first.Offset):][:nd.loc.Stored] - if err := fill.add(ctx, nd.chunk, stored, nd.loc.Codec, nd.loc.Length); err != nil { + data, id := alone(raw, stored, nd.loc) + if err := fill.add(ctx, nd.chunk, data, id, nd.loc.Length); err != nil { return 0, fmt.Errorf("node cache: %w", err) } } diff --git a/pkg/store/push.go b/pkg/store/push.go index b9fd192..4620e11 100644 --- a/pkg/store/push.go +++ b/pkg/store/push.go @@ -48,10 +48,11 @@ type PushStats struct { } type pending struct { - data []byte // stored form - raw uint32 - codec codec.ID - chunk digest.Digest + data []byte // stored form + raw uint32 + codec codec.ID + chunk digest.Digest + members []member // a frame's chunks (#52); empty: data is the one chunk above } type sealed struct { @@ -207,7 +208,13 @@ func (s *Store) Push(ctx context.Context, root string, o PushOptions) (digest.Di continue } h.Write(p.data) - entries = append(entries, indexEntry{Chunk: p.chunk, Loc: Loc{Offset: uint32(size), Stored: uint32(len(p.data)), Length: p.raw, Codec: p.codec}}) + if len(p.members) == 0 { + entries = append(entries, indexEntry{Chunk: p.chunk, Loc: Loc{Offset: uint32(size), Stored: uint32(len(p.data)), Length: p.raw, Codec: p.codec}}) + } + for _, m := range p.members { + entries = append(entries, indexEntry{Chunk: m.chunk, Loc: Loc{Offset: uint32(size), Stored: uint32(len(p.data)), + Length: m.length, Codec: p.codec, Inner: m.inner, Frame: p.raw}}) + } size += int64(len(p.data)) if size >= s.Config.PackSize { seal() @@ -400,9 +407,18 @@ func (s *Store) chunkFile(ctx context.Context, path string, twoLevel bool, mode c := chunker.New(f, s.Config.Chunker) var list []ChunkRef var nb int64 + var fb *frameBuilder // new chunks of a sub-chunked file go out in compression frames (#52) + if twoLevel && s.Config.Frame > 0 { + fb = &frameBuilder{limit: int(s.Config.Frame), mode: mode} + } for { ch, err := c.Next() if errors.Is(err, io.EOF) { + if fb != nil { + if err := fb.flush(ctx, out); err != nil { + return nil, 0, err + } + } return list, nb, nil } if err != nil { @@ -427,6 +443,11 @@ func (s *Store) chunkFile(ctx context.Context, path string, twoLevel bool, mode sd := digest.Of(sub) subs = append(subs, ChunkRef{Digest: sd, Length: uint32(len(sub))}) if sd == d || claim(sd) { // sd == d: the whole chunk is one sub-chunk; this push owns d already + if fb != nil { + n, err := fb.add(ctx, sub, sd, out) + nb += n + return err + } n, err := send(ctx, sub, sd, mode, out) if err != nil { return err diff --git a/pkg/store/source.go b/pkg/store/source.go index f338597..50fc253 100644 --- a/pkg/store/source.go +++ b/pkg/store/source.go @@ -60,14 +60,18 @@ func (s *Store) ReadStored(ctx context.Context, ds []digest.Digest, fn func(d di if items[i].l.Pack != items[j].l.Pack { return items[i].l.Pack < items[j].l.Pack } - return items[i].l.Offset < items[j].l.Offset + if items[i].l.Offset != items[j].l.Offset { + return items[i].l.Offset < items[j].l.Offset + } + return items[i].l.Inner < items[j].l.Inner }) const maxRange = 16 << 20 var buf []byte for i := 0; i < len(items); { j := i + 1 for j < len(items) && items[j].l.Pack == items[i].l.Pack && - items[j].l.Offset+items[j].l.Stored-items[i].l.Offset <= maxRange && items[j].l.Offset >= items[j-1].l.Offset+items[j-1].l.Stored { + items[j].l.Offset+items[j].l.Stored-items[i].l.Offset <= maxRange && + (items[j].l.Offset >= items[j-1].l.Offset+items[j-1].l.Stored || items[j].l.Frame != 0 && items[j].l.Offset == items[j-1].l.Offset) { j++ } first, last := items[i].l, items[j-1].l @@ -85,9 +89,22 @@ func (s *Store) ReadStored(ctx context.Context, ds []digest.Digest, fn func(d di if err != nil { return err } + var fr frameReader for _, it := range items[i:j] { off := it.l.Offset - first.Offset - if err := fn(it.d, b[off:off+it.l.Stored], it.l.Codec, it.l.Length); err != nil { + stored := b[off : off+it.l.Stored] + if it.l.Frame != 0 { // sent on its own: readers decode chunk by chunk (#52) + raw, err := fr.chunk(stored, it.l) + if err != nil { + return err + } + data, id := alone(raw, stored, it.l) + if err := fn(it.d, data, id, it.l.Length); err != nil { + return err + } + continue + } + if err := fn(it.d, stored, it.l.Codec, it.l.Length); err != nil { return err } } diff --git a/pkg/store/store.go b/pkg/store/store.go index b856059..124161f 100644 --- a/pkg/store/store.go +++ b/pkg/store/store.go @@ -51,6 +51,7 @@ type Create struct { Sub *chunker.Params // nil: one level of chunks SubKinds []layout.Kind // file kinds that get sub-chunks; empty: all PackSize int64 + Frame int64 // compression frames for sub-chunked files (#52); used only with Sub } // DefaultSubKinds are the kinds whose compile variants rearrange the same bytes (#15). @@ -83,7 +84,7 @@ func OpenWith(ctx context.Context, b backend.Backend, c Create) (*Store, error) } cfg = Config{Schema: "onggi.store.v1", Chunker: p, PackSize: packSize, Sub: c.Sub, Created: time.Now().UTC()} if c.Sub != nil { - cfg.SubKinds = c.SubKinds + cfg.SubKinds, cfg.Frame = c.SubKinds, c.Frame } data, _ := json.MarshalIndent(cfg, "", " ") if err := b.Put(ctx, configKey, bytes.NewReader(data), int64(len(data))); err != nil {