feat(index): HNSW writer bridge for Pebble crawl path + ops systemd units

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)
+}

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())
+	}
+}

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

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

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

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

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

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

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