Compare commits

..
Author SHA1 Message Date
Emil a055473706 Add final result reduction
coordinator / test (push) Waiting to run
2026-07-24 14:58:18 +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
37 changed files with 1588 additions and 91 deletions
+4 -4
View File
@@ -3,10 +3,10 @@
SciMesh is a scientific-workload framework for molecular datasets. Its public CLI
runs exact similarity search and sparse similarity-graph construction locally in
one Python process; it creates no dense similarity matrix. The Go/PostgreSQL
coordinator and Python worker can run a diagnostic, shard-based
`similarity-search` pipeline locally. Its CSV artifacts are not a global result
until CTX-07--09 add planning and reduction; use the local CLI for scientific
results today. See [`STATUS.md`](STATUS.md).
coordinator and Python worker can run a shard-based `similarity-search`
pipeline locally. After every shard succeeds, the coordinator deterministically
merges its candidates into one final global top-k CSV. See
[`STATUS.md`](STATUS.md).
The ChEMBL TSV database is intentionally not included in this repository. Download it separately and pass its path to the commands below. The expected columns are `chembl_id` and `canonical_smiles`.
+14 -16
View File
@@ -1,7 +1,7 @@
# SciMesh Status
**Updated:** 2026-07-24
**Branch baseline:** `main` at `f953112` (distributed pipeline hardening)
**Branch baseline:** `main` at `6e67daa` (distributed similarity-search)
## Current state
@@ -18,8 +18,10 @@ the reference behaviour for future distributed execution:
The Go coordinator and its PostgreSQL-backed task lifecycle are implemented:
registration, atomic claiming, lease renewal, artifact storage, dataset
chunking, result/failure reporting, and job progress. The Python worker now
uses the live coordinator contract; its HTTP path was exercised against a real
Docker PostgreSQL stack on 2026-07-23.
uses the live coordinator contract. Completed similarity-search shard results
are reduced once into a checksum-protected final CSV, which is downloadable
through the coordinator. The full Go checks (including a fresh migration and
real PostgreSQL smoke test) passed on 2026-07-24.
## Milestone tracker
@@ -32,29 +34,25 @@ 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-09 Reducer and final-result API | Not started | Depends on CTX-07 and CTX-08. |
| 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 | 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. |
| CTX-09 Reducer and final-result API | Implemented | Atomic `reducing` claim, deterministic coordinator-side top-k reducer, sanitized reducer failure, final artifact persistence, `result_uri`, and final CSV download. |
| 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. |
| CTX-11 Dashboard/operator view | Implemented | Protected local view: job/task/worker status, validated similarity-search upload, partial-artifact diagnostics, final-result download, and bounded polling. |
| CTX-12 Reliability, security, CI | In progress | Unit, race, PostgreSQL integration, and smoke checks exist; CI hardening remains. |
## 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-10** to the distributed-science role: implement deterministic
block-pair planning and reduction for `similarity-graph`.
## 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.
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`; the upload UI
currently accepts `query_smiles` only.
- 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.
+4 -3
View File
@@ -76,9 +76,10 @@ UI_AUTH_TOKEN='local-ui-secret' make up
```
The UI is disabled by default and never accepts the worker bearer token.
It shows recent jobs, task/worker state, and the per-job partial artifacts.
Those files are explicitly diagnostic until the CTX-09 reducer creates a final
result; the UI does not present them as final scientific output.
It shows recent jobs, task/worker state, diagnostic shard artifacts, and the
final CSV for completed similarity-search jobs. The coordinator enters
`reducing` after the last shard completes, then exposes the final deterministic
global top-k result when merging succeeds.
`up` starts three services in order: Postgres waits until `pg_isready` passes, a
one-shot `migrate` container applies the schema and exits, and only then does the
+2
View File
@@ -75,11 +75,13 @@ func run() error {
ClaimTask: usecase.NewClaimTask(taskRepo, jobRepo, workerRepo, tx, clk, cfg.LeaseDuration),
RenewLease: usecase.NewRenewLease(taskRepo, workerRepo, tx, clk, cfg.LeaseDuration),
CompleteTask: usecase.NewCompleteTask(taskRepo, jobRepo, artifactRepo, tx, clk),
ReduceJob: usecase.NewReduceJob(jobRepo, taskRepo, artifactRepo, blobStore, tx, clk),
FailTask: usecase.NewFailTask(taskRepo, jobRepo, tx, clk),
GetJobStatus: usecase.NewGetJobStatus(jobRepo, taskRepo),
CancelJob: usecase.NewCancelJob(jobRepo, taskRepo, tx, clk),
UploadArtifact: usecase.NewUploadArtifact(taskRepo, artifactRepo, blobStore, tx, clk),
DownloadArtifact: usecase.NewDownloadArtifact(artifactRepo, blobStore),
GetJobResult: usecase.NewGetJobResult(jobRepo, usecase.NewDownloadArtifact(artifactRepo, blobStore)),
GetTaskInput: usecase.NewGetTaskInput(taskRepo, artifactRepo, blobStore),
Dashboard: usecase.NewDashboard(uiReadRepo),
}
+15 -8
View File
@@ -11,6 +11,7 @@ type JobStatus string
const (
JobPending JobStatus = "pending"
JobRunning JobStatus = "running"
JobReducing JobStatus = "reducing"
JobCompleted JobStatus = "completed"
JobFailed JobStatus = "failed"
JobCancelled JobStatus = "cancelled"
@@ -18,14 +19,18 @@ const (
// Job is one user submission that fans out into one or more tasks.
type Job struct {
ID uuid.UUID
Workload string
InputURI string // external input URI; empty for uploaded datasets
InputArtifactID *uuid.UUID // uploaded input artifact; nil for URI submissions
Parameters map[string]any
Status JobStatus
CreatedAt time.Time
CompletedAt *time.Time
ID uuid.UUID
Workload string
InputURI string // external input URI; empty for uploaded datasets
InputArtifactID *uuid.UUID // uploaded input artifact; nil for URI submissions
ResultArtifactID *uuid.UUID
Parameters map[string]any
Status JobStatus
CreatedAt time.Time
CompletedAt *time.Time
ReducerStartedAt *time.Time
ErrorCode *string
ErrorMessage *string
}
// NewUploadedJob builds a job whose input was uploaded to the coordinator. The
@@ -114,6 +119,8 @@ func (p JobProgress) DeriveStatus() JobStatus {
switch {
case p.Job.Status == JobCancelled:
return JobCancelled
case p.Job.Status == JobReducing:
return JobReducing
case p.Total == 0:
return JobPending
case p.Done == p.Total:
+45
View File
@@ -216,6 +216,51 @@ func (r *JobRepo) UpdateStatus(ctx context.Context, id uuid.UUID, status domain.
return nil
}
func (r *JobRepo) ClaimReduction(_ context.Context, id uuid.UUID, startedAt time.Time) (bool, error) {
r.mu.Lock()
defer r.mu.Unlock()
j, ok := r.jobs[id]
if !ok {
return false, domain.ErrJobNotFound
}
if j.Status != domain.JobReducing || j.ReducerStartedAt != nil {
return false, nil
}
j.ReducerStartedAt = &startedAt
return true, nil
}
func (r *JobRepo) CompleteWithResult(_ context.Context, id, resultArtifactID uuid.UUID, completedAt time.Time) error {
r.mu.Lock()
defer r.mu.Unlock()
j, ok := r.jobs[id]
if !ok {
return domain.ErrJobNotFound
}
j.ResultArtifactID = &resultArtifactID
j.Status = domain.JobCompleted
j.CompletedAt = &completedAt
j.ReducerStartedAt = nil
j.ErrorCode = nil
j.ErrorMessage = nil
return nil
}
func (r *JobRepo) FailReduction(_ context.Context, id uuid.UUID, code, message string, completedAt time.Time) error {
r.mu.Lock()
defer r.mu.Unlock()
j, ok := r.jobs[id]
if !ok {
return domain.ErrJobNotFound
}
j.Status = domain.JobFailed
j.CompletedAt = &completedAt
j.ReducerStartedAt = nil
j.ErrorCode = &code
j.ErrorMessage = &message
return nil
}
// --- WorkerRepo ----------------------------------------------------------
type WorkerRepo struct {
@@ -0,0 +1,195 @@
// Package reducer contains deterministic, coordinator-side result reductions.
package reducer
import (
"bytes"
"encoding/csv"
"errors"
"fmt"
"io"
"math"
"sort"
"strconv"
)
var searchHeader = []string{"rank", "chembl_id", "canonical_smiles", "similarity"}
type similarityMatch struct {
similarity float64
id string
smiles string
}
// ReduceSimilaritySearch streams worker-local top-k CSVs into the exact global
// top-k. Each partial is validated before it can affect the final artifact.
func ReduceSimilaritySearch(partials []io.Reader, parameters map[string]any) ([]byte, error) {
topK, err := positiveInt(parameters["top_k"], 20)
if err != nil {
return nil, err
}
direction, err := thresholdDirection(parameters["threshold_direction"])
if err != nil {
return nil, err
}
h := &matchHeap{direction: direction}
for _, partial := range partials {
if err := readPartial(partial, direction, func(match similarityMatch) {
if len(h.items) < topK {
heapPush(h, match)
return
}
if better(match, h.items[0], direction) {
h.items[0] = match
heapDown(h, 0)
}
}); err != nil {
return nil, err
}
}
matches := append([]similarityMatch(nil), h.items...)
sort.Slice(matches, func(i, j int) bool { return better(matches[i], matches[j], direction) })
var out bytes.Buffer
writer := csv.NewWriter(&out)
if err := writer.Write(searchHeader); err != nil {
return nil, err
}
for index, match := range matches {
if err := writer.Write([]string{
strconv.Itoa(index + 1), match.id, match.smiles, fmt.Sprintf("%.6f", match.similarity),
}); err != nil {
return nil, err
}
}
writer.Flush()
if err := writer.Error(); err != nil {
return nil, err
}
return out.Bytes(), nil
}
func readPartial(input io.Reader, direction string, consume func(similarityMatch)) error {
reader := csv.NewReader(input)
header, err := reader.Read()
if err != nil {
return fmt.Errorf("read partial header: %w", err)
}
if !equalStrings(header, searchHeader) {
return fmt.Errorf("partial result has an invalid CSV header")
}
var previous *similarityMatch
for rank := 1; ; rank++ {
row, err := reader.Read()
if errors.Is(err, io.EOF) {
return nil
}
if err != nil {
return fmt.Errorf("read partial row: %w", err)
}
if len(row) != len(searchHeader) || row[0] != strconv.Itoa(rank) {
return fmt.Errorf("partial result has an invalid rank")
}
score, err := strconv.ParseFloat(row[3], 64)
if err != nil || math.IsNaN(score) || math.IsInf(score, 0) || score < 0 || score > 1 {
return fmt.Errorf("partial result has an invalid similarity")
}
match := similarityMatch{similarity: score, id: row[1], smiles: row[2]}
if previous != nil && better(match, *previous, direction) {
return fmt.Errorf("partial result is not sorted deterministically")
}
previous = &match
consume(match)
}
}
func positiveInt(value any, fallback int) (int, error) {
if value == nil {
return fallback, nil
}
switch n := value.(type) {
case int:
if n > 0 {
return n, nil
}
case int64:
if n > 0 && n <= math.MaxInt {
return int(n), nil
}
case float64:
if n > 0 && n == math.Trunc(n) && n <= math.MaxInt {
return int(n), nil
}
}
return 0, fmt.Errorf("top_k must be a positive integer")
}
func thresholdDirection(value any) (string, error) {
if value == nil {
return "greater", nil
}
direction, ok := value.(string)
if !ok || (direction != "greater" && direction != "less") {
return "", fmt.Errorf("threshold_direction must be greater or less")
}
return direction, nil
}
func better(left, right similarityMatch, direction string) bool {
if left.similarity != right.similarity {
if direction == "less" {
return left.similarity < right.similarity
}
return left.similarity > right.similarity
}
if left.id != right.id {
return left.id < right.id
}
return left.smiles < right.smiles
}
func equalStrings(left, right []string) bool {
if len(left) != len(right) {
return false
}
for index := range left {
if left[index] != right[index] {
return false
}
}
return true
}
// matchHeap keeps the worst retained match at index zero.
type matchHeap struct {
items []similarityMatch
direction string
}
func heapPush(h *matchHeap, value similarityMatch) {
h.items = append(h.items, value)
for child := len(h.items) - 1; child > 0; {
parent := (child - 1) / 2
if !better(h.items[parent], h.items[child], h.direction) {
break
}
h.items[parent], h.items[child] = h.items[child], h.items[parent]
child = parent
}
}
func heapDown(h *matchHeap, parent int) {
for {
child := parent*2 + 1
if child >= len(h.items) {
return
}
if right := child + 1; right < len(h.items) && better(h.items[child], h.items[right], h.direction) {
child = right
}
if !better(h.items[parent], h.items[child], h.direction) {
return
}
h.items[parent], h.items[child] = h.items[child], h.items[parent]
parent = child
}
}
@@ -0,0 +1,41 @@
package reducer
import (
"io"
"strings"
"testing"
)
func TestReduceSimilaritySearchKeepsExactCrossShardRanking(t *testing.T) {
first := strings.NewReader("rank,chembl_id,canonical_smiles,similarity\n1,B,CCC,0.50000048\n2,C,CCCC,0.1\n")
second := strings.NewReader("rank,chembl_id,canonical_smiles,similarity\n1,A,CC,0.50000049\n")
output, err := ReduceSimilaritySearch([]io.Reader{first, second}, map[string]any{"top_k": 2})
if err != nil {
t.Fatal(err)
}
want := "rank,chembl_id,canonical_smiles,similarity\n1,A,CC,0.500000\n2,B,CCC,0.500000\n"
if string(output) != want {
t.Fatalf("output = %q, want %q", output, want)
}
}
func TestReduceSimilaritySearchSupportsLeastSimilarDirection(t *testing.T) {
partial := strings.NewReader("rank,chembl_id,canonical_smiles,similarity\n1,A,CC,0.1\n2,B,CCC,0.8\n")
output, err := ReduceSimilaritySearch([]io.Reader{partial}, map[string]any{
"top_k": 1, "threshold_direction": "less",
})
if err != nil {
t.Fatal(err)
}
if got, want := string(output), "rank,chembl_id,canonical_smiles,similarity\n1,A,CC,0.100000\n"; got != want {
t.Fatalf("output = %q, want %q", got, want)
}
}
func TestReduceSimilaritySearchRejectsMalformedPartial(t *testing.T) {
partial := strings.NewReader("rank,chembl_id,canonical_smiles,similarity\n2,A,CC,0.1\n")
if _, err := ReduceSimilaritySearch([]io.Reader{partial}, nil); err == nil {
t.Fatal("expected malformed rank error")
}
}
@@ -90,6 +90,49 @@ func TestCreateJobPersistsEveryTask(t *testing.T) {
}
}
func TestClaimReductionIsAtomic(t *testing.T) {
pool := testPool(t)
job, _ := seedJob(t, pool, 1)
repo := NewJobRepo(pool)
ctx := context.Background()
if err := repo.UpdateStatus(ctx, job.ID, domain.JobReducing, nil); err != nil {
t.Fatal(err)
}
var (
wg sync.WaitGroup
mu sync.Mutex
claimed int
)
for range 8 {
wg.Add(1)
go func() {
defer wg.Done()
ok, err := repo.ClaimReduction(context.Background(), job.ID, time.Now().UTC())
if err != nil {
t.Errorf("claim reduction: %v", err)
return
}
if ok {
mu.Lock()
claimed++
mu.Unlock()
}
}()
}
wg.Wait()
if claimed != 1 {
t.Fatalf("reducer claims = %d, want 1", claimed)
}
stored, err := repo.Get(ctx, job.ID)
if err != nil {
t.Fatal(err)
}
if stored.Status != domain.JobReducing || stored.ReducerStartedAt == nil {
t.Fatalf("stored reduction state = %+v", stored)
}
}
// A job must land whole or not at all: a half-created job leaves chunks no
// worker could ever complete.
func TestCreateJobRollsBackOnFailure(t *testing.T) {
@@ -25,7 +25,10 @@ func NewJobRepo(pool *pgxpool.Pool) *JobRepo {
var _ usecase.JobRepository = (*JobRepo)(nil)
var jobColumns = []string{"id", "workload", "input_uri", "parameters", "status", "created_at", "completed_at"}
var jobColumns = []string{
"id", "workload", "input_uri", "parameters", "status", "created_at", "completed_at",
"input_artifact_id", "result_artifact_id", "error_code", "error_message", "reducer_started_at",
}
// Insert runs inside the caller's transaction, alongside the job's tasks — that
// is what makes "all tasks or none" hold.
@@ -55,7 +58,8 @@ func (r *JobRepo) Get(ctx context.Context, id uuid.UUID) (*domain.Job, error) {
status string
)
err = conn(ctx, r.pool).QueryRow(ctx, sql, args...).Scan(
&j.ID, &j.Workload, &j.InputURI, &j.Parameters, &status, &j.CreatedAt, &j.CompletedAt)
&j.ID, &j.Workload, &j.InputURI, &j.Parameters, &status, &j.CreatedAt, &j.CompletedAt,
&j.InputArtifactID, &j.ResultArtifactID, &j.ErrorCode, &j.ErrorMessage, &j.ReducerStartedAt)
if errors.Is(err, pgx.ErrNoRows) {
return nil, domain.ErrJobNotFound
}
@@ -66,6 +70,64 @@ func (r *JobRepo) Get(ctx context.Context, id uuid.UUID) (*domain.Job, error) {
return &j, nil
}
func (r *JobRepo) ClaimReduction(ctx context.Context, id uuid.UUID, startedAt time.Time) (bool, error) {
sql, args, err := psql.Update("jobs").
Set("reducer_started_at", startedAt).
Where(sq.Eq{"id": id, "status": string(domain.JobReducing), "reducer_started_at": nil}).
ToSql()
if err != nil {
return false, err
}
tag, err := conn(ctx, r.pool).Exec(ctx, sql, args...)
if err != nil {
return false, err
}
return tag.RowsAffected() == 1, nil
}
func (r *JobRepo) CompleteWithResult(ctx context.Context, id, resultArtifactID uuid.UUID, completedAt time.Time) error {
sql, args, err := psql.Update("jobs").
SetMap(map[string]any{
"status": string(domain.JobCompleted),
"result_artifact_id": resultArtifactID,
"completed_at": completedAt,
"reducer_started_at": nil,
"error_code": nil,
"error_message": nil,
}).
Where(sq.Eq{"id": id, "status": string(domain.JobReducing)}).
ToSql()
if err != nil {
return err
}
tag, err := conn(ctx, r.pool).Exec(ctx, sql, args...)
if err != nil {
return err
}
if tag.RowsAffected() == 0 {
return domain.ErrJobNotFound
}
return nil
}
func (r *JobRepo) FailReduction(ctx context.Context, id uuid.UUID, code, message string, completedAt time.Time) error {
sql, args, err := psql.Update("jobs").
SetMap(map[string]any{
"status": string(domain.JobFailed),
"completed_at": completedAt,
"error_code": code,
"error_message": message,
"reducer_started_at": nil,
}).
Where(sq.Eq{"id": id, "status": string(domain.JobReducing)}).
ToSql()
if err != nil {
return err
}
_, err = conn(ctx, r.pool).Exec(ctx, sql, args...)
return err
}
func (r *JobRepo) UpdateStatus(ctx context.Context, id uuid.UUID,
status domain.JobStatus, completedAt *time.Time) error {
+10 -1
View File
@@ -119,6 +119,8 @@ type jobProgressResponse struct {
Done int `json:"completed"`
Failed int `json:"failed"`
Cancelled int `json:"cancelled"`
ResultURI string `json:"result_uri,omitempty"`
ErrorCode string `json:"error_code,omitempty"`
}
type uploadArtifactResponse struct {
@@ -153,7 +155,7 @@ func toClaimedTaskResponse(c domain.ClaimedTask) claimedTaskResponse {
}
func toJobProgressResponse(p domain.JobProgress) jobProgressResponse {
return jobProgressResponse{
out := jobProgressResponse{
ID: p.Job.ID,
Status: string(p.DeriveStatus()),
Total: p.Total,
@@ -163,4 +165,11 @@ func toJobProgressResponse(p domain.JobProgress) jobProgressResponse {
Failed: p.Failed,
Cancelled: p.Cancelled,
}
if p.Job.ResultArtifactID != nil && out.Status == string(domain.JobCompleted) {
out.ResultURI = "/jobs/" + p.Job.ID.String() + "/result"
}
if p.Job.ErrorCode != nil {
out.ErrorCode = *p.Job.ErrorCode
}
return out
}
@@ -0,0 +1,25 @@
package http
import (
"testing"
"github.com/google/uuid"
"github.com/emil28092005/SciMesh/coordinator/internal/domain"
)
func TestJobProgressResponseExposesFinalResultOnlyWhenCompleted(t *testing.T) {
id := uuid.New()
result := uuid.New()
progress := domain.JobProgress{Job: domain.Job{
ID: id, Status: domain.JobCompleted, ResultArtifactID: &result,
}, Total: 1, Done: 1}
if got, want := toJobProgressResponse(progress).ResultURI, "/jobs/"+id.String()+"/result"; got != want {
t.Fatalf("result URI = %q, want %q", got, want)
}
progress.Job.Status = domain.JobReducing
if got := toJobProgressResponse(progress).ResultURI; got != "" {
t.Fatalf("reducing job exposes result URI %q", got)
}
}
@@ -150,6 +150,12 @@ func (s *Server) handleResult(w http.ResponseWriter, r *http.Request) {
s.writeError(w, r, err)
return
}
if s.uc.ReduceJob != nil {
if err := s.uc.ReduceJob.Execute(ctx, task.JobID); err != nil {
s.writeError(w, r, err)
return
}
}
writeJSON(w, http.StatusOK, taskResponse{ID: task.ID, JobID: task.JobID, Status: string(task.Status)})
}
@@ -399,6 +405,25 @@ func (s *Server) handleGetJob(w http.ResponseWriter, r *http.Request) {
writeJSON(w, http.StatusOK, toJobProgressResponse(progress))
}
func (s *Server) handleGetJobResult(w http.ResponseWriter, r *http.Request) {
ctx, cancel := s.reqCtx(r)
defer cancel()
jobID, ok := s.pathUUID(w, r, "job_id")
if !ok {
return
}
art, body, err := s.uc.GetJobResult.Execute(ctx, jobID)
if err != nil {
s.writeError(w, r, err)
return
}
defer func() { _ = body.Close() }()
w.Header().Set("Content-Type", art.ContentType)
w.Header().Set("Content-Length", strconv.FormatInt(art.SizeBytes, 10))
w.Header().Set("Content-Disposition", fmt.Sprintf("attachment; filename=%q", art.Filename))
_, _ = io.Copy(w, body)
}
// handleCancelJob stops all non-terminal shards for an operator-requested job.
// It is available to both the bearer API and the separately authenticated UI.
func (s *Server) handleCancelJob(w http.ResponseWriter, r *http.Request) {
@@ -22,8 +22,10 @@ type UseCases struct {
ClaimTask *usecase.ClaimTask
RenewLease *usecase.RenewLease
CompleteTask *usecase.CompleteTask
ReduceJob *usecase.ReduceJob
FailTask *usecase.FailTask
GetJobStatus *usecase.GetJobStatus
GetJobResult *usecase.GetJobResult
CancelJob *usecase.CancelJob
UploadArtifact *usecase.UploadArtifact
DownloadArtifact *usecase.DownloadArtifact
@@ -62,6 +64,7 @@ func (s *Server) Handler(token string, uiToken ...string) http.Handler {
protected.HandleFunc("POST /jobs", s.handleCreateJob)
protected.HandleFunc("POST /jobs/upload", s.handleUploadDataset)
protected.HandleFunc("GET /jobs/{job_id}", s.handleGetJob)
protected.HandleFunc("GET /jobs/{job_id}/result", s.handleGetJobResult)
protected.HandleFunc("POST /jobs/{job_id}/cancel", s.handleCancelJob)
protected.HandleFunc("POST /tasks/claim", s.handleClaim)
protected.HandleFunc("GET /tasks/{task_id}/input", s.handleGetTaskInput)
@@ -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
downloadArtifact := usecase.NewDownloadArtifact(arts, blobs)
uc := coordhttp.UseCases{
RegisterWorker: usecase.NewRegisterWorker(work, clk),
@@ -50,11 +51,13 @@ func newEnvWithUIToken(t *testing.T, ready func(context.Context) error, configur
ClaimTask: usecase.NewClaimTask(tasks, jobs, work, tx, clk, lease),
RenewLease: usecase.NewRenewLease(tasks, work, tx, clk, lease),
CompleteTask: usecase.NewCompleteTask(tasks, jobs, arts, tx, clk),
ReduceJob: usecase.NewReduceJob(jobs, tasks, arts, blobs, tx, clk),
FailTask: usecase.NewFailTask(tasks, jobs, tx, clk),
GetJobStatus: usecase.NewGetJobStatus(jobs, tasks),
GetJobResult: usecase.NewGetJobResult(jobs, downloadArtifact),
CancelJob: usecase.NewCancelJob(jobs, tasks, tx, clk),
UploadArtifact: usecase.NewUploadArtifact(tasks, arts, blobs, tx, clk),
DownloadArtifact: usecase.NewDownloadArtifact(arts, blobs),
DownloadArtifact: downloadArtifact,
GetTaskInput: usecase.NewGetTaskInput(tasks, arts, blobs),
Dashboard: usecase.NewDashboard(memstore.NewUIReadRepo(jobs, tasks, work, arts)),
}
@@ -400,6 +403,42 @@ func TestFullLifecycle(t *testing.T) {
}
}
func TestSimilaritySearchLifecyclePublishesFinalResult(t *testing.T) {
e := newEnv(t, healthy)
code, job := e.uploadDataset(t, "similarity-search", 10, "chembl_id\tcanonical_smiles\nA\tCC\n")
if code != http.StatusCreated {
t.Fatalf("upload job: %d", code)
}
jobID := job["job_id"].(string)
code, claim := e.do(t, "POST", "/tasks/claim", `{"worker_id":"w1","capabilities":["similarity-search"]}`)
if code != http.StatusOK {
t.Fatalf("claim: %d", code)
}
taskID := claim["task_id"].(string)
attempt := int(claim["attempt"].(float64))
artifactID := e.putArtifact(t, taskID, "w1", attempt, "rank,chembl_id,canonical_smiles,similarity\n1,A,CC,0.900000\n")
if code, _ := e.do(t, "POST", "/tasks/"+taskID+"/result",
`{"worker_id":"w1","attempt":`+itoa(attempt)+`,"result":{"artifact_id":"`+artifactID+`"}}`); code != http.StatusOK {
t.Fatalf("complete: %d", code)
}
code, progress := e.do(t, "GET", "/jobs/"+jobID, "")
if code != http.StatusOK || progress["status"] != "completed" || progress["result_uri"] != "/jobs/"+jobID+"/result" {
t.Fatalf("progress = (%d, %v)", code, progress)
}
req, _ := http.NewRequestWithContext(context.Background(), "GET", e.ts.URL+progress["result_uri"].(string), nil)
req.Header.Set("Authorization", "Bearer "+token)
resp, err := http.DefaultClient.Do(req)
if err != nil {
t.Fatal(err)
}
defer resp.Body.Close()
body, _ := io.ReadAll(resp.Body)
if resp.StatusCode != http.StatusOK || string(body) != "rank,chembl_id,canonical_smiles,similarity\n1,A,CC,0.900000\n" {
t.Fatalf("final result = (%d, %q)", resp.StatusCode, body)
}
}
func TestForeignArtifactResultConflict(t *testing.T) {
e := newEnv(t, healthy)
e.do(t, "POST", "/jobs", `{"workload":"w","input_uri":"s3://in",
@@ -12,7 +12,7 @@
<body>
<main class="page">
<header class="top"><div><p class="eyebrow">Local coordinator</p><h1>SciMesh operator dashboard</h1><p class="lead">See where a computation is and what should happen next.</p></div><a class="button" href="/ui/jobs/new">Start a check</a></header>
<section class="notice" aria-label="Current pipeline limitation"><strong>This screen currently diagnoses shard jobs.</strong><span>Workers upload partial CSVs to the coordinator. Until a reducer is implemented, those files are not one final scientific result.</span><div class="steps"><div class="step"><b>1. Upload TSV</b>The coordinator splits the file into shard tasks.</div><div class="step"><b>2. Wait for a worker</b>A worker claims a shard, calculates similarity, and returns a CSV.</div><div class="step"><b>3. Inspect artifacts</b>Download a partial result from the job page.</div></div></section>
<section class="notice" aria-label="Pipeline flow"><strong>Similarity-search jobs produce a final CSV.</strong><span>Workers return shard-level candidates; after every shard succeeds, the coordinator deterministically merges them into one global top-k result.</span><div class="steps"><div class="step"><b>1. Upload TSV</b>The coordinator splits the file into shard tasks.</div><div class="step"><b>2. Wait for a worker</b>A worker claims a shard, calculates similarity, and returns a CSV.</div><div class="step"><b>3. Download result</b>When merging finishes, download the final CSV from the job page.</div></div></section>
<h2>Recent jobs</h2>
<div class="table-wrap"><table><tr><th>Computation</th><th>State</th><th>Progress</th><th class="hide-mobile">Created</th></tr>{{range .Jobs}}<tr><td><a href="/ui/jobs/{{.ID}}"><strong>{{workloadLabel .Workload}}</strong></a><br><small>Open job details</small></td><td><span class="status status-{{statusClass .Status}}">{{statusLabel .Status}}</span><br><small>{{statusHint .Status}}</small></td><td class="kicker"><strong>{{.Completed}} / {{.Total}}</strong> complete{{if gt .Failed 0}} · <span style="color:#a31135">failed: {{.Failed}}</span>{{end}}{{if gt .Cancelled 0}} · <span>stopped: {{.Cancelled}}</span>{{end}}<div class="bar"><span style="width:{{progressPercent .Completed .Failed .Cancelled .Total}}%"></span></div></td><td class="hide-mobile"><small>{{time .CreatedAt}}</small></td></tr>{{else}}<tr><td colspan="4" class="empty"><strong>No jobs yet.</strong><br>Click “Start a check”, upload a small TSV, and leave a worker running.</td></tr>{{end}}</table></div>
<h2>Workers</h2>
@@ -13,14 +13,14 @@
<main class="page">
<a class="back" href="/ui">← Back to jobs</a><p class="eyebrow">{{workloadLabel .Workload}}</p><h1>Execution progress</h1>
<section class="summary"><div class="summary-head"><div><span id="status" class="status status-{{statusClass .Status}}">{{statusLabel .Status}}</span><p id="hint" class="hint">{{statusHint .Status}}</p></div><div>{{if cancellable .Status}}<button id="stop-job" class="stop" type="button">Stop job</button><small>This cancels every shard that is not finished yet.</small>{{else}}<small>Summary refreshes automatically every two seconds.</small>{{end}}</div></div><div class="bar" aria-label="Progress"><span id="progress-bar" style="width:{{progressPercent .Completed .Failed .Cancelled .Total}}%"></span></div><p id="progress" class="muted">{{.Completed}} of {{.Total}} tasks complete</p><div class="numbers"><div class="number"><b id="total">{{.Total}}</b><small>total shards</small></div><div class="number"><b id="completed">{{.Completed}}</b><small>complete</small></div><div class="number"><b id="pending">{{.Pending}}</b><small>waiting</small></div><div class="number"><b id="active">{{add .Leased .Running}}</b><small>with workers</small></div><div class="number"><b id="failed">{{.Failed}}</b><small>failed</small></div><div class="number"><b id="cancelled">{{.Cancelled}}</b><small>stopped</small></div></div><details><summary>Technical details</summary><p>Job ID: <code>{{.ID}}</code><br>Workload: <code>{{.Workload}}</code><br>Created: {{time .CreatedAt}}</p></details></section>
<section class="notice"><strong>What can be downloaded now?</strong><br><code>partial_result</code> files come from individual shards. They are useful for checking the pipeline, but are not a merged final CSV because the reducer is not implemented yet.</section>
<section class="notice">{{if .FinalResultAvailable}}<strong>Final result ready</strong><br>Download the <code>final_result</code> CSV below. It is the deterministic global top-k across all completed shards. The <code>partial_result</code> files remain available for diagnostics.{{else if eq .Status "reducing"}}<strong>Merging completed shards</strong><br>Every shard has finished. The coordinator is building one deterministic global CSV; refresh in a moment to download it.{{else}}<strong>What can be downloaded now?</strong><br><code>partial_result</code> files come from individual shards. They are useful for checking the pipeline, but are not a merged final CSV yet.{{end}}</section>
<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>
</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.']};
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.'],reducing:['Merging results','active','All shards are complete. The coordinator is merging their candidates into one final CSV.'],completed:['Completed','success','The final result is ready to download.'],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.']};
const stop=document.querySelector('#stop-job');if(stop)stop.addEventListener('click',async()=>{if(!confirm('Stop this job? Unfinished shards will be cancelled.'))return;stop.disabled=true;const response=await fetch('/ui/api/jobs/'+id+'/cancel',{method:'POST'});if(!response.ok){stop.disabled=false;alert('Unable to stop this job.');return}location.reload()});
const terminal=new Set(['completed','failed','cancelled']);let timer;const poll=async()=>{try{const response=await fetch('/ui/api/jobs/'+id);if(!response.ok)return;const job=await response.json(),info=state[job.status]||[job.status,'waiting','Status reported by the coordinator.'],done=job.completed+job.failed+job.cancelled,percent=job.total?Math.min(100,Math.floor(done*100/job.total)):0,badge=document.querySelector('#status');badge.textContent=info[0];badge.className='status status-'+info[1];document.querySelector('#hint').textContent=info[2];document.querySelector('#progress').textContent=job.completed+' of '+job.total+' tasks complete'+(job.failed?' · failed: '+job.failed:'')+(job.cancelled?' · stopped: '+job.cancelled:'');document.querySelector('#progress-bar').style.width=percent+'%';for(const key of ['total','completed','pending','failed','cancelled'])document.querySelector('#'+key).textContent=job[key];document.querySelector('#active').textContent=job.leased+job.running;if(terminal.has(job.status)&&timer){clearInterval(timer);timer=undefined}}catch(_){}};const start=()=>{if(!timer&&!document.hidden&&!terminal.has(document.querySelector('#status').textContent.toLowerCase()))timer=setInterval(poll,2000)};document.addEventListener('visibilitychange',()=>{if(document.hidden&&timer){clearInterval(timer);timer=undefined}else start()});start();
</script>
+7 -3
View File
@@ -48,8 +48,10 @@ func uiStatusLabel(status string) string {
return "Assigned to a worker"
case "running":
return "Running"
case "reducing":
return "Merging results"
case "completed":
return "Tasks complete"
return "Completed"
case "failed":
return "Needs attention"
case "cancelled":
@@ -67,8 +69,10 @@ func uiStatusHint(status string) string {
return "A worker has claimed the task and should begin processing shortly."
case "running":
return "A worker is reading a shard, calculating fingerprints, and uploading its result through the coordinator."
case "reducing":
return "All shards are complete. The coordinator is merging their candidates into one final CSV."
case "completed":
return "Every shard task is complete. Files below are still partial results."
return "The final result is ready to download."
case "failed":
return "One or more shard tasks failed. Open the task list below for details."
case "cancelled":
@@ -86,7 +90,7 @@ func uiStatusClass(status string) string {
return "danger"
case "cancelled":
return "waiting"
case "running", "leased":
case "running", "leased", "reducing":
return "active"
default:
return "waiting"
@@ -10,7 +10,8 @@ func TestUIStatusPresentation(t *testing.T) {
}{
{"pending", "Waiting for a worker", "waiting"},
{"running", "Running", "active"},
{"completed", "Tasks complete", "success"},
{"reducing", "Merging results", "active"},
{"completed", "Completed", "success"},
{"failed", "Needs attention", "danger"},
}
for _, test := range tests {
+14 -4
View File
@@ -100,7 +100,7 @@ func (uc *CancelJob) Execute(ctx context.Context, jobID uuid.UUID) (int64, error
if job.Status == domain.JobCancelled {
return nil
}
if job.Status == domain.JobCompleted || job.Status == domain.JobFailed {
if job.Status == domain.JobReducing || job.Status == domain.JobCompleted || job.Status == domain.JobFailed {
return domain.ErrJobNotCancellable
}
// The lease reaper can be the transition that exhausted the final task.
@@ -111,7 +111,7 @@ func (uc *CancelJob) Execute(ctx context.Context, jobID uuid.UUID) (int64, error
return err
}
derived := progressFrom(*job, counts).DeriveStatus()
if derived == domain.JobCompleted || derived == domain.JobFailed {
if derived == domain.JobReducing || derived == domain.JobCompleted || derived == domain.JobFailed {
return domain.ErrJobNotCancellable
}
cancelled, err = uc.tasks.CancelByJob(ctx, jobID, now)
@@ -226,10 +226,20 @@ func syncJobStatus(ctx context.Context, jobs JobRepository, tasks TaskRepository
if err != nil {
return err
}
status := progressFrom(domain.Job{}, counts).DeriveStatus()
job, err := jobs.Get(ctx, jobID)
if err != nil {
return err
}
status := progressFrom(*job, counts).DeriveStatus()
// All worker shards being complete means scientific reduction is ready, not
// that the job's final artifact already exists. CTX-09 owns the transition
// from reducing to completed after it persists that artifact.
if status == domain.JobCompleted && job.Workload == "similarity-search" {
status = domain.JobReducing
}
var completedAt *time.Time
if status == domain.JobCompleted || status == domain.JobFailed {
if status == domain.JobFailed {
completedAt = &now
}
return jobs.UpdateStatus(ctx, jobID, status, completedAt)
+3
View File
@@ -70,6 +70,9 @@ type JobRepository interface {
Insert(ctx context.Context, j *domain.Job) error
Get(ctx context.Context, id uuid.UUID) (*domain.Job, error)
UpdateStatus(ctx context.Context, id uuid.UUID, status domain.JobStatus, completedAt *time.Time) error
ClaimReduction(ctx context.Context, id uuid.UUID, startedAt time.Time) (bool, error)
CompleteWithResult(ctx context.Context, id, resultArtifactID uuid.UUID, completedAt time.Time) error
FailReduction(ctx context.Context, id uuid.UUID, code, message string, completedAt time.Time) error
}
// WorkerRepository persists the worker registry.
+133
View File
@@ -0,0 +1,133 @@
package usecase
import (
"bytes"
"context"
"io"
"github.com/google/uuid"
"github.com/emil28092005/SciMesh/coordinator/internal/domain"
"github.com/emil28092005/SciMesh/coordinator/internal/reducer"
)
// ReduceJob turns completed coordinator-owned partial artifacts into one final
// artifact. It performs no worker I/O and never trusts a worker URI or path.
type ReduceJob struct {
jobs JobRepository
tasks TaskRepository
artifacts ArtifactRepository
blobs BlobStore
tx TxManager
clock Clock
}
func NewReduceJob(jobs JobRepository, tasks TaskRepository, artifacts ArtifactRepository,
blobs BlobStore, tx TxManager, clock Clock) *ReduceJob {
return &ReduceJob{jobs: jobs, tasks: tasks, artifacts: artifacts, blobs: blobs, tx: tx, clock: clock}
}
// Execute is idempotent for jobs that are not currently reducing. The worker
// completion path may call it after every result; only the last task changes a
// similarity-search job into reducing state.
func (uc *ReduceJob) Execute(ctx context.Context, jobID uuid.UUID) error {
claimed, err := uc.jobs.ClaimReduction(ctx, jobID, uc.clock.Now())
if err != nil || !claimed {
return err
}
job, err := uc.jobs.Get(ctx, jobID)
if err != nil {
return err
}
if job.Status != domain.JobReducing {
return nil
}
if job.Workload != "similarity-search" {
return uc.fail(ctx, jobID)
}
completed, err := uc.tasks.ListCompleted(ctx, jobID)
if err != nil {
return uc.fail(ctx, jobID)
}
if len(completed) == 0 {
return uc.fail(ctx, jobID)
}
readers := make([]io.Reader, 0, len(completed))
closers := make([]io.Closer, 0, len(completed))
for _, task := range completed {
if task.ResultArtifactID == nil {
closeAll(closers)
return uc.fail(ctx, jobID)
}
artifact, err := uc.artifacts.Get(ctx, *task.ResultArtifactID)
if err != nil || artifact.JobID != jobID || artifact.TaskID == nil || *artifact.TaskID != task.ID || artifact.Kind != domain.ArtifactPartialResult {
closeAll(closers)
return uc.fail(ctx, jobID)
}
body, err := uc.blobs.Open(ctx, artifact.StorageKey)
if err != nil {
closeAll(closers)
return uc.fail(ctx, jobID)
}
readers = append(readers, body)
closers = append(closers, body)
}
output, reduceErr := reducer.ReduceSimilaritySearch(readers, job.Parameters)
closeAll(closers)
if reduceErr != nil {
return uc.fail(ctx, jobID)
}
final, err := domain.NewArtifact(jobID, nil, domain.ArtifactFinalResult, "similarity-search.csv", "text/csv", uc.clock.Now())
if err != nil {
return uc.fail(ctx, jobID)
}
sum, size, err := uc.blobs.Put(ctx, final.StorageKey, bytes.NewReader(output))
if err != nil {
return uc.fail(ctx, jobID)
}
final.SetContent(sum, size)
if err := uc.tx.WithinTx(ctx, func(ctx context.Context) error {
if err := uc.artifacts.Insert(ctx, final); err != nil {
return err
}
return uc.jobs.CompleteWithResult(ctx, jobID, final.ID, uc.clock.Now())
}); err != nil {
_ = uc.blobs.Delete(ctx, final.StorageKey)
return err
}
return nil
}
func (uc *ReduceJob) fail(ctx context.Context, jobID uuid.UUID) error {
// The public state carries a stable sanitized failure, never parser/storage
// internals that may include local paths or implementation details.
return uc.jobs.FailReduction(ctx, jobID, "reducer_failed", "final result reduction failed", uc.clock.Now())
}
func closeAll(closers []io.Closer) {
for _, closer := range closers {
_ = closer.Close()
}
}
type GetJobResult struct {
jobs JobRepository
download *DownloadArtifact
}
func NewGetJobResult(jobs JobRepository, download *DownloadArtifact) *GetJobResult {
return &GetJobResult{jobs: jobs, download: download}
}
func (uc *GetJobResult) Execute(ctx context.Context, jobID uuid.UUID) (*domain.Artifact, io.ReadCloser, error) {
job, err := uc.jobs.Get(ctx, jobID)
if err != nil {
return nil, nil, err
}
if job.Status != domain.JobCompleted || job.ResultArtifactID == nil {
return nil, nil, domain.ErrArtifactNotFound
}
return uc.download.Execute(ctx, *job.ResultArtifactID)
}
@@ -55,6 +55,8 @@ type harness struct {
getInput *usecase.GetTaskInput
expire *usecase.ExpireLeases
cancel *usecase.CancelJob
reduce *usecase.ReduceJob
jobResult *usecase.GetJobResult
}
func newHarness() *harness {
@@ -81,9 +83,88 @@ func newHarness() *harness {
h.getInput = usecase.NewGetTaskInput(h.tasks, h.arts, h.blobs)
h.expire = usecase.NewExpireLeases(h.tasks, h.jobs, tx, h.clk)
h.cancel = usecase.NewCancelJob(h.jobs, h.tasks, tx, h.clk)
h.reduce = usecase.NewReduceJob(h.jobs, h.tasks, h.arts, h.blobs, tx, h.clk)
h.jobResult = usecase.NewGetJobResult(h.jobs, h.downloadArt)
return h
}
func TestSimilaritySearchReductionCreatesFinalArtifact(t *testing.T) {
h := newHarness()
jobID := h.seedJob(t, "similarity-search", 2)
_, err := h.jobs.Get(ctx, jobID)
if err != nil {
t.Fatal(err)
}
if err := h.jobs.UpdateStatus(ctx, jobID, domain.JobRunning, nil); err != nil {
t.Fatal(err)
}
partials := []string{
"rank,chembl_id,canonical_smiles,similarity\n1,B,CCC,0.50000048\n",
"rank,chembl_id,canonical_smiles,similarity\n1,A,CC,0.50000049\n",
}
for _, partial := range partials {
taskID, attempt := h.leaseOne(t, "w1", "similarity-search")
art, err := h.uploadArt.Execute(ctx, usecase.UploadArtifactInput{TaskID: taskID, WorkerID: "w1", Attempt: attempt, Filename: "partial.csv", ContentType: "text/csv", Body: strings.NewReader(partial)})
if err != nil {
t.Fatal(err)
}
if _, err := h.complete.Execute(ctx, usecase.CompleteTaskInput{TaskID: taskID, WorkerID: "w1", Attempt: attempt, ResultArtifactID: art.ID}); err != nil {
t.Fatal(err)
}
}
if err := h.reduce.Execute(ctx, jobID); err != nil {
t.Fatal(err)
}
progress, err := h.status.Execute(ctx, jobID)
if err != nil || progress.Job.Status != domain.JobCompleted {
t.Fatalf("status=%s err=%v", progress.Job.Status, err)
}
art, body, err := h.jobResult.Execute(ctx, jobID)
if err != nil {
t.Fatal(err)
}
defer body.Close()
bytes, _ := io.ReadAll(body)
if art.Kind != domain.ArtifactFinalResult || string(bytes) != "rank,chembl_id,canonical_smiles,similarity\n1,A,CC,0.500000\n2,B,CCC,0.500000\n" {
t.Fatalf("unexpected final %q", bytes)
}
}
func TestSimilaritySearchReductionFailureIsSanitized(t *testing.T) {
h := newHarness()
jobID := h.seedJob(t, "similarity-search", 1)
if err := h.jobs.UpdateStatus(ctx, jobID, domain.JobRunning, nil); err != nil {
t.Fatal(err)
}
taskID, attempt := h.leaseOne(t, "w1", "similarity-search")
art, err := h.uploadArt.Execute(ctx, usecase.UploadArtifactInput{
TaskID: taskID, WorkerID: "w1", Attempt: attempt, Filename: "partial.csv",
ContentType: "text/csv", Body: strings.NewReader("rank,chembl_id,canonical_smiles,similarity\n2,A,CC,0.9\n"),
})
if err != nil {
t.Fatal(err)
}
if _, err := h.complete.Execute(ctx, usecase.CompleteTaskInput{
TaskID: taskID, WorkerID: "w1", Attempt: attempt, ResultArtifactID: art.ID,
}); err != nil {
t.Fatal(err)
}
if err := h.reduce.Execute(ctx, jobID); err != nil {
t.Fatal(err)
}
job, err := h.jobs.Get(ctx, jobID)
if err != nil {
t.Fatal(err)
}
if job.Status != domain.JobFailed || job.ErrorCode == nil || *job.ErrorCode != "reducer_failed" ||
job.ErrorMessage == nil || *job.ErrorMessage != "final result reduction failed" {
t.Fatalf("unexpected failed job: %+v", job)
}
if job.ResultArtifactID != nil {
t.Fatal("failed reduction must not expose a final result")
}
}
// seedJob creates a URI-chunked job with n chunks and returns its id.
func (h *harness) seedJob(t *testing.T, workload string, n int) uuid.UUID {
t.Helper()
@@ -0,0 +1,7 @@
BEGIN;
ALTER TABLE jobs DROP COLUMN IF EXISTS error_message;
ALTER TABLE jobs DROP COLUMN IF EXISTS error_code;
ALTER TABLE jobs DROP COLUMN IF EXISTS reducer_started_at;
COMMIT;
@@ -0,0 +1,7 @@
-- PostgreSQL enum values must be committed before they are used by a later
-- transaction, so this migration intentionally has no BEGIN/COMMIT wrapper.
ALTER TYPE job_status ADD VALUE IF NOT EXISTS 'reducing';
ALTER TABLE jobs ADD COLUMN IF NOT EXISTS error_code text;
ALTER TABLE jobs ADD COLUMN IF NOT EXISTS error_message text;
ALTER TABLE jobs ADD COLUMN IF NOT EXISTS reducer_started_at timestamptz;
+31 -1
View File
@@ -26,7 +26,8 @@ must be updated in the same change as any behaviour it describes.
| `POST /tasks/{id}/heartbeat` | renew lease | ✅ done |
| `POST /tasks/{id}/result` | complete | ✅ done, references `artifact_id` |
| `POST /tasks/{id}/failure` | fail | ✅ done |
| `GET /jobs/{id}` | progress | ✅ done |
| `GET /jobs/{id}` | progress and final-result URI | ✅ done |
| `GET /jobs/{id}/result` | download final CSV | ✅ done |
| `PUT /tasks/{id}/artifacts/{name}` | upload partial | ✅ done |
| `GET /artifacts/{id}/download` | download by id | ✅ done |
| `POST /jobs/upload` | upload dataset, coordinator chunks it | ✅ done |
@@ -67,6 +68,35 @@ artifact. The coordinator splits the selected TSV rows into shard artifacts
Each resulting task's claim response carries `input.uri = /tasks/{id}/input`,
served by §5.4.
## Job progress and final result
```http
GET /jobs/{job_id}
Authorization: Bearer <token>
```
The response contains task counters and a derived status. A successful
`similarity-search` enters `reducing` after the last shard completes, then
becomes `completed` only after the coordinator stores its deterministic final
CSV. At that point `result_uri` is present:
```json
{
"id": "uuid",
"status": "completed",
"total": 3,
"pending": 0,
"leased": 0,
"completed": 3,
"failed": 0,
"cancelled": 0,
"result_uri": "/jobs/uuid/result"
}
```
`GET /jobs/{job_id}/result` downloads that final coordinator-owned CSV. Before
the job is completed (or for a failed/cancelled job), it returns `404`.
## Stop a job
```http
+4 -3
View File
@@ -56,9 +56,10 @@ Response: `{ "worker_id": "<uuid>", "heartbeat_interval_seconds": 15 }`.
- **Keep `worker_id`**. Use it as your identity in every later call. Using the
registered UUID is what lets the coordinator track your liveness (it marks
workers offline after they go silent).
- Current diagnostic uploads use `similarity-search` with `query_smiles`. The
reference worker accepts the legacy `similarity_search` spelling too. Do not
advertise `similarity-graph` until CTX-10 implements cross-shard pair planning.
- Current distributed uploads use `similarity-search` with `query_smiles`; the
coordinator merges completed shard candidates into a final CSV. The reference
worker accepts the legacy `similarity_search` spelling too. Do not advertise
`similarity-graph` until CTX-10 implements cross-shard pair planning.
## 2. Claim a task
+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.
+29 -4
View File
@@ -97,8 +97,9 @@ paths:
multipart/form-data. The text fields (`workload`, `parameters`,
`chunk_rows`, `max_rows`) MUST precede the `file` part: the file is streamed, not
buffered, so the fields have to be parsed before it arrives. Currently
only diagnostic `similarity-search` with `parameters.query_smiles` is
accepted; distributed graph planning is not implemented.
only `similarity-search` with `parameters.query_smiles` is accepted.
When every shard succeeds, the coordinator merges their candidates into
one final CSV; distributed graph planning is not implemented.
requestBody:
required: true
content:
@@ -124,7 +125,9 @@ paths:
- $ref: "#/components/parameters/JobID"
responses:
"200":
description: Progress counts and derived status.
description: >
Progress counts and derived status. A completed similarity-search
response includes `result_uri` for its final CSV.
content:
application/json:
schema: { $ref: "#/components/schemas/JobProgress" }
@@ -132,6 +135,21 @@ paths:
"401": { $ref: "#/components/responses/Unauthorized" }
"404": { $ref: "#/components/responses/NotFound" }
/jobs/{job_id}/result:
get:
tags: [jobs]
summary: Download a completed job's final result
parameters:
- $ref: "#/components/parameters/JobID"
responses:
"200":
description: Final coordinator-owned CSV.
content:
text/csv:
schema: { type: string, format: binary }
"401": { $ref: "#/components/responses/Unauthorized" }
"404": { $ref: "#/components/responses/NotFound" }
/jobs/{job_id}/cancel:
post:
tags: [jobs]
@@ -483,6 +501,13 @@ components:
completed: { type: integer }
failed: { type: integer }
cancelled: { type: integer }
result_uri:
type: string
description: Present only when the final result is available.
example: /jobs/5a4c3a7f-ccfc-47d6-b78d-2d1fa565bafd/result
error_code:
type: string
description: Sanitized terminal reducer failure code, when applicable.
ClaimRequest:
type: object
@@ -577,7 +602,7 @@ components:
JobStatus:
type: string
enum: [pending, running, completed, failed, cancelled]
enum: [pending, running, reducing, completed, failed, cancelled]
TaskStatus:
type: string
+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: