Merge distributed pipeline hardening

This commit is contained in:
Emil
2026-07-24 14:16:42 +03:00
38 changed files with 1417 additions and 181 deletions
+28
View File
@@ -0,0 +1,28 @@
name: python
on:
push:
paths:
- "scimesh/**"
- "tests/**"
- "pyproject.toml"
- ".github/workflows/python.yml"
pull_request:
paths:
- "scimesh/**"
- "tests/**"
- "pyproject.toml"
- ".github/workflows/python.yml"
jobs:
test:
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v4
- uses: actions/setup-python@v5
with:
python-version: "3.11"
cache: pip
- run: python -m pip install --upgrade pip
- run: python -m pip install -e '.[dev]'
- run: pytest -q
+4
View File
@@ -11,3 +11,7 @@ results/
*_similarities.csv
test_results.csv
test_structures/
# Local coordinator-worker execution state
worker-data*/
scimesh-worker-data/
+6 -5
View File
@@ -1,11 +1,12 @@
# SciMesh
SciMesh is a scientific-workload framework for molecular datasets. Its public CLI
currently runs exact similarity search and sparse similarity-graph construction
locally in one Python process; it creates no dense similarity matrix. A Python
Worker client and the planned Go/PostgreSQL coordinator contract are tracked in
the repository, but distributed execution is not available yet; see
[`STATUS.md`](STATUS.md).
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).
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`.
+4 -1
View File
@@ -36,7 +36,7 @@ Docker PostgreSQL stack on 2026-07-23.
| 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-10 Distributed similarity-graph | Not started | Local reference exists. |
| CTX-11 Dashboard/operator view | In progress | `feat/web-interface` adds a protected local view: job/task/worker status, dataset upload, diagnostic partial-artifact download, and polling. Final-result reduction remains CTX-09. |
| 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-12 Reliability, security, CI | In progress | Unit, race, PostgreSQL integration, and smoke checks exist; CI hardening remains. |
## Next recommended assignment
@@ -53,6 +53,9 @@ reduction boundaries before implementing distributed search or graph execution.
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.
- 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.
## Update rule
+4 -4
View File
@@ -71,14 +71,14 @@ func run() error {
useCases := httptransport.UseCases{
RegisterWorker: usecase.NewRegisterWorker(workerRepo, clk),
CreateJob: usecase.NewCreateJob(jobRepo, taskRepo, tx, clk),
SubmitDataset: usecase.NewSubmitDataset(blobStore, artifactRepo, jobRepo, taskRepo, tx, clk),
ClaimTask: usecase.NewClaimTask(taskRepo, clk, cfg.LeaseDuration),
SubmitDataset: usecase.NewSubmitDataset(blobStore, artifactRepo, jobRepo, taskRepo, tx, clk, cfg.DefaultMaxAttempts),
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),
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, clk),
UploadArtifact: usecase.NewUploadArtifact(taskRepo, artifactRepo, blobStore, tx, clk),
DownloadArtifact: usecase.NewDownloadArtifact(artifactRepo, blobStore),
GetTaskInput: usecase.NewGetTaskInput(taskRepo, artifactRepo, blobStore),
Dashboard: usecase.NewDashboard(uiReadRepo),
@@ -87,7 +87,7 @@ func run() error {
// Background reapers are tracked so shutdown can wait for them. Without this
// the process would exit mid-UPDATE, and the deferred pool.Close() would pull
// connections out from under them.
expireLeases := usecase.NewExpireLeases(taskRepo, clk)
expireLeases := usecase.NewExpireLeases(taskRepo, jobRepo, tx, clk)
markOffline := usecase.NewMarkWorkersOffline(workerRepo, clk, cfg.WorkerOfflineAfter)
var wg sync.WaitGroup
+45 -1
View File
@@ -7,12 +7,19 @@ import (
"bytes"
"fmt"
"io"
"strings"
)
// ErrNoRows is returned when the input has a header but no data rows: a job with
// zero tasks could never complete, so it is rejected at the source.
var ErrNoRows = fmt.Errorf("input has no data rows")
// maxShardBytes bounds the coordinator memory used by one in-progress shard.
// The uploaded file may be much larger: it is first stored on disk, then split
// in small bounded pieces. Operators can lower rowsPerShard when this limit is
// reached rather than exhausting the coordinator process.
const maxShardBytes = 64 << 20 // 64 MiB
// SplitTSV reads a header-plus-rows text stream and cuts it into shards of at
// most rowsPerShard data rows. Every shard repeats the header, so a worker can
// parse its shard in isolation. emit is called once per shard, in order, with a
@@ -26,13 +33,25 @@ var ErrNoRows = fmt.Errorf("input has no data rows")
// Only one shard is buffered at a time, so memory is bounded by shard size (a
// worker-sized slice of the data), not by the size of the whole dataset.
func SplitTSV(r io.Reader, rowsPerShard int, emit func(index int, shard io.Reader) error) error {
return SplitTSVLimit(r, rowsPerShard, 0, emit)
return splitTSVLimit(r, rowsPerShard, 0, nil, emit)
}
// SplitTSVLimit behaves like SplitTSV but emits no more than maxRows data rows.
// A maxRows value of zero means unlimited. This lets an operator make a small,
// representative pipeline check without materialising a second dataset file.
func SplitTSVLimit(r io.Reader, rowsPerShard, maxRows int, emit func(index int, shard io.Reader) error) error {
return splitTSVLimit(r, rowsPerShard, maxRows, nil, emit)
}
// SplitChEMBLTSVLimit is the coordinator's scientific-upload splitter. It
// validates the two columns every local SciMesh workload requires before any
// shard task is persisted, while generic SplitTSV remains reusable for future
// non-chemistry workloads.
func SplitChEMBLTSVLimit(r io.Reader, rowsPerShard, maxRows int, emit func(index int, shard io.Reader) error) error {
return splitTSVLimit(r, rowsPerShard, maxRows, validateChEMBLHeader, emit)
}
func splitTSVLimit(r io.Reader, rowsPerShard, maxRows int, validateHeader func([]byte) error, emit func(index int, shard io.Reader) error) error {
if rowsPerShard <= 0 {
return fmt.Errorf("rowsPerShard must be positive, got %d", rowsPerShard)
}
@@ -51,6 +70,11 @@ func SplitTSVLimit(r io.Reader, rowsPerShard, maxRows int, emit func(index int,
return ErrNoRows // completely empty input
}
header := append([]byte(nil), sc.Bytes()...)
if validateHeader != nil {
if err := validateHeader(header); err != nil {
return err
}
}
var (
buf bytes.Buffer
@@ -71,9 +95,15 @@ func SplitTSVLimit(r io.Reader, rowsPerShard, maxRows int, emit func(index int,
for sc.Scan() {
if rows == 0 {
if len(header)+1 > maxShardBytes {
return fmt.Errorf("TSV header exceeds maximum shard size of %d bytes", maxShardBytes)
}
buf.Write(header)
buf.WriteByte('\n')
}
if buf.Len()+len(sc.Bytes())+1 > maxShardBytes {
return fmt.Errorf("shard exceeds maximum size of %d bytes; lower rowsPerShard", maxShardBytes)
}
buf.Write(sc.Bytes())
buf.WriteByte('\n')
rows++
@@ -103,3 +133,17 @@ func SplitTSVLimit(r io.Reader, rowsPerShard, maxRows int, emit func(index int,
}
return nil
}
func validateChEMBLHeader(header []byte) error {
seen := make(map[string]struct{})
for _, field := range strings.Split(strings.TrimPrefix(string(header), "\ufeff"), "\t") {
seen[field] = struct{}{}
}
if _, ok := seen["chembl_id"]; !ok {
return fmt.Errorf("TSV is missing required column chembl_id")
}
if _, ok := seen["canonical_smiles"]; !ok {
return fmt.Errorf("TSV is missing required column canonical_smiles")
}
return nil
}
+8
View File
@@ -119,6 +119,14 @@ func TestSplitLimitUsesOnlyLeadingDataRows(t *testing.T) {
}
}
func TestChEMBLSplitRejectsMissingRequiredColumns(t *testing.T) {
err := SplitChEMBLTSVLimit(strings.NewReader("id\tsmiles\nA\tCC\n"), 1, 0,
func(int, io.Reader) error { return nil })
if err == nil || !strings.Contains(err.Error(), "chembl_id") {
t.Errorf("err = %v, want missing-column error", err)
}
}
// The scanned bytes are reused by bufio; the shard buffer must copy them, or a
// later row would corrupt an earlier one. This guards that copy.
func TestSplitDoesNotAliasScannerBuffer(t *testing.T) {
+3
View File
@@ -129,6 +129,9 @@ func LoadConfig() (Config, error) {
if cfg.DefaultMaxAttempts, err = getEnvInt("DEFAULT_MAX_ATTEMPTS", cfg.DefaultMaxAttempts); err != nil {
return Config{}, err
}
if cfg.DefaultMaxAttempts < 1 {
return Config{}, fmt.Errorf("DEFAULT_MAX_ATTEMPTS must be positive")
}
return cfg, nil
}
+11
View File
@@ -32,3 +32,14 @@ func TestLoadConfigAllowsDistinctUIAndWorkerTokens(t *testing.T) {
t.Fatalf("unexpected tokens: %+v", cfg)
}
}
func TestLoadConfigRejectsNonPositiveDefaultMaxAttempts(t *testing.T) {
t.Setenv("ENV_FILE", filepath.Join(t.TempDir(), "missing.env"))
t.Setenv("DATABASE_URL", "postgres://test")
t.Setenv("DEFAULT_MAX_ATTEMPTS", "0")
_, err := LoadConfig()
if err == nil || !strings.Contains(err.Error(), "DEFAULT_MAX_ATTEMPTS") {
t.Fatalf("LoadConfig error = %v, want default-attempt validation", err)
}
}
+23 -5
View File
@@ -160,17 +160,18 @@ func (r *TaskRepo) CancelByJob(_ context.Context, jobID uuid.UUID, now time.Time
return cancelled, nil
}
func (r *TaskRepo) ExpireLeases(ctx context.Context, now time.Time) (int64, error) {
func (r *TaskRepo) ExpireLeases(ctx context.Context, now time.Time) ([]uuid.UUID, error) {
r.mu.Lock()
defer r.mu.Unlock()
var n int64
affected := make([]uuid.UUID, 0)
for _, t := range r.tasks {
if t.Status == domain.TaskLeased && t.LeaseExpiresAt != nil && t.LeaseExpiresAt.Before(now) {
if (t.Status == domain.TaskLeased || t.Status == domain.TaskRunning) &&
t.LeaseExpiresAt != nil && t.LeaseExpiresAt.Before(now) {
t.ExpireLease(now)
n++
affected = append(affected, t.JobID)
}
}
return n, nil
return affected, nil
}
// --- JobRepo -------------------------------------------------------------
@@ -298,6 +299,23 @@ func (r *ArtifactRepo) Get(ctx context.Context, id uuid.UUID) (*domain.Artifact,
return &cp, nil
}
func (r *ArtifactRepo) FindPartialResult(_ context.Context, taskID uuid.UUID, attempt int) (*domain.Artifact, error) {
r.mu.Lock()
defer r.mu.Unlock()
for _, a := range r.arts {
if a.TaskID != nil && *a.TaskID == taskID && a.Kind == domain.ArtifactPartialResult &&
a.Attempt != nil && *a.Attempt == attempt {
return cloneArtifact(a), nil
}
}
return nil, nil
}
func cloneArtifact(a *domain.Artifact) *domain.Artifact {
cp := *a
return &cp
}
// --- BlobStore -----------------------------------------------------------
type BlobStore struct {
+12
View File
@@ -60,6 +60,18 @@ func (r *UIReadRepo) ListTasksByJob(_ context.Context, jobID uuid.UUID) ([]domai
sort.Slice(out, func(i, j int) bool { return out[i].ChunkIndex < out[j].ChunkIndex })
return out, nil
}
func (r *UIReadRepo) ListTasksByJobs(ctx context.Context, jobIDs []uuid.UUID) (map[uuid.UUID][]domain.Task, error) {
out := make(map[uuid.UUID][]domain.Task, len(jobIDs))
for _, id := range jobIDs {
tasks, err := r.ListTasksByJob(ctx, id)
if err != nil {
return nil, err
}
out[id] = tasks
}
return out, nil
}
func (r *UIReadRepo) ListWorkers(_ context.Context, limit int) ([]domain.Worker, error) {
if limit < 1 || limit > 100 {
return nil, domain.ErrInvalidInput
@@ -70,3 +70,33 @@ func (r *ArtifactRepo) Get(ctx context.Context, id uuid.UUID) (*domain.Artifact,
a.Kind = domain.ArtifactKind(kind)
return &a, nil
}
func (r *ArtifactRepo) FindPartialResult(ctx context.Context, taskID uuid.UUID, attempt int) (*domain.Artifact, error) {
sql, args, err := psql.Select(artifactColumns...).
From("artifacts").
Where(sq.Eq{
"task_id": taskID,
"attempt": attempt,
"kind": string(domain.ArtifactPartialResult),
}).
ToSql()
if err != nil {
return nil, err
}
var (
a domain.Artifact
kind string
)
err = conn(ctx, r.pool).QueryRow(ctx, sql, args...).Scan(
&a.ID, &a.JobID, &a.TaskID, &a.Attempt, &kind, &a.Filename, &a.StorageKey,
&a.ContentType, &a.SizeBytes, &a.SHA256, &a.CreatedAt)
if errors.Is(err, pgx.ErrNoRows) {
return nil, nil
}
if err != nil {
return nil, fmt.Errorf("find partial result: %w", err)
}
a.Kind = domain.ArtifactKind(kind)
return &a, nil
}
@@ -360,12 +360,23 @@ func TestCompleteTaskReplayIsIdempotent(t *testing.T) {
t.Errorf("replay must be idempotent, got %v", err)
}
// A different result artifact for the same task is a genuine conflict.
art2 := seedArtifact(t, pool, job.ID, &claimed.ID, domain.ArtifactPartialResult)
other := in
other.ResultArtifactID = art2.ID
if _, err := uc.Execute(ctx, other); !errors.Is(err, domain.ErrResultConflict) {
t.Errorf("err = %v, want ErrResultConflict", err)
}
func TestPartialResultIsUniquePerTaskAttempt(t *testing.T) {
pool := testPool(t)
ctx := context.Background()
job, tasks := seedJob(t, pool, 1)
taskID := tasks[0].ID
first := seedArtifact(t, pool, job.ID, &taskID, domain.ArtifactPartialResult)
second, err := domain.NewArtifact(job.ID, &taskID, domain.ArtifactPartialResult, "retry.csv", "text/csv", time.Now().UTC())
if err != nil {
t.Fatal(err)
}
attempt := 1
second.Attempt = &attempt
second.SetContent("other-sha", 5)
if err := NewArtifactRepo(pool).Insert(ctx, second); err == nil {
t.Fatalf("second partial artifact for %s/%d was accepted after %s", taskID, attempt, first.ID)
}
}
@@ -337,17 +337,26 @@ SET status = CASE WHEN attempt < max_attempts THEN 'pending'::task_sta
ELSE error_message END,
completed_at = CASE WHEN attempt >= max_attempts THEN $1 ELSE completed_at END,
version = version + 1
WHERE status IN ('leased','running') AND lease_expires_at < $1`
WHERE status IN ('leased','running') AND lease_expires_at < $1
RETURNING job_id`
func (r *TaskRepo) ExpireLeases(ctx context.Context, now time.Time) (int64, error) {
var affected int64
func (r *TaskRepo) ExpireLeases(ctx context.Context, now time.Time) ([]uuid.UUID, error) {
var affected []uuid.UUID
err := withRetry(ctx, func(ctx context.Context) error {
tag, err := conn(ctx, r.pool).Exec(ctx, expireLeasesSQL, now, domain.ErrCodeLeaseExpired)
rows, err := conn(ctx, r.pool).Query(ctx, expireLeasesSQL, now, domain.ErrCodeLeaseExpired)
if err != nil {
return err
}
affected = tag.RowsAffected()
return nil
defer rows.Close()
affected = affected[:0]
for rows.Next() {
var jobID uuid.UUID
if err := rows.Scan(&jobID); err != nil {
return err
}
affected = append(affected, jobID)
}
return rows.Err()
})
return affected, err
}
@@ -75,6 +75,31 @@ func (r *UIReadRepo) ListTasksByJob(ctx context.Context, jobID uuid.UUID) ([]dom
return tasks, rows.Err()
}
func (r *UIReadRepo) ListTasksByJobs(ctx context.Context, jobIDs []uuid.UUID) (map[uuid.UUID][]domain.Task, error) {
out := make(map[uuid.UUID][]domain.Task, len(jobIDs))
if len(jobIDs) == 0 {
return out, nil
}
sql, args, err := psql.Select(taskColumns...).From("tasks").
Where(sq.Eq{"job_id": jobIDs}).OrderBy("job_id ASC", "chunk_index ASC").ToSql()
if err != nil {
return nil, err
}
rows, err := conn(ctx, r.pool).Query(ctx, sql, args...)
if err != nil {
return nil, fmt.Errorf("list tasks by jobs: %w", err)
}
defer rows.Close()
for rows.Next() {
task, err := scanTask(rows)
if err != nil {
return nil, err
}
out[task.JobID] = append(out[task.JobID], *task)
}
return out, rows.Err()
}
func (r *UIReadRepo) ListWorkers(ctx context.Context, limit int) ([]domain.Worker, error) {
if limit < 1 || limit > 100 {
return nil, domain.ErrInvalidInput
@@ -79,6 +79,10 @@ func (s *Server) handleClaim(w http.ResponseWriter, r *http.Request) {
s.writeError(w, r, domain.ErrInvalidInput)
return
}
if _, err := uuid.Parse(req.WorkerID); err != nil {
s.writeError(w, r, domain.ErrInvalidInput)
return
}
claimed, err := s.uc.ClaimTask.Execute(ctx, usecase.ClaimTaskInput{
WorkerID: req.WorkerID,
@@ -23,8 +23,9 @@ const token = "secret"
const uiToken = "ui-secret"
type env struct {
ts *httptest.Server
blobs *memstore.BlobStore
ts *httptest.Server
blobs *memstore.BlobStore
workerID string
}
func newEnv(t *testing.T, ready func(context.Context) error) *env {
@@ -45,22 +46,28 @@ func newEnvWithUIToken(t *testing.T, ready func(context.Context) error, configur
uc := coordhttp.UseCases{
RegisterWorker: usecase.NewRegisterWorker(work, clk),
CreateJob: usecase.NewCreateJob(jobs, tasks, tx, clk),
SubmitDataset: usecase.NewSubmitDataset(blobs, arts, jobs, tasks, tx, clk),
ClaimTask: usecase.NewClaimTask(tasks, clk, lease),
SubmitDataset: usecase.NewSubmitDataset(blobs, arts, jobs, tasks, tx, clk, 3),
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),
FailTask: usecase.NewFailTask(tasks, jobs, tx, clk),
GetJobStatus: usecase.NewGetJobStatus(jobs, tasks),
CancelJob: usecase.NewCancelJob(jobs, tasks, tx, clk),
UploadArtifact: usecase.NewUploadArtifact(tasks, arts, blobs, clk),
UploadArtifact: usecase.NewUploadArtifact(tasks, arts, blobs, tx, clk),
DownloadArtifact: usecase.NewDownloadArtifact(arts, blobs),
GetTaskInput: usecase.NewGetTaskInput(tasks, arts, blobs),
Dashboard: usecase.NewDashboard(memstore.NewUIReadRepo(jobs, tasks, work, arts)),
}
worker, err := uc.RegisterWorker.Execute(context.Background(), usecase.RegisterWorkerInput{
Name: "test-worker", Capabilities: []string{"w", "similarity-search"},
})
if err != nil {
t.Fatalf("register test worker: %v", err)
}
srv := coordhttp.NewServer(uc, slog.New(slog.NewTextHandler(io.Discard, nil)), 5*time.Second, 15*time.Second, 1<<30, ready)
ts := httptest.NewServer(srv.Handler(token, configuredUIToken))
t.Cleanup(ts.Close)
return &env{ts: ts, blobs: blobs}
return &env{ts: ts, blobs: blobs, workerID: worker.ID.String()}
}
func healthy(context.Context) error { return nil }
@@ -68,6 +75,7 @@ func healthy(context.Context) error { return nil }
// do sends an authenticated JSON request and returns status + decoded body.
func (e *env) do(t *testing.T, method, path, body string) (int, map[string]any) {
t.Helper()
body = strings.ReplaceAll(body, `"worker_id":"w1"`, `"worker_id":"`+e.workerID+`"`)
req, _ := http.NewRequestWithContext(context.Background(), method, e.ts.URL+path, strings.NewReader(body))
req.Header.Set("Authorization", "Bearer "+token)
if body != "" {
@@ -330,6 +338,26 @@ func TestRegisterRejectsNoCapabilities(t *testing.T) {
}
}
func TestClaimRequiresRegisteredWorkerAndUsesStoredCapabilities(t *testing.T) {
e := newEnv(t, healthy)
if code, _ := e.do(t, "POST", "/tasks/claim", `{"worker_id":"not-a-uuid"}`); code != http.StatusBadRequest {
t.Fatalf("invalid worker id claim = %d, want 400", code)
}
if code, _ := e.do(t, "POST", "/tasks/claim", `{"worker_id":"11111111-1111-4111-8111-111111111111"}`); code != http.StatusNotFound {
t.Fatalf("unregistered worker claim = %d, want 404", code)
}
if code, _ := e.do(t, "POST", "/jobs", `{"workload":"w","input_uri":"s3://in","chunks":[{"chunk_index":0,"input_uri":"s3://c","input_sha256":"sha"}]}`); code != http.StatusCreated {
t.Fatalf("create job = %d", code)
}
code, worker := e.do(t, "POST", "/workers/register", `{"name":"search-only","capabilities":["similarity-search"]}`)
if code != http.StatusCreated {
t.Fatalf("register = %d", code)
}
if code, _ := e.do(t, "POST", "/tasks/claim", `{"worker_id":"`+worker["worker_id"].(string)+`","capabilities":["w"]}`); code != http.StatusNoContent {
t.Fatalf("forged capability claim = %d, want 204", code)
}
}
func TestFullLifecycle(t *testing.T) {
e := newEnv(t, healthy)
@@ -393,9 +421,9 @@ func TestForeignArtifactResultConflict(t *testing.T) {
func TestUploadDatasetChunksAndServesInput(t *testing.T) {
e := newEnv(t, healthy)
tsv := "id\tsmiles\nA\tCC\nB\tCCC\nC\tCCCC\nD\tCCCCC\nE\tCCCCCC\n"
tsv := "chembl_id\tcanonical_smiles\nA\tCC\nB\tCCC\nC\tCCCC\nD\tCCCCC\nE\tCCCCCC\n"
code, body := e.uploadDataset(t, "w", 2, tsv)
code, body := e.uploadDataset(t, "similarity-search", 2, tsv)
if code != 201 {
t.Fatalf("upload: status = %d", code)
}
@@ -421,7 +449,7 @@ func TestUploadDatasetChunksAndServesInput(t *testing.T) {
t.Fatalf("get input: status = %d", resp.StatusCode)
}
shard, _ := io.ReadAll(resp.Body)
if !strings.HasPrefix(string(shard), "id\tsmiles\n") {
if !strings.HasPrefix(string(shard), "chembl_id\tcanonical_smiles\n") {
t.Errorf("shard missing header: %q", shard)
}
}
@@ -430,11 +458,12 @@ func TestUploadDatasetLimitsRows(t *testing.T) {
e := newEnv(t, healthy)
var buf bytes.Buffer
mw := multipart.NewWriter(&buf)
_ = mw.WriteField("workload", "w")
_ = mw.WriteField("workload", "similarity-search")
_ = mw.WriteField("parameters", `{"query_smiles":"CCO"}`)
_ = mw.WriteField("chunk_rows", "2")
_ = mw.WriteField("max_rows", "3")
fw, _ := mw.CreateFormFile("file", "chembl.tsv")
_, _ = io.Copy(fw, strings.NewReader("id\tsmiles\nA\tCC\nB\tCCC\nC\tCCCC\nD\tCCCCC\n"))
_, _ = io.Copy(fw, strings.NewReader("chembl_id\tcanonical_smiles\nA\tCC\nB\tCCC\nC\tCCCC\nD\tCCCCC\n"))
_ = mw.Close()
req, _ := http.NewRequestWithContext(context.Background(), "POST", e.ts.URL+"/jobs/upload", &buf)
req.Header.Set("Authorization", "Bearer "+token)
@@ -451,6 +480,29 @@ func TestUploadDatasetLimitsRows(t *testing.T) {
}
}
func TestUploadDatasetRejectsMissingChEMBLColumns(t *testing.T) {
e := newEnv(t, healthy)
var buf bytes.Buffer
mw := multipart.NewWriter(&buf)
_ = mw.WriteField("workload", "similarity-search")
_ = mw.WriteField("parameters", `{"query_smiles":"CCO"}`)
_ = mw.WriteField("chunk_rows", "2")
fw, _ := mw.CreateFormFile("file", "not-chembl.tsv")
_, _ = io.Copy(fw, strings.NewReader("id\tsmiles\nA\tCC\n"))
_ = mw.Close()
req, _ := http.NewRequestWithContext(context.Background(), "POST", e.ts.URL+"/jobs/upload", &buf)
req.Header.Set("Authorization", "Bearer "+token)
req.Header.Set("Content-Type", mw.FormDataContentType())
resp, err := http.DefaultClient.Do(req)
if err != nil {
t.Fatal(err)
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusBadRequest {
t.Errorf("missing ChEMBL columns = %d, want 400", resp.StatusCode)
}
}
func TestErrorMappings(t *testing.T) {
e := newEnv(t, healthy)
zero := "00000000-0000-0000-0000-000000000000"
@@ -478,10 +530,11 @@ func TestUploadDatasetRejectsAmbiguousMultipartInput(t *testing.T) {
e := newEnv(t, healthy)
var buf bytes.Buffer
mw := multipart.NewWriter(&buf)
_ = mw.WriteField("workload", "w")
_ = mw.WriteField("workload", "similarity-search")
_ = mw.WriteField("parameters", `{"query_smiles":"CCO"}`)
_ = mw.WriteField("chunk_rows", "not-a-number")
fw, _ := mw.CreateFormFile("file", "chembl.tsv")
_, _ = io.Copy(fw, strings.NewReader("id\tsmiles\nA\tCC\n"))
_, _ = io.Copy(fw, strings.NewReader("chembl_id\tcanonical_smiles\nA\tCC\n"))
_ = mw.Close()
req, _ := http.NewRequestWithContext(context.Background(), "POST", e.ts.URL+"/jobs/upload", &buf)
@@ -505,6 +558,9 @@ func (e *env) putArtifact(t *testing.T, taskID, worker string, attempt int, data
e.ts.URL+"/tasks/"+taskID+"/artifacts/r.csv", strings.NewReader(data))
req.Header.Set("Authorization", "Bearer "+token)
req.Header.Set("Content-Type", "text/csv")
if worker == "w1" {
worker = e.workerID
}
req.Header.Set("X-Worker-ID", worker)
req.Header.Set("X-Task-Attempt", itoa(attempt))
resp, err := http.DefaultClient.Do(req)
@@ -526,6 +582,7 @@ func (e *env) uploadDataset(t *testing.T, workload string, rows int, tsv string)
var buf bytes.Buffer
mw := multipart.NewWriter(&buf)
_ = mw.WriteField("workload", workload)
_ = mw.WriteField("parameters", `{"query_smiles":"CCO"}`)
_ = mw.WriteField("chunk_rows", itoa(rows))
fw, _ := mw.CreateFormFile("file", "chembl.tsv")
_, _ = io.Copy(fw, strings.NewReader(tsv))
@@ -22,7 +22,7 @@
<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 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()});
setInterval(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}catch(_){}} ,2000);
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>
</body>
</html>
+41 -15
View File
@@ -14,12 +14,13 @@ type UploadArtifact struct {
tasks TaskRepository
artifacts ArtifactRepository
blobs BlobStore
tx TxManager
clk Clock
}
func NewUploadArtifact(tasks TaskRepository, artifacts ArtifactRepository,
blobs BlobStore, clk Clock) *UploadArtifact {
return &UploadArtifact{tasks: tasks, artifacts: artifacts, blobs: blobs, clk: clk}
blobs BlobStore, tx TxManager, clk Clock) *UploadArtifact {
return &UploadArtifact{tasks: tasks, artifacts: artifacts, blobs: blobs, tx: tx, clk: clk}
}
func (uc *UploadArtifact) Execute(ctx context.Context, in UploadArtifactInput) (*domain.Artifact, error) {
@@ -32,6 +33,15 @@ func (uc *UploadArtifact) Execute(ctx context.Context, in UploadArtifactInput) (
if !task.IsLeaseHeldBy(in.WorkerID, in.Attempt, uc.clk.Now()) {
return nil, domain.ErrLeaseConflict
}
// A client can retry a PUT after losing the response. Return the one durable
// result for this lease attempt instead of storing duplicate artifacts.
existing, err := uc.artifacts.FindPartialResult(ctx, in.TaskID, in.Attempt)
if err != nil {
return nil, err
}
if existing != nil {
return existing, nil
}
taskID := task.ID
art, err := domain.NewArtifact(task.JobID, &taskID, domain.ArtifactPartialResult,
@@ -50,25 +60,41 @@ func (uc *UploadArtifact) Execute(ctx context.Context, in UploadArtifactInput) (
}
art.SetContent(sum, size)
// The stream may take longer than the lease. Re-check after it finishes so
// an expired worker cannot leave a durable result record behind. Completion
// performs the same ownership check under its transaction.
current, err := uc.tasks.Get(ctx, in.TaskID)
// The stream may take longer than the lease. Lock the task while re-checking
// ownership and inserting metadata: completion or another upload cannot race
// this final decision. The database unique index is a second line of defence.
var durable *domain.Artifact
err = uc.tx.WithinTx(ctx, func(ctx context.Context) error {
current, err := uc.tasks.GetForUpdate(ctx, in.TaskID)
if err != nil {
return err
}
if !current.IsLeaseHeldBy(in.WorkerID, in.Attempt, uc.clk.Now()) {
return domain.ErrLeaseConflict
}
existing, err := uc.artifacts.FindPartialResult(ctx, in.TaskID, in.Attempt)
if err != nil {
return err
}
if existing != nil {
durable = existing
return nil
}
if err := uc.artifacts.Insert(ctx, art); err != nil {
return err
}
durable = art
return nil
})
if err != nil {
_ = uc.blobs.Delete(ctx, art.StorageKey)
return nil, err
}
if !current.IsLeaseHeldBy(in.WorkerID, in.Attempt, uc.clk.Now()) {
if durable != art {
// Another request won the race while this stream was being written.
_ = uc.blobs.Delete(ctx, art.StorageKey)
return nil, domain.ErrLeaseConflict
}
// Persist the record. If that fails the blob would be an orphan, so remove it.
if err := uc.artifacts.Insert(ctx, art); err != nil {
_ = uc.blobs.Delete(ctx, art.StorageKey)
return nil, err
}
return art, nil
return durable, nil
}
// DownloadArtifact returns an artifact's metadata together with a reader over
+37
View File
@@ -33,6 +33,17 @@ func NewCreateJob(jobs JobRepository, tasks TaskRepository, tx TxManager, clock
// The all-or-none guarantee comes from TxManager: a half-created job would
// leave chunks no worker could ever complete.
func (uc *CreateJob) Execute(ctx context.Context, in CreateJobInput) (*domain.Job, error) {
if in.Workload == "similarity-graph" || in.Workload == "similarity_graph" {
// CTX-10 must plan triangular block pairs; ordinary independent input
// chunks would silently omit every cross-chunk molecular pair.
return nil, domain.ErrInvalidInput
}
if (in.Workload == "similarity-search" || in.Workload == "similarity_search") &&
len(in.Chunks) > 1 && in.Parameters["query_id"] != nil {
// Resolving once against the source dataset belongs to CTX-07. Letting
// each shard resolve it would make most tasks fail or use inconsistent data.
return nil, domain.ErrInvalidInput
}
chunks := make([]domain.ChunkSpec, 0, len(in.Chunks))
for _, c := range in.Chunks {
chunks = append(chunks, domain.ChunkSpec(c))
@@ -92,6 +103,17 @@ func (uc *CancelJob) Execute(ctx context.Context, jobID uuid.UUID) (int64, error
if job.Status == domain.JobCompleted || job.Status == domain.JobFailed {
return domain.ErrJobNotCancellable
}
// The lease reaper can be the transition that exhausted the final task.
// Check the authoritative task histogram as well as the cached job status,
// so a stale status can never turn a failed/completed job into cancelled.
counts, err := uc.tasks.CountByStatus(ctx, jobID)
if err != nil {
return err
}
derived := progressFrom(*job, counts).DeriveStatus()
if derived == domain.JobCompleted || derived == domain.JobFailed {
return domain.ErrJobNotCancellable
}
cancelled, err = uc.tasks.CancelByJob(ctx, jobID, now)
if err != nil {
return err
@@ -212,3 +234,18 @@ func syncJobStatus(ctx context.Context, jobs JobRepository, tasks TaskRepository
}
return jobs.UpdateStatus(ctx, jobID, status, completedAt)
}
func syncExpiredJobStatuses(ctx context.Context, jobs JobRepository, tasks TaskRepository,
jobIDs []uuid.UUID, now time.Time) error {
seen := make(map[uuid.UUID]struct{}, len(jobIDs))
for _, jobID := range jobIDs {
if _, duplicate := seen[jobID]; duplicate {
continue
}
seen[jobID] = struct{}{}
if err := syncJobStatus(ctx, jobs, tasks, jobID, now); err != nil {
return err
}
}
return nil
}
+6 -3
View File
@@ -60,9 +60,9 @@ type TaskRepository interface {
// lease. It returns how many tasks changed.
CancelByJob(ctx context.Context, jobID uuid.UUID, now time.Time) (int64, error)
// ExpireLeases applies the lease-expiry rule to every elapsed task and
// reports how many were affected.
ExpireLeases(ctx context.Context, now time.Time) (int64, error)
// ExpireLeases applies the lease-expiry rule to every elapsed task and returns
// the distinct jobs whose aggregate status may have changed.
ExpireLeases(ctx context.Context, now time.Time) ([]uuid.UUID, error)
}
// JobRepository persists jobs.
@@ -89,6 +89,9 @@ type WorkerRepository interface {
type ArtifactRepository interface {
Insert(ctx context.Context, a *domain.Artifact) error
Get(ctx context.Context, id uuid.UUID) (*domain.Artifact, error)
// FindPartialResult returns the durable result already uploaded for one task
// attempt. A nil artifact means the attempt has not uploaded one yet.
FindPartialResult(ctx context.Context, taskID uuid.UUID, attempt int) (*domain.Artifact, error)
}
// BlobStore holds artifact bytes, addressed by an opaque storage key. It streams
+55 -21
View File
@@ -21,12 +21,15 @@ import (
type ClaimTask struct {
tasks TaskRepository
jobs JobRepository
workers WorkerRepository
tx TxManager
clock Clock
leaseDuration time.Duration
}
func NewClaimTask(tasks TaskRepository, clock Clock, leaseDuration time.Duration) *ClaimTask {
return &ClaimTask{tasks: tasks, clock: clock, leaseDuration: leaseDuration}
func NewClaimTask(tasks TaskRepository, jobs JobRepository, workers WorkerRepository, tx TxManager, clock Clock, leaseDuration time.Duration) *ClaimTask {
return &ClaimTask{tasks: tasks, jobs: jobs, workers: workers, tx: tx, clock: clock, leaseDuration: leaseDuration}
}
// Execute reclaims elapsed leases first, then hands out one task.
@@ -42,27 +45,46 @@ func (uc *ClaimTask) Execute(ctx context.Context, in ClaimTaskInput) (*domain.Cl
if in.WorkerID == "" {
return nil, domain.ErrInvalidInput
}
now := uc.clock.Now()
if _, err := uc.tasks.ExpireLeases(ctx, now); err != nil {
return nil, err
workloads := in.Workloads
if workerID, err := uuid.Parse(in.WorkerID); err == nil {
worker, err := uc.workers.Get(ctx, workerID)
if err != nil {
return nil, err
}
// Never trust caller-supplied capabilities: registration is the durable
// worker identity and its allowlist.
workloads = worker.Capabilities
}
var claimed *domain.ClaimedTask
err := uc.tx.WithinTx(ctx, func(ctx context.Context) error {
now := uc.clock.Now()
affectedJobs, err := uc.tasks.ExpireLeases(ctx, now)
if err != nil {
return err
}
if err := syncExpiredJobStatuses(ctx, uc.jobs, uc.tasks, affectedJobs, now); err != nil {
return err
}
task, err := uc.tasks.ClaimNext(ctx, ClaimFilter{
Workloads: in.Workloads,
Owner: in.WorkerID,
Now: now,
LeaseUntil: now.Add(uc.leaseDuration),
task, err := uc.tasks.ClaimNext(ctx, ClaimFilter{
Workloads: workloads,
Owner: in.WorkerID,
Now: now,
LeaseUntil: now.Add(uc.leaseDuration),
})
if err != nil {
return err
}
if task != nil {
value := task.AsClaimed()
claimed = &value
}
return nil
})
if err != nil {
return nil, err
}
if task == nil {
return nil, nil // empty queue is a normal state, not an error
}
claimed := task.AsClaimed()
return &claimed, nil
return claimed, nil // nil means an empty queue
}
// --- RenewLease ----------------------------------------------------------
@@ -230,18 +252,30 @@ func (uc *FailTask) Execute(ctx context.Context, in FailTaskInput) (*domain.Task
type ExpireLeases struct {
tasks TaskRepository
jobs JobRepository
tx TxManager
clock Clock
}
func NewExpireLeases(tasks TaskRepository, clock Clock) *ExpireLeases {
return &ExpireLeases{tasks: tasks, clock: clock}
func NewExpireLeases(tasks TaskRepository, jobs JobRepository, tx TxManager, clock Clock) *ExpireLeases {
return &ExpireLeases{tasks: tasks, jobs: jobs, tx: tx, clock: clock}
}
// Execute reports how many tasks were reclaimed.
// Execute reclaims elapsed tasks and persists the state of every affected job.
//
// The sweep is one set-based statement rather than a load-decide-save loop:
// several coordinators run it concurrently, and a single atomic UPDATE makes
// the duplicate work harmless — the loser simply updates 0 rows.
func (uc *ExpireLeases) Execute(ctx context.Context) (int64, error) {
return uc.tasks.ExpireLeases(ctx, uc.clock.Now())
var affected []uuid.UUID
err := uc.tx.WithinTx(ctx, func(ctx context.Context) error {
now := uc.clock.Now()
var err error
affected, err = uc.tasks.ExpireLeases(ctx, now)
if err != nil {
return err
}
return syncExpiredJobStatuses(ctx, uc.jobs, uc.tasks, affected, now)
})
return int64(len(affected)), err
}
+10 -5
View File
@@ -15,6 +15,7 @@ type UIReadRepository interface {
GetJob(ctx context.Context, jobID uuid.UUID) (*domain.Job, error)
ListJobs(ctx context.Context, limit int) ([]domain.Job, error)
ListTasksByJob(ctx context.Context, jobID uuid.UUID) ([]domain.Task, error)
ListTasksByJobs(ctx context.Context, jobIDs []uuid.UUID) (map[uuid.UUID][]domain.Task, error)
ListWorkers(ctx context.Context, limit int) ([]domain.Worker, error)
ListArtifactsByJob(ctx context.Context, jobID uuid.UUID) ([]domain.Artifact, error)
}
@@ -88,12 +89,16 @@ func (d *Dashboard) Overview(ctx context.Context, limit int) (DashboardView, err
return DashboardView{}, err
}
out := DashboardView{Jobs: make([]JobCard, 0, len(jobs)), Workers: make([]WorkerCard, 0, len(workers))}
jobIDs := make([]uuid.UUID, 0, len(jobs))
for _, job := range jobs {
tasks, err := d.read.ListTasksByJob(ctx, job.ID)
if err != nil {
return DashboardView{}, err
}
out.Jobs = append(out.Jobs, jobCard(job, tasks))
jobIDs = append(jobIDs, job.ID)
}
tasksByJob, err := d.read.ListTasksByJobs(ctx, jobIDs)
if err != nil {
return DashboardView{}, err
}
for _, job := range jobs {
out.Jobs = append(out.Jobs, jobCard(job, tasksByJob[job.ID]))
}
for _, worker := range workers {
out.Workers = append(out.Workers, WorkerCard{ID: worker.ID.String(), Name: worker.Name, Status: string(worker.Status), Capabilities: worker.Capabilities, LastHeartbeatAt: worker.LastHeartbeatAt})
+90 -11
View File
@@ -4,6 +4,7 @@ import (
"context"
"fmt"
"io"
"math"
"github.com/google/uuid"
@@ -15,20 +16,27 @@ import (
// creates the job with one task per shard — the coordinator-side counterpart of
// a client submitting pre-chunked URIs.
type SubmitDataset struct {
blobs BlobStore
artifacts ArtifactRepository
jobs JobRepository
tasks TaskRepository
tx TxManager
clk Clock
blobs BlobStore
artifacts ArtifactRepository
jobs JobRepository
tasks TaskRepository
tx TxManager
clk Clock
maxAttempts int
}
func NewSubmitDataset(blobs BlobStore, artifacts ArtifactRepository, jobs JobRepository,
tasks TaskRepository, tx TxManager, clk Clock) *SubmitDataset {
return &SubmitDataset{blobs: blobs, artifacts: artifacts, jobs: jobs, tasks: tasks, tx: tx, clk: clk}
tasks TaskRepository, tx TxManager, clk Clock, maxAttempts int) *SubmitDataset {
return &SubmitDataset{blobs: blobs, artifacts: artifacts, jobs: jobs, tasks: tasks, tx: tx, clk: clk, maxAttempts: maxAttempts}
}
func (uc *SubmitDataset) Execute(ctx context.Context, in SubmitDatasetInput) (SubmitDatasetResult, error) {
if err := validateUploadedWorkload(in.Workload, in.Parameters); err != nil {
return SubmitDatasetResult{}, err
}
if uc.maxAttempts < 1 {
return SubmitDatasetResult{}, domain.ErrInvalidInput
}
now := uc.clk.Now()
job, err := domain.NewUploadedJob(in.Workload, in.Parameters, now)
@@ -64,7 +72,7 @@ func (uc *SubmitDataset) Execute(ctx context.Context, in SubmitDatasetInput) (Su
cleanup()
return SubmitDatasetResult{}, err
}
splitErr := chunk.SplitTSVLimit(rc, in.RowsPerShard, in.MaxRows, func(index int, shard io.Reader) error {
splitErr := chunk.SplitChEMBLTSVLimit(rc, in.RowsPerShard, in.MaxRows, func(index int, shard io.Reader) error {
art, err := domain.NewArtifact(job.ID, nil, domain.ArtifactShard,
fmt.Sprintf("shard-%d.tsv", index), in.ContentType, now)
if err != nil {
@@ -77,7 +85,7 @@ func (uc *SubmitDataset) Execute(ctx context.Context, in SubmitDatasetInput) (Su
putKeys = append(putKeys, art.StorageKey)
art.SetContent(ssum, ssize)
task, err := domain.NewShardTask(job.ID, index, in.Workload, art.ID, ssum, in.Parameters, 0, now)
task, err := domain.NewShardTask(job.ID, index, in.Workload, art.ID, ssum, in.Parameters, uc.maxAttempts, now)
if err != nil {
return err
}
@@ -88,7 +96,8 @@ func (uc *SubmitDataset) Execute(ctx context.Context, in SubmitDatasetInput) (Su
_ = rc.Close()
if splitErr != nil {
cleanup()
return SubmitDatasetResult{}, splitErr
// Dataset shape is caller input, not an internal coordinator failure.
return SubmitDatasetResult{}, domain.ErrInvalidInput
}
// 3. Persist job + all artifacts + all tasks atomically.
@@ -118,6 +127,76 @@ func (uc *SubmitDataset) Execute(ctx context.Context, in SubmitDatasetInput) (Su
}, nil
}
// validateUploadedWorkload is deliberately narrow until CTX-07/08/10 adds a
// typed distributed-workload registry. In particular, running similarity-graph
// independently per TSV shard is scientifically wrong: cross-shard pairs would
// be absent from the apparent graph.
func validateUploadedWorkload(workload string, parameters map[string]any) error {
if workload != "similarity-search" {
return domain.ErrInvalidInput
}
allowed := map[string]struct{}{
"query_smiles": {}, "top_k": {}, "threshold": {},
"threshold_direction": {}, "progress_every": {},
}
for key := range parameters {
if _, ok := allowed[key]; !ok {
return domain.ErrInvalidInput
}
}
query, ok := parameters["query_smiles"].(string)
if !ok || query == "" || len(query) > 200 {
return domain.ErrInvalidInput
}
if value, ok := parameters["top_k"]; ok && !isPositiveJSONInteger(value) {
return domain.ErrInvalidInput
}
if value, ok := parameters["progress_every"]; ok && !isNonNegativeJSONInteger(value) {
return domain.ErrInvalidInput
}
if value, ok := parameters["threshold"]; ok && !isUnitIntervalNumber(value) {
return domain.ErrInvalidInput
}
if value, ok := parameters["threshold_direction"]; ok && value != "greater" && value != "less" {
return domain.ErrInvalidInput
}
return nil
}
func isPositiveJSONInteger(value any) bool { return isJSONInteger(value, false) }
func isNonNegativeJSONInteger(value any) bool { return isJSONInteger(value, true) }
func isJSONInteger(value any, allowZero bool) bool {
var n int64
switch v := value.(type) {
case int:
n = int64(v)
case int64:
n = v
case float64:
if math.Trunc(v) != v || v > math.MaxInt64 || v < math.MinInt64 {
return false
}
n = int64(v)
default:
return false
}
return n >= 0 && (allowZero || n > 0)
}
func isUnitIntervalNumber(value any) bool {
switch v := value.(type) {
case float64:
return !math.IsNaN(v) && !math.IsInf(v, 0) && v >= 0 && v <= 1
case int:
return v >= 0 && v <= 1
case int64:
return v >= 0 && v <= 1
default:
return false
}
}
// GetTaskInput resolves a task's input shard and opens it for streaming. The
// caller closes the reader.
type GetTaskInput struct {
+104 -10
View File
@@ -68,18 +68,18 @@ func newHarness() *harness {
}
tx := memstore.Tx{}
h.createJob = usecase.NewCreateJob(h.jobs, h.tasks, tx, h.clk)
h.submit = usecase.NewSubmitDataset(h.blobs, h.arts, h.jobs, h.tasks, tx, h.clk)
h.claim = usecase.NewClaimTask(h.tasks, h.clk, lease)
h.submit = usecase.NewSubmitDataset(h.blobs, h.arts, h.jobs, h.tasks, tx, h.clk, 3)
h.claim = usecase.NewClaimTask(h.tasks, h.jobs, h.work, tx, h.clk, lease)
h.renew = usecase.NewRenewLease(h.tasks, h.work, tx, h.clk, lease)
h.complete = usecase.NewCompleteTask(h.tasks, h.jobs, h.arts, tx, h.clk)
h.fail = usecase.NewFailTask(h.tasks, h.jobs, tx, h.clk)
h.status = usecase.NewGetJobStatus(h.jobs, h.tasks)
h.results = usecase.NewListResults(h.tasks)
h.register = usecase.NewRegisterWorker(h.work, h.clk)
h.uploadArt = usecase.NewUploadArtifact(h.tasks, h.arts, h.blobs, h.clk)
h.uploadArt = usecase.NewUploadArtifact(h.tasks, h.arts, h.blobs, tx, h.clk)
h.downloadArt = usecase.NewDownloadArtifact(h.arts, h.blobs)
h.getInput = usecase.NewGetTaskInput(h.tasks, h.arts, h.blobs)
h.expire = usecase.NewExpireLeases(h.tasks, h.clk)
h.expire = usecase.NewExpireLeases(h.tasks, h.jobs, tx, h.clk)
h.cancel = usecase.NewCancelJob(h.jobs, h.tasks, tx, h.clk)
return h
}
@@ -138,6 +138,44 @@ func TestClaimLeasesAndAdvancesAttempt(t *testing.T) {
}
}
func TestRegisteredWorkerCannotBroadenItsCapabilitiesAtClaim(t *testing.T) {
h := newHarness()
h.seedJob(t, "restricted", 1)
worker, err := h.register.Execute(ctx, usecase.RegisterWorkerInput{
Name: "search-only", Capabilities: []string{"similarity-search"},
})
if err != nil {
t.Fatal(err)
}
claimed, err := h.claim.Execute(ctx, usecase.ClaimTaskInput{
WorkerID: worker.ID.String(), Workloads: []string{"restricted"},
})
if err != nil || claimed != nil {
t.Fatalf("claim = (%v, %v), want no compatible task", claimed, err)
}
}
func TestCreateJobRejectsUnsafeDistributedScientificPlans(t *testing.T) {
h := newHarness()
_, err := h.createJob.Execute(ctx, usecase.CreateJobInput{
Workload: "similarity-graph", InputURI: "s3://input",
Chunks: []usecase.ChunkInput{{ChunkIndex: 0, InputURI: "s3://chunk", InputSHA256: "sha"}},
})
if !errors.Is(err, domain.ErrInvalidInput) {
t.Errorf("graph job err = %v, want ErrInvalidInput", err)
}
_, err = h.createJob.Execute(ctx, usecase.CreateJobInput{
Workload: "similarity-search", InputURI: "s3://input", Parameters: map[string]any{"query_id": "CHEMBL1"},
Chunks: []usecase.ChunkInput{
{ChunkIndex: 0, InputURI: "s3://chunk0", InputSHA256: "sha"},
{ChunkIndex: 1, InputURI: "s3://chunk1", InputSHA256: "sha"},
},
})
if !errors.Is(err, domain.ErrInvalidInput) {
t.Errorf("sharded query_id job err = %v, want ErrInvalidInput", err)
}
}
func TestClaimEmptyQueueReturnsNil(t *testing.T) {
h := newHarness()
c, err := h.claim.Execute(ctx, usecase.ClaimTaskInput{WorkerID: "w1", Workloads: []string{"w"}})
@@ -281,7 +319,7 @@ func TestUploadRejectsLeaseThatExpiresDuringStreaming(t *testing.T) {
h.seedJob(t, "w", 1)
taskID, attempt := h.leaseOne(t, "w1", "w")
h.uploadArt = usecase.NewUploadArtifact(
h.tasks, h.arts, expiringBlobStore{BlobStore: h.blobs, clock: h.clk}, h.clk,
h.tasks, h.arts, expiringBlobStore{BlobStore: h.blobs, clock: h.clk}, memstore.Tx{}, h.clk,
)
_, err := h.uploadArt.Execute(ctx, usecase.UploadArtifactInput{
@@ -414,6 +452,23 @@ func TestUploadArtifactRejectsForeignWorker(t *testing.T) {
}
}
func TestUploadArtifactIsIdempotentPerTaskAttempt(t *testing.T) {
h := newHarness()
h.seedJob(t, "w", 1)
taskID, attempt := h.leaseOne(t, "w1", "w")
first := h.uploadResult(t, taskID, "w1", attempt)
second, err := h.uploadArt.Execute(ctx, usecase.UploadArtifactInput{
TaskID: taskID, WorkerID: "w1", Attempt: attempt,
Filename: "retry.csv", ContentType: "text/csv", Body: strings.NewReader("different bytes"),
})
if err != nil {
t.Fatal(err)
}
if second.ID != first {
t.Errorf("retry artifact = %s, want existing %s", second.ID, first)
}
}
func TestDownloadArtifactRoundTrips(t *testing.T) {
h := newHarness()
h.seedJob(t, "w", 1)
@@ -434,10 +489,10 @@ func TestDownloadArtifactRoundTrips(t *testing.T) {
func TestSubmitDatasetChunksAndServesInput(t *testing.T) {
h := newHarness()
tsv := "id\tsmiles\nA\tCC\nB\tCCC\nC\tCCCC\nD\tCCCCC\nE\tCCCCCC\n"
tsv := "chembl_id\tcanonical_smiles\nA\tCC\nB\tCCC\nC\tCCCC\nD\tCCCCC\nE\tCCCCCC\n"
res, err := h.submit.Execute(ctx, usecase.SubmitDatasetInput{
Workload: "w", RowsPerShard: 2, Filename: "chembl.tsv",
Workload: "similarity-search", Parameters: map[string]any{"query_smiles": "CCO"}, RowsPerShard: 2, Filename: "chembl.tsv",
ContentType: "text/tab-separated-values", Body: strings.NewReader(tsv),
})
if err != nil {
@@ -453,7 +508,7 @@ func TestSubmitDatasetChunksAndServesInput(t *testing.T) {
t.Errorf("job total = %d, want 3", prog.Total)
}
c, err := h.claim.Execute(ctx, usecase.ClaimTaskInput{WorkerID: "w1", Workloads: []string{"w"}})
c, err := h.claim.Execute(ctx, usecase.ClaimTaskInput{WorkerID: "w1", Workloads: []string{"similarity-search"}})
if err != nil || c == nil {
t.Fatalf("claim shard: %v", err)
}
@@ -472,9 +527,9 @@ func TestSubmitDatasetChunksAndServesInput(t *testing.T) {
func TestSubmitDatasetLimitsRowsBeforeCreatingShards(t *testing.T) {
h := newHarness()
tsv := "id\tsmiles\nA\tCC\nB\tCCC\nC\tCCCC\nD\tCCCCC\nE\tCCCCCC\n"
tsv := "chembl_id\tcanonical_smiles\nA\tCC\nB\tCCC\nC\tCCCC\nD\tCCCCC\nE\tCCCCCC\n"
res, err := h.submit.Execute(ctx, usecase.SubmitDatasetInput{
Workload: "w", RowsPerShard: 2, MaxRows: 3, Filename: "chembl.tsv",
Workload: "similarity-search", Parameters: map[string]any{"query_smiles": "CCO"}, RowsPerShard: 2, MaxRows: 3, Filename: "chembl.tsv",
ContentType: "text/tab-separated-values", Body: strings.NewReader(tsv),
})
if err != nil {
@@ -485,6 +540,26 @@ func TestSubmitDatasetLimitsRowsBeforeCreatingShards(t *testing.T) {
}
}
func TestSubmitDatasetRejectsUnsupportedDistributedWorkloads(t *testing.T) {
h := newHarness()
_, err := h.submit.Execute(ctx, usecase.SubmitDatasetInput{
Workload: "similarity-graph", Parameters: map[string]any{"threshold": 0.7}, RowsPerShard: 2,
Filename: "chembl.tsv", ContentType: "text/tab-separated-values",
Body: strings.NewReader("chembl_id\tcanonical_smiles\nA\tCC\n"),
})
if !errors.Is(err, domain.ErrInvalidInput) {
t.Errorf("graph submission err = %v, want ErrInvalidInput", err)
}
_, err = h.submit.Execute(ctx, usecase.SubmitDatasetInput{
Workload: "similarity-search", Parameters: map[string]any{"query_id": "CHEMBL1"}, RowsPerShard: 2,
Filename: "chembl.tsv", ContentType: "text/tab-separated-values",
Body: strings.NewReader("chembl_id\tcanonical_smiles\nA\tCC\n"),
})
if !errors.Is(err, domain.ErrInvalidInput) {
t.Errorf("query_id submission err = %v, want ErrInvalidInput", err)
}
}
func TestCancelJobInvalidatesClaimedAndPendingTasks(t *testing.T) {
h := newHarness()
jobID := h.seedJob(t, "w", 3)
@@ -531,3 +606,22 @@ func TestExpireLeasesReclaims(t *testing.T) {
t.Errorf("expire = (%d, %v), want (1, nil)", n, err)
}
}
func TestFinalLeaseExpiryPersistsFailedJobAndCannotBeCancelled(t *testing.T) {
h := newHarness()
jobID := h.seedJob(t, "w", 1)
for attempt := 1; attempt <= domain.DefaultMaxAttempts; attempt++ {
h.leaseOne(t, "w1", "w")
h.clk.Advance(lease + time.Second)
if _, err := h.expire.Execute(ctx); err != nil {
t.Fatalf("expire attempt %d: %v", attempt, err)
}
}
progress, err := h.status.Execute(ctx, jobID)
if err != nil || progress.Job.Status != domain.JobFailed || progress.DeriveStatus() != domain.JobFailed {
t.Fatalf("progress = %+v, err = %v; want persisted failed job", progress, err)
}
if _, err := h.cancel.Execute(ctx, jobID); !errors.Is(err, domain.ErrJobNotCancellable) {
t.Errorf("cancel terminal lease failure = %v, want ErrJobNotCancellable", err)
}
}
@@ -0,0 +1,5 @@
BEGIN;
DROP INDEX IF EXISTS uq_partial_result_task_attempt;
COMMIT;
@@ -0,0 +1,26 @@
BEGIN;
-- Old deployments can contain more than one partial result because earlier
-- versions accepted repeated PUTs. Preserve the one referenced by a completed
-- task and discard stale rows; unfinished tasks must upload again after a
-- deploy, just as they do after a lost lease.
DELETE FROM artifacts AS a
USING tasks AS t
WHERE a.task_id = t.id
AND a.kind = 'partial_result'::artifact_kind
AND t.status <> 'completed'::task_status;
DELETE FROM artifacts AS a
USING tasks AS t
WHERE a.task_id = t.id
AND a.kind = 'partial_result'::artifact_kind
AND t.status = 'completed'::task_status
AND a.id <> t.result_artifact_id;
-- One lease attempt has one durable partial result. This makes an upload retry
-- idempotent and prevents repeated uploads from accumulating orphan artifacts.
CREATE UNIQUE INDEX uq_partial_result_task_attempt
ON artifacts (task_id, attempt)
WHERE kind = 'partial_result'::artifact_kind;
COMMIT;
+27 -16
View File
@@ -51,6 +51,13 @@ check "register worker" 201 -X POST "${HOST}/workers/register" "
-d '{"name":"smoke-worker","capabilities":["similarity_search"],"cpu_count":4,"memory_mb":8192}'
check "register without capabilities → 400" 400 -X POST "${HOST}/workers/register" "${auth[@]}" \
-d '{"name":"bad"}'
registration=$(curl -sS "${auth[@]}" -X POST "${HOST}/workers/register" \
-d '{"name":"smoke-worker-active","capabilities":["similarity_search","similarity-search"],"cpu_count":4}')
worker_id=$(printf '%s' "$registration" | python3 -c 'import json,sys;print(json.load(sys.stdin)["worker_id"])' 2>/dev/null)
if [[ -z "${worker_id:-}" ]]; then
echo " ✗ could not register active worker: $registration"
exit 1
fi
echo
echo "job lifecycle"
@@ -75,7 +82,7 @@ declare -A our_chunks
task_id=""
attempt=""
for _ in $(seq 1 40); do
claim=$(curl -sS "${auth[@]}" -X POST "${HOST}/tasks/claim" -d '{"worker_id":"w1"}')
claim=$(curl -sS "${auth[@]}" -X POST "${HOST}/tasks/claim" -d "{\"worker_id\":\"${worker_id}\"}")
[[ -z "$claim" ]] && break # 204: queue drained
read -r c_job c_task c_chunk c_attempt < <(printf '%s' "$claim" |
@@ -100,7 +107,7 @@ else
fi
check "heartbeat" 200 -X POST "${HOST}/tasks/${task_id}/heartbeat" "${auth[@]}" \
-d "{\"worker_id\":\"w1\",\"attempt\":${attempt}}"
-d "{\"worker_id\":\"${worker_id}\",\"attempt\":${attempt}}"
# --- artifacts + result (uploads happen while the task is still leased) ---
bearer=(-H "Authorization: Bearer ${TOKEN}")
@@ -108,43 +115,47 @@ bearer=(-H "Authorization: Bearer ${TOKEN}")
# upload <filename> -> prints the artifact_id
upload() {
curl -sS -X PUT "${HOST}/tasks/${task_id}/artifacts/$1" "${bearer[@]}" \
-H 'Content-Type: text/csv' -H 'X-Worker-ID: w1' -H "X-Task-Attempt: ${attempt}" \
-H 'Content-Type: text/csv' -H "X-Worker-ID: ${worker_id}" -H "X-Task-Attempt: ${attempt}" \
--data-binary $'query,match,score\nA,B,0.9\n' |
python3 -c 'import json,sys;print(json.load(sys.stdin)["artifact_id"])' 2>/dev/null
}
check "upload artifact" 200 -X PUT "${HOST}/tasks/${task_id}/artifacts/result.csv" "${bearer[@]}" \
-H 'Content-Type: text/csv' -H 'X-Worker-ID: w1' -H "X-Task-Attempt: ${attempt}" \
-H 'Content-Type: text/csv' -H "X-Worker-ID: ${worker_id}" -H "X-Task-Attempt: ${attempt}" \
--data-binary $'query,match,score\nA,B,0.9\n'
check "foreign worker upload → 409" 409 -X PUT "${HOST}/tasks/${task_id}/artifacts/x.csv" "${bearer[@]}" \
-H 'Content-Type: text/csv' -H 'X-Worker-ID: impostor' -H "X-Task-Attempt: ${attempt}" \
--data-binary 'x'
# Two result artifacts, uploaded now while the lease is held: one to complete
# with, a second to prove a different manifest is rejected after completion.
# A retry of a PUT returns the same durable artifact for the task attempt.
art_id=$(upload primary.csv)
art_id2=$(upload secondary.csv)
if [[ "$art_id" == "$art_id2" ]]; then
printf ' \033[32m✓\033[0m %-46s %s\n' "duplicate upload is idempotent" "$art_id"
pass=$((pass + 1))
else
printf ' \033[31m✗\033[0m %-46s got %s and %s\n' "duplicate upload is idempotent" "$art_id" "$art_id2"
fail=$((fail + 1))
fi
check "download artifact" 200 "${HOST}/artifacts/${art_id}/download" "${bearer[@]}"
check "foreign worker submits → 409" 409 -X POST "${HOST}/tasks/${task_id}/result" "${auth[@]}" \
-d "{\"worker_id\":\"impostor\",\"attempt\":${attempt},\"result\":{\"artifact_id\":\"${art_id}\"}}"
check "submit result" 200 -X POST "${HOST}/tasks/${task_id}/result" "${auth[@]}" \
-d "{\"worker_id\":\"w1\",\"attempt\":${attempt},\"result\":{\"artifact_id\":\"${art_id}\"}}"
-d "{\"worker_id\":\"${worker_id}\",\"attempt\":${attempt},\"result\":{\"artifact_id\":\"${art_id}\"}}"
check "replay same result → idempotent" 200 -X POST "${HOST}/tasks/${task_id}/result" "${auth[@]}" \
-d "{\"worker_id\":\"w1\",\"attempt\":${attempt},\"result\":{\"artifact_id\":\"${art_id}\"}}"
check "different result → 409" 409 -X POST "${HOST}/tasks/${task_id}/result" "${auth[@]}" \
-d "{\"worker_id\":\"w1\",\"attempt\":${attempt},\"result\":{\"artifact_id\":\"${art_id2}\"}}"
-d "{\"worker_id\":\"${worker_id}\",\"attempt\":${attempt},\"result\":{\"artifact_id\":\"${art_id}\"}}"
check "GET /jobs/{id}" 200 "${HOST}/jobs/${job_id}" "${auth[@]}"
echo
echo "input validation"
check "malformed uuid → 400" 400 -X POST "${HOST}/tasks/not-a-uuid/result" "${auth[@]}" \
-d '{"worker_id":"w1","attempt":1,"result":{"artifact_id":"00000000-0000-0000-0000-000000000000"}}'
-d "{\"worker_id\":\"${worker_id}\",\"attempt\":1,\"result\":{\"artifact_id\":\"00000000-0000-0000-0000-000000000000\"}}"
# Note: Go's encoding/json matches field names case-insensitively, so
# "worker_ID" would be accepted as "worker_id". Only a genuinely unknown key
# trips DisallowUnknownFields.
check "unknown json field → 400" 400 -X POST "${HOST}/tasks/claim" "${auth[@]}" \
-d '{"worker_id":"w1","totally_unknown":1}'
-d "{\"worker_id\":\"${worker_id}\",\"totally_unknown\":1}"
check "unknown job → 404" 404 "${HOST}/jobs/00000000-0000-0000-0000-000000000000" "${auth[@]}"
echo
@@ -152,11 +163,11 @@ echo "dataset upload → chunking"
# Upload a 5-row TSV split at 2 rows/shard → expect 3 shard tasks. The text
# fields precede the file part, which the coordinator streams.
up=$(curl -sS "${bearer[@]}" -X POST "${HOST}/jobs/upload" \
-F 'workload=similarity_search' \
-F 'parameters={"top_k":10}' \
-F 'workload=similarity-search' \
-F 'parameters={"query_smiles":"CCO","top_k":10}' \
-F 'chunk_rows=2' \
-F 'file=@-;filename=chembl.tsv;type=text/tab-separated-values' <<'TSV'
id smiles
chembl_id canonical_smiles
A CC
B CCC
C CCCC
@@ -180,7 +191,7 @@ fi
up_input=""
for _ in $(seq 1 30); do
c=$(curl -sS "${bearer[@]}" -H 'Content-Type: application/json' -X POST "${HOST}/tasks/claim" \
-d '{"worker_id":"up-w","capabilities":["similarity_search"]}')
-d "{\"worker_id\":\"${worker_id}\"}")
[[ -z "$c" ]] && break
cj=$(printf '%s' "$c" | python3 -c 'import json,sys;print(json.load(sys.stdin)["job_id"])' 2>/dev/null)
[[ "$cj" != "$up_job" ]] && continue
+4 -2
View File
@@ -107,7 +107,8 @@ Content-Type: application/json
```
`cpu_count`/`memory_mb` are accepted for forward compatibility and not yet
persisted. `capabilities` must be non-empty (an allowlisted workload set).
persisted. `capabilities` must be non-empty. A claim uses the capabilities
stored at registration; the request cannot broaden them.
## Claim task
@@ -120,7 +121,8 @@ Content-Type: application/json
```
- `204 No Content`: no compatible task.
- `200 OK`: a task is leased atomically.
- `200 OK`: a task is leased atomically. `worker_id` must be a registered UUID;
its persisted capabilities, rather than this request field, decide eligibility.
```json
{
+35 -8
View File
@@ -46,24 +46,25 @@ coordinator was started with. Never log it, never send it in an error body.
```http
POST /workers/register
{ "name": "lab-worker-01", "capabilities": ["similarity_search"] }
{ "name": "lab-worker-01", "capabilities": ["similarity-search"] }
```
Response: `{ "worker_id": "<uuid>", "heartbeat_interval_seconds": 15 }`.
- `capabilities` are the workload names you can run — the coordinator only hands
you matching tasks.
- `capabilities` are fixed at registration — the coordinator only hands you
matching tasks and a later claim cannot broaden that set.
- **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 coordinator jobs use `similarity_search` / `similarity_graph`; the
reference Python worker also accepts the public CLI spellings with hyphens.
- 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.
## 2. Claim a task
```http
POST /tasks/claim
{ "worker_id": "<uuid>", "capabilities": ["similarity_search"] }
{ "worker_id": "<uuid>", "capabilities": ["similarity-search"] }
```
- `200` → a leased task (below).
@@ -74,9 +75,9 @@ POST /tasks/claim
"task_id": "<uuid>",
"attempt": 1,
"lease_expires_at": "2026-07-22T12:05:00Z",
"workload": "similarity_search",
"workload": "similarity-search",
"input": { "uri": "/tasks/<uuid>/input", "sha256": "<hex>" },
"parameters": { "query_id": "CHEMBL939", "top_k": 20 }
"parameters": { "query_smiles": "CCO", "top_k": 20 }
}
```
@@ -193,6 +194,32 @@ Per the worker contract, at minimum:
- poll interval and request timeout
- a working directory for downloaded inputs and generated outputs
## Run the reference worker locally
Use one terminal per worker and a distinct work directory for each process:
```sh
SCIMESH_COORDINATOR_URL=http://localhost:8080 \
SCIMESH_BEARER_TOKEN=dev-token \
SCIMESH_WORKER_NAME=worker-1 \
scimesh-worker --work-dir "$PWD/worker-data-1"
```
For a bounded manual check, use one of these lifecycle modes:
```sh
# Make exactly one claim; exit immediately when no task is available.
scimesh-worker --work-dir "$PWD/worker-data-check" --once
# Keep polling until two tasks complete successfully, then exit.
scimesh-worker --work-dir "$PWD/worker-data-check" --max-tasks 2
```
`SCIMESH_MAX_TASKS` provides the same limit through the environment. Pressing
`Ctrl+C` stops the reference worker cleanly. If it interrupts an active task,
the worker reports a sanitized retriable failure first, emits no traceback, and
exits with status `130`.
## Generate a client from the spec
Instead of hand-writing request code, generate it:
@@ -0,0 +1,207 @@
# CTX-07: distributed workload protocol and planner contract
## Status and scope
This document is the implementation contract for CTX-07. It does not implement
a planner, reducer, API endpoint, database migration, or final artifact. Until
CTX-08 and CTX-09 are 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.
The Go coordinator owns durable artifacts, transactions, task rows, leases, and
HTTP. A Python workload must never access PostgreSQL or call the coordinator.
Read `PLAN.md`, `.agents/workloads.md`, and `docs/api-contract.md` before
implementing this CTX.
## Canonical vocabulary
- External workload names are lowercase hyphenated names: `similarity-search`
and, later, `similarity-graph`.
- The existing underscore spellings are a temporary compatibility alias at the
Python worker boundary only. Planners, persisted job/task payloads, and new
API examples use the canonical hyphenated spelling.
- A **plan** contains only JSON-compatible values and coordinator artifact
references. It contains no local filesystem path, worker URI, presigned URL,
database connection, or callable.
- `chunk_index` is a non-negative integer, unique within a plan, and sorted
ascending whenever results are enumerated.
## Python boundary
CTX-07 adds a small `DistributedWorkload` protocol under `scimesh/distributed/`
and a registry separate from the local CLI registry. Names below are proposed
public types; keep concrete implementation details minimal.
```python
class DistributedWorkload(Protocol):
name: str
def validate_job(self, parameters: Mapping[str, object]) -> None: ...
def plan(
self,
input_path: Path,
input_artifact_id: str,
parameters: Mapping[str, object],
shard_rows: int,
workspace: Path,
) -> DistributedPlan: ...
def reduce(
self,
partial_results: Sequence[CompletedPartial],
parameters: Mapping[str, object],
workspace: Path,
) -> FinalResult: ...
```
`input_path` and `workspace` are temporary files supplied by the coordinator
bridge. They are never serialized. `plan()` returns only a `DistributedPlan`;
the bridge validates it, persists artifact/task rows in one coordinator
transaction, and removes its temporary workspace. If validation or planning
fails, no job or task may be written.
## JSON models
All objects below are schema version `1`. Future incompatible changes require a
new version; never infer a schema from missing fields.
### Artifact reference
```json
{
"artifact_id": "c4273293-f8b4-4ecb-99df-3b9f5a32b6a6",
"sha256": "3b2d...64-lowercase-hex-characters",
"content_type": "text/tab-separated-values"
}
```
The artifact ID is coordinator-owned. The checksum is included so planning and
tests can assert exactly which immutable input was used. A worker receives the
coordinator-generated download URI only through `POST /tasks/claim`.
### Distributed plan
```json
{
"schema_version": 1,
"workload": "similarity-search",
"resolved_parameters": {
"query_smiles": "COc1ccc(Nc2ncnc3cc(OCCCN4CCOCC4)c(OC)c23)cc1",
"query_source": {"kind": "chembl_id", "value": "CHEMBL939"},
"top_k": 20,
"threshold": 0.7,
"threshold_direction": "greater",
"fingerprint": {"algorithm": "morgan", "radius": 2, "fp_size": 2048}
},
"tasks": [
{
"chunk_index": 0,
"input_artifact": {
"artifact_id": "69e41105-d9fb-4c7f-a2db-7dd9e3ba2c76",
"sha256": "4c92...64-lowercase-hex-characters",
"content_type": "text/tab-separated-values"
},
"parameters": {
"query_smiles": "COc1ccc(Nc2ncnc3cc(OCCCN4CCOCC4)c(OC)c23)cc1",
"top_k": 20,
"threshold": 0.7,
"threshold_direction": "greater",
"fingerprint": {"algorithm": "morgan", "radius": 2, "fp_size": 2048}
}
}
]
}
```
`resolved_parameters` are immutable job metadata. A task copies only the
values required by its worker runner. The coordinator may add its own durable
task ID and generated input URI; it must not alter scientific parameters.
## Similarity-search planning rules
1. Accept exactly one of `query_id` and `query_smiles` at the public boundary.
2. Validate a supplied SMILES once. For `query_id`, find and validate that
molecule once against the original uploaded TSV **before** creating shards.
3. Persist the resolved canonical query SMILES and the original query source in
`resolved_parameters`. Workers receive `query_smiles`, never `query_id`.
4. Fingerprint settings are fixed to Morgan radius `2` and `fp_size` `2048`.
Reject a request that tries to override them rather than silently changing
scientific semantics.
5. Split source rows in input order. Every shard includes the original TSV
header and has a contiguous, zero-based `chunk_index`.
6. Each shard uses the global `top_k`, not a smaller local limit. A global
reducer cannot recover a candidate discarded by every shard.
7. Preserve `threshold`, `threshold_direction`, and valid `max_rows` semantics
in the resolved plan. A job-level row limit is applied before sharding, not
independently by every worker.
Invalid row SMILES are not planner failures. They remain shard data and are
counted by the worker exactly as the local workload does. An invalid query is a
planning failure.
## Partial-result contract
A completed similarity-search task owns exactly one coordinator-uploaded CSV
artifact with content type `text/csv` and these columns, in this order:
```csv
rank,chembl_id,canonical_smiles,similarity
1,CHEMBL123,CCO,0.875000
```
- `rank` is one-based local rank.
- `similarity` uses the local CLI's six-decimal formatting.
- Rows are sorted by `(-similarity, chembl_id, canonical_smiles)` for
`threshold_direction=greater`, or `(similarity, chembl_id,
canonical_smiles)` for `less`.
- The query molecule and every candidate with the same canonical query SMILES
are excluded using the existing local-workload definition.
- Empty valid result files still include the header.
The worker completion metrics must include JSON numbers for `scanned_rows`,
`valid_molecules`, `invalid_smiles`, `matches_emitted`, and
`elapsed_seconds`. Metrics are observability data; the reducer derives final
scientific output exclusively from coordinator-owned partial artifacts.
## Reduction boundary
CTX-09 invokes the registered reducer only after every task is completed. It
passes `CompletedPartial` values ordered by `chunk_index`, each containing its
coordinator artifact reference and validated metrics.
For similarity-search the reducer:
1. reads partial CSVs in `chunk_index` order;
2. validates header, row shape, rank, finite similarity in `[0, 1]`, and sort
order;
3. retains a bounded heap of at most the global `top_k` candidates using the
exact local ranking key;
4. writes the same header and deterministic rank numbering as the local CLI.
It must not deduplicate ordinary records: the local reference keeps input-row
multiplicity. Reduction is independent of worker completion order and uses
`O(top_k + shard_rows)` memory apart from CSV streaming buffers.
## Required tests for the CTX-07 implementation
- unknown workload is rejected before any coordinator job/task write;
- invalid public parameters and invalid `query_id` produce no partial plan;
- `query_id` resolution occurs once, before shard construction;
- the same input, parameters, and shard size generate byte-equivalent
JSON plans and identical shard order;
- every task payload is JSON-serializable and contains only artifact references
and validated scalar/object values;
- a two-shard dummy workload proves coordinator transaction rollback on planner
validation failure;
- completed partial artifacts reach the reducer ordered by `chunk_index`, even
when workers finish in a different order;
- the protocol registry never imports the Go coordinator or database code.
## 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.
+9 -7
View File
@@ -96,7 +96,9 @@ paths:
description: >
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.
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.
requestBody:
required: true
content:
@@ -388,7 +390,7 @@ components:
type: array
minItems: 1
items: { type: string }
example: [similarity_search, similarity_graph]
example: [similarity-search]
cpu_count:
type: integer
description: Accepted for forward compatibility; not yet persisted.
@@ -419,7 +421,7 @@ components:
type: object
required: [workload, input_uri, chunks]
properties:
workload: { type: string, example: similarity_search }
workload: { type: string, example: similarity-search }
input_uri: { type: string }
parameters: { type: object, additionalProperties: true }
chunks:
@@ -437,11 +439,11 @@ components:
type: object
required: [workload, file]
properties:
workload: { type: string, example: similarity_search }
workload: { type: string, enum: [similarity-search], example: similarity-search }
parameters:
type: string
description: JSON object, sent as a string form field.
example: '{"top_k":10}'
example: '{"query_smiles":"CCO","top_k":10}'
chunk_rows:
type: integer
description: Data rows per shard. Default 1000.
@@ -486,11 +488,11 @@ components:
type: object
required: [worker_id]
properties:
worker_id: { type: string }
worker_id: { type: string, format: uuid, description: Registered worker identity. }
capabilities:
type: array
items: { type: string }
description: Workloads this worker can run. Empty means "any".
description: Accepted for compatibility only; registration capabilities decide eligibility.
max_concurrency:
type: integer
description: Accepted; the coordinator leases one task per call.
+4 -4
View File
@@ -59,7 +59,7 @@ Content-Type: application/json
{
"name": "lab-worker-01",
"capabilities": ["similarity-search", "similarity-graph"],
"capabilities": ["similarity-search"],
"cpu_count": 8,
"memory_mb": 16384
}
@@ -74,8 +74,8 @@ POST /tasks/claim
Content-Type: application/json
{
"worker_id": "worker-01",
"capabilities": ["similarity-search", "similarity-graph"],
"worker_id": "<registered-uuid>",
"capabilities": ["similarity-search"],
"max_concurrency": 1
}
```
@@ -95,7 +95,7 @@ When a task is available, it returns `200 OK`:
"sha256": "..."
},
"parameters": {
"query_id": "CHEMBL939",
"query_smiles": "CCO",
"top_k": 20
}
}
+31 -6
View File
@@ -13,15 +13,17 @@ from .daemon import WorkerDaemon
from .runners import SciMeshRunner
def main(argv: list[str] | None = None) -> int:
def build_parser() -> argparse.ArgumentParser:
"""Build the worker CLI parser for command-line use and focused tests."""
parser = argparse.ArgumentParser(
prog="scimesh-worker",
epilog=(
"Environment: SCIMESH_COORDINATOR_URL, SCIMESH_WORK_DIR, "
"SCIMESH_WORKER_NAME, SCIMESH_CPU_COUNT, SCIMESH_MEMORY_MB, "
"SCIMESH_POLL_INTERVAL, SCIMESH_REQUEST_TIMEOUT, "
"SCIMESH_HEARTBEAT_INTERVAL, SCIMESH_CLEANUP_AFTER_SECONDS, and "
"SCIMESH_BEARER_TOKEN. SCIMESH_WORKER_ID is a legacy/test override."
"SCIMESH_HEARTBEAT_INTERVAL, SCIMESH_CLEANUP_AFTER_SECONDS, "
"SCIMESH_MAX_TASKS, and SCIMESH_BEARER_TOKEN. "
"SCIMESH_WORKER_ID is a legacy/test override."
),
)
parser.add_argument("--coordinator-url")
@@ -34,8 +36,31 @@ def main(argv: list[str] | None = None) -> int:
parser.add_argument("--request-timeout", type=float)
parser.add_argument("--heartbeat-interval", type=float)
parser.add_argument("--cleanup-after-seconds", type=float)
lifecycle = parser.add_mutually_exclusive_group()
lifecycle.add_argument(
"--once",
action="store_true",
help="Claim at most one task, then exit; exit immediately when the queue is empty",
)
lifecycle.add_argument(
"--max-tasks",
type=int,
help="Process this many claimed tasks, then exit",
)
return parser
def main(argv: list[str] | None = None) -> int:
parser = build_parser()
args = parser.parse_args(argv)
overrides = {key: value for key, value in vars(args).items() if value is not None}
overrides = {
key: value
for key, value in vars(args).items()
if value is not None and key != "once"
}
if args.once:
overrides["max_tasks"] = 1
overrides["exit_when_idle"] = True
if "work_dir" in overrides:
overrides["work_dir"] = Path(overrides["work_dir"])
try:
@@ -44,13 +69,13 @@ def main(argv: list[str] | None = None) -> int:
parser.error(str(error))
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
client = HttpCoordinatorClient(config.coordinator_url, config.request_timeout, config.bearer_token)
WorkerDaemon(
completed_without_interruption = WorkerDaemon(
config,
client,
HttpArtifactClient(config.coordinator_url, config.request_timeout, config.bearer_token),
SciMeshRunner(),
).run_forever()
return 0
return 0 if completed_without_interruption else 130
if __name__ == "__main__":
+18 -3
View File
@@ -36,13 +36,16 @@ class WorkerConfig:
heartbeat_interval: float = 15.0
bearer_token: str | None = None
cleanup_after_seconds: float | None = None
max_tasks: int | None = None
exit_when_idle: bool = False
# Distributed similarity-graph requires triangular block-pair planning and
# is deliberately not advertised until CTX-10. A normal worker must never
# make a multi-shard graph job appear scientifically complete.
# The local CLI uses hyphens; the first coordinator contract used
# underscores. Advertise both stable spellings while jobs are migrated.
# underscores, so retain the search alias during migration.
capabilities: tuple[str, ...] = (
"similarity-search",
"similarity-graph",
"similarity_search",
"similarity_graph",
)
def __post_init__(self) -> None:
@@ -66,6 +69,15 @@ class WorkerConfig:
_positive_number(self.heartbeat_interval, "heartbeat_interval")
if self.cleanup_after_seconds is not None:
_positive_number(self.cleanup_after_seconds, "cleanup_after_seconds", allow_zero=True)
if self.max_tasks is not None:
if (
isinstance(self.max_tasks, bool)
or not isinstance(self.max_tasks, int)
or self.max_tasks < 1
):
raise ValueError("max_tasks must be positive when set")
if not isinstance(self.exit_when_idle, bool):
raise ValueError("exit_when_idle must be a boolean")
if not self.capabilities:
raise ValueError("capabilities cannot be empty")
# Runner subprocesses use a task directory as their cwd. Keep the
@@ -90,6 +102,7 @@ class WorkerConfig:
cleanup = value("cleanup_after_seconds", "SCIMESH_CLEANUP_AFTER_SECONDS")
cpu_count = value("cpu_count", "SCIMESH_CPU_COUNT", os.cpu_count() or 1)
memory_mb = value("memory_mb", "SCIMESH_MEMORY_MB")
max_tasks = value("max_tasks", "SCIMESH_MAX_TASKS")
return cls(
coordinator_url=url.rstrip("/"),
worker_id=value("worker_id", "SCIMESH_WORKER_ID"),
@@ -102,4 +115,6 @@ class WorkerConfig:
heartbeat_interval=float(value("heartbeat_interval", "SCIMESH_HEARTBEAT_INTERVAL", "15")),
bearer_token=value("bearer_token", "SCIMESH_BEARER_TOKEN"),
cleanup_after_seconds=float(cleanup) if cleanup else None,
max_tasks=int(max_tasks) if max_tasks is not None else None,
exit_when_idle=bool(values.get("exit_when_idle", False)),
)
+94 -23
View File
@@ -4,8 +4,10 @@ from __future__ import annotations
import logging
from dataclasses import replace
from dataclasses import dataclass
from pathlib import Path
import random
import re
import shutil
import threading
import time
@@ -72,6 +74,14 @@ class LeaseHeartbeat:
return seconds
@dataclass(frozen=True)
class RunOnceOutcome:
"""Whether a claim was made and whether that claimed task completed."""
claimed: bool
completed: bool
class WorkerDaemon:
def __init__(self, config: WorkerConfig, coordinator: CoordinatorClient, artifacts: ArtifactClient, runner: Runner) -> None:
self.config, self.coordinator, self.artifacts, self.runner = config, coordinator, artifacts, runner
@@ -79,32 +89,68 @@ class WorkerDaemon:
self._registered = False
self.log = logging.getLogger("scimesh.worker")
def run_forever(self) -> None:
def run_forever(self) -> bool:
"""Run until stopped; return false only when interrupted by the operator."""
failures = 0
while True:
try:
if not self._registered:
self._register_worker()
self._cleanup_expired_directories()
claimed = self.run_once()
failures = 0
if not claimed:
self._sleep(self.config.poll_interval)
except CoordinatorTransientError as error:
failures += 1
self._log("failed", error_type=type(error).__name__)
self._sleep(min(self.config.poll_interval * 2 ** min(failures, 6), 60.0))
completed_tasks = 0
self._log(
"started",
max_tasks=self.config.max_tasks,
exit_when_idle=self.config.exit_when_idle,
)
try:
while True:
try:
if not self._registered:
self._register_worker()
self._cleanup_expired_directories()
outcome = self.run_once()
failures = 0
if outcome.claimed:
if outcome.completed:
completed_tasks += 1
if self.config.exit_when_idle:
self._log(
"stopped",
reason="one_claim_processed",
completed_tasks=completed_tasks,
)
return True
if (
outcome.completed
and self.config.max_tasks is not None
and completed_tasks >= self.config.max_tasks
):
self._log(
"stopped",
reason="max_tasks_reached",
completed_tasks=completed_tasks,
)
return True
elif self.config.exit_when_idle:
self._log("stopped", reason="queue_empty", completed_tasks=completed_tasks)
return True
else:
self._sleep(self.config.poll_interval)
except CoordinatorTransientError as error:
failures += 1
self._log("failed", error_type=type(error).__name__)
self._sleep(min(self.config.poll_interval * 2 ** min(failures, 6), 60.0))
except KeyboardInterrupt:
self._log("stopped", reason="interrupted", completed_tasks=completed_tasks)
return False
def run_once(self) -> bool:
def run_once(self) -> RunOnceOutcome:
worker_id = self._worker_id()
self._log("claiming")
self._log("claiming", log_level=logging.DEBUG)
task = self.coordinator.claim(worker_id, self.config.capabilities)
if task is None:
self._log("idle")
return False
self._log("idle", log_level=logging.DEBUG)
return RunOnceOutcome(claimed=False, completed=False)
started = time.monotonic()
task_dir = self.config.work_dir / task.task_id / str(task.attempt)
heartbeat = LeaseHeartbeat(task, self.coordinator, self.config)
completed = False
try:
task_dir.mkdir(parents=True, exist_ok=False)
heartbeat.start()
@@ -135,7 +181,15 @@ class WorkerDaemon:
},
},
)
self._log("idle", task, elapsed_seconds=round(time.monotonic() - started, 3))
completed = True
self._log("completed", task, elapsed_seconds=round(time.monotonic() - started, 3))
except KeyboardInterrupt:
self._log("interrupted", task)
try:
self._report_failure(task, InterruptedError("worker interrupted by operator"))
except CoordinatorTransientError:
self._log("failed", task, error_type="FailureReportError")
raise
except CoordinatorConflictError as error:
self._log("lease_lost", task, error_type=type(error).__name__)
except Exception as error:
@@ -143,10 +197,10 @@ class WorkerDaemon:
self._report_failure(task, error)
finally:
heartbeat.stop()
return True
return RunOnceOutcome(claimed=True, completed=completed)
def _report_failure(self, task: ClaimedTask, error: Exception) -> None:
message = str(error).replace(str(self.config.work_dir), "<worker-dir>")[:300]
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})
except CoordinatorTransientError:
@@ -154,6 +208,16 @@ class WorkerDaemon:
except Exception:
self._log("failed", task, error_type="FailureReportError")
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>")
# CalledProcessError includes the complete argv, including sys.executable
# outside work_dir. Replace POSIX and Windows absolute paths before the
# message reaches the coordinator database or operator UI.
message = re.sub(r"(?<![\w:])[A-Za-z]:\\[^\s'\"\],)]+", "<path>", message)
message = re.sub(r"(?<![\w:])/(?:[^\s'\"\],)]+)", "<path>", message)
return message[:300]
def _register_worker(self) -> None:
registered = self.coordinator.register(
self.config.worker_name,
@@ -180,9 +244,16 @@ class WorkerDaemon:
"""Keep completion payload exact: coordinator owns all artifact metadata."""
return {"artifact_id": uploaded.artifact_id}
def _log(self, state: str, task: ClaimedTask | None = None, **extra: object) -> None:
def _log(
self,
state: str,
task: ClaimedTask | None = None,
*,
log_level: int = logging.INFO,
**extra: object,
) -> None:
fields = {"worker_id": self.config.worker_id, "task_id": task.task_id if task else None, "attempt": task.attempt if task else None, "state": state, **extra}
self.log.info("worker_event %s", fields)
self.log.log(log_level, "worker_event %s", fields)
def _cleanup_expired_directories(self) -> None:
"""Remove only old task attempt directories when retention was configured."""
+165
View File
@@ -0,0 +1,165 @@
#!/usr/bin/env bash
#
# End-to-end check: two Python workers process separate coordinator shards.
#
# Requires Docker, curl, python3, and an installed scimesh-worker (normally
# from this repository's .venv). The test uses its own Compose project, ports,
# volumes, and temporary worker directories, leaving a developer stack alone.
set -euo pipefail
ROOT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)
COORDINATOR_DIR="$ROOT_DIR/coordinator"
COMPOSE_PROJECT=${COMPOSE_PROJECT:-scimesh-two-worker-smoke}
COORDINATOR_PORT=${COORDINATOR_PORT:-18081}
POSTGRES_PORT=${POSTGRES_PORT:-55434}
HOST="http://127.0.0.1:${COORDINATOR_PORT}"
TOKEN=${SCIMESH_SMOKE_TOKEN:-two-worker-smoke-token}
WORK_DIR=$(mktemp -d "${TMPDIR:-/tmp}/scimesh-two-worker-smoke.XXXXXX")
WORKER_ONE_PID=""
WORKER_TWO_PID=""
WORKER_PYTHON=${SCIMESH_WORKER_PYTHON:-"$ROOT_DIR/.venv/bin/python"}
cleanup() {
local exit_code=$?
if [[ "$exit_code" -ne 0 ]]; then
printf '\nTwo-worker smoke failed; worker logs follow.\n' >&2
sed -n '1,200p' "$WORK_DIR/worker-a.log" >&2 || true
sed -n '1,200p' "$WORK_DIR/worker-b.log" >&2 || true
(
cd "$COORDINATOR_DIR"
POSTGRES_PORT="$POSTGRES_PORT" COORDINATOR_PORT="$COORDINATOR_PORT" \
docker compose -p "$COMPOSE_PROJECT" logs coordinator >&2 || true
)
fi
if [[ -n "$WORKER_ONE_PID" ]]; then kill "$WORKER_ONE_PID" 2>/dev/null || true; fi
if [[ -n "$WORKER_TWO_PID" ]]; then kill "$WORKER_TWO_PID" 2>/dev/null || true; fi
if [[ -n "$WORKER_ONE_PID" ]]; then wait "$WORKER_ONE_PID" 2>/dev/null || true; fi
if [[ -n "$WORKER_TWO_PID" ]]; then wait "$WORKER_TWO_PID" 2>/dev/null || true; fi
(
cd "$COORDINATOR_DIR"
POSTGRES_PORT="$POSTGRES_PORT" COORDINATOR_PORT="$COORDINATOR_PORT" \
docker compose -p "$COMPOSE_PROJECT" down -v --remove-orphans >/dev/null 2>&1 || true
)
rm -rf "$WORK_DIR"
exit "$exit_code"
}
trap cleanup EXIT INT TERM
require() {
command -v "$1" >/dev/null || {
printf 'missing required command: %s\n' "$1" >&2
exit 2
}
}
for command in docker curl python3; do require "$command"; done
[[ -x "$WORKER_PYTHON" ]] || {
printf 'worker Python is not executable: %s\n' "$WORKER_PYTHON" >&2
printf 'Set SCIMESH_WORKER_PYTHON to a Python environment with SciMesh and RDKit.\n' >&2
exit 2
}
printf 'Starting isolated coordinator on %s (project %s)\n' "$HOST" "$COMPOSE_PROJECT"
(
cd "$COORDINATOR_DIR"
POSTGRES_PORT="$POSTGRES_PORT" COORDINATOR_PORT="$COORDINATOR_PORT" \
WORKER_AUTH_TOKEN="$TOKEN" UI_AUTH_TOKEN= \
docker compose -p "$COMPOSE_PROJECT" up -d --build
)
for _ in $(seq 1 45); do
if curl -fsS "$HOST/health" >/dev/null; then break; fi
sleep 1
done
curl -fsS "$HOST/health" >/dev/null || {
printf 'coordinator did not become healthy\n' >&2
exit 1
}
start_worker() {
local worker_name=$1
local worker_dir=$2
SCIMESH_COORDINATOR_URL="$HOST" \
SCIMESH_BEARER_TOKEN="$TOKEN" \
SCIMESH_WORKER_NAME="$worker_name" \
SCIMESH_POLL_INTERVAL=0.2 \
"$WORKER_PYTHON" -m scimesh.worker.cli --work-dir "$worker_dir" --max-tasks 2 >"$worker_dir.log" 2>&1 &
STARTED_WORKER_PID=$!
}
start_worker two-worker-smoke-a "$WORK_DIR/worker-a"
WORKER_ONE_PID=$STARTED_WORKER_PID
start_worker two-worker-smoke-b "$WORK_DIR/worker-b"
WORKER_TWO_PID=$STARTED_WORKER_PID
for _ in $(seq 1 30); do
registered=$(docker compose -p "$COMPOSE_PROJECT" -f "$COORDINATOR_DIR/docker-compose.yml" \
exec -T postgres psql -U scimesh -d scimesh -Atc "SELECT count(*) FROM workers" 2>/dev/null || printf '0')
[[ "$registered" == "2" ]] && break
sleep 1
done
[[ "${registered:-0}" == "2" ]] || {
printf 'workers did not register; logs follow\n' >&2
sed -n '1,160p' "$WORK_DIR/worker-a.log" >&2 || true
sed -n '1,160p' "$WORK_DIR/worker-b.log" >&2 || true
exit 1
}
DATASET="$WORK_DIR/fixture.tsv"
printf '%s\n' \
$'chembl_id\tcanonical_smiles' \
$'TEST001\tCC' $'TEST002\tCCC' $'TEST003\tCCCC' $'TEST004\tCCCO' $'TEST005\tCCN' \
$'TEST006\tCCCl' $'TEST007\tCCBr' $'TEST008\tCCF' $'TEST009\tCC=O' $'TEST010\tCC#N' \
$'TEST011\tCO' $'TEST012\tCOC' $'TEST013\tCOCC' $'TEST014\tCN' $'TEST015\tCNC' \
$'TEST016\tO=C=O' $'TEST017\tC1CC1' $'TEST018\tc1ccccc1' $'TEST019\tCC(C)O' $'TEST020\tCC(C)N' \
>"$DATASET"
response=$(curl -fsS -H "Authorization: Bearer $TOKEN" -X POST "$HOST/jobs/upload" \
-F 'workload=similarity-search' \
-F 'parameters={"query_smiles":"CCO","top_k":5,"progress_every":0}' \
-F 'chunk_rows=5' \
-F 'max_rows=20' \
-F "file=@${DATASET};type=text/tab-separated-values")
job_id=$(printf '%s' "$response" | python3 -c 'import json,sys; print(json.load(sys.stdin)["job_id"])')
task_count=$(printf '%s' "$response" | python3 -c 'import json,sys; print(json.load(sys.stdin)["task_count"])')
[[ "$task_count" == "4" ]] || {
printf 'expected four shards, got response: %s\n' "$response" >&2
exit 1
}
printf 'Submitted job %s with four shards\n' "$job_id"
for _ in $(seq 1 90); do
job=$(curl -fsS -H "Authorization: Bearer $TOKEN" "$HOST/jobs/$job_id")
status=$(printf '%s' "$job" | python3 -c 'import json,sys; print(json.load(sys.stdin)["status"])')
[[ "$status" == "completed" || "$status" == "failed" || "$status" == "cancelled" ]] && break
sleep 1
done
printf '%s' "$job" | python3 -c '
import json, sys
job = json.load(sys.stdin)
assert job["status"] == "completed", job
assert job["total"] == 4, job
assert job["completed"] == 4, job
assert job["failed"] == 0, job
'
task_check=$(docker compose -p "$COMPOSE_PROJECT" -f "$COORDINATOR_DIR/docker-compose.yml" \
exec -T postgres psql -U scimesh -d scimesh -Atc \
"SELECT count(*) FROM tasks WHERE job_id = '$job_id'::uuid AND status = 'completed' AND attempt = 1 AND result_artifact_id IS NOT NULL")
[[ "$task_check" == "4" ]] || {
printf 'expected four first-attempt tasks with coordinator artifacts, got %s\n' "$task_check" >&2
exit 1
}
worker_one_results=$(find "$WORK_DIR/worker-a" -name result.csv -type f | wc -l | tr -d ' ')
worker_two_results=$(find "$WORK_DIR/worker-b" -name result.csv -type f | wc -l | tr -d ' ')
[[ "$worker_one_results" -ge 1 && "$worker_two_results" -ge 1 ]] || {
printf 'both workers must process at least one shard (a=%s, b=%s)\n' \
"$worker_one_results" "$worker_two_results" >&2
exit 1
}
printf 'PASS: 4/4 shards completed; worker-a=%s, worker-b=%s\n' \
"$worker_one_results" "$worker_two_results"
+139 -5
View File
@@ -1,7 +1,9 @@
from __future__ import annotations
import hashlib
import logging
from pathlib import Path
import subprocess
import time
from datetime import datetime, timedelta, timezone
from urllib.request import Request
@@ -9,8 +11,10 @@ from urllib.request import Request
import pytest
from scimesh.worker.config import WorkerConfig
from scimesh.worker import cli as worker_cli
from scimesh.worker.cli import build_parser
from scimesh.worker.coordinator import CoordinatorTransientError
from scimesh.worker.daemon import LeaseHeartbeat, WorkerDaemon
from scimesh.worker.daemon import LeaseHeartbeat, RunOnceOutcome, WorkerDaemon
from scimesh.worker.models import (
ClaimedTask,
InputArtifact,
@@ -92,7 +96,7 @@ def daemon(tmp_path: Path, task: ClaimedTask | None, content: bytes):
def test_claims_runs_uploads_and_submits_csv(tmp_path: Path) -> None:
content = b"input fixture"
worker, coordinator, artifacts, runner, _ = daemon(tmp_path, make_task(content), content)
assert worker.run_once() is True
assert worker.run_once() == RunOnceOutcome(claimed=True, completed=True)
assert runner.calls == 1
assert len(artifacts.uploaded) == 1
assert coordinator.heartbeats == [("task-1", 1, "worker-1")]
@@ -104,24 +108,153 @@ def test_claims_runs_uploads_and_submits_csv(tmp_path: Path) -> None:
def test_no_task_does_not_create_directory(tmp_path: Path) -> None:
worker, _, _, runner, config = daemon(tmp_path, None, b"")
assert worker.run_once() is False
assert worker.run_once() == RunOnceOutcome(claimed=False, completed=False)
assert runner.calls == 0
assert not config.work_dir.exists()
def test_once_worker_exits_after_an_empty_claim(tmp_path: Path, caplog: pytest.LogCaptureFixture) -> None:
caplog.set_level(logging.INFO, logger="scimesh.worker")
worker, _, _, runner, _ = daemon(tmp_path, None, b"")
worker.config = WorkerConfig(**{**worker.config.__dict__, "exit_when_idle": True, "max_tasks": 1})
assert worker.run_forever() is True
assert runner.calls == 0
assert "queue_empty" in caplog.text
def test_worker_stops_after_the_configured_number_of_claims(tmp_path: Path, caplog: pytest.LogCaptureFixture) -> None:
caplog.set_level(logging.INFO, logger="scimesh.worker")
content = b"input fixture"
worker, _, _, runner, _ = daemon(tmp_path, make_task(content), content)
worker.config = WorkerConfig(**{**worker.config.__dict__, "max_tasks": 1})
assert worker.run_forever() is True
assert runner.calls == 1
assert "max_tasks_reached" in caplog.text
def test_keyboard_interrupt_stops_worker_without_propagating(tmp_path: Path, caplog: pytest.LogCaptureFixture) -> None:
caplog.set_level(logging.INFO, logger="scimesh.worker")
class InterruptingCoordinator(FakeCoordinator):
def claim(self, worker_id: str, capabilities: tuple[str, ...]) -> ClaimedTask | None:
raise KeyboardInterrupt
worker, _, _, _, _ = daemon(tmp_path, None, b"")
worker.coordinator = InterruptingCoordinator(None)
assert worker.run_forever() is False
assert "interrupted" in caplog.text
def test_interrupting_an_active_task_reports_a_sanitized_failure(tmp_path: Path) -> None:
content = b"input fixture"
worker, coordinator, _, _, _ = daemon(tmp_path, make_task(content), content)
class InterruptingRunner(FakeRunner):
def run(self, task: ClaimedTask, task_dir: Path) -> RunResult:
raise KeyboardInterrupt
worker.runner = InterruptingRunner()
with pytest.raises(KeyboardInterrupt):
worker.run_once()
assert coordinator.failures == [
{
"worker_id": "worker-1",
"attempt": 1,
"error_code": "InterruptedError",
"error_message": "worker interrupted by operator",
}
]
def test_max_tasks_counts_successes_not_failed_claims(tmp_path: Path) -> None:
successful_content = b"successful input"
class SequencedCoordinator(FakeCoordinator):
def __init__(self) -> None:
super().__init__(None)
self.tasks = [
make_task(b"bad input", "wrong-checksum"),
ClaimedTask(
"task-2",
1,
(datetime.now(timezone.utc) + timedelta(seconds=60)).isoformat(),
"similarity-search",
InputArtifact(
"https://example.test/input",
hashlib.sha256(successful_content).hexdigest(),
),
{"query_id": "CHEMBL1"},
),
]
def claim(self, worker_id: str, capabilities: tuple[str, ...]) -> ClaimedTask | None:
return self.tasks.pop(0) if self.tasks else None
coordinator = SequencedCoordinator()
artifacts, runner = FakeArtifacts(successful_content), FakeRunner()
config = WorkerConfig("https://example.test", "worker-1", tmp_path / "work", max_tasks=1)
worker = WorkerDaemon(config, coordinator, artifacts, runner)
assert worker.run_forever() is True
assert len(coordinator.failures) == 1
assert len(coordinator.submissions) == 1
assert runner.calls == 1
def test_worker_cli_lifecycle_options_are_explicit_and_exclusive() -> None:
parser = build_parser()
assert parser.parse_args(["--once"]).once is True
assert parser.parse_args(["--max-tasks", "2"]).max_tasks == 2
with pytest.raises(SystemExit):
parser.parse_args(["--once", "--max-tasks", "2"])
def test_worker_cli_uses_a_nonzero_exit_code_for_interruption(
monkeypatch: pytest.MonkeyPatch, tmp_path: Path
) -> None:
class InterruptedDaemon:
def __init__(self, *_: object) -> None:
pass
def run_forever(self) -> bool:
return False
monkeypatch.setattr(worker_cli, "WorkerDaemon", InterruptedDaemon)
assert worker_cli.main(
["--coordinator-url", "https://example.test", "--work-dir", str(tmp_path)]
) == 130
@pytest.mark.parametrize("value", [0, -1, True])
def test_max_tasks_must_be_positive(value: object, tmp_path: Path) -> None:
with pytest.raises(ValueError, match="max_tasks"):
WorkerConfig("https://example.test", None, tmp_path, max_tasks=value) # type: ignore[arg-type]
def test_bad_checksum_reports_failure_without_running(tmp_path: Path) -> None:
worker, coordinator, _, runner, _ = daemon(tmp_path, make_task(b"actual", "not-the-hash"), b"actual")
assert worker.run_once() is True
assert worker.run_once() == RunOnceOutcome(claimed=True, completed=False)
assert runner.calls == 0
assert coordinator.failures[0]["error_code"] == "ValueError"
assert not coordinator.submissions
def test_failure_reporting_removes_paths_outside_the_worker_directory(tmp_path: Path) -> None:
worker, coordinator, _, _, _ = daemon(tmp_path, make_task(b"input"), b"input")
error = subprocess.CalledProcessError(
1,
["/home/alice/.venv/bin/python", "-m", "scimesh.cli", "/private/input.tsv"],
)
worker._report_failure(make_task(b"input"), error)
message = coordinator.failures[0]["error_message"]
assert "/home/alice" not in message
assert "/private/input.tsv" not in message
assert "<path>" in message
def test_directory_creation_failure_is_reported(tmp_path: Path) -> None:
content = b"input fixture"
worker, coordinator, _, _, config = daemon(tmp_path, make_task(content), content)
(config.work_dir / "task-1" / "1").mkdir(parents=True)
assert worker.run_once() is True
assert worker.run_once() == RunOnceOutcome(claimed=True, completed=False)
assert coordinator.failures[0]["error_code"] == "FileExistsError"
@@ -312,6 +445,7 @@ def test_environment_overrides_allow_cli_only_configuration(monkeypatch: pytest.
assert config.worker_id is None
assert "similarity-search" in config.capabilities
assert "similarity_search" in config.capabilities
assert "similarity-graph" not in config.capabilities
def test_relative_work_dir_is_normalized_for_runner_subprocesses(