Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
a055473706 | ||
|
|
6e67daa9eb | ||
|
|
0f3a2d92d8 | ||
|
|
0bef7604fd |
@@ -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`.
|
||||
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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),
|
||||
}
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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 {
|
||||
|
||||
|
||||
@@ -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>
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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",
|
||||
]
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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()
|
||||
@@ -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>")
|
||||
|
||||
@@ -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)]
|
||||
|
||||
@@ -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
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user