Compare commits

..
Author SHA1 Message Date
reran4ik fd62763313 Add safe CSV artifact preview to job detail UI
coordinator / test (push) Waiting to run
Adds a Preview action next to eligible partial/final CSV artifacts on
the job detail page. Reads at most 64 KiB and 30 rows via a coordinator-
owned blob open, verifying job ownership and the same downloadable rule
as the existing download proxy so an artifact ID from another job is
never disclosed. Non-CSV and malformed/empty content fail safely with a
sanitized message instead of being rendered as text; all cell values go
through html/template escaping.
2026-07-24 15:32:53 +03:00
Emil 6e67daa9eb Merge distributed similarity search 2026-07-24 14:43:57 +03:00
Emil 0f3a2d92d8 Add distributed similarity search
python / test (push) Waiting to run
2026-07-24 14:38:11 +03:00
Emil 0bef7604fd Merge distributed workload protocol 2026-07-24 14:26:58 +03:00
18 changed files with 1314 additions and 49 deletions
+12 -10
View File
@@ -32,8 +32,8 @@ Docker PostgreSQL stack on 2026-07-23.
| CTX-04 Worker registry and HTTP API | Implemented | Registration, claim, heartbeat, result, failure, and status endpoints. |
| CTX-05 Artifact storage | Implemented | Coordinator-owned inputs/results, checksum verification, and upload flow. |
| CTX-06 Python Worker live-contract alignment | Implemented | Worker completed a real uploaded shard via HTTP on 2026-07-23. |
| CTX-07 Distributed workload protocol | Implemented | Versioned Python contract models, registry, strict plan validation, and deterministic reduction ordering are in `scimesh/distributed/`. The concrete molecular planner/reducer remains CTX-08/09. |
| CTX-08 Distributed similarity-search | Not started | Local reference exists. |
| CTX-07 Distributed workload protocol | Implemented | Versioned Python contract models, registry, strict plan validation, and deterministic reduction ordering are in `scimesh/distributed/`. |
| CTX-08 Distributed similarity-search | Implemented (scientific layer) | Python planner resolves `query_id` once, creates deterministic shard plans, worker adapter emits exact partial top-k CSVs/metrics, and reducer matches the local reference. Coordinator persistence/orchestration remains CTX-09. |
| CTX-09 Reducer and final-result API | Not started | Depends on CTX-07 and CTX-08. |
| CTX-10 Distributed similarity-graph | Not started | Local reference exists. |
| CTX-11 Dashboard/operator view | Implemented (diagnostic scope) | Protected local view: job/task/worker status, validated similarity-search upload, diagnostic partial-artifact download, and bounded polling. Final-result reduction remains CTX-09. |
@@ -41,20 +41,22 @@ Docker PostgreSQL stack on 2026-07-23.
## Next recommended assignment
Assign **CTX-08** to the workload role: implement the molecular
`similarity-search` planner and worker adapter on top of the accepted CTX-07
contract.
Assign **CTX-09** to the coordinator role: materialize planned shards,
persist them transactionally, invoke the registered reducer once, and expose a
durable final artifact.
## Known constraints
- The CTX-07 protocol is implemented, but no concrete molecular planner or
reducer is registered yet; the operator UI labels `partial_result` files as
diagnostic and cannot present them as final output.
- The Python `similarity-search` planner/reducer is implemented, but the Go
coordinator does not yet invoke it or persist its final artifact. The
operator UI labels `partial_result` files as diagnostic and cannot present
them as final output.
Use the local `scimesh` CLI for complete workload results.
- The worker/coordinator flow currently accepts both underscore API workload
names and hyphenated CLI names while the contract is consolidated.
- A real-stack worker test uses a small `query_smiles` shard. Resolving a
`query_id` once and sharing it across shards belongs to CTX-07.
- A real-stack worker test uses a small `query_smiles` shard. The Python
planner resolves `query_id` once and shares `query_smiles`; connecting that
planner to uploaded coordinator jobs belongs to CTX-09.
- The coordinator accepts uploaded distributed jobs only for
`similarity-search` with `query_smiles`. It rejects `similarity-graph` until
CTX-10 supplies cross-shard pair planning.
+1
View File
@@ -82,6 +82,7 @@ func run() error {
DownloadArtifact: usecase.NewDownloadArtifact(artifactRepo, blobStore),
GetTaskInput: usecase.NewGetTaskInput(taskRepo, artifactRepo, blobStore),
Dashboard: usecase.NewDashboard(uiReadRepo),
PreviewArtifact: usecase.NewPreviewArtifact(uiReadRepo, blobStore),
}
// Background reapers are tracked so shutdown can wait for them. Without this
@@ -29,6 +29,7 @@ type UseCases struct {
DownloadArtifact *usecase.DownloadArtifact
GetTaskInput *usecase.GetTaskInput
Dashboard *usecase.Dashboard
PreviewArtifact *usecase.PreviewArtifact
}
type Server struct {
@@ -82,6 +83,7 @@ func (s *Server) Handler(token string, uiToken ...string) http.Handler {
ui.HandleFunc("POST /ui/api/jobs/{job_id}/cancel", s.handleCancelJob)
ui.HandleFunc("POST /ui/api/jobs/upload", s.handleUploadDataset)
ui.HandleFunc("GET /ui/jobs/{job_id}/artifacts/{artifact_id}", s.handleUIArtifactDownload)
ui.HandleFunc("GET /ui/jobs/{job_id}/artifacts/{artifact_id}/preview", s.handleUIArtifactPreview)
mux.Handle("/ui", chain(ui, withRequestID, withAccessLog(s.log), withBasicAuth(uiToken[0]), withSameOrigin))
mux.Handle("/ui/", chain(ui, withRequestID, withAccessLog(s.log), withBasicAuth(uiToken[0]), withSameOrigin))
} else {
@@ -42,6 +42,7 @@ func newEnvWithUIToken(t *testing.T, ready func(context.Context) error, configur
clk := memstore.NewClock(time.Date(2026, 7, 21, 12, 0, 0, 0, time.UTC))
tx := memstore.Tx{}
lease := 2 * time.Minute
uiRead := memstore.NewUIReadRepo(jobs, tasks, work, arts)
uc := coordhttp.UseCases{
RegisterWorker: usecase.NewRegisterWorker(work, clk),
@@ -56,7 +57,8 @@ func newEnvWithUIToken(t *testing.T, ready func(context.Context) error, configur
UploadArtifact: usecase.NewUploadArtifact(tasks, arts, blobs, tx, clk),
DownloadArtifact: usecase.NewDownloadArtifact(arts, blobs),
GetTaskInput: usecase.NewGetTaskInput(tasks, arts, blobs),
Dashboard: usecase.NewDashboard(memstore.NewUIReadRepo(jobs, tasks, work, arts)),
Dashboard: usecase.NewDashboard(uiRead),
PreviewArtifact: usecase.NewPreviewArtifact(uiRead, blobs),
}
worker, err := uc.RegisterWorker.Execute(context.Background(), usecase.RegisterWorkerInput{
Name: "test-worker", Capabilities: []string{"w", "similarity-search"},
@@ -287,6 +289,134 @@ func TestUIArtifactDownloadRejectsAnotherJobsArtifact(t *testing.T) {
}
}
func TestUIArtifactPreviewRequiresAuth(t *testing.T) {
e := newEnv(t, healthy)
code, job := e.do(t, "POST", "/jobs", `{"workload":"w","input_uri":"s3://in","chunks":[{"chunk_index":0,"input_uri":"s3://c","input_sha256":"sha"}]}`)
if code != http.StatusCreated {
t.Fatalf("create: %d", code)
}
_, claim := e.do(t, "POST", "/tasks/claim", `{"worker_id":"w1","capabilities":["w"]}`)
artifactID := e.putArtifact(t, claim["task_id"].(string), "w1", int(claim["attempt"].(float64)), "a,b\n1,2\n")
req, _ := http.NewRequestWithContext(context.Background(), "GET",
e.ts.URL+"/ui/jobs/"+job["id"].(string)+"/artifacts/"+artifactID+"/preview", nil)
resp, err := http.DefaultClient.Do(req)
if err != nil {
t.Fatal(err)
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusUnauthorized {
t.Errorf("unauthenticated preview = %d, want 401", resp.StatusCode)
}
}
func TestUIArtifactPreviewRendersEscapedCSVRows(t *testing.T) {
e := newEnv(t, healthy)
code, job := e.do(t, "POST", "/jobs", `{"workload":"w","input_uri":"s3://in","chunks":[{"chunk_index":0,"input_uri":"s3://c","input_sha256":"sha"}]}`)
if code != http.StatusCreated {
t.Fatalf("create: %d", code)
}
_, claim := e.do(t, "POST", "/tasks/claim", `{"worker_id":"w1","capabilities":["w"]}`)
csv := "chembl_id,note\nCHEMBL1,<script>alert(1)</script>\n"
artifactID := e.putArtifact(t, claim["task_id"].(string), "w1", int(claim["attempt"].(float64)), csv)
req, _ := http.NewRequestWithContext(context.Background(), "GET",
e.ts.URL+"/ui/jobs/"+job["id"].(string)+"/artifacts/"+artifactID+"/preview", nil)
req.SetBasicAuth("operator", uiToken)
resp, err := http.DefaultClient.Do(req)
if err != nil {
t.Fatal(err)
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
t.Fatalf("preview: %d", resp.StatusCode)
}
body, _ := io.ReadAll(resp.Body)
if strings.Contains(string(body), "<script>alert(1)</script>") {
t.Error("preview must escape HTML-like CSV values, found raw <script> tag")
}
if !strings.Contains(string(body), "&lt;script&gt;") {
t.Errorf("expected escaped script tag in preview body: %s", body)
}
if !strings.Contains(string(body), "CHEMBL1") {
t.Error("preview missing expected cell value")
}
}
func TestUIArtifactPreviewRejectsAnotherJobsArtifact(t *testing.T) {
e := newEnv(t, healthy)
code, _ := e.do(t, "POST", "/jobs", `{"workload":"w","input_uri":"s3://in","chunks":[{"chunk_index":0,"input_uri":"s3://c","input_sha256":"sha"}]}`)
if code != http.StatusCreated {
t.Fatalf("first job: %d", code)
}
_, claim := e.do(t, "POST", "/tasks/claim", `{"worker_id":"w1","capabilities":["w"]}`)
artifactID := e.putArtifact(t, claim["task_id"].(string), "w1", int(claim["attempt"].(float64)), "a,b\n1,2\n")
code, second := e.do(t, "POST", "/jobs", `{"workload":"w","input_uri":"s3://in","chunks":[{"chunk_index":0,"input_uri":"s3://c","input_sha256":"sha"}]}`)
if code != http.StatusCreated {
t.Fatalf("second job: %d", code)
}
req, _ := http.NewRequestWithContext(context.Background(), "GET",
e.ts.URL+"/ui/jobs/"+second["id"].(string)+"/artifacts/"+artifactID+"/preview", nil)
req.SetBasicAuth("operator", uiToken)
resp, err := http.DefaultClient.Do(req)
if err != nil {
t.Fatal(err)
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusNotFound {
t.Errorf("cross-job preview = %d, want 404", resp.StatusCode)
}
}
func TestUIArtifactPreviewIsFriendlyForNonCSV(t *testing.T) {
e := newEnv(t, healthy)
code, job := e.do(t, "POST", "/jobs", `{"workload":"w","input_uri":"s3://in","chunks":[{"chunk_index":0,"input_uri":"s3://c","input_sha256":"sha"}]}`)
if code != http.StatusCreated {
t.Fatalf("create: %d", code)
}
_, claim := e.do(t, "POST", "/tasks/claim", `{"worker_id":"w1","capabilities":["w"]}`)
taskID := claim["task_id"].(string)
attempt := int(claim["attempt"].(float64))
req, _ := http.NewRequestWithContext(context.Background(), "PUT",
e.ts.URL+"/tasks/"+taskID+"/artifacts/notes.bin", strings.NewReader("\x00\x01binary garbage"))
req.Header.Set("Authorization", "Bearer "+token)
req.Header.Set("Content-Type", "application/octet-stream")
req.Header.Set("X-Worker-ID", e.workerID)
req.Header.Set("X-Task-Attempt", strconv.Itoa(attempt))
putResp, err := http.DefaultClient.Do(req)
if err != nil {
t.Fatal(err)
}
defer putResp.Body.Close()
if putResp.StatusCode != http.StatusOK {
t.Fatalf("put non-csv artifact: %d", putResp.StatusCode)
}
var m map[string]any
b, _ := io.ReadAll(putResp.Body)
_ = json.Unmarshal(b, &m)
artifactID := m["artifact_id"].(string)
previewReq, _ := http.NewRequestWithContext(context.Background(), "GET",
e.ts.URL+"/ui/jobs/"+job["id"].(string)+"/artifacts/"+artifactID+"/preview", nil)
previewReq.SetBasicAuth("operator", uiToken)
resp, err := http.DefaultClient.Do(previewReq)
if err != nil {
t.Fatal(err)
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
t.Fatalf("preview status: %d", resp.StatusCode)
}
body, _ := io.ReadAll(resp.Body)
if strings.Contains(string(body), "binary garbage") {
t.Error("non-CSV bytes must not be rendered as text")
}
if !strings.Contains(string(body), "not a CSV file") {
t.Errorf("expected a friendly non-CSV explanation, got: %s", body)
}
}
func TestHealthUnavailableWhenDBDown(t *testing.T) {
e := newEnv(t, func(context.Context) error { return context.DeadlineExceeded })
resp := e.get(t, "/health")
@@ -0,0 +1,31 @@
{{define "artifact-preview.html"}}
<!doctype html>
<html lang="en">
<head>
<meta charset="utf-8">
<meta name="viewport" content="width=device-width, initial-scale=1">
<title>SciMesh artifact preview</title>
<style>
:root{color:#172033;background:#f6f8fc;font:16px/1.5 system-ui,sans-serif}body{margin:0}.page{max-width:1180px;margin:auto;padding:32px 20px 56px}a{color:#174ecf}.back{text-decoration:none}h1{margin:18px 0 4px;font-size:1.6rem;word-break:break-word}.muted{color:#68758b}.notice{margin:16px 0;padding:15px 17px;border:1px solid #f2cb72;border-radius:10px;background:#fff8e6}.table-wrap{overflow-x:auto;border:1px solid #dfe5f0;border-radius:10px;background:#fff;margin-top:16px}table{width:100%;border-collapse:collapse}td,th{padding:10px 12px;border-bottom:1px solid #e8ecf4;text-align:left;vertical-align:top;white-space:pre-wrap;word-break:break-word}th{color:#50617d;font-size:.78rem;text-transform:uppercase;letter-spacing:.06em;background:#f6f8fc}tr:last-child td{border:0}.empty{padding:24px;text-align:center;color:#68758b}
</style>
</head>
<body>
<main class="page">
<a class="back" href="/ui/jobs/{{.JobID}}">← Back to job</a>
<h1>Preview: {{.Filename}}</h1>
<p class="muted">Diagnostic preview only — a partial shard result, not a final molecular-search answer. At most {{.RowLimit}} rows and {{.ByteLimit}} bytes are read from storage.</p>
{{if not .Previewable}}
<div class="notice">{{.Reason}}</div>
{{else}}
{{if .Truncated}}<div class="notice">Truncated: showing at most the first {{.RowLimit}} rows or {{.ByteLimit}} bytes of this artifact. Download it for the full contents.</div>{{end}}
<div class="table-wrap">
<table>
<tr>{{range .Headers}}<th>{{.}}</th>{{end}}</tr>
{{range .Rows}}<tr>{{range .}}<td>{{.}}</td>{{end}}</tr>{{else}}<tr><td class="empty" colspan="99">No data rows.</td></tr>{{end}}
</table>
</div>
{{end}}
</main>
</body>
</html>
{{end}}
@@ -17,7 +17,7 @@
<h2>Shard tasks</h2><p class="muted">If a task fails, its code and message appear here. Refresh the page to update the detailed rows.</p>
<div class="table-wrap"><table><tr><th>Shard</th><th>State</th><th>Attempt</th><th>Worker / lease</th><th>Error</th></tr>{{range .Tasks}}<tr><td>#{{.ChunkIndex}}</td><td><span class="status status-{{statusClass .Status}}">{{statusLabel .Status}}</span></td><td>{{.Attempt}} / {{.MaxAttempts}}</td><td>{{if .LeaseOwner}}<code>{{.LeaseOwner}}</code>{{if .LeaseExpiresAt}}<br><small>until {{time .LeaseExpiresAt}}</small>{{end}}{{else}}<span class="muted"></span>{{end}}</td><td class="error">{{if .ErrorCode}}<strong>{{taskErrorLabel .ErrorCode}}</strong><br><small>{{taskErrorHint .ErrorCode}}</small>{{else}}<span class="muted"></span>{{end}}</td></tr>{{else}}<tr><td colspan="5" class="empty">No tasks have appeared yet.</td></tr>{{end}}</table></div>
<h2>Coordinator artifacts</h2>
<div class="table-wrap"><table><tr><th>Type</th><th>File</th><th>Size</th><th>Integrity check</th><th></th></tr>{{range .Artifacts}}<tr><td>{{if .Diagnostic}}<strong>Partial result</strong><br><small>diagnostic</small>{{else}}{{.Kind}}{{end}}</td><td>{{.Filename}}</td><td>{{bytes .SizeBytes}}</td><td><code>{{.SHA256}}</code></td><td>{{if .Downloadable}}<a class="download" href="/ui/jobs/{{$.ID}}/artifacts/{{.ID}}">Download CSV</a>{{else}}<span class="muted">Unavailable</span>{{end}}</td></tr>{{else}}<tr><td colspan="5" class="empty">No artifacts yet. The worker uploads a CSV after it completes a shard.</td></tr>{{end}}</table></div>
<div class="table-wrap"><table><tr><th>Type</th><th>File</th><th>Size</th><th>Integrity check</th><th></th></tr>{{range .Artifacts}}<tr><td>{{if .Diagnostic}}<strong>Partial result</strong><br><small>diagnostic</small>{{else}}{{.Kind}}{{end}}</td><td>{{.Filename}}</td><td>{{bytes .SizeBytes}}</td><td><code>{{.SHA256}}</code></td><td>{{if .Downloadable}}<a class="download" href="/ui/jobs/{{$.ID}}/artifacts/{{.ID}}">Download CSV</a> <a class="download" href="/ui/jobs/{{$.ID}}/artifacts/{{.ID}}/preview">Preview</a>{{else}}<span class="muted">Unavailable</span>{{end}}</td></tr>{{else}}<tr><td colspan="5" class="empty">No artifacts yet. The worker uploads a CSV after it completes a shard.</td></tr>{{end}}</table></div>
</main>
<script>
const id={{printf "%q" .ID}},state={pending:['Waiting for a worker','waiting','Waiting for an available worker with the required capability.'],leased:['Assigned to a worker','active','A worker has claimed the task and should begin processing shortly.'],running:['Running','active','A worker is reading a shard, calculating fingerprints, and uploading its result through the coordinator.'],completed:['Tasks complete','success','Every shard task is complete. Files below are still partial results.'],failed:['Needs attention','danger','One or more shard tasks failed. Open the task list below for details.'],cancelled:['Stopped','waiting','The operator stopped this job. No new shards can be claimed.']};
+23
View File
@@ -279,3 +279,26 @@ func (s *Server) handleUIArtifactDownload(w http.ResponseWriter, r *http.Request
w.Header().Set("X-Checksum-SHA256", art.SHA256)
_, _ = io.Copy(w, body)
}
// handleUIArtifactPreview renders a bounded, job-scoped CSV preview. The use
// case enforces the same ownership and downloadable rule as the download
// proxy above; nothing here trusts the artifact ID beyond that check.
func (s *Server) handleUIArtifactPreview(w http.ResponseWriter, r *http.Request) {
jobID, ok := s.uiJobID(w, r)
if !ok {
return
}
artifactID, err := uuid.Parse(r.PathValue("artifact_id"))
if err != nil {
s.writeError(w, r, domain.ErrInvalidInput)
return
}
ctx, cancel := s.reqCtx(r)
defer cancel()
view, err := s.uc.PreviewArtifact.Execute(ctx, jobID, artifactID)
if err != nil {
s.writeError(w, r, err)
return
}
s.renderUI(w, "artifact-preview.html", view)
}
+150
View File
@@ -0,0 +1,150 @@
package usecase
import (
"context"
"encoding/csv"
"errors"
"io"
"strings"
"github.com/google/uuid"
"github.com/emil28092005/SciMesh/coordinator/internal/domain"
)
// previewMaxRows and previewMaxBytes bound how much of an artifact the
// diagnostic preview ever reads or renders: a partial shard CSV can be large,
// and this is a diagnostic aid, not a viewer for the full file.
const (
previewMaxRows = 30
previewMaxBytes = 64 * 1024
)
// ArtifactPreviewView is what the UI renders for a diagnostic CSV preview. It
// never carries a storage path, database error, or worker-local detail.
type ArtifactPreviewView struct {
JobID string
ArtifactID string
Filename string
Previewable bool
Reason string
Headers []string
Rows [][]string
Truncated bool
RowLimit int
ByteLimit int64
}
// PreviewArtifact renders at most the first previewMaxRows rows of a CSV
// artifact, reading at most previewMaxBytes from storage. It reuses the same
// job-scoped, downloadable-artifact rule as the download proxy so an artifact
// ID from another job is never previewable.
type PreviewArtifact struct {
read UIReadRepository
blobs BlobStore
}
func NewPreviewArtifact(read UIReadRepository, blobs BlobStore) *PreviewArtifact {
return &PreviewArtifact{read: read, blobs: blobs}
}
func (p *PreviewArtifact) Execute(ctx context.Context, jobID, artifactID uuid.UUID) (ArtifactPreviewView, error) {
job, err := p.read.GetJob(ctx, jobID)
if err != nil {
return ArtifactPreviewView{}, err
}
tasks, err := p.read.ListTasksByJob(ctx, jobID)
if err != nil {
return ArtifactPreviewView{}, err
}
// Same status derivation the dashboard uses, so a final artifact previews
// exactly when it would also be offered for download.
status := jobCard(*job, tasks).Status
artifacts, err := p.read.ListArtifactsByJob(ctx, jobID)
if err != nil {
return ArtifactPreviewView{}, err
}
var art *domain.Artifact
for i := range artifacts {
if artifacts[i].ID == artifactID {
art = &artifacts[i]
break
}
}
if art == nil {
return ArtifactPreviewView{}, domain.ErrArtifactNotFound
}
downloadable := art.Kind == domain.ArtifactPartialResult ||
(art.Kind == domain.ArtifactFinalResult && status == string(domain.JobCompleted))
if !downloadable {
return ArtifactPreviewView{}, domain.ErrArtifactNotFound
}
view := ArtifactPreviewView{
JobID: jobID.String(),
ArtifactID: art.ID.String(),
Filename: art.Filename,
RowLimit: previewMaxRows,
ByteLimit: previewMaxBytes,
}
if !isCSVArtifact(art) {
view.Reason = "This artifact is not a CSV file, so it cannot be shown as text here. Download it instead."
return view, nil
}
if art.SizeBytes == 0 {
view.Reason = "This artifact is empty."
return view, nil
}
rc, err := p.blobs.Open(ctx, art.StorageKey)
if err != nil {
return ArtifactPreviewView{}, err
}
defer func() { _ = rc.Close() }()
// LimitedReader caps the bytes read from storage regardless of how many
// rows are found within that window — the artifact is never loaded whole.
limited := &io.LimitedReader{R: rc, N: previewMaxBytes}
reader := csv.NewReader(limited)
reader.FieldsPerRecord = -1 // a byte-limited cut mid-row must not look like a schema error
header, err := reader.Read()
if err != nil {
view.Reason = "This artifact could not be read as CSV."
return view, nil
}
view.Headers = append([]string(nil), header...)
rows := make([][]string, 0, previewMaxRows)
for len(rows) < previewMaxRows {
record, err := reader.Read()
if err != nil {
if !errors.Is(err, io.EOF) {
// Malformed content further into the stream: keep what parsed
// cleanly and say the preview stopped early.
view.Truncated = true
}
break
}
rows = append(rows, append([]string(nil), record...))
}
view.Rows = rows
if art.SizeBytes > previewMaxBytes {
view.Truncated = true
} else if len(rows) == previewMaxRows {
if _, err := reader.Read(); err == nil {
view.Truncated = true
}
}
view.Previewable = true
return view, nil
}
func isCSVArtifact(a *domain.Artifact) bool {
if a.ContentType == "text/csv" {
return true
}
return strings.HasSuffix(strings.ToLower(a.Filename), ".csv")
}
@@ -0,0 +1,226 @@
package usecase_test
import (
"context"
"errors"
"strconv"
"strings"
"testing"
"time"
"github.com/google/uuid"
"github.com/emil28092005/SciMesh/coordinator/internal/domain"
"github.com/emil28092005/SciMesh/coordinator/internal/memstore"
"github.com/emil28092005/SciMesh/coordinator/internal/usecase"
)
func newPreviewHarness() (*usecase.PreviewArtifact, *memstore.JobRepo, *memstore.TaskRepo, *memstore.ArtifactRepo, *memstore.BlobStore) {
jobs := memstore.NewJobRepo()
tasks := memstore.NewTaskRepo()
work := memstore.NewWorkerRepo()
arts := memstore.NewArtifactRepo()
blobs := memstore.NewBlobStore()
read := memstore.NewUIReadRepo(jobs, tasks, work, arts)
return usecase.NewPreviewArtifact(read, blobs), jobs, tasks, arts, blobs
}
func mustInsertJob(t *testing.T, jobs *memstore.JobRepo, status domain.JobStatus) uuid.UUID {
t.Helper()
job := &domain.Job{ID: uuid.New(), Workload: "similarity-search", Status: status, CreatedAt: time.Now()}
if err := jobs.Insert(context.Background(), job); err != nil {
t.Fatalf("insert job: %v", err)
}
return job.ID
}
func mustCompleteJob(t *testing.T, jobs *memstore.JobRepo, tasks *memstore.TaskRepo) uuid.UUID {
t.Helper()
jobID := mustInsertJob(t, jobs, domain.JobRunning)
task := &domain.Task{
ID: uuid.New(), JobID: jobID, ChunkIndex: 0, Workload: "similarity-search",
Status: domain.TaskCompleted, MaxAttempts: 3, CreatedAt: time.Now(),
}
if err := tasks.InsertBatch(context.Background(), []*domain.Task{task}); err != nil {
t.Fatalf("insert task: %v", err)
}
return jobID
}
func mustInsertArtifact(t *testing.T, arts *memstore.ArtifactRepo, blobs *memstore.BlobStore,
jobID uuid.UUID, kind domain.ArtifactKind, filename, contentType, body string) uuid.UUID {
t.Helper()
id := uuid.New()
sha, size, err := blobs.Put(context.Background(), id.String(), strings.NewReader(body))
if err != nil {
t.Fatalf("put blob: %v", err)
}
art := &domain.Artifact{
ID: id, JobID: jobID, Kind: kind, Filename: filename,
StorageKey: id.String(), ContentType: contentType,
SizeBytes: size, SHA256: sha, CreatedAt: time.Now(),
}
if err := arts.Insert(context.Background(), art); err != nil {
t.Fatalf("insert artifact: %v", err)
}
return id
}
func TestPreviewArtifactRendersCSVRows(t *testing.T) {
preview, jobs, _, arts, blobs := newPreviewHarness()
jobID := mustInsertJob(t, jobs, domain.JobRunning)
artID := mustInsertArtifact(t, arts, blobs, jobID, domain.ArtifactPartialResult, "result.csv", "text/csv",
"chembl_id,score\nCHEMBL1,0.9\nCHEMBL2,0.8\n")
view, err := preview.Execute(context.Background(), jobID, artID)
if err != nil {
t.Fatalf("Execute: %v", err)
}
if !view.Previewable {
t.Fatalf("expected previewable, reason=%q", view.Reason)
}
if view.Truncated {
t.Error("small CSV should not be truncated")
}
if len(view.Headers) != 2 || view.Headers[0] != "chembl_id" {
t.Errorf("headers = %v", view.Headers)
}
if len(view.Rows) != 2 || view.Rows[0][0] != "CHEMBL1" {
t.Errorf("rows = %v", view.Rows)
}
}
func TestPreviewArtifactTruncatesAt30Rows(t *testing.T) {
preview, jobs, _, arts, blobs := newPreviewHarness()
jobID := mustInsertJob(t, jobs, domain.JobRunning)
var sb strings.Builder
sb.WriteString("id,value\n")
for i := 0; i < 40; i++ {
sb.WriteString("R" + strconv.Itoa(i) + ",v\n")
}
artID := mustInsertArtifact(t, arts, blobs, jobID, domain.ArtifactPartialResult, "result.csv", "text/csv", sb.String())
view, err := preview.Execute(context.Background(), jobID, artID)
if err != nil {
t.Fatalf("Execute: %v", err)
}
if len(view.Rows) != 30 {
t.Fatalf("rows = %d, want 30", len(view.Rows))
}
if !view.Truncated {
t.Error("expected truncated for more than 30 data rows")
}
}
func TestPreviewArtifactTruncatesAt64KiB(t *testing.T) {
preview, jobs, _, arts, blobs := newPreviewHarness()
jobID := mustInsertJob(t, jobs, domain.JobRunning)
var sb strings.Builder
sb.WriteString("id,value\n")
row := "row," + strings.Repeat("x", 200) + "\n"
for sb.Len() < 70*1024 {
sb.WriteString(row)
}
artID := mustInsertArtifact(t, arts, blobs, jobID, domain.ArtifactPartialResult, "result.csv", "text/csv", sb.String())
view, err := preview.Execute(context.Background(), jobID, artID)
if err != nil {
t.Fatalf("Execute: %v", err)
}
if !view.Truncated {
t.Error("expected truncated for an artifact bigger than 64KiB")
}
if len(view.Rows) > 30 {
t.Errorf("rows = %d, want <= 30", len(view.Rows))
}
}
func TestPreviewArtifactRejectsNonCSV(t *testing.T) {
preview, jobs, _, arts, blobs := newPreviewHarness()
jobID := mustInsertJob(t, jobs, domain.JobRunning)
artID := mustInsertArtifact(t, arts, blobs, jobID, domain.ArtifactPartialResult,
"shard-0.tsv", "text/tab-separated-values", "chembl_id\tcanonical_smiles\nCHEMBL1\tCCO\n")
view, err := preview.Execute(context.Background(), jobID, artID)
if err != nil {
t.Fatalf("Execute: %v", err)
}
if view.Previewable {
t.Error("non-CSV artifact must not be previewable as text")
}
if view.Reason == "" {
t.Error("expected a friendly reason")
}
}
func TestPreviewArtifactFailsSafelyOnEmptyArtifact(t *testing.T) {
preview, jobs, _, arts, blobs := newPreviewHarness()
jobID := mustInsertJob(t, jobs, domain.JobRunning)
artID := mustInsertArtifact(t, arts, blobs, jobID, domain.ArtifactPartialResult, "result.csv", "text/csv", "")
view, err := preview.Execute(context.Background(), jobID, artID)
if err != nil {
t.Fatalf("Execute: %v", err)
}
if view.Previewable {
t.Error("empty artifact must not be previewable")
}
if view.Reason == "" {
t.Error("expected a friendly reason")
}
}
func TestPreviewArtifactFailsSafelyOnMalformedCSV(t *testing.T) {
preview, jobs, _, arts, blobs := newPreviewHarness()
jobID := mustInsertJob(t, jobs, domain.JobRunning)
// An unterminated quote makes even the header row unparsable.
artID := mustInsertArtifact(t, arts, blobs, jobID, domain.ArtifactPartialResult, "result.csv", "text/csv", `"unterminated`)
view, err := preview.Execute(context.Background(), jobID, artID)
if err != nil {
t.Fatalf("Execute: %v", err)
}
if view.Previewable {
t.Error("malformed CSV must not be previewable")
}
if view.Reason == "" {
t.Error("expected a friendly reason")
}
}
func TestPreviewArtifactRejectsCrossJobArtifact(t *testing.T) {
preview, jobs, _, arts, blobs := newPreviewHarness()
jobA := mustInsertJob(t, jobs, domain.JobRunning)
jobB := mustInsertJob(t, jobs, domain.JobRunning)
artID := mustInsertArtifact(t, arts, blobs, jobA, domain.ArtifactPartialResult, "result.csv", "text/csv", "a,b\n1,2\n")
if _, err := preview.Execute(context.Background(), jobB, artID); !errors.Is(err, domain.ErrArtifactNotFound) {
t.Fatalf("err = %v, want ErrArtifactNotFound", err)
}
}
func TestPreviewArtifactRejectsUncompletedFinalResult(t *testing.T) {
preview, jobs, _, arts, blobs := newPreviewHarness()
jobID := mustInsertJob(t, jobs, domain.JobRunning)
artID := mustInsertArtifact(t, arts, blobs, jobID, domain.ArtifactFinalResult, "final.csv", "text/csv", "a,b\n1,2\n")
if _, err := preview.Execute(context.Background(), jobID, artID); !errors.Is(err, domain.ErrArtifactNotFound) {
t.Fatalf("err = %v, want ErrArtifactNotFound", err)
}
}
func TestPreviewArtifactAllowsFinalResultOnceJobIsCompleted(t *testing.T) {
preview, jobs, tasks, arts, blobs := newPreviewHarness()
jobID := mustCompleteJob(t, jobs, tasks)
artID := mustInsertArtifact(t, arts, blobs, jobID, domain.ArtifactFinalResult, "final.csv", "text/csv", "a,b\n1,2\n")
view, err := preview.Execute(context.Background(), jobID, artID)
if err != nil {
t.Fatalf("Execute: %v", err)
}
if !view.Previewable {
t.Fatalf("expected previewable, reason=%q", view.Reason)
}
}
+13 -8
View File
@@ -4,9 +4,11 @@
This document is the implementation contract for CTX-07. Its generic protocol,
registry, strict JSON models, and deterministic reduction ordering are
implemented in `scimesh/distributed/`. It does not implement a molecular
planner, reducer, API endpoint, database migration, or final artifact. Until
CTX-08 and CTX-09 are complete, shard CSVs remain diagnostic partial results.
implemented in `scimesh/distributed/`. CTX-08 implements the molecular
similarity-search planner, worker adapter, and pure reducer on top of it. This
document does not implement a coordinator API endpoint, database migration, or
durable final artifact. Until CTX-09 is complete, shard CSVs remain diagnostic
partial results.
The protocol gives local scientific workloads a coordinator-independent way to
validate a job, plan artifact-backed tasks, and later reduce completed outputs.
@@ -154,7 +156,10 @@ rank,chembl_id,canonical_smiles,similarity
```
- `rank` is one-based local rank.
- `similarity` uses the local CLI's six-decimal formatting.
- `similarity` uses a round-trip decimal representation of the computed float
(for Python, `repr(similarity)`). This preserves exact cross-shard ranking;
the reducer writes the user-facing final CSV with the local CLI's six-decimal
display formatting.
- Rows are sorted by `(-similarity, chembl_id, canonical_smiles)` for
`threshold_direction=greater`, or `(similarity, chembl_id,
canonical_smiles)` for `less`.
@@ -203,7 +208,7 @@ multiplicity. Reduction is independent of worker completion order and uses
## Deferred work
CTX-08 implements the similarity-search planner, runner adapter, reducer, and
comparison against the local CLI. CTX-09 persists the final artifact and job
state. CTX-10 defines graph-specific triangular block plans; it must not reuse
the search shard scheme without its pair-coverage invariants.
CTX-09 materializes the planned shard files as coordinator artifacts, invokes
the registered reducer once, and persists its final artifact/job state. CTX-10
defines graph-specific triangular block plans; it must not reuse the search
shard scheme without its pair-coverage invariants.
+9 -3
View File
@@ -122,7 +122,10 @@ Content-Type: application/json
},
"metrics": {
"elapsed_seconds": 12.4,
"processed_rows": 10000
"scanned_rows": 10000,
"valid_molecules": 9876,
"invalid_smiles": 124,
"matches_emitted": 20
}
}
```
@@ -189,8 +192,11 @@ class Runner(Protocol):
"""Run one task and return output artifacts plus safe metrics."""
```
`SciMeshRunner` should map `workload` and validated parameters to the existing
SciMesh CLI. For example, a `similarity-search` task invokes:
`SciMeshRunner` maps an allowlisted workload and validated parameters to the
local SciMesh reference functions. A planned `similarity-search` task contains
a resolved `query_smiles` (never `query_id`) and writes one exact local top-k
partial CSV plus the metrics above. Legacy single-shard tasks may still use the
CLI compatibility path:
```text
scimesh similarity-search <local-input> --query-id ... --output <task-dir>/result.csv
+9 -1
View File
@@ -12,17 +12,25 @@ from .models import (
FinalResult,
PlannedTask,
)
from .registry import DistributedWorkloadRegistry, PlanningService, WorkloadDescription
from .registry import (
DistributedWorkloadRegistry,
PlanningService,
WorkloadDescription,
default_distributed_registry,
)
from .similarity_search import SimilaritySearchDistributedWorkload
from .workload import DistributedWorkload
__all__ = [
"ArtifactReference",
"CompletedPartial",
"default_distributed_registry",
"DistributedPlan",
"DistributedWorkload",
"DistributedWorkloadRegistry",
"FinalResult",
"PlannedTask",
"PlanningService",
"SimilaritySearchDistributedWorkload",
"WorkloadDescription",
]
+13 -1
View File
@@ -50,7 +50,8 @@ class PlanningService:
It writes neither jobs nor artifacts. A Go coordinator bridge can therefore
validate and produce a plan before opening its own all-or-nothing persistence
transaction; CTX-08/09 will implement that concrete bridge and reducers.
transaction; CTX-09 will implement that concrete bridge and durable result
orchestration.
"""
def __init__(self, registry: DistributedWorkloadRegistry) -> None:
@@ -94,3 +95,14 @@ class PlanningService:
if not isinstance(result, FinalResult):
raise ValueError("distributed reducer must return a FinalResult")
return result
def default_distributed_registry() -> DistributedWorkloadRegistry:
"""Return the currently supported distributed scientific workloads."""
# Delayed import keeps the generic registry independent of concrete RDKit
# workloads and avoids making the contract layer import application setup.
from .similarity_search import SimilaritySearchDistributedWorkload
registry = DistributedWorkloadRegistry()
registry.register(SimilaritySearchDistributedWorkload())
return registry
+380
View File
@@ -0,0 +1,380 @@
"""Distributed planning and reduction for exact molecular similarity search."""
from __future__ import annotations
import csv
import hashlib
import heapq
import math
from pathlib import Path
from typing import Any, Iterator, Mapping, Sequence
from uuid import UUID, uuid5
from rdkit import Chem
from scimesh.chemistry.dataset import MoleculeRecord, find_molecule_by_id, parse_smiles
from scimesh.chemistry.fingerprints import FP_RADIUS, FP_SIZE
from scimesh.workloads.similarity_search import (
SimilarityMatch,
_HeapEntry,
search_similar,
write_search_results,
)
from .models import ArtifactReference, CompletedPartial, DistributedPlan, FinalResult, PlannedTask
_TSV_CONTENT_TYPE = "text/tab-separated-values"
_CSV_CONTENT_TYPE = "text/csv"
_SEARCH_COLUMNS = ("rank", "chembl_id", "canonical_smiles", "similarity")
_REQUIRED_COLUMNS = {"chembl_id", "canonical_smiles"}
def write_similarity_search_partial(output_path: Path, matches: Sequence[SimilarityMatch]) -> None:
"""Write a worker partial with a round-trip score, not display rounding.
The public final CSV continues to use the local CLI's six-decimal display.
A reducer needs the full binary float representation to rank candidates
from separate shards exactly as the single-process reference does.
"""
output_path.parent.mkdir(parents=True, exist_ok=True)
with output_path.open("w", encoding="utf-8", newline="") as destination:
writer = csv.DictWriter(destination, fieldnames=_SEARCH_COLUMNS)
writer.writeheader()
for rank, match in enumerate(matches, start=1):
writer.writerow({
"rank": rank,
"chembl_id": match.molecule_id,
"canonical_smiles": match.smiles,
"similarity": repr(match.similarity),
})
class SimilaritySearchDistributedWorkload:
"""Planner/reducer for exact global top-k Tanimoto similarity search."""
name = "similarity-search"
description = "Exact top-k molecular similarity search over deterministic TSV shards."
def validate_job(self, parameters: Mapping[str, object]) -> None:
allowed = {
"query_id", "query_smiles", "top_k", "threshold",
"threshold_direction", "max_rows", "progress_every",
}
unknown = set(parameters) - allowed
if unknown:
raise ValueError(f"unsupported similarity-search parameters: {', '.join(sorted(unknown))}")
query_id = parameters.get("query_id")
query_smiles = parameters.get("query_smiles")
if (query_id is None) == (query_smiles is None):
raise ValueError("exactly one of query_id or query_smiles is required")
if query_id is not None:
self._string(query_id, "query_id")
if query_smiles is not None:
self._string(query_smiles, "query_smiles")
self._positive_int(parameters.get("top_k", 20), "top_k")
if "max_rows" in parameters:
self._positive_int(parameters["max_rows"], "max_rows")
if "progress_every" in parameters:
self._nonnegative_int(parameters["progress_every"], "progress_every")
if "threshold" in parameters:
self._unit_interval(parameters["threshold"], "threshold")
if "threshold_direction" in parameters and parameters["threshold_direction"] not in {"greater", "less"}:
raise ValueError("threshold_direction must be 'greater' or 'less'")
def plan(
self,
input_path: Path,
input_artifact_id: str,
parameters: Mapping[str, object],
shard_rows: int,
workspace: Path,
) -> DistributedPlan:
self.validate_job(parameters)
if not input_path.is_file():
raise ValueError("input_path must be a readable dataset file")
if isinstance(shard_rows, bool) or not isinstance(shard_rows, int) or shard_rows < 1:
raise ValueError("shard_rows must be a positive integer")
try:
input_id = UUID(input_artifact_id)
except ValueError as error:
raise ValueError("input_artifact_id must be a UUID") from error
query_smiles, query_source = self._resolve_query(input_path, parameters)
resolved = self._resolved_parameters(parameters, query_smiles, query_source)
workspace.mkdir(parents=True, exist_ok=True)
shard_paths: list[Path] = []
try:
shard_paths = self._write_shards(input_path, workspace, shard_rows, resolved.get("max_rows"))
tasks = tuple(
PlannedTask(
chunk_index=index,
input_artifact=ArtifactReference(
artifact_id=str(uuid5(input_id, f"scimesh:similarity-search:shard:{index}")),
sha256=_sha256_file(path),
content_type=_TSV_CONTENT_TYPE,
),
parameters=self._task_parameters(resolved),
)
for index, path in enumerate(shard_paths)
)
except Exception:
for path in shard_paths:
path.unlink(missing_ok=True)
raise
return DistributedPlan(self.name, resolved, tasks)
def reduce(
self,
partial_results: Sequence[CompletedPartial],
parameters: Mapping[str, object],
workspace: Path,
) -> FinalResult:
"""Merge materialized partial CSVs into one deterministic final CSV.
The coordinator bridge materializes each downloaded artifact at
``workspace / artifact_id`` before it calls this method. Those local
paths are an ephemeral bridge detail, never present in the plan or task
payload. CTX-09 owns the durable final-artifact upload and job state.
"""
if not partial_results:
raise ValueError("at least one partial result is required")
resolved = self._validate_resolved_parameters(parameters)
top_k = resolved["top_k"]
direction = resolved["threshold_direction"]
heap: list[_HeapEntry] = []
ordered_partials = tuple(sorted(partial_results, key=lambda partial: partial.chunk_index))
indexes = [partial.chunk_index for partial in ordered_partials]
if len(indexes) != len(set(indexes)):
raise ValueError("partial results must have unique chunk_index values")
for partial in ordered_partials:
path = workspace / partial.artifact.artifact_id
if not path.is_file():
raise ValueError("materialized partial result is missing")
if _sha256_file(path) != partial.artifact.sha256:
raise ValueError("materialized partial result checksum does not match its artifact reference")
for match in self._read_partial(path, direction):
rank_key = match.sort_key(direction)
entry = _HeapEntry(match, rank_key)
if len(heap) < top_k:
heapq.heappush(heap, entry)
elif rank_key < heap[0].rank_key:
heapq.heapreplace(heap, entry)
matches = sorted((entry.match for entry in heap), key=lambda match: match.sort_key(direction))
output = workspace / "result.csv"
write_search_results(output, matches)
final_id = uuid5(
UUID(ordered_partials[0].artifact.artifact_id),
"scimesh:similarity-search:final:" + ",".join(
partial.artifact.artifact_id for partial in ordered_partials
),
)
return FinalResult(
ArtifactReference(str(final_id), _sha256_file(output), _CSV_CONTENT_TYPE),
{"matches_emitted": len(matches), "partial_count": len(ordered_partials)},
)
def _resolve_query(
self, input_path: Path, parameters: Mapping[str, object]
) -> tuple[str, dict[str, str]]:
query_id = parameters.get("query_id")
if isinstance(query_id, str):
record = find_molecule_by_id(input_path, query_id)
return Chem.MolToSmiles(record.molecule, canonical=True), {"kind": "chembl_id", "value": query_id}
supplied = parameters["query_smiles"]
assert isinstance(supplied, str) # checked by validate_job
molecule = parse_smiles(supplied)
if molecule is None:
raise ValueError("query_smiles is invalid")
return Chem.MolToSmiles(molecule, canonical=True), {"kind": "smiles", "value": supplied}
def _resolved_parameters(
self, parameters: Mapping[str, object], query_smiles: str, query_source: Mapping[str, str]
) -> dict[str, object]:
resolved: dict[str, object] = {
"query_smiles": query_smiles,
"query_source": dict(query_source),
"top_k": self._positive_int(parameters.get("top_k", 20), "top_k"),
"threshold_direction": parameters.get("threshold_direction", "greater"),
"fingerprint": {"algorithm": "morgan", "radius": FP_RADIUS, "fp_size": FP_SIZE},
}
if "threshold" in parameters:
resolved["threshold"] = self._unit_interval(parameters["threshold"], "threshold")
if "max_rows" in parameters:
resolved["max_rows"] = self._positive_int(parameters["max_rows"], "max_rows")
if "progress_every" in parameters:
resolved["progress_every"] = self._nonnegative_int(parameters["progress_every"], "progress_every")
return resolved
def _validate_resolved_parameters(self, parameters: Mapping[str, object]) -> dict[str, object]:
query_smiles = self._string(parameters.get("query_smiles"), "query_smiles")
if parse_smiles(query_smiles) is None:
raise ValueError("query_smiles is invalid")
resolved = self._resolved_parameters(
parameters,
Chem.MolToSmiles(parse_smiles(query_smiles), canonical=True),
{"kind": "resolved", "value": query_smiles},
)
# A reducer receives immutable plan metadata, whose query source and
# fixed fingerprint are observational context rather than worker input.
if "fingerprint" in parameters:
fingerprint = parameters["fingerprint"]
if fingerprint != {"algorithm": "morgan", "radius": FP_RADIUS, "fp_size": FP_SIZE}:
raise ValueError("resolved fingerprint does not match SciMesh defaults")
return resolved
@staticmethod
def _task_parameters(resolved: Mapping[str, object]) -> dict[str, object]:
# max_rows is applied before sharding. Passing it to each task would
# silently scan N rows per shard instead of the requested global prefix.
return {
key: value for key, value in resolved.items()
if key in {"query_smiles", "top_k", "threshold", "threshold_direction", "progress_every"}
}
def _write_shards(
self, input_path: Path, workspace: Path, shard_rows: int, max_rows: object
) -> list[Path]:
limit = int(max_rows) if isinstance(max_rows, int) else None
paths: list[Path] = []
current: Path | None = None
destination = None
rows_in_shard = 0
seen_rows = 0
try:
with input_path.open("r", encoding="utf-8", newline="") as source:
reader = csv.DictReader(source, delimiter="\t")
fieldnames = reader.fieldnames or []
if not _REQUIRED_COLUMNS.issubset(set(fieldnames)):
missing = sorted(_REQUIRED_COLUMNS - set(fieldnames))
raise ValueError(f"dataset is missing required columns: {', '.join(missing)}")
for row in reader:
if limit is not None and seen_rows >= limit:
break
if destination is None or rows_in_shard == shard_rows:
if destination is not None:
destination.close()
current = workspace / f"shard-{len(paths)}.tsv"
destination = current.open("w", encoding="utf-8", newline="")
writer = csv.DictWriter(destination, fieldnames=fieldnames, delimiter="\t", lineterminator="\n")
writer.writeheader()
paths.append(current)
rows_in_shard = 0
writer.writerow(row)
rows_in_shard += 1
seen_rows += 1
finally:
if destination is not None:
destination.close()
if not paths:
raise ValueError("dataset has no data rows")
return paths
@staticmethod
def _read_partial(path: Path, direction: object) -> Iterator[SimilarityMatch]:
if not path.is_file():
raise ValueError("materialized partial result is missing")
if direction not in {"greater", "less"}:
raise ValueError("threshold_direction must be 'greater' or 'less'")
previous_key: tuple[float, str, str] | None = None
with path.open("r", encoding="utf-8", newline="") as source:
reader = csv.DictReader(source)
if tuple(reader.fieldnames or ()) != _SEARCH_COLUMNS:
raise ValueError("partial result has an invalid CSV header")
for expected_rank, row in enumerate(reader, start=1):
if set(row) != set(_SEARCH_COLUMNS) or row["rank"] != str(expected_rank):
raise ValueError("partial result has an invalid rank")
try:
similarity = float(row["similarity"])
except (TypeError, ValueError) as error:
raise ValueError("partial result has an invalid similarity") from error
if not math.isfinite(similarity) or not 0 <= similarity <= 1:
raise ValueError("partial result has an invalid similarity")
match = SimilarityMatch(similarity, row["chembl_id"], row["canonical_smiles"])
key = match.sort_key(direction)
if previous_key is not None and key < previous_key:
raise ValueError("partial result is not sorted deterministically")
previous_key = key
yield match
@staticmethod
def _string(value: object, name: str) -> str:
if not isinstance(value, str) or not value.strip() or len(value) > 200:
raise ValueError(f"{name} must be a non-empty string")
return value
@staticmethod
def _positive_int(value: object, name: str) -> int:
if isinstance(value, bool) or not isinstance(value, int) or value < 1:
raise ValueError(f"{name} must be a positive integer")
return value
@staticmethod
def _nonnegative_int(value: object, name: str) -> int:
if isinstance(value, bool) or not isinstance(value, int) or value < 0:
raise ValueError(f"{name} must be a non-negative integer")
return value
@staticmethod
def _unit_interval(value: object, name: str) -> float:
if isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(value) or not 0 <= value <= 1:
raise ValueError(f"{name} must be a number between 0 and 1")
return float(value)
def run_similarity_search_shard(
input_path: Path, parameters: Mapping[str, object], output_path: Path
) -> dict[str, int]:
"""Run one planned shard using the local reference implementation.
This is the worker adapter used by CTX-08. It deliberately accepts only
resolved ``query_smiles``: resolving an identifier independently in each
shard would make the distributed search scientifically invalid.
"""
allowed = {"query_smiles", "top_k", "threshold", "threshold_direction", "progress_every"}
unknown = set(parameters) - allowed
if unknown:
raise ValueError(f"unsupported similarity-search parameters: {', '.join(sorted(unknown))}")
query_smiles = parameters.get("query_smiles")
if not isinstance(query_smiles, str) or not query_smiles.strip():
raise ValueError("query_smiles is required for a distributed shard")
molecule = parse_smiles(query_smiles)
if molecule is None:
raise ValueError("query_smiles is invalid")
top_k = SimilaritySearchDistributedWorkload._positive_int(parameters.get("top_k", 20), "top_k")
threshold = None
if "threshold" in parameters:
threshold = SimilaritySearchDistributedWorkload._unit_interval(parameters["threshold"], "threshold")
direction = parameters.get("threshold_direction", "greater")
if direction not in {"greater", "less"}:
raise ValueError("threshold_direction must be 'greater' or 'less'")
progress_every = 0
if "progress_every" in parameters:
progress_every = SimilaritySearchDistributedWorkload._nonnegative_int(
parameters["progress_every"], "progress_every"
)
result = search_similar(
input_path,
MoleculeRecord("query", query_smiles, molecule),
top_k=top_k,
progress_every=progress_every,
threshold=threshold,
threshold_direction=direction,
)
write_similarity_search_partial(output_path, result.matches)
return {
"scanned_rows": result.stats.scanned,
"valid_molecules": result.stats.valid,
"invalid_smiles": result.stats.invalid,
"matches_emitted": len(result.matches),
}
def _sha256_file(path: Path) -> str:
digest = hashlib.sha256()
with path.open("rb") as source:
for block in iter(lambda: source.read(1024 * 1024), b""):
digest.update(block)
return digest.hexdigest()
+13 -1
View File
@@ -9,6 +9,7 @@ from pathlib import Path
import random
import re
import shutil
import subprocess
import threading
import time
from datetime import datetime, timezone
@@ -202,12 +203,23 @@ class WorkerDaemon:
def _report_failure(self, task: ClaimedTask, error: Exception) -> None:
message = self._sanitize_error_message(error)
try:
self.coordinator.fail(task, {"worker_id": self._worker_id(), "attempt": task.attempt, "error_code": type(error).__name__, "error_message": message})
self.coordinator.fail(task, {
"worker_id": self._worker_id(),
"attempt": task.attempt,
"error_code": type(error).__name__,
"error_message": message,
"retryable": self._is_retryable(error),
})
except CoordinatorTransientError:
raise
except Exception:
self._log("failed", task, error_type="FailureReportError")
@staticmethod
def _is_retryable(error: Exception) -> bool:
"""Retry transient worker/transport failures, never invalid scientific input."""
return not isinstance(error, (ValueError, FileNotFoundError, subprocess.CalledProcessError))
def _sanitize_error_message(self, error: Exception) -> str:
"""Keep coordinator-visible failures useful without exposing local paths."""
message = str(error).replace(str(self.config.work_dir), "<worker-dir>")
+8
View File
@@ -7,6 +7,8 @@ import subprocess
import sys
from typing import Protocol
from scimesh.distributed.similarity_search import run_similarity_search_shard
from .models import ClaimedTask, ProducedArtifact, RunResult
@@ -34,6 +36,12 @@ class SciMeshRunner:
query_id, query_smiles = params.get("query_id"), params.get("query_smiles")
if (query_id is None) == (query_smiles is None):
raise ValueError("exactly one of query_id or query_smiles is required")
if query_smiles is not None and "max_rows" not in params:
metrics = run_similarity_search_shard(input_path, params, output_path)
return RunResult((ProducedArtifact(output_path, "text/csv"),), metrics)
# Legacy URI jobs may still use query_id or an explicitly task-local
# max_rows value. CTX-08 plans never create those payloads; retain
# CLI execution only for backwards compatibility at this boundary.
top_k = self._positive_int(params, "top_k", default=20)
command += ["--query-id", self._string(params, "query_id")] if query_id is not None else ["--query-smiles", self._string(params, "query_smiles")]
command += ["--top-k", str(top_k)]
+184
View File
@@ -0,0 +1,184 @@
"""Scientific reference tests for the CTX-08 distributed search workload."""
from __future__ import annotations
import csv
import hashlib
from pathlib import Path
from uuid import NAMESPACE_URL, uuid5
import pytest
from scimesh.chemistry.dataset import find_molecule_by_id
from scimesh.distributed import (
ArtifactReference,
CompletedPartial,
PlanningService,
default_distributed_registry,
)
from scimesh.distributed.registry import DistributedWorkloadRegistry
from scimesh.distributed.similarity_search import (
SimilaritySearchDistributedWorkload,
run_similarity_search_shard,
write_similarity_search_partial,
)
from scimesh.workloads.similarity_search import search_similar, write_search_results
def make_dataset(path: Path) -> None:
path.write_text(
"chembl_id\tcanonical_smiles\textra\n"
"CHEMBL_QUERY\tCCO\tquery\n"
"CHEMBL_A\tCCCO\ta\n"
"CHEMBL_B\tCCCC\tb\n"
"CHEMBL_INVALID\tnot-a-smiles\tbad\n"
"CHEMBL_DUPLICATE\tCCO\tduplicate\n"
"CHEMBL_C\tCCN\tc\n",
encoding="utf-8",
)
def planner() -> tuple[PlanningService, SimilaritySearchDistributedWorkload]:
workload = SimilaritySearchDistributedWorkload()
registry = DistributedWorkloadRegistry()
registry.register(workload)
return PlanningService(registry), workload
def test_default_registry_exposes_only_supported_distributed_search() -> None:
assert [item.name for item in default_distributed_registry().descriptions()] == ["similarity-search"]
def checksum(path: Path) -> str:
return hashlib.sha256(path.read_bytes()).hexdigest()
def test_query_id_is_resolved_once_before_deterministic_shards(
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
dataset = tmp_path / "chembl.tsv"
workspace = tmp_path / "workspace"
make_dataset(dataset)
service, _ = planner()
calls = 0
real_find = find_molecule_by_id
def count_find(path: Path, query_id: str) -> MoleculeRecord:
nonlocal calls
calls += 1
return real_find(path, query_id)
monkeypatch.setattr("scimesh.distributed.similarity_search.find_molecule_by_id", count_find)
input_id = str(uuid5(NAMESPACE_URL, "dataset"))
plan = service.plan(
"similarity-search", dataset, input_id,
{"query_id": "CHEMBL_QUERY", "top_k": 3, "max_rows": 5, "progress_every": 0},
2, workspace,
)
assert calls == 1
assert plan.resolved_parameters["query_smiles"] == "CCO"
assert plan.resolved_parameters["query_source"] == {"kind": "chembl_id", "value": "CHEMBL_QUERY"}
assert [task.chunk_index for task in plan.tasks] == [0, 1, 2]
assert all("query_id" not in task.parameters for task in plan.tasks)
assert all("max_rows" not in task.parameters for task in plan.tasks)
assert all(task.parameters["query_smiles"] == "CCO" for task in plan.tasks)
assert [
sum(1 for _ in path.open(encoding="utf-8")) - 1
for path in sorted(workspace.glob("shard-*.tsv"))
] == [2, 2, 1]
def test_distributed_reduction_matches_single_process_reference(tmp_path: Path) -> None:
dataset = tmp_path / "chembl.tsv"
workspace = tmp_path / "workspace"
make_dataset(dataset)
service, workload = planner()
plan = service.plan(
"similarity-search", dataset, str(uuid5(NAMESPACE_URL, "dataset")),
{"query_smiles": "CCO", "top_k": 3, "threshold": 0.0}, 2, workspace,
)
partials: list[CompletedPartial] = []
# Worker two finishes the latter shards first. Worker one loses its first
# attempt for shard zero, then retries it last. The reducer must remain
# independent of both completion and retry order.
for task in reversed(plan.tasks):
shard = workspace / f"shard-{task.chunk_index}.tsv"
temporary_partial = workspace / f"worker-output-{task.chunk_index}.csv"
metrics = run_similarity_search_shard(shard, task.parameters, temporary_partial)
partial_id = str(uuid5(NAMESPACE_URL, f"partial:{task.chunk_index}"))
partials.append(
CompletedPartial(
task.chunk_index,
ArtifactReference(
partial_id, checksum(temporary_partial), "text/csv",
),
metrics,
)
)
# The reducer materializes result files under their own coordinator IDs,
# not shard input IDs. Keep this fixture faithful to that boundary.
temporary_partial.rename(workspace / partial_id)
final = workload.reduce(tuple(partials), plan.resolved_parameters, workspace)
reference = tmp_path / "reference.csv"
query_record = find_molecule_by_id(dataset, "CHEMBL_QUERY")
write_search_results(reference, search_similar(dataset, query_record, top_k=3, threshold=0.0).matches)
assert (workspace / "result.csv").read_bytes() == reference.read_bytes()
assert final.metrics == {"matches_emitted": 3, "partial_count": 3}
rows = list(csv.DictReader((workspace / "result.csv").open(encoding="utf-8")))
assert {row["chembl_id"] for row in rows}.isdisjoint({"CHEMBL_QUERY", "CHEMBL_DUPLICATE"})
def test_reducer_rejects_unsorted_or_invalid_partial_csv(tmp_path: Path) -> None:
workspace = tmp_path / "workspace"
workspace.mkdir()
workload = SimilaritySearchDistributedWorkload()
artifact_id = str(uuid5(NAMESPACE_URL, "bad"))
partial_path = workspace / artifact_id
partial_path.write_text(
"rank,chembl_id,canonical_smiles,similarity\n1,A,CC,0.1\n2,B,CCC,0.9\n",
encoding="utf-8",
)
artifact = ArtifactReference(artifact_id, checksum(partial_path), "text/csv")
with pytest.raises(ValueError, match="not sorted"):
workload.reduce(
(CompletedPartial(0, artifact, {"scanned_rows": 2}),),
{"query_smiles": "CCO", "top_k": 2, "threshold_direction": "greater", "fingerprint": {"algorithm": "morgan", "radius": 2, "fp_size": 2048}},
workspace,
)
def test_partial_csv_preserves_exact_scores_for_global_ranking(tmp_path: Path) -> None:
partial = tmp_path / "partial.csv"
# Both values look identical in a six-decimal final CSV. The exact value
# must survive shard transport so the global reducer can still rank them.
from scimesh.workloads.similarity_search import SimilarityMatch
write_similarity_search_partial(
partial,
[SimilarityMatch(0.50000049, "A", "CC"), SimilarityMatch(0.50000048, "B", "CCC")],
)
values = list(csv.DictReader(partial.open(encoding="utf-8")))
assert values[0]["similarity"] == repr(0.50000049)
assert values[1]["similarity"] == repr(0.50000048)
def test_planner_rejects_fingerprint_override_and_invalid_query(tmp_path: Path) -> None:
dataset = tmp_path / "chembl.tsv"
make_dataset(dataset)
service, _ = planner()
with pytest.raises(ValueError, match="unsupported similarity-search parameters"):
service.plan(
"similarity-search", dataset, str(uuid5(NAMESPACE_URL, "dataset")),
{"query_smiles": "CCO", "fingerprint": {"radius": 1}}, 2, tmp_path / "workspace",
)
with pytest.raises(ValueError, match="query_smiles is invalid"):
service.plan(
"similarity-search", dataset, str(uuid5(NAMESPACE_URL, "dataset")),
{"query_smiles": "invalid"}, 2, tmp_path / "workspace",
)
+108 -23
View File
@@ -106,6 +106,90 @@ def test_claims_runs_uploads_and_submits_csv(tmp_path: Path) -> None:
}
def test_worker_executes_a_resolved_similarity_search_shard(tmp_path: Path) -> None:
content = b"chembl_id\tcanonical_smiles\nQUERY\tCCO\nMATCH\tCCCO\nINVALID\tnot-a-smiles\n"
task = make_task(content)
task = ClaimedTask(
task.task_id, task.attempt, task.lease_expires_at, task.workload, task.input,
{"query_smiles": "CCO", "top_k": 5, "progress_every": 0},
)
worker, coordinator, artifacts, _, _ = daemon(tmp_path, task, content)
worker.runner = SciMeshRunner()
assert worker.run_once() == RunOnceOutcome(claimed=True, completed=True)
output = artifacts.uploaded[0][2].read_text(encoding="utf-8")
assert output.startswith("rank,chembl_id,canonical_smiles,similarity\n")
metrics = coordinator.submissions[0]["metrics"]
assert metrics["scanned_rows"] == 3
assert metrics["valid_molecules"] == 2
assert metrics["invalid_smiles"] == 1
assert metrics["matches_emitted"] == 1
assert isinstance(metrics["elapsed_seconds"], float)
def test_two_workers_complete_resolved_shards_after_one_retry(tmp_path: Path) -> None:
content = b"chembl_id\tcanonical_smiles\nQUERY\tCCO\nMATCH\tCCCO\n"
first = make_task(content)
first = ClaimedTask(
"retry-task", 1, first.lease_expires_at, "similarity-search", first.input,
{"query_smiles": "CCO", "top_k": 5},
)
second = ClaimedTask(
"other-task", 1, first.lease_expires_at, "similarity-search", first.input,
{"query_smiles": "CCO", "top_k": 5},
)
class RetryCoordinator(FakeCoordinator):
def __init__(self) -> None:
super().__init__(None)
self.queue = [first, second]
self.claimants: list[str] = []
def claim(self, worker_id: str, capabilities: tuple[str, ...]) -> ClaimedTask | None:
self.claimants.append(worker_id)
return self.queue.pop(0) if self.queue else None
def fail(self, task: ClaimedTask, payload: dict) -> None:
self.failures.append(payload)
if task.task_id == "retry-task" and task.attempt == 1 and payload["retryable"]:
self.queue.append(
ClaimedTask(
task.task_id, 2, task.lease_expires_at, task.workload, task.input,
task.parameters,
)
)
class FailFirstAttempt:
def __init__(self) -> None:
self.calls = 0
self.delegate = SciMeshRunner()
def run(self, task: ClaimedTask, task_dir: Path) -> RunResult:
self.calls += 1
if self.calls == 1:
raise RuntimeError("simulated retryable shard failure")
return self.delegate.run(task, task_dir)
coordinator = RetryCoordinator()
artifacts = FakeArtifacts(content)
worker_a = WorkerDaemon(
WorkerConfig("https://example.test", "worker-a", tmp_path / "worker-a"),
coordinator, artifacts, FailFirstAttempt(),
)
worker_b = WorkerDaemon(
WorkerConfig("https://example.test", "worker-b", tmp_path / "worker-b"),
coordinator, artifacts, SciMeshRunner(),
)
assert worker_a.run_once() == RunOnceOutcome(claimed=True, completed=False)
assert worker_b.run_once() == RunOnceOutcome(claimed=True, completed=True)
assert worker_a.run_once() == RunOnceOutcome(claimed=True, completed=True)
assert coordinator.claimants == ["worker-a", "worker-b", "worker-a"]
assert len(coordinator.failures) == 1
assert coordinator.failures[0]["retryable"] is True
assert len(coordinator.submissions) == 2
def test_no_task_does_not_create_directory(tmp_path: Path) -> None:
worker, _, _, runner, config = daemon(tmp_path, None, b"")
assert worker.run_once() == RunOnceOutcome(claimed=False, completed=False)
@@ -161,6 +245,7 @@ def test_interrupting_an_active_task_reports_a_sanitized_failure(tmp_path: Path)
"attempt": 1,
"error_code": "InterruptedError",
"error_message": "worker interrupted by operator",
"retryable": True,
}
]
@@ -234,6 +319,7 @@ def test_bad_checksum_reports_failure_without_running(tmp_path: Path) -> None:
assert worker.run_once() == RunOnceOutcome(claimed=True, completed=False)
assert runner.calls == 0
assert coordinator.failures[0]["error_code"] == "ValueError"
assert coordinator.failures[0]["retryable"] is False
assert not coordinator.submissions
@@ -356,30 +442,35 @@ def test_runner_maps_graph_and_smiles_search_parameters(tmp_path: Path, monkeypa
runner = SciMeshRunner()
graph = ClaimedTask("graph", 1, "2026-07-30T00:00:00Z", "similarity-graph", InputArtifact("https://example/input", "x"), {"threshold": 0.2, "threshold_direction": "less", "block_size": 42, "max_rows": 7, "progress_every": 0})
search = ClaimedTask("search", 1, "2026-07-30T00:00:00Z", "similarity-search", InputArtifact("https://example/input", "x"), {"query_smiles": "CCO", "top_k": 3})
search_dir = tmp_path / "search"
search_dir.mkdir()
(search_dir / "input").write_text(
"chembl_id\tcanonical_smiles\nA\tCCO\nB\tCCCO\n", encoding="utf-8"
)
runner.run(graph, tmp_path / "graph")
runner.run(search, tmp_path / "search")
result = runner.run(search, search_dir)
assert "--threshold-direction" in commands[0] and "less" in commands[0]
assert "--block-size" in commands[0] and "42" in commands[0]
assert "--max-rows" in commands[0] and "7" in commands[0]
assert "--query-smiles" in commands[1] and "CCO" in commands[1]
assert len(commands) == 1
assert result.metrics == {
"scanned_rows": 2, "valid_molecules": 2, "invalid_smiles": 0, "matches_emitted": 1,
}
def test_runner_accepts_coordinator_workload_names(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
commands: list[list[str]] = []
def fake_run(command: list[str], **_: object) -> None:
commands.append(command)
output = Path(command[command.index("--output") + 1])
output.parent.mkdir(parents=True, exist_ok=True)
output.write_text("a,b\\n", encoding="utf-8")
monkeypatch.setattr("scimesh.worker.runners.subprocess.run", fake_run)
task_dir = tmp_path / "search"
task_dir.mkdir()
(task_dir / "input").write_text(
"chembl_id\tcanonical_smiles\nA\tCCO\nB\tCCCO\n", encoding="utf-8"
)
task = ClaimedTask(
"search", 1, "2026-07-30T00:00:00Z", "similarity_search",
InputArtifact("https://example/input", "a" * 64), {"query_smiles": "CCO"},
)
SciMeshRunner().run(task, tmp_path / "search")
assert commands[0][3] == "similarity-search"
result = SciMeshRunner().run(task, task_dir)
assert result.metrics["matches_emitted"] == 1
assert (task_dir / "result.csv").is_file()
def test_claimed_task_rejects_path_traversal_and_invalid_metadata() -> None:
@@ -457,21 +548,15 @@ def test_relative_work_dir_is_normalized_for_runner_subprocesses(
task_dir = config.work_dir / "task" / "1"
task_dir.mkdir(parents=True)
(task_dir / "input").write_text("fixture", encoding="utf-8")
command: list[str] = []
def fake_run(args: list[str], **_: object) -> None:
command.extend(args)
output = Path(args[args.index("--output") + 1])
output.write_text("id,score\n", encoding="utf-8")
monkeypatch.setattr("scimesh.worker.runners.subprocess.run", fake_run)
(task_dir / "input").write_text(
"chembl_id\tcanonical_smiles\nA\tCCO\nB\tCCCO\n", encoding="utf-8"
)
task = ClaimedTask(
"task", 1, "2026-07-30T00:00:00Z", "similarity-search",
InputArtifact("https://example.test/input", "a" * 64), {"query_smiles": "CCO"},
)
SciMeshRunner().run(task, task_dir)
assert command[4] == str(task_dir / "input")
assert (task_dir / "result.csv").is_file()
def test_worker_registration_sets_returned_identity(tmp_path: Path) -> None: