From 79dc7a356e30288f7ec260254ddb577d20acf650 Mon Sep 17 00:00:00 2001 From: Teodor Calin Date: Thu, 23 Jul 2026 23:55:55 +0300 Subject: [PATCH] feat(index): add HNSW writer bridge for Pebble crawl path + ops systemd units MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit internal/index/hnsw_writer.go implements PassageWriter against a Pebble store + in-memory HNSW graph, with periodic checkpointing and a manual Flush for shutdown. Covered by hnsw_writer_test.go (bridges passages into the graph, and verifies the auto-flush-at-threshold contract). Not yet wired into cmd_crawl/cmd_serve — those already use their own hnswPassageWriter in cmd/cosift; this is an alternate, independently tested implementation living in internal/index for review. Also adds the systemd unit/timer files used to run cosift-crawl, cosift-serve, cosift-snapshot, and the ollama service(s) in production. No secrets included — cosift-snapshot.service loads its admin token via EnvironmentFile from /etc/cosift/snapshot.env, not inline. Co-Authored-By: Claude Opus 4.8 --- internal/index/hnsw_writer.go | 105 +++++++++++++++++++++++ internal/index/hnsw_writer_test.go | 126 ++++++++++++++++++++++++++++ scripts/cosift-crawl.service | 22 +++++ scripts/cosift-crawl.timer | 14 ++++ scripts/cosift-serve.service | 30 +++++++ scripts/cosift-snapshot.service | 19 +++++ scripts/cosift-snapshot.timer | 13 +++ scripts/ollama-replica.service.tmpl | 18 ++++ scripts/ollama.service | 16 ++++ 9 files changed, 363 insertions(+) create mode 100644 internal/index/hnsw_writer.go create mode 100644 internal/index/hnsw_writer_test.go create mode 100644 scripts/cosift-crawl.service create mode 100644 scripts/cosift-crawl.timer create mode 100644 scripts/cosift-serve.service create mode 100644 scripts/cosift-snapshot.service create mode 100644 scripts/cosift-snapshot.timer create mode 100644 scripts/ollama-replica.service.tmpl create mode 100644 scripts/ollama.service diff --git a/internal/index/hnsw_writer.go b/internal/index/hnsw_writer.go new file mode 100644 index 0000000..53d59e5 --- /dev/null +++ b/internal/index/hnsw_writer.go @@ -0,0 +1,105 @@ +// HNSW write bridge for the Pebble-backed crawler path. +// +// The PassageWriter interface lets the crawler stream passage +// vectors to whatever backend the operator wires. On SQLite, *store.Store +// satisfies it directly (writes go to the passages table). On Pebble there +// was no implementation — vector indexing during crawl was silently +// skipped, which made the Pebble path BM25-only. +// +// HNSWWriter bridges the gap: passages flow into an in-memory HNSW graph +// and get periodically flushed to the PebbleStore via HNSW.Persist. +// Cheap per-passage cost (one graph insert + one +// GetDocMeta lookup); the heavy on-disk write is amortized across +// many inserts. + +package index + +import ( + "context" + "fmt" + "sync" + + "github.com/pilot-protocol/cosift/internal/store" +) + +// HNSWWriter satisfies the crawler's PassageWriter interface against a +// PebbleStore + HNSW pair. Passage URLs/titles are resolved from the +// 'i' family (cheap, no full Document gob decode). +// +// Persistence is checkpointed every PersistEvery passages. Setting +// PersistEvery to 0 (the default if you don't configure it) means +// "never auto-persist" — callers can call Persist manually at safe +// points (e.g. crawler shutdown). +type HNSWWriter struct { + hnsw *HNSW + store *store.PebbleStore + persistEvery int + + mu sync.Mutex + sinceLastFlush int +} + +// NewHNSWWriter returns a PassageWriter that adds each incoming passage +// to the HNSW graph and persists to the PebbleStore every persistEvery +// passages. persistEvery=0 disables auto-persist (caller drives it via +// HNSWWriter.Flush). +func NewHNSWWriter(h *HNSW, ps *store.PebbleStore, persistEvery int) *HNSWWriter { + return &HNSWWriter{hnsw: h, store: ps, persistEvery: persistEvery} +} + +// UpsertPassage satisfies crawler.PassageWriter. Resolves the doc's URL + +// title via PebbleStore.GetDocMeta, adds the passage to HNSW, and +// triggers a checkpoint Persist when the threshold is reached. +// +// store.Passage carries Embedding + DocID + Offset/Length + Model; +// HNSW.AddPassage needs URL + title for hit-display. +func (w *HNSWWriter) UpsertPassage(ctx context.Context, p *store.Passage) error { + if err := ctx.Err(); err != nil { + return err + } + url, title, ok, err := w.store.GetDocMeta(ctx, p.DocID) + if err != nil { + return fmt.Errorf("lookup doc %d: %w", p.DocID, err) + } + if !ok { + // Doc was deleted between Upsert and Passage write. Skip silently — + // no point indexing a vector pointing at nothing. + return nil + } + + w.mu.Lock() + defer w.mu.Unlock() + w.hnsw.AddPassage(url, title, p.Offset, p.Length, p.Embedding) + w.sinceLastFlush++ + if w.persistEvery > 0 && w.sinceLastFlush >= w.persistEvery { + w.sinceLastFlush = 0 + // Persist holds the HNSW RLock; we release our own writer mutex + // briefly first? Actually no — w.mu protects sinceLastFlush, not + // HNSW. HNSW.Persist takes its own RLock independent of ours. + if err := w.hnsw.Persist(ctx, w.store); err != nil { + return fmt.Errorf("hnsw persist: %w", err) + } + } + return nil +} + +// MarkURLInvalid zeros out every HNSW node whose url matches. Returns +// the count zeroed. Lets the crawler tell us "this URL's prior passages +// are stale" before pushing the new chunk batch. Satisfies the optional +// crawler.URLInvalidator interface. +func (w *HNSWWriter) MarkURLInvalid(ctx context.Context, url string) (int, error) { + if err := ctx.Err(); err != nil { + return 0, err + } + return w.hnsw.MarkURLPassagesInvalid(url), nil +} + +// Flush forces an immediate Persist of the current HNSW state. Crawler +// callers should invoke this at shutdown so the last partial batch of +// in-memory inserts isn't lost on next startup. +func (w *HNSWWriter) Flush(ctx context.Context) error { + w.mu.Lock() + defer w.mu.Unlock() + w.sinceLastFlush = 0 + return w.hnsw.Persist(ctx, w.store) +} diff --git a/internal/index/hnsw_writer_test.go b/internal/index/hnsw_writer_test.go new file mode 100644 index 0000000..56a0922 --- /dev/null +++ b/internal/index/hnsw_writer_test.go @@ -0,0 +1,126 @@ +package index + +import ( + "context" + "path/filepath" + "testing" + "time" + + "github.com/pilot-protocol/cosift/internal/store" +) + +// TestHNSWWriterBridgesPassagesIntoGraph — The PassageWriter +// satisfies the crawler interface and routes incoming passages into the +// HNSW graph + the PebbleStore. +func TestHNSWWriterBridgesPassagesIntoGraph(t *testing.T) { + dir := filepath.Join(t.TempDir(), "pebble") + ps, err := store.OpenPebble(dir) + if err != nil { + t.Fatalf("OpenPebble: %v", err) + } + defer ps.Close() + + // Need a Document in the store so GetDocMeta resolves URL+title. + ctx := context.Background() + id, err := ps.UpsertDocument(ctx, &store.Document{ + URL: "https://x/doc1", Title: "Doc One", FetchedAt: time.Now(), + }) + if err != nil { + t.Fatalf("upsert: %v", err) + } + + h := NewHNSW(4) + w := NewHNSWWriter(h, ps, 0) // no auto-flush; test drives Flush manually + + vec := []float32{0.7, 0.1, 0.1, 0.7} + if err := w.UpsertPassage(ctx, &store.Passage{ + DocID: id, Offset: 0, Length: 50, Model: "test", Embedding: vec, + }); err != nil { + t.Fatalf("upsert passage: %v", err) + } + + // HNSW now holds one entry; search returns the right URL. + hits := h.Search(ctx, vec, 5) + if len(hits) != 1 { + t.Fatalf("want 1 hit, got %d", len(hits)) + } + if hits[0].URL != "https://x/doc1" { + t.Errorf("hit URL: want https://x/doc1, got %s", hits[0].URL) + } + if hits[0].Title != "Doc One" { + t.Errorf("hit Title: got %q", hits[0].Title) + } + + // Persist + reload — confirm the bridge produces a recoverable graph. + if err := w.Flush(ctx); err != nil { + t.Fatalf("flush: %v", err) + } + loaded, ok, err := LoadHNSW(ctx, ps) + if err != nil || !ok { + t.Fatalf("load: ok=%v err=%v", ok, err) + } + loadedHits := loaded.Search(ctx, vec, 5) + if len(loadedHits) != 1 || loadedHits[0].URL != "https://x/doc1" { + t.Errorf("post-reload hits: %+v", loadedHits) + } +} + +// TestHNSWWriterAutoFlushOnThreshold — when persistEvery is set, the +// writer triggers Persist automatically after N inserts. Locks in the +// auto-checkpoint contract so callers don't have to poll Flush manually. +func TestHNSWWriterAutoFlushOnThreshold(t *testing.T) { + dir := filepath.Join(t.TempDir(), "pebble") + ps, err := store.OpenPebble(dir) + if err != nil { + t.Fatalf("OpenPebble: %v", err) + } + defer ps.Close() + + ctx := context.Background() + // Three docs to write three passages for. + for i := 1; i <= 3; i++ { + _, err := ps.UpsertDocument(ctx, &store.Document{ + URL: "https://x/" + string(rune('a'+i-1)), + Title: "doc", + FetchedAt: time.Now(), + }) + if err != nil { + t.Fatalf("upsert %d: %v", i, err) + } + } + + h := NewHNSW(4) + w := NewHNSWWriter(h, ps, 2) // auto-persist every 2 passages + + vec := []float32{1, 0, 0, 0} + for i := int64(1); i <= 3; i++ { + if err := w.UpsertPassage(ctx, &store.Passage{ + DocID: i, Offset: 0, Length: 50, Model: "test", Embedding: vec, + }); err != nil { + t.Fatalf("passage %d: %v", i, err) + } + } + + // After 3 passages and persistEvery=2, exactly one auto-persist + // happened (after the 2nd). The 3rd is in memory but not yet on disk — + // load should see 2 nodes, not 3. + loaded, ok, err := LoadHNSW(ctx, ps) + if err != nil || !ok { + t.Fatalf("load: ok=%v err=%v", ok, err) + } + // Loaded must have at least the auto-flushed batch; in-memory writer + // has all three. After explicit Flush, persisted count catches up. + if loaded.Len() < 2 { + t.Errorf("auto-flush at threshold 2: want loaded ≥2 nodes, got %d", loaded.Len()) + } + if err := w.Flush(ctx); err != nil { + t.Fatalf("manual flush: %v", err) + } + loaded2, ok, err := LoadHNSW(ctx, ps) + if err != nil || !ok { + t.Fatalf("re-load: ok=%v err=%v", ok, err) + } + if loaded2.Len() != 3 { + t.Errorf("post-flush count: want 3, got %d", loaded2.Len()) + } +} diff --git a/scripts/cosift-crawl.service b/scripts/cosift-crawl.service new file mode 100644 index 0000000..5ef3af0 --- /dev/null +++ b/scripts/cosift-crawl.service @@ -0,0 +1,22 @@ +[Unit] +Description=cosift crawl (oneshot — stops cosift-serve, crawls, restarts serve) +# Both services want the Pebble writer lock; declare mutual exclusion so +# starting one cleanly stops the other. +Conflicts=cosift-serve.service + +[Service] +Type=oneshot +User=ubuntu +Group=ubuntu +WorkingDirectory=/home/ubuntu +Environment=COSIFT_HNSW_CHECKPOINT_SEC=60 +# Run for up to 30 minutes per cycle. Periodic HNSW checkpoint means we +# keep all vectors even if SIGTERM cuts us short. +ExecStart=/home/ubuntu/cosift -config /home/ubuntu/cosift.json crawl -backend pebble -seeds-file /home/ubuntu/seeds.txt -duration 30m +ExecStartPost=/bin/systemctl start cosift-serve.service +TimeoutStartSec=2400 +StandardOutput=journal +StandardError=journal + +[Install] +WantedBy=multi-user.target diff --git a/scripts/cosift-crawl.timer b/scripts/cosift-crawl.timer new file mode 100644 index 0000000..13dfab4 --- /dev/null +++ b/scripts/cosift-crawl.timer @@ -0,0 +1,14 @@ +[Unit] +Description=cosift crawl — cycle every 6h +# Triggers cosift-crawl.service, which stops cosift-serve.service via the +# Conflicts= directive, runs the crawl, then ExecStartPost restarts serve. + +[Timer] +# 4 cycles per day. Crawl runs 30 min, serve runs the other ~85 min before +# the next trigger. Tune to taste. +OnCalendar=*-*-* 00,06,12,18:00:00 +Persistent=true +Unit=cosift-crawl.service + +[Install] +WantedBy=timers.target diff --git a/scripts/cosift-serve.service b/scripts/cosift-serve.service new file mode 100644 index 0000000..9b107ab --- /dev/null +++ b/scripts/cosift-serve.service @@ -0,0 +1,30 @@ +[Unit] +Description=cosift pebble-serve (with in-process continuous crawler) +After=network-online.target ollama.service +Wants=network-online.target + +[Service] +Type=simple +User=ubuntu +Group=ubuntu +WorkingDirectory=/home/ubuntu +Environment=COSIFT_LOAD_HNSW=true +Environment=COSIFT_RATELIMIT_RPM=240 +Environment=COSIFT_RATELIMIT_BURST=80 +Environment=COSIFT_RATELIMIT_WHITELIST=104.28.216.88,127.0.0.1 +Environment=COSIFT_PPROF_ADDR=127.0.0.1:6060 +Environment=COSIFT_MUTEX_PROFILE_FRACTION=100 +Environment=COSIFT_DISABLE_PQ=true +Environment=COSIFT_HNSW_EF_SEARCH=200 +Environment=COSIFT_CRAWL_EMBED_CONCURRENCY=16 +Environment=COSIFT_PEBBLE_CACHE_MB=32768 +ExecStart=/home/ubuntu/cosift -config /home/ubuntu/cosift.json pebble-serve -dir /home/ubuntu/cosift-data/pebble -addr 127.0.0.1:7777 -crawl-seeds-file /home/ubuntu/seeds.txt -crawl-checkpoint 60s +Restart=on-failure +RestartSec=5 +LimitNOFILE=131072 +TimeoutStopSec=30 +MemoryMax=262G +MemoryHigh=240G + +[Install] +WantedBy=multi-user.target diff --git a/scripts/cosift-snapshot.service b/scripts/cosift-snapshot.service new file mode 100644 index 0000000..2f86294 --- /dev/null +++ b/scripts/cosift-snapshot.service @@ -0,0 +1,19 @@ +[Unit] +Description=cosift snapshot — tarball pebble + config, upload to GCS +After=network-online.target +Wants=network-online.target + +[Service] +Type=oneshot +User=ubuntu +Group=ubuntu +WorkingDirectory=/home/ubuntu +Environment=COSIFT_ADMIN_URL=http://127.0.0.1:7777 +Environment=COSIFT_CONFIG=/home/ubuntu/cosift.json +Environment=COSIFT_GCS_BUCKET=gs://pilot-cosift-index +Environment=COSIFT_KEEP=14 +# COSIFT_ADMIN_TOKEN — set via /etc/cosift/snapshot.env if cluster.peer_auth_token is configured +EnvironmentFile=-/etc/cosift/snapshot.env +ExecStart=/usr/bin/bash /home/ubuntu/snapshot.sh +StandardOutput=journal +StandardError=journal diff --git a/scripts/cosift-snapshot.timer b/scripts/cosift-snapshot.timer new file mode 100644 index 0000000..834d86d --- /dev/null +++ b/scripts/cosift-snapshot.timer @@ -0,0 +1,13 @@ +[Unit] +Description=cosift snapshot — every 4h +# Triggers cosift-snapshot.service. 4-hour cadence × KEEP=14 = ~56 hours +# of recovery history. Tune by changing OnCalendar + COSIFT_KEEP env. + +[Timer] +OnCalendar=*-*-* 00,04,08,12,16,20:00:00 +Persistent=true +RandomizedDelaySec=60 +Unit=cosift-snapshot.service + +[Install] +WantedBy=timers.target diff --git a/scripts/ollama-replica.service.tmpl b/scripts/ollama-replica.service.tmpl new file mode 100644 index 0000000..b01cc72 --- /dev/null +++ b/scripts/ollama-replica.service.tmpl @@ -0,0 +1,18 @@ +[Unit] +Description=Ollama Service (replica 1, port 11435) +After=network-online.target + +[Service] +ExecStart=/usr/local/bin/ollama serve +User=ollama +Group=ollama +Restart=always +RestartSec=3 +Environment="PATH=/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin" +Environment="OLLAMA_HOST=127.0.0.1:11435" +Environment="OLLAMA_NUM_PARALLEL=32" +Environment="OLLAMA_KEEP_ALIVE=24h" +Environment="HOME=/usr/share/ollama" + +[Install] +WantedBy=default.target diff --git a/scripts/ollama.service b/scripts/ollama.service new file mode 100644 index 0000000..875e97a --- /dev/null +++ b/scripts/ollama.service @@ -0,0 +1,16 @@ +[Unit] +Description=Ollama Service +After=network-online.target + +[Service] +ExecStart=/usr/local/bin/ollama serve +Environment="OLLAMA_NUM_PARALLEL=10" +Environment="OLLAMA_KEEP_ALIVE=24h" +User=ollama +Group=ollama +Restart=always +RestartSec=3 +Environment="PATH=/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:/usr/games:/usr/local/games:/snap/bin" + +[Install] +WantedBy=default.target