Merge distributed pipeline hardening
This commit is contained in:
@@ -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
|
||||||
@@ -11,3 +11,7 @@ results/
|
|||||||
*_similarities.csv
|
*_similarities.csv
|
||||||
test_results.csv
|
test_results.csv
|
||||||
test_structures/
|
test_structures/
|
||||||
|
|
||||||
|
# Local coordinator-worker execution state
|
||||||
|
worker-data*/
|
||||||
|
scimesh-worker-data/
|
||||||
|
|||||||
@@ -1,11 +1,12 @@
|
|||||||
# SciMesh
|
# SciMesh
|
||||||
|
|
||||||
SciMesh is a scientific-workload framework for molecular datasets. Its public CLI
|
SciMesh is a scientific-workload framework for molecular datasets. Its public CLI
|
||||||
currently runs exact similarity search and sparse similarity-graph construction
|
runs exact similarity search and sparse similarity-graph construction locally in
|
||||||
locally in one Python process; it creates no dense similarity matrix. A Python
|
one Python process; it creates no dense similarity matrix. The Go/PostgreSQL
|
||||||
Worker client and the planned Go/PostgreSQL coordinator contract are tracked in
|
coordinator and Python worker can run a diagnostic, shard-based
|
||||||
the repository, but distributed execution is not available yet; see
|
`similarity-search` pipeline locally. Its CSV artifacts are not a global result
|
||||||
[`STATUS.md`](STATUS.md).
|
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`.
|
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`.
|
||||||
|
|
||||||
|
|||||||
@@ -36,7 +36,7 @@ Docker PostgreSQL stack on 2026-07-23.
|
|||||||
| CTX-08 Distributed similarity-search | Not started | Local reference exists. |
|
| 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-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-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. |
|
| CTX-12 Reliability, security, CI | In progress | Unit, race, PostgreSQL integration, and smoke checks exist; CI hardening remains. |
|
||||||
|
|
||||||
## Next recommended assignment
|
## 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.
|
names and hyphenated CLI names while the contract is consolidated.
|
||||||
- A real-stack worker test uses a small `query_smiles` shard. Resolving a
|
- 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.
|
`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
|
## Update rule
|
||||||
|
|
||||||
|
|||||||
@@ -71,14 +71,14 @@ func run() error {
|
|||||||
useCases := httptransport.UseCases{
|
useCases := httptransport.UseCases{
|
||||||
RegisterWorker: usecase.NewRegisterWorker(workerRepo, clk),
|
RegisterWorker: usecase.NewRegisterWorker(workerRepo, clk),
|
||||||
CreateJob: usecase.NewCreateJob(jobRepo, taskRepo, tx, clk),
|
CreateJob: usecase.NewCreateJob(jobRepo, taskRepo, tx, clk),
|
||||||
SubmitDataset: usecase.NewSubmitDataset(blobStore, artifactRepo, jobRepo, taskRepo, tx, clk),
|
SubmitDataset: usecase.NewSubmitDataset(blobStore, artifactRepo, jobRepo, taskRepo, tx, clk, cfg.DefaultMaxAttempts),
|
||||||
ClaimTask: usecase.NewClaimTask(taskRepo, clk, cfg.LeaseDuration),
|
ClaimTask: usecase.NewClaimTask(taskRepo, jobRepo, workerRepo, tx, clk, cfg.LeaseDuration),
|
||||||
RenewLease: usecase.NewRenewLease(taskRepo, workerRepo, tx, clk, cfg.LeaseDuration),
|
RenewLease: usecase.NewRenewLease(taskRepo, workerRepo, tx, clk, cfg.LeaseDuration),
|
||||||
CompleteTask: usecase.NewCompleteTask(taskRepo, jobRepo, artifactRepo, tx, clk),
|
CompleteTask: usecase.NewCompleteTask(taskRepo, jobRepo, artifactRepo, tx, clk),
|
||||||
FailTask: usecase.NewFailTask(taskRepo, jobRepo, tx, clk),
|
FailTask: usecase.NewFailTask(taskRepo, jobRepo, tx, clk),
|
||||||
GetJobStatus: usecase.NewGetJobStatus(jobRepo, taskRepo),
|
GetJobStatus: usecase.NewGetJobStatus(jobRepo, taskRepo),
|
||||||
CancelJob: usecase.NewCancelJob(jobRepo, taskRepo, tx, clk),
|
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),
|
DownloadArtifact: usecase.NewDownloadArtifact(artifactRepo, blobStore),
|
||||||
GetTaskInput: usecase.NewGetTaskInput(taskRepo, artifactRepo, blobStore),
|
GetTaskInput: usecase.NewGetTaskInput(taskRepo, artifactRepo, blobStore),
|
||||||
Dashboard: usecase.NewDashboard(uiReadRepo),
|
Dashboard: usecase.NewDashboard(uiReadRepo),
|
||||||
@@ -87,7 +87,7 @@ func run() error {
|
|||||||
// Background reapers are tracked so shutdown can wait for them. Without this
|
// 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
|
// the process would exit mid-UPDATE, and the deferred pool.Close() would pull
|
||||||
// connections out from under them.
|
// connections out from under them.
|
||||||
expireLeases := usecase.NewExpireLeases(taskRepo, clk)
|
expireLeases := usecase.NewExpireLeases(taskRepo, jobRepo, tx, clk)
|
||||||
markOffline := usecase.NewMarkWorkersOffline(workerRepo, clk, cfg.WorkerOfflineAfter)
|
markOffline := usecase.NewMarkWorkersOffline(workerRepo, clk, cfg.WorkerOfflineAfter)
|
||||||
|
|
||||||
var wg sync.WaitGroup
|
var wg sync.WaitGroup
|
||||||
|
|||||||
@@ -7,12 +7,19 @@ import (
|
|||||||
"bytes"
|
"bytes"
|
||||||
"fmt"
|
"fmt"
|
||||||
"io"
|
"io"
|
||||||
|
"strings"
|
||||||
)
|
)
|
||||||
|
|
||||||
// ErrNoRows is returned when the input has a header but no data rows: a job with
|
// 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.
|
// zero tasks could never complete, so it is rejected at the source.
|
||||||
var ErrNoRows = fmt.Errorf("input has no data rows")
|
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
|
// 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
|
// 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
|
// 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
|
// 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.
|
// 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 {
|
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.
|
// 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,
|
// A maxRows value of zero means unlimited. This lets an operator make a small,
|
||||||
// representative pipeline check without materialising a second dataset file.
|
// 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 {
|
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 {
|
if rowsPerShard <= 0 {
|
||||||
return fmt.Errorf("rowsPerShard must be positive, got %d", rowsPerShard)
|
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
|
return ErrNoRows // completely empty input
|
||||||
}
|
}
|
||||||
header := append([]byte(nil), sc.Bytes()...)
|
header := append([]byte(nil), sc.Bytes()...)
|
||||||
|
if validateHeader != nil {
|
||||||
|
if err := validateHeader(header); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
var (
|
var (
|
||||||
buf bytes.Buffer
|
buf bytes.Buffer
|
||||||
@@ -71,9 +95,15 @@ func SplitTSVLimit(r io.Reader, rowsPerShard, maxRows int, emit func(index int,
|
|||||||
|
|
||||||
for sc.Scan() {
|
for sc.Scan() {
|
||||||
if rows == 0 {
|
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.Write(header)
|
||||||
buf.WriteByte('\n')
|
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.Write(sc.Bytes())
|
||||||
buf.WriteByte('\n')
|
buf.WriteByte('\n')
|
||||||
rows++
|
rows++
|
||||||
@@ -103,3 +133,17 @@ func SplitTSVLimit(r io.Reader, rowsPerShard, maxRows int, emit func(index int,
|
|||||||
}
|
}
|
||||||
return nil
|
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
|
||||||
|
}
|
||||||
|
|||||||
@@ -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
|
// 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.
|
// later row would corrupt an earlier one. This guards that copy.
|
||||||
func TestSplitDoesNotAliasScannerBuffer(t *testing.T) {
|
func TestSplitDoesNotAliasScannerBuffer(t *testing.T) {
|
||||||
|
|||||||
@@ -129,6 +129,9 @@ func LoadConfig() (Config, error) {
|
|||||||
if cfg.DefaultMaxAttempts, err = getEnvInt("DEFAULT_MAX_ATTEMPTS", cfg.DefaultMaxAttempts); err != nil {
|
if cfg.DefaultMaxAttempts, err = getEnvInt("DEFAULT_MAX_ATTEMPTS", cfg.DefaultMaxAttempts); err != nil {
|
||||||
return Config{}, err
|
return Config{}, err
|
||||||
}
|
}
|
||||||
|
if cfg.DefaultMaxAttempts < 1 {
|
||||||
|
return Config{}, fmt.Errorf("DEFAULT_MAX_ATTEMPTS must be positive")
|
||||||
|
}
|
||||||
|
|
||||||
return cfg, nil
|
return cfg, nil
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -32,3 +32,14 @@ func TestLoadConfigAllowsDistinctUIAndWorkerTokens(t *testing.T) {
|
|||||||
t.Fatalf("unexpected tokens: %+v", cfg)
|
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)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -160,17 +160,18 @@ func (r *TaskRepo) CancelByJob(_ context.Context, jobID uuid.UUID, now time.Time
|
|||||||
return cancelled, nil
|
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()
|
r.mu.Lock()
|
||||||
defer r.mu.Unlock()
|
defer r.mu.Unlock()
|
||||||
var n int64
|
affected := make([]uuid.UUID, 0)
|
||||||
for _, t := range r.tasks {
|
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)
|
t.ExpireLease(now)
|
||||||
n++
|
affected = append(affected, t.JobID)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
return n, nil
|
return affected, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
// --- JobRepo -------------------------------------------------------------
|
// --- JobRepo -------------------------------------------------------------
|
||||||
@@ -298,6 +299,23 @@ func (r *ArtifactRepo) Get(ctx context.Context, id uuid.UUID) (*domain.Artifact,
|
|||||||
return &cp, nil
|
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 -----------------------------------------------------------
|
// --- BlobStore -----------------------------------------------------------
|
||||||
|
|
||||||
type BlobStore struct {
|
type BlobStore struct {
|
||||||
|
|||||||
@@ -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 })
|
sort.Slice(out, func(i, j int) bool { return out[i].ChunkIndex < out[j].ChunkIndex })
|
||||||
return out, nil
|
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) {
|
func (r *UIReadRepo) ListWorkers(_ context.Context, limit int) ([]domain.Worker, error) {
|
||||||
if limit < 1 || limit > 100 {
|
if limit < 1 || limit > 100 {
|
||||||
return nil, domain.ErrInvalidInput
|
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)
|
a.Kind = domain.ArtifactKind(kind)
|
||||||
return &a, nil
|
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)
|
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
|
func TestPartialResultIsUniquePerTaskAttempt(t *testing.T) {
|
||||||
other.ResultArtifactID = art2.ID
|
pool := testPool(t)
|
||||||
if _, err := uc.Execute(ctx, other); !errors.Is(err, domain.ErrResultConflict) {
|
ctx := context.Background()
|
||||||
t.Errorf("err = %v, want ErrResultConflict", err)
|
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,
|
ELSE error_message END,
|
||||||
completed_at = CASE WHEN attempt >= max_attempts THEN $1 ELSE completed_at END,
|
completed_at = CASE WHEN attempt >= max_attempts THEN $1 ELSE completed_at END,
|
||||||
version = version + 1
|
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) {
|
func (r *TaskRepo) ExpireLeases(ctx context.Context, now time.Time) ([]uuid.UUID, error) {
|
||||||
var affected int64
|
var affected []uuid.UUID
|
||||||
err := withRetry(ctx, func(ctx context.Context) error {
|
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 {
|
if err != nil {
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
affected = tag.RowsAffected()
|
defer rows.Close()
|
||||||
return nil
|
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
|
return affected, err
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -75,6 +75,31 @@ func (r *UIReadRepo) ListTasksByJob(ctx context.Context, jobID uuid.UUID) ([]dom
|
|||||||
return tasks, rows.Err()
|
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) {
|
func (r *UIReadRepo) ListWorkers(ctx context.Context, limit int) ([]domain.Worker, error) {
|
||||||
if limit < 1 || limit > 100 {
|
if limit < 1 || limit > 100 {
|
||||||
return nil, domain.ErrInvalidInput
|
return nil, domain.ErrInvalidInput
|
||||||
|
|||||||
@@ -79,6 +79,10 @@ func (s *Server) handleClaim(w http.ResponseWriter, r *http.Request) {
|
|||||||
s.writeError(w, r, domain.ErrInvalidInput)
|
s.writeError(w, r, domain.ErrInvalidInput)
|
||||||
return
|
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{
|
claimed, err := s.uc.ClaimTask.Execute(ctx, usecase.ClaimTaskInput{
|
||||||
WorkerID: req.WorkerID,
|
WorkerID: req.WorkerID,
|
||||||
|
|||||||
@@ -23,8 +23,9 @@ const token = "secret"
|
|||||||
const uiToken = "ui-secret"
|
const uiToken = "ui-secret"
|
||||||
|
|
||||||
type env struct {
|
type env struct {
|
||||||
ts *httptest.Server
|
ts *httptest.Server
|
||||||
blobs *memstore.BlobStore
|
blobs *memstore.BlobStore
|
||||||
|
workerID string
|
||||||
}
|
}
|
||||||
|
|
||||||
func newEnv(t *testing.T, ready func(context.Context) error) *env {
|
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{
|
uc := coordhttp.UseCases{
|
||||||
RegisterWorker: usecase.NewRegisterWorker(work, clk),
|
RegisterWorker: usecase.NewRegisterWorker(work, clk),
|
||||||
CreateJob: usecase.NewCreateJob(jobs, tasks, tx, clk),
|
CreateJob: usecase.NewCreateJob(jobs, tasks, tx, clk),
|
||||||
SubmitDataset: usecase.NewSubmitDataset(blobs, arts, jobs, tasks, tx, clk),
|
SubmitDataset: usecase.NewSubmitDataset(blobs, arts, jobs, tasks, tx, clk, 3),
|
||||||
ClaimTask: usecase.NewClaimTask(tasks, clk, lease),
|
ClaimTask: usecase.NewClaimTask(tasks, jobs, work, tx, clk, lease),
|
||||||
RenewLease: usecase.NewRenewLease(tasks, work, tx, clk, lease),
|
RenewLease: usecase.NewRenewLease(tasks, work, tx, clk, lease),
|
||||||
CompleteTask: usecase.NewCompleteTask(tasks, jobs, arts, tx, clk),
|
CompleteTask: usecase.NewCompleteTask(tasks, jobs, arts, tx, clk),
|
||||||
FailTask: usecase.NewFailTask(tasks, jobs, tx, clk),
|
FailTask: usecase.NewFailTask(tasks, jobs, tx, clk),
|
||||||
GetJobStatus: usecase.NewGetJobStatus(jobs, tasks),
|
GetJobStatus: usecase.NewGetJobStatus(jobs, tasks),
|
||||||
CancelJob: usecase.NewCancelJob(jobs, tasks, tx, clk),
|
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),
|
DownloadArtifact: usecase.NewDownloadArtifact(arts, blobs),
|
||||||
GetTaskInput: usecase.NewGetTaskInput(tasks, arts, blobs),
|
GetTaskInput: usecase.NewGetTaskInput(tasks, arts, blobs),
|
||||||
Dashboard: usecase.NewDashboard(memstore.NewUIReadRepo(jobs, tasks, work, arts)),
|
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)
|
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))
|
ts := httptest.NewServer(srv.Handler(token, configuredUIToken))
|
||||||
t.Cleanup(ts.Close)
|
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 }
|
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.
|
// 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) {
|
func (e *env) do(t *testing.T, method, path, body string) (int, map[string]any) {
|
||||||
t.Helper()
|
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, _ := http.NewRequestWithContext(context.Background(), method, e.ts.URL+path, strings.NewReader(body))
|
||||||
req.Header.Set("Authorization", "Bearer "+token)
|
req.Header.Set("Authorization", "Bearer "+token)
|
||||||
if body != "" {
|
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) {
|
func TestFullLifecycle(t *testing.T) {
|
||||||
e := newEnv(t, healthy)
|
e := newEnv(t, healthy)
|
||||||
|
|
||||||
@@ -393,9 +421,9 @@ func TestForeignArtifactResultConflict(t *testing.T) {
|
|||||||
|
|
||||||
func TestUploadDatasetChunksAndServesInput(t *testing.T) {
|
func TestUploadDatasetChunksAndServesInput(t *testing.T) {
|
||||||
e := newEnv(t, healthy)
|
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 {
|
if code != 201 {
|
||||||
t.Fatalf("upload: status = %d", code)
|
t.Fatalf("upload: status = %d", code)
|
||||||
}
|
}
|
||||||
@@ -421,7 +449,7 @@ func TestUploadDatasetChunksAndServesInput(t *testing.T) {
|
|||||||
t.Fatalf("get input: status = %d", resp.StatusCode)
|
t.Fatalf("get input: status = %d", resp.StatusCode)
|
||||||
}
|
}
|
||||||
shard, _ := io.ReadAll(resp.Body)
|
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)
|
t.Errorf("shard missing header: %q", shard)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -430,11 +458,12 @@ func TestUploadDatasetLimitsRows(t *testing.T) {
|
|||||||
e := newEnv(t, healthy)
|
e := newEnv(t, healthy)
|
||||||
var buf bytes.Buffer
|
var buf bytes.Buffer
|
||||||
mw := multipart.NewWriter(&buf)
|
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("chunk_rows", "2")
|
||||||
_ = mw.WriteField("max_rows", "3")
|
_ = mw.WriteField("max_rows", "3")
|
||||||
fw, _ := mw.CreateFormFile("file", "chembl.tsv")
|
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()
|
_ = mw.Close()
|
||||||
req, _ := http.NewRequestWithContext(context.Background(), "POST", e.ts.URL+"/jobs/upload", &buf)
|
req, _ := http.NewRequestWithContext(context.Background(), "POST", e.ts.URL+"/jobs/upload", &buf)
|
||||||
req.Header.Set("Authorization", "Bearer "+token)
|
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) {
|
func TestErrorMappings(t *testing.T) {
|
||||||
e := newEnv(t, healthy)
|
e := newEnv(t, healthy)
|
||||||
zero := "00000000-0000-0000-0000-000000000000"
|
zero := "00000000-0000-0000-0000-000000000000"
|
||||||
@@ -478,10 +530,11 @@ func TestUploadDatasetRejectsAmbiguousMultipartInput(t *testing.T) {
|
|||||||
e := newEnv(t, healthy)
|
e := newEnv(t, healthy)
|
||||||
var buf bytes.Buffer
|
var buf bytes.Buffer
|
||||||
mw := multipart.NewWriter(&buf)
|
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")
|
_ = mw.WriteField("chunk_rows", "not-a-number")
|
||||||
fw, _ := mw.CreateFormFile("file", "chembl.tsv")
|
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()
|
_ = mw.Close()
|
||||||
|
|
||||||
req, _ := http.NewRequestWithContext(context.Background(), "POST", e.ts.URL+"/jobs/upload", &buf)
|
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))
|
e.ts.URL+"/tasks/"+taskID+"/artifacts/r.csv", strings.NewReader(data))
|
||||||
req.Header.Set("Authorization", "Bearer "+token)
|
req.Header.Set("Authorization", "Bearer "+token)
|
||||||
req.Header.Set("Content-Type", "text/csv")
|
req.Header.Set("Content-Type", "text/csv")
|
||||||
|
if worker == "w1" {
|
||||||
|
worker = e.workerID
|
||||||
|
}
|
||||||
req.Header.Set("X-Worker-ID", worker)
|
req.Header.Set("X-Worker-ID", worker)
|
||||||
req.Header.Set("X-Task-Attempt", itoa(attempt))
|
req.Header.Set("X-Task-Attempt", itoa(attempt))
|
||||||
resp, err := http.DefaultClient.Do(req)
|
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
|
var buf bytes.Buffer
|
||||||
mw := multipart.NewWriter(&buf)
|
mw := multipart.NewWriter(&buf)
|
||||||
_ = mw.WriteField("workload", workload)
|
_ = mw.WriteField("workload", workload)
|
||||||
|
_ = mw.WriteField("parameters", `{"query_smiles":"CCO"}`)
|
||||||
_ = mw.WriteField("chunk_rows", itoa(rows))
|
_ = mw.WriteField("chunk_rows", itoa(rows))
|
||||||
fw, _ := mw.CreateFormFile("file", "chembl.tsv")
|
fw, _ := mw.CreateFormFile("file", "chembl.tsv")
|
||||||
_, _ = io.Copy(fw, strings.NewReader(tsv))
|
_, _ = io.Copy(fw, strings.NewReader(tsv))
|
||||||
|
|||||||
@@ -22,7 +22,7 @@
|
|||||||
<script>
|
<script>
|
||||||
const id={{printf "%q" .ID}},state={pending:['Waiting for a worker','waiting','Waiting for an available worker with the required capability.'],leased:['Assigned to a worker','active','A worker has claimed the task and should begin processing shortly.'],running:['Running','active','A worker is reading a shard, calculating fingerprints, and uploading its result through the coordinator.'],completed:['Tasks complete','success','Every shard task is complete. Files below are still partial results.'],failed:['Needs attention','danger','One or more shard tasks failed. Open the task list below for details.'],cancelled:['Stopped','waiting','The operator stopped this job. No new shards can be claimed.']};
|
const id={{printf "%q" .ID}},state={pending:['Waiting for a worker','waiting','Waiting for an available worker with the required capability.'],leased:['Assigned to a worker','active','A worker has claimed the task and should begin processing shortly.'],running:['Running','active','A worker is reading a shard, calculating fingerprints, and uploading its result through the coordinator.'],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()});
|
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>
|
</script>
|
||||||
</body>
|
</body>
|
||||||
</html>
|
</html>
|
||||||
|
|||||||
@@ -14,12 +14,13 @@ type UploadArtifact struct {
|
|||||||
tasks TaskRepository
|
tasks TaskRepository
|
||||||
artifacts ArtifactRepository
|
artifacts ArtifactRepository
|
||||||
blobs BlobStore
|
blobs BlobStore
|
||||||
|
tx TxManager
|
||||||
clk Clock
|
clk Clock
|
||||||
}
|
}
|
||||||
|
|
||||||
func NewUploadArtifact(tasks TaskRepository, artifacts ArtifactRepository,
|
func NewUploadArtifact(tasks TaskRepository, artifacts ArtifactRepository,
|
||||||
blobs BlobStore, clk Clock) *UploadArtifact {
|
blobs BlobStore, tx TxManager, clk Clock) *UploadArtifact {
|
||||||
return &UploadArtifact{tasks: tasks, artifacts: artifacts, blobs: blobs, clk: clk}
|
return &UploadArtifact{tasks: tasks, artifacts: artifacts, blobs: blobs, tx: tx, clk: clk}
|
||||||
}
|
}
|
||||||
|
|
||||||
func (uc *UploadArtifact) Execute(ctx context.Context, in UploadArtifactInput) (*domain.Artifact, error) {
|
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()) {
|
if !task.IsLeaseHeldBy(in.WorkerID, in.Attempt, uc.clk.Now()) {
|
||||||
return nil, domain.ErrLeaseConflict
|
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
|
taskID := task.ID
|
||||||
art, err := domain.NewArtifact(task.JobID, &taskID, domain.ArtifactPartialResult,
|
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)
|
art.SetContent(sum, size)
|
||||||
|
|
||||||
// The stream may take longer than the lease. Re-check after it finishes so
|
// The stream may take longer than the lease. Lock the task while re-checking
|
||||||
// an expired worker cannot leave a durable result record behind. Completion
|
// ownership and inserting metadata: completion or another upload cannot race
|
||||||
// performs the same ownership check under its transaction.
|
// this final decision. The database unique index is a second line of defence.
|
||||||
current, err := uc.tasks.Get(ctx, in.TaskID)
|
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 {
|
if err != nil {
|
||||||
_ = uc.blobs.Delete(ctx, art.StorageKey)
|
_ = uc.blobs.Delete(ctx, art.StorageKey)
|
||||||
return nil, err
|
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)
|
_ = uc.blobs.Delete(ctx, art.StorageKey)
|
||||||
return nil, domain.ErrLeaseConflict
|
|
||||||
}
|
}
|
||||||
|
return durable, nil
|
||||||
// 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
|
|
||||||
}
|
}
|
||||||
|
|
||||||
// DownloadArtifact returns an artifact's metadata together with a reader over
|
// DownloadArtifact returns an artifact's metadata together with a reader over
|
||||||
|
|||||||
@@ -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
|
// The all-or-none guarantee comes from TxManager: a half-created job would
|
||||||
// leave chunks no worker could ever complete.
|
// leave chunks no worker could ever complete.
|
||||||
func (uc *CreateJob) Execute(ctx context.Context, in CreateJobInput) (*domain.Job, error) {
|
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))
|
chunks := make([]domain.ChunkSpec, 0, len(in.Chunks))
|
||||||
for _, c := range in.Chunks {
|
for _, c := range in.Chunks {
|
||||||
chunks = append(chunks, domain.ChunkSpec(c))
|
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 {
|
if job.Status == domain.JobCompleted || job.Status == domain.JobFailed {
|
||||||
return domain.ErrJobNotCancellable
|
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)
|
cancelled, err = uc.tasks.CancelByJob(ctx, jobID, now)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return err
|
return err
|
||||||
@@ -212,3 +234,18 @@ func syncJobStatus(ctx context.Context, jobs JobRepository, tasks TaskRepository
|
|||||||
}
|
}
|
||||||
return jobs.UpdateStatus(ctx, jobID, status, completedAt)
|
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
|
||||||
|
}
|
||||||
|
|||||||
@@ -60,9 +60,9 @@ type TaskRepository interface {
|
|||||||
// lease. It returns how many tasks changed.
|
// lease. It returns how many tasks changed.
|
||||||
CancelByJob(ctx context.Context, jobID uuid.UUID, now time.Time) (int64, error)
|
CancelByJob(ctx context.Context, jobID uuid.UUID, now time.Time) (int64, error)
|
||||||
|
|
||||||
// ExpireLeases applies the lease-expiry rule to every elapsed task and
|
// ExpireLeases applies the lease-expiry rule to every elapsed task and returns
|
||||||
// reports how many were affected.
|
// the distinct jobs whose aggregate status may have changed.
|
||||||
ExpireLeases(ctx context.Context, now time.Time) (int64, error)
|
ExpireLeases(ctx context.Context, now time.Time) ([]uuid.UUID, error)
|
||||||
}
|
}
|
||||||
|
|
||||||
// JobRepository persists jobs.
|
// JobRepository persists jobs.
|
||||||
@@ -89,6 +89,9 @@ type WorkerRepository interface {
|
|||||||
type ArtifactRepository interface {
|
type ArtifactRepository interface {
|
||||||
Insert(ctx context.Context, a *domain.Artifact) error
|
Insert(ctx context.Context, a *domain.Artifact) error
|
||||||
Get(ctx context.Context, id uuid.UUID) (*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
|
// BlobStore holds artifact bytes, addressed by an opaque storage key. It streams
|
||||||
|
|||||||
@@ -21,12 +21,15 @@ import (
|
|||||||
|
|
||||||
type ClaimTask struct {
|
type ClaimTask struct {
|
||||||
tasks TaskRepository
|
tasks TaskRepository
|
||||||
|
jobs JobRepository
|
||||||
|
workers WorkerRepository
|
||||||
|
tx TxManager
|
||||||
clock Clock
|
clock Clock
|
||||||
leaseDuration time.Duration
|
leaseDuration time.Duration
|
||||||
}
|
}
|
||||||
|
|
||||||
func NewClaimTask(tasks TaskRepository, clock Clock, leaseDuration time.Duration) *ClaimTask {
|
func NewClaimTask(tasks TaskRepository, jobs JobRepository, workers WorkerRepository, tx TxManager, clock Clock, leaseDuration time.Duration) *ClaimTask {
|
||||||
return &ClaimTask{tasks: tasks, clock: clock, leaseDuration: leaseDuration}
|
return &ClaimTask{tasks: tasks, jobs: jobs, workers: workers, tx: tx, clock: clock, leaseDuration: leaseDuration}
|
||||||
}
|
}
|
||||||
|
|
||||||
// Execute reclaims elapsed leases first, then hands out one task.
|
// 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 == "" {
|
if in.WorkerID == "" {
|
||||||
return nil, domain.ErrInvalidInput
|
return nil, domain.ErrInvalidInput
|
||||||
}
|
}
|
||||||
now := uc.clock.Now()
|
workloads := in.Workloads
|
||||||
|
if workerID, err := uuid.Parse(in.WorkerID); err == nil {
|
||||||
if _, err := uc.tasks.ExpireLeases(ctx, now); err != nil {
|
worker, err := uc.workers.Get(ctx, workerID)
|
||||||
return nil, err
|
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{
|
task, err := uc.tasks.ClaimNext(ctx, ClaimFilter{
|
||||||
Workloads: in.Workloads,
|
Workloads: workloads,
|
||||||
Owner: in.WorkerID,
|
Owner: in.WorkerID,
|
||||||
Now: now,
|
Now: now,
|
||||||
LeaseUntil: now.Add(uc.leaseDuration),
|
LeaseUntil: now.Add(uc.leaseDuration),
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
if task != nil {
|
||||||
|
value := task.AsClaimed()
|
||||||
|
claimed = &value
|
||||||
|
}
|
||||||
|
return nil
|
||||||
})
|
})
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return nil, err
|
return nil, err
|
||||||
}
|
}
|
||||||
if task == nil {
|
return claimed, nil // nil means an empty queue
|
||||||
return nil, nil // empty queue is a normal state, not an error
|
|
||||||
}
|
|
||||||
|
|
||||||
claimed := task.AsClaimed()
|
|
||||||
return &claimed, nil
|
|
||||||
}
|
}
|
||||||
|
|
||||||
// --- RenewLease ----------------------------------------------------------
|
// --- RenewLease ----------------------------------------------------------
|
||||||
@@ -230,18 +252,30 @@ func (uc *FailTask) Execute(ctx context.Context, in FailTaskInput) (*domain.Task
|
|||||||
|
|
||||||
type ExpireLeases struct {
|
type ExpireLeases struct {
|
||||||
tasks TaskRepository
|
tasks TaskRepository
|
||||||
|
jobs JobRepository
|
||||||
|
tx TxManager
|
||||||
clock Clock
|
clock Clock
|
||||||
}
|
}
|
||||||
|
|
||||||
func NewExpireLeases(tasks TaskRepository, clock Clock) *ExpireLeases {
|
func NewExpireLeases(tasks TaskRepository, jobs JobRepository, tx TxManager, clock Clock) *ExpireLeases {
|
||||||
return &ExpireLeases{tasks: tasks, clock: clock}
|
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:
|
// 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
|
// several coordinators run it concurrently, and a single atomic UPDATE makes
|
||||||
// the duplicate work harmless — the loser simply updates 0 rows.
|
// the duplicate work harmless — the loser simply updates 0 rows.
|
||||||
func (uc *ExpireLeases) Execute(ctx context.Context) (int64, error) {
|
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
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -15,6 +15,7 @@ type UIReadRepository interface {
|
|||||||
GetJob(ctx context.Context, jobID uuid.UUID) (*domain.Job, error)
|
GetJob(ctx context.Context, jobID uuid.UUID) (*domain.Job, error)
|
||||||
ListJobs(ctx context.Context, limit int) ([]domain.Job, error)
|
ListJobs(ctx context.Context, limit int) ([]domain.Job, error)
|
||||||
ListTasksByJob(ctx context.Context, jobID uuid.UUID) ([]domain.Task, 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)
|
ListWorkers(ctx context.Context, limit int) ([]domain.Worker, error)
|
||||||
ListArtifactsByJob(ctx context.Context, jobID uuid.UUID) ([]domain.Artifact, 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
|
return DashboardView{}, err
|
||||||
}
|
}
|
||||||
out := DashboardView{Jobs: make([]JobCard, 0, len(jobs)), Workers: make([]WorkerCard, 0, len(workers))}
|
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 {
|
for _, job := range jobs {
|
||||||
tasks, err := d.read.ListTasksByJob(ctx, job.ID)
|
jobIDs = append(jobIDs, job.ID)
|
||||||
if err != nil {
|
}
|
||||||
return DashboardView{}, err
|
tasksByJob, err := d.read.ListTasksByJobs(ctx, jobIDs)
|
||||||
}
|
if err != nil {
|
||||||
out.Jobs = append(out.Jobs, jobCard(job, tasks))
|
return DashboardView{}, err
|
||||||
|
}
|
||||||
|
for _, job := range jobs {
|
||||||
|
out.Jobs = append(out.Jobs, jobCard(job, tasksByJob[job.ID]))
|
||||||
}
|
}
|
||||||
for _, worker := range workers {
|
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})
|
out.Workers = append(out.Workers, WorkerCard{ID: worker.ID.String(), Name: worker.Name, Status: string(worker.Status), Capabilities: worker.Capabilities, LastHeartbeatAt: worker.LastHeartbeatAt})
|
||||||
|
|||||||
@@ -4,6 +4,7 @@ import (
|
|||||||
"context"
|
"context"
|
||||||
"fmt"
|
"fmt"
|
||||||
"io"
|
"io"
|
||||||
|
"math"
|
||||||
|
|
||||||
"github.com/google/uuid"
|
"github.com/google/uuid"
|
||||||
|
|
||||||
@@ -15,20 +16,27 @@ import (
|
|||||||
// creates the job with one task per shard — the coordinator-side counterpart of
|
// creates the job with one task per shard — the coordinator-side counterpart of
|
||||||
// a client submitting pre-chunked URIs.
|
// a client submitting pre-chunked URIs.
|
||||||
type SubmitDataset struct {
|
type SubmitDataset struct {
|
||||||
blobs BlobStore
|
blobs BlobStore
|
||||||
artifacts ArtifactRepository
|
artifacts ArtifactRepository
|
||||||
jobs JobRepository
|
jobs JobRepository
|
||||||
tasks TaskRepository
|
tasks TaskRepository
|
||||||
tx TxManager
|
tx TxManager
|
||||||
clk Clock
|
clk Clock
|
||||||
|
maxAttempts int
|
||||||
}
|
}
|
||||||
|
|
||||||
func NewSubmitDataset(blobs BlobStore, artifacts ArtifactRepository, jobs JobRepository,
|
func NewSubmitDataset(blobs BlobStore, artifacts ArtifactRepository, jobs JobRepository,
|
||||||
tasks TaskRepository, tx TxManager, clk Clock) *SubmitDataset {
|
tasks TaskRepository, tx TxManager, clk Clock, maxAttempts int) *SubmitDataset {
|
||||||
return &SubmitDataset{blobs: blobs, artifacts: artifacts, jobs: jobs, tasks: tasks, tx: tx, clk: clk}
|
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) {
|
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()
|
now := uc.clk.Now()
|
||||||
|
|
||||||
job, err := domain.NewUploadedJob(in.Workload, in.Parameters, 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()
|
cleanup()
|
||||||
return SubmitDatasetResult{}, err
|
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,
|
art, err := domain.NewArtifact(job.ID, nil, domain.ArtifactShard,
|
||||||
fmt.Sprintf("shard-%d.tsv", index), in.ContentType, now)
|
fmt.Sprintf("shard-%d.tsv", index), in.ContentType, now)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
@@ -77,7 +85,7 @@ func (uc *SubmitDataset) Execute(ctx context.Context, in SubmitDatasetInput) (Su
|
|||||||
putKeys = append(putKeys, art.StorageKey)
|
putKeys = append(putKeys, art.StorageKey)
|
||||||
art.SetContent(ssum, ssize)
|
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 {
|
if err != nil {
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
@@ -88,7 +96,8 @@ func (uc *SubmitDataset) Execute(ctx context.Context, in SubmitDatasetInput) (Su
|
|||||||
_ = rc.Close()
|
_ = rc.Close()
|
||||||
if splitErr != nil {
|
if splitErr != nil {
|
||||||
cleanup()
|
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.
|
// 3. Persist job + all artifacts + all tasks atomically.
|
||||||
@@ -118,6 +127,76 @@ func (uc *SubmitDataset) Execute(ctx context.Context, in SubmitDatasetInput) (Su
|
|||||||
}, nil
|
}, 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
|
// GetTaskInput resolves a task's input shard and opens it for streaming. The
|
||||||
// caller closes the reader.
|
// caller closes the reader.
|
||||||
type GetTaskInput struct {
|
type GetTaskInput struct {
|
||||||
|
|||||||
@@ -68,18 +68,18 @@ func newHarness() *harness {
|
|||||||
}
|
}
|
||||||
tx := memstore.Tx{}
|
tx := memstore.Tx{}
|
||||||
h.createJob = usecase.NewCreateJob(h.jobs, h.tasks, tx, h.clk)
|
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.submit = usecase.NewSubmitDataset(h.blobs, h.arts, h.jobs, h.tasks, tx, h.clk, 3)
|
||||||
h.claim = usecase.NewClaimTask(h.tasks, h.clk, lease)
|
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.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.complete = usecase.NewCompleteTask(h.tasks, h.jobs, h.arts, tx, h.clk)
|
||||||
h.fail = usecase.NewFailTask(h.tasks, h.jobs, tx, h.clk)
|
h.fail = usecase.NewFailTask(h.tasks, h.jobs, tx, h.clk)
|
||||||
h.status = usecase.NewGetJobStatus(h.jobs, h.tasks)
|
h.status = usecase.NewGetJobStatus(h.jobs, h.tasks)
|
||||||
h.results = usecase.NewListResults(h.tasks)
|
h.results = usecase.NewListResults(h.tasks)
|
||||||
h.register = usecase.NewRegisterWorker(h.work, h.clk)
|
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.downloadArt = usecase.NewDownloadArtifact(h.arts, h.blobs)
|
||||||
h.getInput = usecase.NewGetTaskInput(h.tasks, 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)
|
h.cancel = usecase.NewCancelJob(h.jobs, h.tasks, tx, h.clk)
|
||||||
return h
|
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) {
|
func TestClaimEmptyQueueReturnsNil(t *testing.T) {
|
||||||
h := newHarness()
|
h := newHarness()
|
||||||
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{"w"}})
|
||||||
@@ -281,7 +319,7 @@ func TestUploadRejectsLeaseThatExpiresDuringStreaming(t *testing.T) {
|
|||||||
h.seedJob(t, "w", 1)
|
h.seedJob(t, "w", 1)
|
||||||
taskID, attempt := h.leaseOne(t, "w1", "w")
|
taskID, attempt := h.leaseOne(t, "w1", "w")
|
||||||
h.uploadArt = usecase.NewUploadArtifact(
|
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{
|
_, 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) {
|
func TestDownloadArtifactRoundTrips(t *testing.T) {
|
||||||
h := newHarness()
|
h := newHarness()
|
||||||
h.seedJob(t, "w", 1)
|
h.seedJob(t, "w", 1)
|
||||||
@@ -434,10 +489,10 @@ func TestDownloadArtifactRoundTrips(t *testing.T) {
|
|||||||
|
|
||||||
func TestSubmitDatasetChunksAndServesInput(t *testing.T) {
|
func TestSubmitDatasetChunksAndServesInput(t *testing.T) {
|
||||||
h := newHarness()
|
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{
|
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),
|
ContentType: "text/tab-separated-values", Body: strings.NewReader(tsv),
|
||||||
})
|
})
|
||||||
if err != nil {
|
if err != nil {
|
||||||
@@ -453,7 +508,7 @@ func TestSubmitDatasetChunksAndServesInput(t *testing.T) {
|
|||||||
t.Errorf("job total = %d, want 3", prog.Total)
|
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 {
|
if err != nil || c == nil {
|
||||||
t.Fatalf("claim shard: %v", err)
|
t.Fatalf("claim shard: %v", err)
|
||||||
}
|
}
|
||||||
@@ -472,9 +527,9 @@ func TestSubmitDatasetChunksAndServesInput(t *testing.T) {
|
|||||||
|
|
||||||
func TestSubmitDatasetLimitsRowsBeforeCreatingShards(t *testing.T) {
|
func TestSubmitDatasetLimitsRowsBeforeCreatingShards(t *testing.T) {
|
||||||
h := newHarness()
|
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{
|
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),
|
ContentType: "text/tab-separated-values", Body: strings.NewReader(tsv),
|
||||||
})
|
})
|
||||||
if err != nil {
|
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) {
|
func TestCancelJobInvalidatesClaimedAndPendingTasks(t *testing.T) {
|
||||||
h := newHarness()
|
h := newHarness()
|
||||||
jobID := h.seedJob(t, "w", 3)
|
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)
|
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;
|
||||||
@@ -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}'
|
-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[@]}" \
|
check "register without capabilities → 400" 400 -X POST "${HOST}/workers/register" "${auth[@]}" \
|
||||||
-d '{"name":"bad"}'
|
-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
|
||||||
echo "job lifecycle"
|
echo "job lifecycle"
|
||||||
@@ -75,7 +82,7 @@ declare -A our_chunks
|
|||||||
task_id=""
|
task_id=""
|
||||||
attempt=""
|
attempt=""
|
||||||
for _ in $(seq 1 40); do
|
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
|
[[ -z "$claim" ]] && break # 204: queue drained
|
||||||
|
|
||||||
read -r c_job c_task c_chunk c_attempt < <(printf '%s' "$claim" |
|
read -r c_job c_task c_chunk c_attempt < <(printf '%s' "$claim" |
|
||||||
@@ -100,7 +107,7 @@ else
|
|||||||
fi
|
fi
|
||||||
|
|
||||||
check "heartbeat" 200 -X POST "${HOST}/tasks/${task_id}/heartbeat" "${auth[@]}" \
|
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) ---
|
# --- artifacts + result (uploads happen while the task is still leased) ---
|
||||||
bearer=(-H "Authorization: Bearer ${TOKEN}")
|
bearer=(-H "Authorization: Bearer ${TOKEN}")
|
||||||
@@ -108,43 +115,47 @@ bearer=(-H "Authorization: Bearer ${TOKEN}")
|
|||||||
# upload <filename> -> prints the artifact_id
|
# upload <filename> -> prints the artifact_id
|
||||||
upload() {
|
upload() {
|
||||||
curl -sS -X PUT "${HOST}/tasks/${task_id}/artifacts/$1" "${bearer[@]}" \
|
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' |
|
--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
|
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[@]}" \
|
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'
|
--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[@]}" \
|
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}" \
|
-H 'Content-Type: text/csv' -H 'X-Worker-ID: impostor' -H "X-Task-Attempt: ${attempt}" \
|
||||||
--data-binary 'x'
|
--data-binary 'x'
|
||||||
|
|
||||||
# Two result artifacts, uploaded now while the lease is held: one to complete
|
# A retry of a PUT returns the same durable artifact for the task attempt.
|
||||||
# with, a second to prove a different manifest is rejected after completion.
|
|
||||||
art_id=$(upload primary.csv)
|
art_id=$(upload primary.csv)
|
||||||
art_id2=$(upload secondary.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 "download artifact" 200 "${HOST}/artifacts/${art_id}/download" "${bearer[@]}"
|
||||||
|
|
||||||
check "foreign worker submits → 409" 409 -X POST "${HOST}/tasks/${task_id}/result" "${auth[@]}" \
|
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}\"}}"
|
-d "{\"worker_id\":\"impostor\",\"attempt\":${attempt},\"result\":{\"artifact_id\":\"${art_id}\"}}"
|
||||||
check "submit result" 200 -X POST "${HOST}/tasks/${task_id}/result" "${auth[@]}" \
|
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[@]}" \
|
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}\"}}"
|
-d "{\"worker_id\":\"${worker_id}\",\"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}\"}}"
|
|
||||||
check "GET /jobs/{id}" 200 "${HOST}/jobs/${job_id}" "${auth[@]}"
|
check "GET /jobs/{id}" 200 "${HOST}/jobs/${job_id}" "${auth[@]}"
|
||||||
|
|
||||||
echo
|
echo
|
||||||
echo "input validation"
|
echo "input validation"
|
||||||
check "malformed uuid → 400" 400 -X POST "${HOST}/tasks/not-a-uuid/result" "${auth[@]}" \
|
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
|
# Note: Go's encoding/json matches field names case-insensitively, so
|
||||||
# "worker_ID" would be accepted as "worker_id". Only a genuinely unknown key
|
# "worker_ID" would be accepted as "worker_id". Only a genuinely unknown key
|
||||||
# trips DisallowUnknownFields.
|
# trips DisallowUnknownFields.
|
||||||
check "unknown json field → 400" 400 -X POST "${HOST}/tasks/claim" "${auth[@]}" \
|
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[@]}"
|
check "unknown job → 404" 404 "${HOST}/jobs/00000000-0000-0000-0000-000000000000" "${auth[@]}"
|
||||||
|
|
||||||
echo
|
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
|
# 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.
|
# fields precede the file part, which the coordinator streams.
|
||||||
up=$(curl -sS "${bearer[@]}" -X POST "${HOST}/jobs/upload" \
|
up=$(curl -sS "${bearer[@]}" -X POST "${HOST}/jobs/upload" \
|
||||||
-F 'workload=similarity_search' \
|
-F 'workload=similarity-search' \
|
||||||
-F 'parameters={"top_k":10}' \
|
-F 'parameters={"query_smiles":"CCO","top_k":10}' \
|
||||||
-F 'chunk_rows=2' \
|
-F 'chunk_rows=2' \
|
||||||
-F 'file=@-;filename=chembl.tsv;type=text/tab-separated-values' <<'TSV'
|
-F 'file=@-;filename=chembl.tsv;type=text/tab-separated-values' <<'TSV'
|
||||||
id smiles
|
chembl_id canonical_smiles
|
||||||
A CC
|
A CC
|
||||||
B CCC
|
B CCC
|
||||||
C CCCC
|
C CCCC
|
||||||
@@ -180,7 +191,7 @@ fi
|
|||||||
up_input=""
|
up_input=""
|
||||||
for _ in $(seq 1 30); do
|
for _ in $(seq 1 30); do
|
||||||
c=$(curl -sS "${bearer[@]}" -H 'Content-Type: application/json' -X POST "${HOST}/tasks/claim" \
|
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
|
[[ -z "$c" ]] && break
|
||||||
cj=$(printf '%s' "$c" | python3 -c 'import json,sys;print(json.load(sys.stdin)["job_id"])' 2>/dev/null)
|
cj=$(printf '%s' "$c" | python3 -c 'import json,sys;print(json.load(sys.stdin)["job_id"])' 2>/dev/null)
|
||||||
[[ "$cj" != "$up_job" ]] && continue
|
[[ "$cj" != "$up_job" ]] && continue
|
||||||
|
|||||||
@@ -107,7 +107,8 @@ Content-Type: application/json
|
|||||||
```
|
```
|
||||||
|
|
||||||
`cpu_count`/`memory_mb` are accepted for forward compatibility and not yet
|
`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
|
## Claim task
|
||||||
|
|
||||||
@@ -120,7 +121,8 @@ Content-Type: application/json
|
|||||||
```
|
```
|
||||||
|
|
||||||
- `204 No Content`: no compatible task.
|
- `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
|
```json
|
||||||
{
|
{
|
||||||
|
|||||||
@@ -46,24 +46,25 @@ coordinator was started with. Never log it, never send it in an error body.
|
|||||||
|
|
||||||
```http
|
```http
|
||||||
POST /workers/register
|
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 }`.
|
Response: `{ "worker_id": "<uuid>", "heartbeat_interval_seconds": 15 }`.
|
||||||
|
|
||||||
- `capabilities` are the workload names you can run — the coordinator only hands
|
- `capabilities` are fixed at registration — the coordinator only hands you
|
||||||
you matching tasks.
|
matching tasks and a later claim cannot broaden that set.
|
||||||
- **Keep `worker_id`**. Use it as your identity in every later call. Using the
|
- **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
|
registered UUID is what lets the coordinator track your liveness (it marks
|
||||||
workers offline after they go silent).
|
workers offline after they go silent).
|
||||||
- Current coordinator jobs use `similarity_search` / `similarity_graph`; the
|
- Current diagnostic uploads use `similarity-search` with `query_smiles`. The
|
||||||
reference Python worker also accepts the public CLI spellings with hyphens.
|
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
|
## 2. Claim a task
|
||||||
|
|
||||||
```http
|
```http
|
||||||
POST /tasks/claim
|
POST /tasks/claim
|
||||||
{ "worker_id": "<uuid>", "capabilities": ["similarity_search"] }
|
{ "worker_id": "<uuid>", "capabilities": ["similarity-search"] }
|
||||||
```
|
```
|
||||||
|
|
||||||
- `200` → a leased task (below).
|
- `200` → a leased task (below).
|
||||||
@@ -74,9 +75,9 @@ POST /tasks/claim
|
|||||||
"task_id": "<uuid>",
|
"task_id": "<uuid>",
|
||||||
"attempt": 1,
|
"attempt": 1,
|
||||||
"lease_expires_at": "2026-07-22T12:05:00Z",
|
"lease_expires_at": "2026-07-22T12:05:00Z",
|
||||||
"workload": "similarity_search",
|
"workload": "similarity-search",
|
||||||
"input": { "uri": "/tasks/<uuid>/input", "sha256": "<hex>" },
|
"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
|
- poll interval and request timeout
|
||||||
- a working directory for downloaded inputs and generated outputs
|
- 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
|
## Generate a client from the spec
|
||||||
|
|
||||||
Instead of hand-writing request code, generate it:
|
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
@@ -96,7 +96,9 @@ paths:
|
|||||||
description: >
|
description: >
|
||||||
multipart/form-data. The text fields (`workload`, `parameters`,
|
multipart/form-data. The text fields (`workload`, `parameters`,
|
||||||
`chunk_rows`, `max_rows`) MUST precede the `file` part: the file is streamed, not
|
`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:
|
requestBody:
|
||||||
required: true
|
required: true
|
||||||
content:
|
content:
|
||||||
@@ -388,7 +390,7 @@ components:
|
|||||||
type: array
|
type: array
|
||||||
minItems: 1
|
minItems: 1
|
||||||
items: { type: string }
|
items: { type: string }
|
||||||
example: [similarity_search, similarity_graph]
|
example: [similarity-search]
|
||||||
cpu_count:
|
cpu_count:
|
||||||
type: integer
|
type: integer
|
||||||
description: Accepted for forward compatibility; not yet persisted.
|
description: Accepted for forward compatibility; not yet persisted.
|
||||||
@@ -419,7 +421,7 @@ components:
|
|||||||
type: object
|
type: object
|
||||||
required: [workload, input_uri, chunks]
|
required: [workload, input_uri, chunks]
|
||||||
properties:
|
properties:
|
||||||
workload: { type: string, example: similarity_search }
|
workload: { type: string, example: similarity-search }
|
||||||
input_uri: { type: string }
|
input_uri: { type: string }
|
||||||
parameters: { type: object, additionalProperties: true }
|
parameters: { type: object, additionalProperties: true }
|
||||||
chunks:
|
chunks:
|
||||||
@@ -437,11 +439,11 @@ components:
|
|||||||
type: object
|
type: object
|
||||||
required: [workload, file]
|
required: [workload, file]
|
||||||
properties:
|
properties:
|
||||||
workload: { type: string, example: similarity_search }
|
workload: { type: string, enum: [similarity-search], example: similarity-search }
|
||||||
parameters:
|
parameters:
|
||||||
type: string
|
type: string
|
||||||
description: JSON object, sent as a string form field.
|
description: JSON object, sent as a string form field.
|
||||||
example: '{"top_k":10}'
|
example: '{"query_smiles":"CCO","top_k":10}'
|
||||||
chunk_rows:
|
chunk_rows:
|
||||||
type: integer
|
type: integer
|
||||||
description: Data rows per shard. Default 1000.
|
description: Data rows per shard. Default 1000.
|
||||||
@@ -486,11 +488,11 @@ components:
|
|||||||
type: object
|
type: object
|
||||||
required: [worker_id]
|
required: [worker_id]
|
||||||
properties:
|
properties:
|
||||||
worker_id: { type: string }
|
worker_id: { type: string, format: uuid, description: Registered worker identity. }
|
||||||
capabilities:
|
capabilities:
|
||||||
type: array
|
type: array
|
||||||
items: { type: string }
|
items: { type: string }
|
||||||
description: Workloads this worker can run. Empty means "any".
|
description: Accepted for compatibility only; registration capabilities decide eligibility.
|
||||||
max_concurrency:
|
max_concurrency:
|
||||||
type: integer
|
type: integer
|
||||||
description: Accepted; the coordinator leases one task per call.
|
description: Accepted; the coordinator leases one task per call.
|
||||||
|
|||||||
@@ -59,7 +59,7 @@ Content-Type: application/json
|
|||||||
|
|
||||||
{
|
{
|
||||||
"name": "lab-worker-01",
|
"name": "lab-worker-01",
|
||||||
"capabilities": ["similarity-search", "similarity-graph"],
|
"capabilities": ["similarity-search"],
|
||||||
"cpu_count": 8,
|
"cpu_count": 8,
|
||||||
"memory_mb": 16384
|
"memory_mb": 16384
|
||||||
}
|
}
|
||||||
@@ -74,8 +74,8 @@ POST /tasks/claim
|
|||||||
Content-Type: application/json
|
Content-Type: application/json
|
||||||
|
|
||||||
{
|
{
|
||||||
"worker_id": "worker-01",
|
"worker_id": "<registered-uuid>",
|
||||||
"capabilities": ["similarity-search", "similarity-graph"],
|
"capabilities": ["similarity-search"],
|
||||||
"max_concurrency": 1
|
"max_concurrency": 1
|
||||||
}
|
}
|
||||||
```
|
```
|
||||||
@@ -95,7 +95,7 @@ When a task is available, it returns `200 OK`:
|
|||||||
"sha256": "..."
|
"sha256": "..."
|
||||||
},
|
},
|
||||||
"parameters": {
|
"parameters": {
|
||||||
"query_id": "CHEMBL939",
|
"query_smiles": "CCO",
|
||||||
"top_k": 20
|
"top_k": 20
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
+31
-6
@@ -13,15 +13,17 @@ from .daemon import WorkerDaemon
|
|||||||
from .runners import SciMeshRunner
|
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(
|
parser = argparse.ArgumentParser(
|
||||||
prog="scimesh-worker",
|
prog="scimesh-worker",
|
||||||
epilog=(
|
epilog=(
|
||||||
"Environment: SCIMESH_COORDINATOR_URL, SCIMESH_WORK_DIR, "
|
"Environment: SCIMESH_COORDINATOR_URL, SCIMESH_WORK_DIR, "
|
||||||
"SCIMESH_WORKER_NAME, SCIMESH_CPU_COUNT, SCIMESH_MEMORY_MB, "
|
"SCIMESH_WORKER_NAME, SCIMESH_CPU_COUNT, SCIMESH_MEMORY_MB, "
|
||||||
"SCIMESH_POLL_INTERVAL, SCIMESH_REQUEST_TIMEOUT, "
|
"SCIMESH_POLL_INTERVAL, SCIMESH_REQUEST_TIMEOUT, "
|
||||||
"SCIMESH_HEARTBEAT_INTERVAL, SCIMESH_CLEANUP_AFTER_SECONDS, and "
|
"SCIMESH_HEARTBEAT_INTERVAL, SCIMESH_CLEANUP_AFTER_SECONDS, "
|
||||||
"SCIMESH_BEARER_TOKEN. SCIMESH_WORKER_ID is a legacy/test override."
|
"SCIMESH_MAX_TASKS, and SCIMESH_BEARER_TOKEN. "
|
||||||
|
"SCIMESH_WORKER_ID is a legacy/test override."
|
||||||
),
|
),
|
||||||
)
|
)
|
||||||
parser.add_argument("--coordinator-url")
|
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("--request-timeout", type=float)
|
||||||
parser.add_argument("--heartbeat-interval", type=float)
|
parser.add_argument("--heartbeat-interval", type=float)
|
||||||
parser.add_argument("--cleanup-after-seconds", 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)
|
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:
|
if "work_dir" in overrides:
|
||||||
overrides["work_dir"] = Path(overrides["work_dir"])
|
overrides["work_dir"] = Path(overrides["work_dir"])
|
||||||
try:
|
try:
|
||||||
@@ -44,13 +69,13 @@ def main(argv: list[str] | None = None) -> int:
|
|||||||
parser.error(str(error))
|
parser.error(str(error))
|
||||||
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
|
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
|
||||||
client = HttpCoordinatorClient(config.coordinator_url, config.request_timeout, config.bearer_token)
|
client = HttpCoordinatorClient(config.coordinator_url, config.request_timeout, config.bearer_token)
|
||||||
WorkerDaemon(
|
completed_without_interruption = WorkerDaemon(
|
||||||
config,
|
config,
|
||||||
client,
|
client,
|
||||||
HttpArtifactClient(config.coordinator_url, config.request_timeout, config.bearer_token),
|
HttpArtifactClient(config.coordinator_url, config.request_timeout, config.bearer_token),
|
||||||
SciMeshRunner(),
|
SciMeshRunner(),
|
||||||
).run_forever()
|
).run_forever()
|
||||||
return 0
|
return 0 if completed_without_interruption else 130
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
|
|||||||
@@ -36,13 +36,16 @@ class WorkerConfig:
|
|||||||
heartbeat_interval: float = 15.0
|
heartbeat_interval: float = 15.0
|
||||||
bearer_token: str | None = None
|
bearer_token: str | None = None
|
||||||
cleanup_after_seconds: float | 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
|
# 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, ...] = (
|
capabilities: tuple[str, ...] = (
|
||||||
"similarity-search",
|
"similarity-search",
|
||||||
"similarity-graph",
|
|
||||||
"similarity_search",
|
"similarity_search",
|
||||||
"similarity_graph",
|
|
||||||
)
|
)
|
||||||
|
|
||||||
def __post_init__(self) -> None:
|
def __post_init__(self) -> None:
|
||||||
@@ -66,6 +69,15 @@ class WorkerConfig:
|
|||||||
_positive_number(self.heartbeat_interval, "heartbeat_interval")
|
_positive_number(self.heartbeat_interval, "heartbeat_interval")
|
||||||
if self.cleanup_after_seconds is not None:
|
if self.cleanup_after_seconds is not None:
|
||||||
_positive_number(self.cleanup_after_seconds, "cleanup_after_seconds", allow_zero=True)
|
_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:
|
if not self.capabilities:
|
||||||
raise ValueError("capabilities cannot be empty")
|
raise ValueError("capabilities cannot be empty")
|
||||||
# Runner subprocesses use a task directory as their cwd. Keep the
|
# 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")
|
cleanup = value("cleanup_after_seconds", "SCIMESH_CLEANUP_AFTER_SECONDS")
|
||||||
cpu_count = value("cpu_count", "SCIMESH_CPU_COUNT", os.cpu_count() or 1)
|
cpu_count = value("cpu_count", "SCIMESH_CPU_COUNT", os.cpu_count() or 1)
|
||||||
memory_mb = value("memory_mb", "SCIMESH_MEMORY_MB")
|
memory_mb = value("memory_mb", "SCIMESH_MEMORY_MB")
|
||||||
|
max_tasks = value("max_tasks", "SCIMESH_MAX_TASKS")
|
||||||
return cls(
|
return cls(
|
||||||
coordinator_url=url.rstrip("/"),
|
coordinator_url=url.rstrip("/"),
|
||||||
worker_id=value("worker_id", "SCIMESH_WORKER_ID"),
|
worker_id=value("worker_id", "SCIMESH_WORKER_ID"),
|
||||||
@@ -102,4 +115,6 @@ class WorkerConfig:
|
|||||||
heartbeat_interval=float(value("heartbeat_interval", "SCIMESH_HEARTBEAT_INTERVAL", "15")),
|
heartbeat_interval=float(value("heartbeat_interval", "SCIMESH_HEARTBEAT_INTERVAL", "15")),
|
||||||
bearer_token=value("bearer_token", "SCIMESH_BEARER_TOKEN"),
|
bearer_token=value("bearer_token", "SCIMESH_BEARER_TOKEN"),
|
||||||
cleanup_after_seconds=float(cleanup) if cleanup else None,
|
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
@@ -4,8 +4,10 @@ from __future__ import annotations
|
|||||||
|
|
||||||
import logging
|
import logging
|
||||||
from dataclasses import replace
|
from dataclasses import replace
|
||||||
|
from dataclasses import dataclass
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
import random
|
import random
|
||||||
|
import re
|
||||||
import shutil
|
import shutil
|
||||||
import threading
|
import threading
|
||||||
import time
|
import time
|
||||||
@@ -72,6 +74,14 @@ class LeaseHeartbeat:
|
|||||||
return seconds
|
return seconds
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass(frozen=True)
|
||||||
|
class RunOnceOutcome:
|
||||||
|
"""Whether a claim was made and whether that claimed task completed."""
|
||||||
|
|
||||||
|
claimed: bool
|
||||||
|
completed: bool
|
||||||
|
|
||||||
|
|
||||||
class WorkerDaemon:
|
class WorkerDaemon:
|
||||||
def __init__(self, config: WorkerConfig, coordinator: CoordinatorClient, artifacts: ArtifactClient, runner: Runner) -> None:
|
def __init__(self, config: WorkerConfig, coordinator: CoordinatorClient, artifacts: ArtifactClient, runner: Runner) -> None:
|
||||||
self.config, self.coordinator, self.artifacts, self.runner = config, coordinator, artifacts, runner
|
self.config, self.coordinator, self.artifacts, self.runner = config, coordinator, artifacts, runner
|
||||||
@@ -79,32 +89,68 @@ class WorkerDaemon:
|
|||||||
self._registered = False
|
self._registered = False
|
||||||
self.log = logging.getLogger("scimesh.worker")
|
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
|
failures = 0
|
||||||
while True:
|
completed_tasks = 0
|
||||||
try:
|
self._log(
|
||||||
if not self._registered:
|
"started",
|
||||||
self._register_worker()
|
max_tasks=self.config.max_tasks,
|
||||||
self._cleanup_expired_directories()
|
exit_when_idle=self.config.exit_when_idle,
|
||||||
claimed = self.run_once()
|
)
|
||||||
failures = 0
|
try:
|
||||||
if not claimed:
|
while True:
|
||||||
self._sleep(self.config.poll_interval)
|
try:
|
||||||
except CoordinatorTransientError as error:
|
if not self._registered:
|
||||||
failures += 1
|
self._register_worker()
|
||||||
self._log("failed", error_type=type(error).__name__)
|
self._cleanup_expired_directories()
|
||||||
self._sleep(min(self.config.poll_interval * 2 ** min(failures, 6), 60.0))
|
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()
|
worker_id = self._worker_id()
|
||||||
self._log("claiming")
|
self._log("claiming", log_level=logging.DEBUG)
|
||||||
task = self.coordinator.claim(worker_id, self.config.capabilities)
|
task = self.coordinator.claim(worker_id, self.config.capabilities)
|
||||||
if task is None:
|
if task is None:
|
||||||
self._log("idle")
|
self._log("idle", log_level=logging.DEBUG)
|
||||||
return False
|
return RunOnceOutcome(claimed=False, completed=False)
|
||||||
started = time.monotonic()
|
started = time.monotonic()
|
||||||
task_dir = self.config.work_dir / task.task_id / str(task.attempt)
|
task_dir = self.config.work_dir / task.task_id / str(task.attempt)
|
||||||
heartbeat = LeaseHeartbeat(task, self.coordinator, self.config)
|
heartbeat = LeaseHeartbeat(task, self.coordinator, self.config)
|
||||||
|
completed = False
|
||||||
try:
|
try:
|
||||||
task_dir.mkdir(parents=True, exist_ok=False)
|
task_dir.mkdir(parents=True, exist_ok=False)
|
||||||
heartbeat.start()
|
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:
|
except CoordinatorConflictError as error:
|
||||||
self._log("lease_lost", task, error_type=type(error).__name__)
|
self._log("lease_lost", task, error_type=type(error).__name__)
|
||||||
except Exception as error:
|
except Exception as error:
|
||||||
@@ -143,10 +197,10 @@ class WorkerDaemon:
|
|||||||
self._report_failure(task, error)
|
self._report_failure(task, error)
|
||||||
finally:
|
finally:
|
||||||
heartbeat.stop()
|
heartbeat.stop()
|
||||||
return True
|
return RunOnceOutcome(claimed=True, completed=completed)
|
||||||
|
|
||||||
def _report_failure(self, task: ClaimedTask, error: Exception) -> None:
|
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:
|
try:
|
||||||
self.coordinator.fail(task, {"worker_id": self._worker_id(), "attempt": task.attempt, "error_code": type(error).__name__, "error_message": message})
|
self.coordinator.fail(task, {"worker_id": self._worker_id(), "attempt": task.attempt, "error_code": type(error).__name__, "error_message": message})
|
||||||
except CoordinatorTransientError:
|
except CoordinatorTransientError:
|
||||||
@@ -154,6 +208,16 @@ class WorkerDaemon:
|
|||||||
except Exception:
|
except Exception:
|
||||||
self._log("failed", task, error_type="FailureReportError")
|
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:
|
def _register_worker(self) -> None:
|
||||||
registered = self.coordinator.register(
|
registered = self.coordinator.register(
|
||||||
self.config.worker_name,
|
self.config.worker_name,
|
||||||
@@ -180,9 +244,16 @@ class WorkerDaemon:
|
|||||||
"""Keep completion payload exact: coordinator owns all artifact metadata."""
|
"""Keep completion payload exact: coordinator owns all artifact metadata."""
|
||||||
return {"artifact_id": uploaded.artifact_id}
|
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}
|
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:
|
def _cleanup_expired_directories(self) -> None:
|
||||||
"""Remove only old task attempt directories when retention was configured."""
|
"""Remove only old task attempt directories when retention was configured."""
|
||||||
|
|||||||
Executable
+165
@@ -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
@@ -1,7 +1,9 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
import hashlib
|
import hashlib
|
||||||
|
import logging
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
import subprocess
|
||||||
import time
|
import time
|
||||||
from datetime import datetime, timedelta, timezone
|
from datetime import datetime, timedelta, timezone
|
||||||
from urllib.request import Request
|
from urllib.request import Request
|
||||||
@@ -9,8 +11,10 @@ from urllib.request import Request
|
|||||||
import pytest
|
import pytest
|
||||||
|
|
||||||
from scimesh.worker.config import WorkerConfig
|
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.coordinator import CoordinatorTransientError
|
||||||
from scimesh.worker.daemon import LeaseHeartbeat, WorkerDaemon
|
from scimesh.worker.daemon import LeaseHeartbeat, RunOnceOutcome, WorkerDaemon
|
||||||
from scimesh.worker.models import (
|
from scimesh.worker.models import (
|
||||||
ClaimedTask,
|
ClaimedTask,
|
||||||
InputArtifact,
|
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:
|
def test_claims_runs_uploads_and_submits_csv(tmp_path: Path) -> None:
|
||||||
content = b"input fixture"
|
content = b"input fixture"
|
||||||
worker, coordinator, artifacts, runner, _ = daemon(tmp_path, make_task(content), content)
|
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 runner.calls == 1
|
||||||
assert len(artifacts.uploaded) == 1
|
assert len(artifacts.uploaded) == 1
|
||||||
assert coordinator.heartbeats == [("task-1", 1, "worker-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:
|
def test_no_task_does_not_create_directory(tmp_path: Path) -> None:
|
||||||
worker, _, _, runner, config = daemon(tmp_path, None, b"")
|
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 runner.calls == 0
|
||||||
assert not config.work_dir.exists()
|
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:
|
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")
|
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 runner.calls == 0
|
||||||
assert coordinator.failures[0]["error_code"] == "ValueError"
|
assert coordinator.failures[0]["error_code"] == "ValueError"
|
||||||
assert not coordinator.submissions
|
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:
|
def test_directory_creation_failure_is_reported(tmp_path: Path) -> None:
|
||||||
content = b"input fixture"
|
content = b"input fixture"
|
||||||
worker, coordinator, _, _, config = daemon(tmp_path, make_task(content), content)
|
worker, coordinator, _, _, config = daemon(tmp_path, make_task(content), content)
|
||||||
(config.work_dir / "task-1" / "1").mkdir(parents=True)
|
(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"
|
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 config.worker_id is None
|
||||||
assert "similarity-search" in config.capabilities
|
assert "similarity-search" in config.capabilities
|
||||||
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(
|
def test_relative_work_dir_is_normalized_for_runner_subprocesses(
|
||||||
|
|||||||
Reference in New Issue
Block a user