Compare commits
16
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
6ef92908a1 | ||
|
|
f953112cfd | ||
|
|
19cbf7f113 | ||
|
|
9ec8f50313 | ||
|
|
08f5478a66 | ||
|
|
bde6cdb4ba | ||
|
|
43ceec1f77 | ||
|
|
f5b16b057f | ||
|
|
f8de0b2b9d | ||
|
|
7547a30bde | ||
|
|
6bac7dad3c | ||
|
|
c7956c4683 | ||
|
|
d648beede2 | ||
|
|
ac9b921401 | ||
|
|
5be87ad762 | ||
|
|
e83e0b5e1f |
@@ -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
|
||||
test_results.csv
|
||||
test_structures/
|
||||
|
||||
# Local coordinator-worker execution state
|
||||
worker-data*/
|
||||
scimesh-worker-data/
|
||||
|
||||
@@ -1,11 +1,12 @@
|
||||
# SciMesh
|
||||
|
||||
SciMesh is a scientific-workload framework for molecular datasets. Its public CLI
|
||||
currently runs exact similarity search and sparse similarity-graph construction
|
||||
locally in one Python process; it creates no dense similarity matrix. A Python
|
||||
Worker client and the planned Go/PostgreSQL coordinator contract are tracked in
|
||||
the repository, but distributed execution is not available yet; see
|
||||
[`STATUS.md`](STATUS.md).
|
||||
runs exact similarity search and sparse similarity-graph construction locally in
|
||||
one Python process; it creates no dense similarity matrix. The Go/PostgreSQL
|
||||
coordinator and Python worker can run a diagnostic, shard-based
|
||||
`similarity-search` pipeline locally. Its CSV artifacts are not a global result
|
||||
until CTX-07--09 add planning and reduction; use the local CLI for scientific
|
||||
results today. See [`STATUS.md`](STATUS.md).
|
||||
|
||||
The ChEMBL TSV database is intentionally not included in this repository. Download it separately and pass its path to the commands below. The expected columns are `chembl_id` and `canonical_smiles`.
|
||||
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
# SciMesh Status
|
||||
|
||||
**Updated:** 2026-07-23
|
||||
**Branch baseline:** `main` at `b4a89dd` (coordinator merge)
|
||||
**Updated:** 2026-07-24
|
||||
**Branch baseline:** `main` at `f953112` (distributed pipeline hardening)
|
||||
|
||||
## Current state
|
||||
|
||||
@@ -32,26 +32,32 @@ Docker PostgreSQL stack on 2026-07-23.
|
||||
| CTX-04 Worker registry and HTTP API | Implemented | Registration, claim, heartbeat, result, failure, and status endpoints. |
|
||||
| CTX-05 Artifact storage | Implemented | Coordinator-owned inputs/results, checksum verification, and upload flow. |
|
||||
| CTX-06 Python Worker live-contract alignment | Implemented | Worker completed a real uploaded shard via HTTP on 2026-07-23. |
|
||||
| CTX-07 Distributed workload protocol | Not started | Depends on artifact and Worker contracts. |
|
||||
| CTX-07 Distributed workload protocol | Implemented | Versioned Python contract models, registry, strict plan validation, and deterministic reduction ordering are in `scimesh/distributed/`. The concrete molecular planner/reducer remains CTX-08/09. |
|
||||
| CTX-08 Distributed similarity-search | Not started | Local reference exists. |
|
||||
| CTX-09 Reducer and final-result API | Not started | Depends on CTX-07 and CTX-08. |
|
||||
| CTX-10 Distributed similarity-graph | Not started | Local reference exists. |
|
||||
| CTX-11 Dashboard/operator view | Not started | Deferred until API and reducer work. |
|
||||
| CTX-11 Dashboard/operator view | Implemented (diagnostic scope) | Protected local view: job/task/worker status, validated similarity-search upload, diagnostic partial-artifact download, and bounded polling. Final-result reduction remains CTX-09. |
|
||||
| CTX-12 Reliability, security, CI | In progress | Unit, race, PostgreSQL integration, and smoke checks exist; CI hardening remains. |
|
||||
|
||||
## Next recommended assignment
|
||||
|
||||
Assign **CTX-07** to the workload role: define distributed job planning and
|
||||
reduction boundaries before implementing distributed search or graph execution.
|
||||
Assign **CTX-08** to the workload role: implement the molecular
|
||||
`similarity-search` planner and worker adapter on top of the accepted CTX-07
|
||||
contract.
|
||||
|
||||
## Known constraints
|
||||
|
||||
- Planner/reducer semantics are not implemented; use the local `scimesh` CLI
|
||||
for complete workload results.
|
||||
- The CTX-07 protocol is implemented, but no concrete molecular planner or
|
||||
reducer is registered yet; the operator UI labels `partial_result` files as
|
||||
diagnostic and cannot present them as final output.
|
||||
Use the local `scimesh` CLI for complete workload results.
|
||||
- The worker/coordinator flow currently accepts both underscore API workload
|
||||
names and hyphenated CLI names while the contract is consolidated.
|
||||
- A real-stack worker test uses a small `query_smiles` shard. Resolving a
|
||||
`query_id` once and sharing it across shards belongs to CTX-07.
|
||||
- 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
|
||||
|
||||
|
||||
@@ -6,6 +6,10 @@ DATABASE_URL=postgres://scimesh:scimesh@localhost:5432/scimesh?sslmode=disable
|
||||
# Shared bearer token every worker must present. Leave empty to disable auth (dev only).
|
||||
WORKER_AUTH_TOKEN=change-me
|
||||
|
||||
# Optional local operator UI. Use a separate value; never reuse the worker token.
|
||||
# When empty, /ui is disabled.
|
||||
UI_AUTH_TOKEN=
|
||||
|
||||
# Logging. LOG_LEVEL: debug|info|warn|error. LOG_FILE empty = stdout only;
|
||||
# set a path to also write a size-rotated file (kept across restarts).
|
||||
LOG_LEVEL=info
|
||||
|
||||
@@ -68,6 +68,18 @@ make logs # follow the coordinator
|
||||
make down # stop (add down-clean to drop the DB volume)
|
||||
```
|
||||
|
||||
To enable the local operator UI, set a separate credential before starting:
|
||||
|
||||
```sh
|
||||
UI_AUTH_TOKEN='local-ui-secret' make up
|
||||
# Open http://localhost:8080/ui and use any username with this value as password.
|
||||
```
|
||||
|
||||
The UI is disabled by default and never accepts the worker bearer token.
|
||||
It shows recent jobs, task/worker state, and the per-job partial artifacts.
|
||||
Those files are explicitly diagnostic until the CTX-09 reducer creates a final
|
||||
result; the UI does not present them as final scientific output.
|
||||
|
||||
`up` starts three services in order: Postgres waits until `pg_isready` passes, a
|
||||
one-shot `migrate` container applies the schema and exits, and only then does the
|
||||
coordinator start — so it never queries a database that has no tables.
|
||||
|
||||
@@ -65,26 +65,29 @@ func run() error {
|
||||
jobRepo = postgres.NewJobRepo(pool)
|
||||
workerRepo = postgres.NewWorkerRepo(pool)
|
||||
artifactRepo = postgres.NewArtifactRepo(pool)
|
||||
uiReadRepo = postgres.NewUIReadRepo(pool)
|
||||
)
|
||||
|
||||
useCases := httptransport.UseCases{
|
||||
RegisterWorker: usecase.NewRegisterWorker(workerRepo, clk),
|
||||
CreateJob: usecase.NewCreateJob(jobRepo, taskRepo, tx, clk),
|
||||
SubmitDataset: usecase.NewSubmitDataset(blobStore, artifactRepo, jobRepo, taskRepo, tx, clk),
|
||||
ClaimTask: usecase.NewClaimTask(taskRepo, clk, cfg.LeaseDuration),
|
||||
SubmitDataset: usecase.NewSubmitDataset(blobStore, artifactRepo, jobRepo, taskRepo, tx, clk, cfg.DefaultMaxAttempts),
|
||||
ClaimTask: usecase.NewClaimTask(taskRepo, jobRepo, workerRepo, tx, clk, cfg.LeaseDuration),
|
||||
RenewLease: usecase.NewRenewLease(taskRepo, workerRepo, tx, clk, cfg.LeaseDuration),
|
||||
CompleteTask: usecase.NewCompleteTask(taskRepo, jobRepo, artifactRepo, tx, clk),
|
||||
FailTask: usecase.NewFailTask(taskRepo, jobRepo, tx, clk),
|
||||
GetJobStatus: usecase.NewGetJobStatus(jobRepo, taskRepo),
|
||||
UploadArtifact: usecase.NewUploadArtifact(taskRepo, artifactRepo, blobStore, clk),
|
||||
CancelJob: usecase.NewCancelJob(jobRepo, taskRepo, tx, clk),
|
||||
UploadArtifact: usecase.NewUploadArtifact(taskRepo, artifactRepo, blobStore, tx, clk),
|
||||
DownloadArtifact: usecase.NewDownloadArtifact(artifactRepo, blobStore),
|
||||
GetTaskInput: usecase.NewGetTaskInput(taskRepo, artifactRepo, blobStore),
|
||||
Dashboard: usecase.NewDashboard(uiReadRepo),
|
||||
}
|
||||
|
||||
// Background reapers are tracked so shutdown can wait for them. Without this
|
||||
// the process would exit mid-UPDATE, and the deferred pool.Close() would pull
|
||||
// connections out from under them.
|
||||
expireLeases := usecase.NewExpireLeases(taskRepo, clk)
|
||||
expireLeases := usecase.NewExpireLeases(taskRepo, jobRepo, tx, clk)
|
||||
markOffline := usecase.NewMarkWorkersOffline(workerRepo, clk, cfg.WorkerOfflineAfter)
|
||||
|
||||
var wg sync.WaitGroup
|
||||
@@ -105,7 +108,7 @@ func run() error {
|
||||
// pool.Ping backs /health: readiness means the database answers, not just
|
||||
// that the process is alive.
|
||||
api := httptransport.NewServer(useCases, log, cfg.RequestTimeout, cfg.HeartbeatInterval, cfg.MaxUploadBytes, pool.Ping)
|
||||
err = infra.RunServer(ctx, log, cfg.Addr, api.Handler(cfg.Token))
|
||||
err = infra.RunServer(ctx, log, cfg.Addr, api.Handler(cfg.Token, cfg.UIToken))
|
||||
|
||||
// Shutdown order matters, and defers alone cannot express it (they run
|
||||
// LIFO, so the deferred stop() would fire *after* the wait below).
|
||||
|
||||
@@ -49,6 +49,8 @@ services:
|
||||
# Host is the service name: compose resolves it on the project network.
|
||||
DATABASE_URL: postgres://${POSTGRES_USER:-scimesh}:${POSTGRES_PASSWORD:-scimesh}@postgres:5432/${POSTGRES_DB:-scimesh}?sslmode=disable
|
||||
WORKER_AUTH_TOKEN: ${WORKER_AUTH_TOKEN:-dev-token}
|
||||
# Empty disables /ui. Set this separately from the worker token.
|
||||
UI_AUTH_TOKEN: ${UI_AUTH_TOKEN:-}
|
||||
DB_MAX_CONNS: "10"
|
||||
REQUEST_TIMEOUT: "15s"
|
||||
LEASE_DURATION: "2m"
|
||||
|
||||
@@ -7,12 +7,19 @@ import (
|
||||
"bytes"
|
||||
"fmt"
|
||||
"io"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// ErrNoRows is returned when the input has a header but no data rows: a job with
|
||||
// zero tasks could never complete, so it is rejected at the source.
|
||||
var ErrNoRows = fmt.Errorf("input has no data rows")
|
||||
|
||||
// maxShardBytes bounds the coordinator memory used by one in-progress shard.
|
||||
// The uploaded file may be much larger: it is first stored on disk, then split
|
||||
// in small bounded pieces. Operators can lower rowsPerShard when this limit is
|
||||
// reached rather than exhausting the coordinator process.
|
||||
const maxShardBytes = 64 << 20 // 64 MiB
|
||||
|
||||
// SplitTSV reads a header-plus-rows text stream and cuts it into shards of at
|
||||
// most rowsPerShard data rows. Every shard repeats the header, so a worker can
|
||||
// parse its shard in isolation. emit is called once per shard, in order, with a
|
||||
@@ -26,9 +33,31 @@ var ErrNoRows = fmt.Errorf("input has no data rows")
|
||||
// Only one shard is buffered at a time, so memory is bounded by shard size (a
|
||||
// worker-sized slice of the data), not by the size of the whole dataset.
|
||||
func SplitTSV(r io.Reader, rowsPerShard int, emit func(index int, shard io.Reader) error) error {
|
||||
return splitTSVLimit(r, rowsPerShard, 0, nil, emit)
|
||||
}
|
||||
|
||||
// SplitTSVLimit behaves like SplitTSV but emits no more than maxRows data rows.
|
||||
// A maxRows value of zero means unlimited. This lets an operator make a small,
|
||||
// representative pipeline check without materialising a second dataset file.
|
||||
func SplitTSVLimit(r io.Reader, rowsPerShard, maxRows int, emit func(index int, shard io.Reader) error) error {
|
||||
return splitTSVLimit(r, rowsPerShard, maxRows, nil, emit)
|
||||
}
|
||||
|
||||
// SplitChEMBLTSVLimit is the coordinator's scientific-upload splitter. It
|
||||
// validates the two columns every local SciMesh workload requires before any
|
||||
// shard task is persisted, while generic SplitTSV remains reusable for future
|
||||
// non-chemistry workloads.
|
||||
func SplitChEMBLTSVLimit(r io.Reader, rowsPerShard, maxRows int, emit func(index int, shard io.Reader) error) error {
|
||||
return splitTSVLimit(r, rowsPerShard, maxRows, validateChEMBLHeader, emit)
|
||||
}
|
||||
|
||||
func splitTSVLimit(r io.Reader, rowsPerShard, maxRows int, validateHeader func([]byte) error, emit func(index int, shard io.Reader) error) error {
|
||||
if rowsPerShard <= 0 {
|
||||
return fmt.Errorf("rowsPerShard must be positive, got %d", rowsPerShard)
|
||||
}
|
||||
if maxRows < 0 {
|
||||
return fmt.Errorf("maxRows must be non-negative, got %d", maxRows)
|
||||
}
|
||||
|
||||
sc := bufio.NewScanner(r)
|
||||
// Allow long lines: a SMILES row can be far wider than bufio's 64 KB default.
|
||||
@@ -41,6 +70,11 @@ func SplitTSV(r io.Reader, rowsPerShard int, emit func(index int, shard io.Reade
|
||||
return ErrNoRows // completely empty input
|
||||
}
|
||||
header := append([]byte(nil), sc.Bytes()...)
|
||||
if validateHeader != nil {
|
||||
if err := validateHeader(header); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
var (
|
||||
buf bytes.Buffer
|
||||
@@ -61,9 +95,15 @@ func SplitTSV(r io.Reader, rowsPerShard int, emit func(index int, shard io.Reade
|
||||
|
||||
for sc.Scan() {
|
||||
if rows == 0 {
|
||||
if len(header)+1 > maxShardBytes {
|
||||
return fmt.Errorf("TSV header exceeds maximum shard size of %d bytes", maxShardBytes)
|
||||
}
|
||||
buf.Write(header)
|
||||
buf.WriteByte('\n')
|
||||
}
|
||||
if buf.Len()+len(sc.Bytes())+1 > maxShardBytes {
|
||||
return fmt.Errorf("shard exceeds maximum size of %d bytes; lower rowsPerShard", maxShardBytes)
|
||||
}
|
||||
buf.Write(sc.Bytes())
|
||||
buf.WriteByte('\n')
|
||||
rows++
|
||||
@@ -73,6 +113,9 @@ func SplitTSV(r io.Reader, rowsPerShard int, emit func(index int, shard io.Reade
|
||||
return err
|
||||
}
|
||||
}
|
||||
if maxRows > 0 && index*rowsPerShard+rows == maxRows {
|
||||
break
|
||||
}
|
||||
}
|
||||
if err := sc.Err(); err != nil {
|
||||
return fmt.Errorf("read rows: %w", err)
|
||||
@@ -90,3 +133,17 @@ func SplitTSV(r io.Reader, rowsPerShard int, emit func(index int, shard io.Reade
|
||||
}
|
||||
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
|
||||
}
|
||||
|
||||
@@ -103,6 +103,30 @@ func TestSplitSingleShardWhenSizeExceedsRows(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestSplitLimitUsesOnlyLeadingDataRows(t *testing.T) {
|
||||
input := "h\nr1\nr2\nr3\nr4\nr5\n"
|
||||
var shards []string
|
||||
err := SplitTSVLimit(strings.NewReader(input), 2, 3, func(_ int, shard io.Reader) error {
|
||||
b, _ := io.ReadAll(shard)
|
||||
shards = append(shards, string(b))
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got, want := strings.Join(shards, ""), "h\nr1\nr2\nh\nr3\n"; got != want {
|
||||
t.Errorf("limited shards = %q, want %q", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestChEMBLSplitRejectsMissingRequiredColumns(t *testing.T) {
|
||||
err := SplitChEMBLTSVLimit(strings.NewReader("id\tsmiles\nA\tCC\n"), 1, 0,
|
||||
func(int, io.Reader) error { return nil })
|
||||
if err == nil || !strings.Contains(err.Error(), "chembl_id") {
|
||||
t.Errorf("err = %v, want missing-column error", err)
|
||||
}
|
||||
}
|
||||
|
||||
// The scanned bytes are reused by bufio; the shard buffer must copy them, or a
|
||||
// later row would corrupt an earlier one. This guards that copy.
|
||||
func TestSplitDoesNotAliasScannerBuffer(t *testing.T) {
|
||||
|
||||
@@ -8,13 +8,14 @@ import "errors"
|
||||
//
|
||||
// Always compare with errors.Is — outer layers may wrap these with %w.
|
||||
var (
|
||||
ErrJobNotFound = errors.New("job not found")
|
||||
ErrTaskNotFound = errors.New("task not found")
|
||||
ErrWorkerNotFound = errors.New("worker not found")
|
||||
ErrArtifactNotFound = errors.New("artifact not found")
|
||||
ErrLeaseConflict = errors.New("task leased to another worker")
|
||||
ErrStaleAttempt = errors.New("attempt does not match lease")
|
||||
ErrResultConflict = errors.New("different result already recorded")
|
||||
ErrInvalidInput = errors.New("invalid input")
|
||||
ErrTaskNotLeased = errors.New("task is not currently leased")
|
||||
ErrJobNotFound = errors.New("job not found")
|
||||
ErrTaskNotFound = errors.New("task not found")
|
||||
ErrWorkerNotFound = errors.New("worker not found")
|
||||
ErrArtifactNotFound = errors.New("artifact not found")
|
||||
ErrJobNotCancellable = errors.New("job cannot be cancelled")
|
||||
ErrLeaseConflict = errors.New("task leased to another worker")
|
||||
ErrStaleAttempt = errors.New("attempt does not match lease")
|
||||
ErrResultConflict = errors.New("different result already recorded")
|
||||
ErrInvalidInput = errors.New("invalid input")
|
||||
ErrTaskNotLeased = errors.New("task is not currently leased")
|
||||
)
|
||||
|
||||
@@ -99,18 +99,21 @@ func NewJobWithTasks(workload, inputURI string, params map[string]any,
|
||||
|
||||
// JobProgress is the aggregate view of a job and the state of its tasks.
|
||||
type JobProgress struct {
|
||||
Job Job
|
||||
Total int
|
||||
Pending int
|
||||
Leased int
|
||||
Done int
|
||||
Failed int
|
||||
Job Job
|
||||
Total int
|
||||
Pending int
|
||||
Leased int
|
||||
Done int
|
||||
Failed int
|
||||
Cancelled int
|
||||
}
|
||||
|
||||
// DeriveStatus computes what the job's status should be from its task counts,
|
||||
// so the rule lives here rather than in a SQL trigger or a handler.
|
||||
func (p JobProgress) DeriveStatus() JobStatus {
|
||||
switch {
|
||||
case p.Job.Status == JobCancelled:
|
||||
return JobCancelled
|
||||
case p.Total == 0:
|
||||
return JobPending
|
||||
case p.Done == p.Total:
|
||||
|
||||
@@ -81,6 +81,7 @@ func TestDeriveStatus(t *testing.T) {
|
||||
{"all done", JobProgress{Total: 3, Done: 3}, JobCompleted},
|
||||
{"done and failed", JobProgress{Total: 3, Done: 2, Failed: 1}, JobFailed},
|
||||
{"failed but work remains", JobProgress{Total: 3, Pending: 1, Failed: 2}, JobRunning},
|
||||
{"cancelled job wins over task histogram", JobProgress{Job: Job{Status: JobCancelled}, Total: 3, Done: 1, Cancelled: 2}, JobCancelled},
|
||||
}
|
||||
for _, c := range cases {
|
||||
t.Run(c.name, func(t *testing.T) {
|
||||
|
||||
@@ -259,6 +259,24 @@ func (t *Task) ExpireLease(now time.Time) {
|
||||
t.CompletedAt = &now
|
||||
}
|
||||
|
||||
// Cancel prevents any further worker transition for a task that has not
|
||||
// reached a terminal result. A cancelled lease deliberately becomes invalid:
|
||||
// a worker still running locally must not upload or complete after its job was
|
||||
// stopped by the operator.
|
||||
func (t *Task) Cancel(now time.Time) bool {
|
||||
if t.Status == TaskCompleted || t.Status == TaskFailed || t.Status == TaskCancelled {
|
||||
return false
|
||||
}
|
||||
t.Status = TaskCancelled
|
||||
t.LeaseOwner = nil
|
||||
t.LeaseExpiresAt = nil
|
||||
t.ErrorCode = nil
|
||||
t.ErrorMessage = nil
|
||||
t.CompletedAt = &now
|
||||
t.Version++
|
||||
return true
|
||||
}
|
||||
|
||||
// ClaimedTask is the worker-facing projection of a leased task. Input is either
|
||||
// an external URI or a coordinator-stored shard (InputArtifactID set); the
|
||||
// transport turns the latter into a coordinator download URL.
|
||||
|
||||
@@ -168,6 +168,23 @@ func TestExpireLeaseIgnoresUnleasedTasks(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestCancelInvalidatesLeaseButPreservesTerminalTask(t *testing.T) {
|
||||
task := leasedTask(1, 3)
|
||||
if !task.Cancel(testNow) {
|
||||
t.Fatal("leased task should be cancelled")
|
||||
}
|
||||
if task.Status != TaskCancelled || task.LeaseOwner != nil || task.LeaseExpiresAt != nil {
|
||||
t.Errorf("cancelled task = %+v", task)
|
||||
}
|
||||
if task.Cancel(testLater) {
|
||||
t.Error("cancelled task must not be changed twice")
|
||||
}
|
||||
completed := &Task{Status: TaskCompleted}
|
||||
if completed.Cancel(testNow) {
|
||||
t.Error("completed task must remain terminal")
|
||||
}
|
||||
}
|
||||
|
||||
func TestFirstHeartbeatMovesLeasedToRunning(t *testing.T) {
|
||||
task := leasedTask(1, 3)
|
||||
until := testLater.Add(time.Hour)
|
||||
|
||||
@@ -24,6 +24,8 @@ type Config struct {
|
||||
DatabaseURL string
|
||||
// Shared bearer token workers must present. Empty disables auth (dev only).
|
||||
Token string
|
||||
// Local operator UI credential. Empty disables the embedded UI entirely.
|
||||
UIToken string
|
||||
|
||||
// Minimum log level: debug, info, warn, error.
|
||||
LogLevel string
|
||||
@@ -77,6 +79,7 @@ func LoadConfig() (Config, error) {
|
||||
// COORDINATOR_TOKEN is the contract name; WORKER_AUTH_TOKEN is the
|
||||
// former name, still honoured so existing .env files keep working.
|
||||
Token: getEnv("COORDINATOR_TOKEN", os.Getenv("WORKER_AUTH_TOKEN")),
|
||||
UIToken: os.Getenv("UI_AUTH_TOKEN"),
|
||||
LogLevel: getEnv("LOG_LEVEL", "info"),
|
||||
LogFile: os.Getenv("LOG_FILE"),
|
||||
StorageDir: getEnv("COORDINATOR_STORAGE_DIR", "./data"),
|
||||
@@ -94,6 +97,9 @@ func LoadConfig() (Config, error) {
|
||||
if cfg.DatabaseURL == "" {
|
||||
return Config{}, fmt.Errorf("DATABASE_URL is required")
|
||||
}
|
||||
if cfg.UIToken != "" && cfg.Token != "" && cfg.UIToken == cfg.Token {
|
||||
return Config{}, fmt.Errorf("UI_AUTH_TOKEN must differ from the worker auth token")
|
||||
}
|
||||
|
||||
var err error
|
||||
if cfg.DBMaxConns, err = getEnvInt32("DB_MAX_CONNS", cfg.DBMaxConns); err != nil {
|
||||
@@ -123,6 +129,9 @@ func LoadConfig() (Config, error) {
|
||||
if cfg.DefaultMaxAttempts, err = getEnvInt("DEFAULT_MAX_ATTEMPTS", cfg.DefaultMaxAttempts); err != nil {
|
||||
return Config{}, err
|
||||
}
|
||||
if cfg.DefaultMaxAttempts < 1 {
|
||||
return Config{}, fmt.Errorf("DEFAULT_MAX_ATTEMPTS must be positive")
|
||||
}
|
||||
|
||||
return cfg, nil
|
||||
}
|
||||
|
||||
@@ -0,0 +1,45 @@
|
||||
package infra
|
||||
|
||||
import (
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestLoadConfigRejectsSharedUIAndWorkerToken(t *testing.T) {
|
||||
t.Setenv("ENV_FILE", filepath.Join(t.TempDir(), "missing.env"))
|
||||
t.Setenv("DATABASE_URL", "postgres://test")
|
||||
t.Setenv("COORDINATOR_TOKEN", "shared-secret")
|
||||
t.Setenv("UI_AUTH_TOKEN", "shared-secret")
|
||||
|
||||
_, err := LoadConfig()
|
||||
if err == nil || !strings.Contains(err.Error(), "must differ") {
|
||||
t.Fatalf("LoadConfig error = %v, want distinct-token error", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoadConfigAllowsDistinctUIAndWorkerTokens(t *testing.T) {
|
||||
t.Setenv("ENV_FILE", filepath.Join(t.TempDir(), "missing.env"))
|
||||
t.Setenv("DATABASE_URL", "postgres://test")
|
||||
t.Setenv("COORDINATOR_TOKEN", "worker-secret")
|
||||
t.Setenv("UI_AUTH_TOKEN", "ui-secret")
|
||||
|
||||
cfg, err := LoadConfig()
|
||||
if err != nil {
|
||||
t.Fatalf("LoadConfig: %v", err)
|
||||
}
|
||||
if cfg.Token != "worker-secret" || cfg.UIToken != "ui-secret" {
|
||||
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)
|
||||
}
|
||||
}
|
||||
@@ -148,17 +148,30 @@ func (r *TaskRepo) CountByStatus(ctx context.Context, jobID uuid.UUID) (map[doma
|
||||
return counts, nil
|
||||
}
|
||||
|
||||
func (r *TaskRepo) ExpireLeases(ctx context.Context, now time.Time) (int64, error) {
|
||||
func (r *TaskRepo) CancelByJob(_ context.Context, jobID uuid.UUID, now time.Time) (int64, error) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
var n int64
|
||||
for _, t := range r.tasks {
|
||||
if t.Status == domain.TaskLeased && t.LeaseExpiresAt != nil && t.LeaseExpiresAt.Before(now) {
|
||||
t.ExpireLease(now)
|
||||
n++
|
||||
var cancelled int64
|
||||
for _, task := range r.tasks {
|
||||
if task.JobID == jobID && task.Cancel(now) {
|
||||
cancelled++
|
||||
}
|
||||
}
|
||||
return n, nil
|
||||
return cancelled, nil
|
||||
}
|
||||
|
||||
func (r *TaskRepo) ExpireLeases(ctx context.Context, now time.Time) ([]uuid.UUID, error) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
affected := make([]uuid.UUID, 0)
|
||||
for _, t := range r.tasks {
|
||||
if (t.Status == domain.TaskLeased || t.Status == domain.TaskRunning) &&
|
||||
t.LeaseExpiresAt != nil && t.LeaseExpiresAt.Before(now) {
|
||||
t.ExpireLease(now)
|
||||
affected = append(affected, t.JobID)
|
||||
}
|
||||
}
|
||||
return affected, nil
|
||||
}
|
||||
|
||||
// --- JobRepo -------------------------------------------------------------
|
||||
@@ -286,6 +299,23 @@ func (r *ArtifactRepo) Get(ctx context.Context, id uuid.UUID) (*domain.Artifact,
|
||||
return &cp, nil
|
||||
}
|
||||
|
||||
func (r *ArtifactRepo) FindPartialResult(_ context.Context, taskID uuid.UUID, attempt int) (*domain.Artifact, error) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
for _, a := range r.arts {
|
||||
if a.TaskID != nil && *a.TaskID == taskID && a.Kind == domain.ArtifactPartialResult &&
|
||||
a.Attempt != nil && *a.Attempt == attempt {
|
||||
return cloneArtifact(a), nil
|
||||
}
|
||||
}
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
func cloneArtifact(a *domain.Artifact) *domain.Artifact {
|
||||
cp := *a
|
||||
return &cp
|
||||
}
|
||||
|
||||
// --- BlobStore -----------------------------------------------------------
|
||||
|
||||
type BlobStore struct {
|
||||
|
||||
@@ -0,0 +1,114 @@
|
||||
package memstore
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sort"
|
||||
|
||||
"github.com/google/uuid"
|
||||
|
||||
"github.com/emil28092005/SciMesh/coordinator/internal/domain"
|
||||
"github.com/emil28092005/SciMesh/coordinator/internal/usecase"
|
||||
)
|
||||
|
||||
// UIReadRepo is the in-memory read projection used by HTTP/UI tests.
|
||||
type UIReadRepo struct {
|
||||
jobs *JobRepo
|
||||
tasks *TaskRepo
|
||||
workers *WorkerRepo
|
||||
artifacts *ArtifactRepo
|
||||
}
|
||||
|
||||
func NewUIReadRepo(j *JobRepo, t *TaskRepo, w *WorkerRepo, a *ArtifactRepo) *UIReadRepo {
|
||||
return &UIReadRepo{j, t, w, a}
|
||||
}
|
||||
|
||||
var _ usecase.UIReadRepository = (*UIReadRepo)(nil)
|
||||
|
||||
func (r *UIReadRepo) GetJob(ctx context.Context, id uuid.UUID) (*domain.Job, error) {
|
||||
return r.jobs.Get(ctx, id)
|
||||
}
|
||||
func (r *UIReadRepo) ListJobs(_ context.Context, limit int) ([]domain.Job, error) {
|
||||
if limit < 1 || limit > 100 {
|
||||
return nil, domain.ErrInvalidInput
|
||||
}
|
||||
r.jobs.mu.Lock()
|
||||
defer r.jobs.mu.Unlock()
|
||||
out := make([]domain.Job, 0, len(r.jobs.jobs))
|
||||
for _, job := range r.jobs.jobs {
|
||||
out = append(out, *job)
|
||||
}
|
||||
sort.Slice(out, func(i, j int) bool {
|
||||
if out[i].CreatedAt.Equal(out[j].CreatedAt) {
|
||||
return out[i].ID.String() > out[j].ID.String()
|
||||
}
|
||||
return out[i].CreatedAt.After(out[j].CreatedAt)
|
||||
})
|
||||
if len(out) > limit {
|
||||
out = out[:limit]
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
func (r *UIReadRepo) ListTasksByJob(_ context.Context, jobID uuid.UUID) ([]domain.Task, error) {
|
||||
r.tasks.mu.Lock()
|
||||
defer r.tasks.mu.Unlock()
|
||||
out := []domain.Task{}
|
||||
for _, task := range r.tasks.tasks {
|
||||
if task.JobID == jobID {
|
||||
out = append(out, *clone(task))
|
||||
}
|
||||
}
|
||||
sort.Slice(out, func(i, j int) bool { return out[i].ChunkIndex < out[j].ChunkIndex })
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (r *UIReadRepo) ListTasksByJobs(ctx context.Context, jobIDs []uuid.UUID) (map[uuid.UUID][]domain.Task, error) {
|
||||
out := make(map[uuid.UUID][]domain.Task, len(jobIDs))
|
||||
for _, id := range jobIDs {
|
||||
tasks, err := r.ListTasksByJob(ctx, id)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out[id] = tasks
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
func (r *UIReadRepo) ListWorkers(_ context.Context, limit int) ([]domain.Worker, error) {
|
||||
if limit < 1 || limit > 100 {
|
||||
return nil, domain.ErrInvalidInput
|
||||
}
|
||||
r.workers.mu.Lock()
|
||||
defer r.workers.mu.Unlock()
|
||||
out := []domain.Worker{}
|
||||
for _, worker := range r.workers.workers {
|
||||
copy := *worker
|
||||
copy.Capabilities = append([]string(nil), worker.Capabilities...)
|
||||
out = append(out, copy)
|
||||
}
|
||||
sort.Slice(out, func(i, j int) bool {
|
||||
if out[i].LastHeartbeatAt.Equal(out[j].LastHeartbeatAt) {
|
||||
return out[i].ID.String() > out[j].ID.String()
|
||||
}
|
||||
return out[i].LastHeartbeatAt.After(out[j].LastHeartbeatAt)
|
||||
})
|
||||
if len(out) > limit {
|
||||
out = out[:limit]
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
func (r *UIReadRepo) ListArtifactsByJob(_ context.Context, jobID uuid.UUID) ([]domain.Artifact, error) {
|
||||
r.artifacts.mu.Lock()
|
||||
defer r.artifacts.mu.Unlock()
|
||||
out := []domain.Artifact{}
|
||||
for _, artifact := range r.artifacts.arts {
|
||||
if artifact.JobID == jobID {
|
||||
out = append(out, *artifact)
|
||||
}
|
||||
}
|
||||
sort.Slice(out, func(i, j int) bool {
|
||||
if out[i].CreatedAt.Equal(out[j].CreatedAt) {
|
||||
return out[i].ID.String() < out[j].ID.String()
|
||||
}
|
||||
return out[i].CreatedAt.Before(out[j].CreatedAt)
|
||||
})
|
||||
return out, nil
|
||||
}
|
||||
@@ -70,3 +70,33 @@ func (r *ArtifactRepo) Get(ctx context.Context, id uuid.UUID) (*domain.Artifact,
|
||||
a.Kind = domain.ArtifactKind(kind)
|
||||
return &a, nil
|
||||
}
|
||||
|
||||
func (r *ArtifactRepo) FindPartialResult(ctx context.Context, taskID uuid.UUID, attempt int) (*domain.Artifact, error) {
|
||||
sql, args, err := psql.Select(artifactColumns...).
|
||||
From("artifacts").
|
||||
Where(sq.Eq{
|
||||
"task_id": taskID,
|
||||
"attempt": attempt,
|
||||
"kind": string(domain.ArtifactPartialResult),
|
||||
}).
|
||||
ToSql()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
var (
|
||||
a domain.Artifact
|
||||
kind string
|
||||
)
|
||||
err = conn(ctx, r.pool).QueryRow(ctx, sql, args...).Scan(
|
||||
&a.ID, &a.JobID, &a.TaskID, &a.Attempt, &kind, &a.Filename, &a.StorageKey,
|
||||
&a.ContentType, &a.SizeBytes, &a.SHA256, &a.CreatedAt)
|
||||
if errors.Is(err, pgx.ErrNoRows) {
|
||||
return nil, nil
|
||||
}
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("find partial result: %w", err)
|
||||
}
|
||||
a.Kind = domain.ArtifactKind(kind)
|
||||
return &a, nil
|
||||
}
|
||||
|
||||
@@ -137,30 +137,37 @@ func TestConcurrentClaimGivesEachTaskToExactlyOneWorker(t *testing.T) {
|
||||
claimed = make(map[uuid.UUID]string)
|
||||
wg sync.WaitGroup
|
||||
)
|
||||
// More workers than tasks, so the surplus must come back empty rather than
|
||||
// steal an already-leased row.
|
||||
// More workers than tasks. With SKIP LOCKED, a concurrent caller can
|
||||
// transiently see no eligible row while every remaining row is locked by a
|
||||
// different claim statement. Poll briefly, as a real worker does, before
|
||||
// treating the queue as empty. This verifies the actual contract: tasks are
|
||||
// unique and all eventually become claimable without lock contention.
|
||||
for i := 0; i < tasks*2; i++ {
|
||||
wg.Add(1)
|
||||
go func(n int) {
|
||||
defer wg.Done()
|
||||
task, err := repo.ClaimNext(context.Background(), usecase.ClaimFilter{
|
||||
Owner: fmt.Sprintf("worker-%d", n),
|
||||
Now: now,
|
||||
LeaseUntil: now.Add(time.Minute),
|
||||
})
|
||||
if err != nil {
|
||||
t.Errorf("claim: %v", err)
|
||||
for attempt := 0; attempt < 20; attempt++ {
|
||||
task, err := repo.ClaimNext(context.Background(), usecase.ClaimFilter{
|
||||
Owner: fmt.Sprintf("worker-%d", n),
|
||||
Now: now,
|
||||
LeaseUntil: now.Add(time.Minute),
|
||||
})
|
||||
if err != nil {
|
||||
t.Errorf("claim: %v", err)
|
||||
return
|
||||
}
|
||||
if task == nil || task.JobID != job.ID {
|
||||
time.Sleep(time.Millisecond)
|
||||
continue
|
||||
}
|
||||
mu.Lock()
|
||||
if prev, dup := claimed[task.ID]; dup {
|
||||
t.Errorf("task %s handed to both %s and worker-%d", task.ID, prev, n)
|
||||
}
|
||||
claimed[task.ID] = fmt.Sprintf("worker-%d", n)
|
||||
mu.Unlock()
|
||||
return
|
||||
}
|
||||
if task == nil || task.JobID != job.ID {
|
||||
return // empty queue, or a task from another test's job
|
||||
}
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
if prev, dup := claimed[task.ID]; dup {
|
||||
t.Errorf("task %s handed to both %s and worker-%d", task.ID, prev, n)
|
||||
}
|
||||
claimed[task.ID] = fmt.Sprintf("worker-%d", n)
|
||||
}(i)
|
||||
}
|
||||
wg.Wait()
|
||||
@@ -199,6 +206,27 @@ func TestClaimNextReturnsNilOnEmptyQueue(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestCancelJobCancelsEveryUnfinishedTask(t *testing.T) {
|
||||
pool := testPool(t)
|
||||
ctx := context.Background()
|
||||
job, _ := seedJob(t, pool, 3)
|
||||
clk := fixedClock{now: time.Now().UTC()}
|
||||
uc := usecase.NewCancelJob(NewJobRepo(pool), NewTaskRepo(pool), NewTxManager(pool), clk)
|
||||
|
||||
cancelled, err := uc.Execute(ctx, job.ID)
|
||||
if err != nil || cancelled != 3 {
|
||||
t.Fatalf("cancel = (%d, %v), want (3, nil)", cancelled, err)
|
||||
}
|
||||
stored, err := NewJobRepo(pool).Get(ctx, job.ID)
|
||||
if err != nil || stored.Status != domain.JobCancelled {
|
||||
t.Fatalf("job after cancel = (%+v, %v)", stored, err)
|
||||
}
|
||||
counts, err := NewTaskRepo(pool).CountByStatus(ctx, job.ID)
|
||||
if err != nil || counts[domain.TaskCancelled] != 3 {
|
||||
t.Fatalf("cancelled tasks = %d, err = %v", counts[domain.TaskCancelled], err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUpdateRejectsStaleVersion(t *testing.T) {
|
||||
pool := testPool(t)
|
||||
ctx := context.Background()
|
||||
@@ -332,12 +360,23 @@ func TestCompleteTaskReplayIsIdempotent(t *testing.T) {
|
||||
t.Errorf("replay must be idempotent, got %v", err)
|
||||
}
|
||||
|
||||
// A different result artifact for the same task is a genuine conflict.
|
||||
art2 := seedArtifact(t, pool, job.ID, &claimed.ID, domain.ArtifactPartialResult)
|
||||
other := in
|
||||
other.ResultArtifactID = art2.ID
|
||||
if _, err := uc.Execute(ctx, other); !errors.Is(err, domain.ErrResultConflict) {
|
||||
t.Errorf("err = %v, want ErrResultConflict", err)
|
||||
}
|
||||
|
||||
func TestPartialResultIsUniquePerTaskAttempt(t *testing.T) {
|
||||
pool := testPool(t)
|
||||
ctx := context.Background()
|
||||
job, tasks := seedJob(t, pool, 1)
|
||||
taskID := tasks[0].ID
|
||||
first := seedArtifact(t, pool, job.ID, &taskID, domain.ArtifactPartialResult)
|
||||
second, err := domain.NewArtifact(job.ID, &taskID, domain.ArtifactPartialResult, "retry.csv", "text/csv", time.Now().UTC())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
attempt := 1
|
||||
second.Attempt = &attempt
|
||||
second.SetContent("other-sha", 5)
|
||||
if err := NewArtifactRepo(pool).Insert(ctx, second); err == nil {
|
||||
t.Fatalf("second partial artifact for %s/%d was accepted after %s", taskID, attempt, first.ID)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -295,6 +295,29 @@ func (r *TaskRepo) CountByStatus(ctx context.Context, jobID uuid.UUID) (map[doma
|
||||
return counts, rows.Err()
|
||||
}
|
||||
|
||||
// cancelByJobSQL mirrors domain.Task.Cancel in one set-based update. It runs in
|
||||
// the same transaction as the job-status update, so no claimable shard remains
|
||||
// after an operator receives a successful cancellation response.
|
||||
const cancelByJobSQL = `
|
||||
UPDATE tasks
|
||||
SET status = 'cancelled'::task_status,
|
||||
lease_owner = NULL,
|
||||
lease_expires_at = NULL,
|
||||
error_code = NULL,
|
||||
error_message = NULL,
|
||||
completed_at = $2,
|
||||
version = version + 1
|
||||
WHERE job_id = $1
|
||||
AND status IN ('pending','leased','running')`
|
||||
|
||||
func (r *TaskRepo) CancelByJob(ctx context.Context, jobID uuid.UUID, now time.Time) (int64, error) {
|
||||
tag, err := conn(ctx, r.pool).Exec(ctx, cancelByJobSQL, jobID, now)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return tag.RowsAffected(), nil
|
||||
}
|
||||
|
||||
// expireLeasesSQL applies the lease-expiry rule set-based, mirroring
|
||||
// domain.Task.ExpireLease: requeue while attempts remain, otherwise fail.
|
||||
//
|
||||
@@ -314,17 +337,26 @@ SET status = CASE WHEN attempt < max_attempts THEN 'pending'::task_sta
|
||||
ELSE error_message END,
|
||||
completed_at = CASE WHEN attempt >= max_attempts THEN $1 ELSE completed_at END,
|
||||
version = version + 1
|
||||
WHERE status IN ('leased','running') AND lease_expires_at < $1`
|
||||
WHERE status IN ('leased','running') AND lease_expires_at < $1
|
||||
RETURNING job_id`
|
||||
|
||||
func (r *TaskRepo) ExpireLeases(ctx context.Context, now time.Time) (int64, error) {
|
||||
var affected int64
|
||||
func (r *TaskRepo) ExpireLeases(ctx context.Context, now time.Time) ([]uuid.UUID, error) {
|
||||
var affected []uuid.UUID
|
||||
err := withRetry(ctx, func(ctx context.Context) error {
|
||||
tag, err := conn(ctx, r.pool).Exec(ctx, expireLeasesSQL, now, domain.ErrCodeLeaseExpired)
|
||||
rows, err := conn(ctx, r.pool).Query(ctx, expireLeasesSQL, now, domain.ErrCodeLeaseExpired)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
affected = tag.RowsAffected()
|
||||
return nil
|
||||
defer rows.Close()
|
||||
affected = affected[:0]
|
||||
for rows.Next() {
|
||||
var jobID uuid.UUID
|
||||
if err := rows.Scan(&jobID); err != nil {
|
||||
return err
|
||||
}
|
||||
affected = append(affected, jobID)
|
||||
}
|
||||
return rows.Err()
|
||||
})
|
||||
return affected, err
|
||||
}
|
||||
|
||||
@@ -0,0 +1,148 @@
|
||||
package postgres
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
|
||||
sq "github.com/Masterminds/squirrel"
|
||||
"github.com/google/uuid"
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
|
||||
"github.com/emil28092005/SciMesh/coordinator/internal/domain"
|
||||
"github.com/emil28092005/SciMesh/coordinator/internal/usecase"
|
||||
)
|
||||
|
||||
// UIReadRepo contains bounded, deterministic read queries for the operator UI.
|
||||
type UIReadRepo struct{ pool *pgxpool.Pool }
|
||||
|
||||
func NewUIReadRepo(pool *pgxpool.Pool) *UIReadRepo { return &UIReadRepo{pool: pool} }
|
||||
|
||||
var _ usecase.UIReadRepository = (*UIReadRepo)(nil)
|
||||
|
||||
func (r *UIReadRepo) GetJob(ctx context.Context, id uuid.UUID) (*domain.Job, error) {
|
||||
job, err := NewJobRepo(r.pool).Get(ctx, id)
|
||||
return job, err
|
||||
}
|
||||
|
||||
func (r *UIReadRepo) ListJobs(ctx context.Context, limit int) ([]domain.Job, error) {
|
||||
if limit < 1 || limit > 100 {
|
||||
return nil, domain.ErrInvalidInput
|
||||
}
|
||||
sql, args, err := psql.Select(jobColumns...).From("jobs").OrderBy("created_at DESC", "id DESC").Limit(uint64(limit)).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 jobs: %w", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
jobs := make([]domain.Job, 0)
|
||||
for rows.Next() {
|
||||
var j domain.Job
|
||||
var status string
|
||||
var inputURI *string
|
||||
if err := rows.Scan(&j.ID, &j.Workload, &inputURI, &j.Parameters, &status, &j.CreatedAt, &j.CompletedAt); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if inputURI != nil {
|
||||
j.InputURI = *inputURI
|
||||
}
|
||||
j.Status = domain.JobStatus(status)
|
||||
jobs = append(jobs, j)
|
||||
}
|
||||
return jobs, rows.Err()
|
||||
}
|
||||
|
||||
func (r *UIReadRepo) ListTasksByJob(ctx context.Context, jobID uuid.UUID) ([]domain.Task, error) {
|
||||
sql, args, err := psql.Select(taskColumns...).From("tasks").Where(sq.Eq{"job_id": jobID}).OrderBy("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: %w", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
tasks := make([]domain.Task, 0)
|
||||
for rows.Next() {
|
||||
task, err := scanTask(rows)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
tasks = append(tasks, *task)
|
||||
}
|
||||
return tasks, rows.Err()
|
||||
}
|
||||
|
||||
func (r *UIReadRepo) ListTasksByJobs(ctx context.Context, jobIDs []uuid.UUID) (map[uuid.UUID][]domain.Task, error) {
|
||||
out := make(map[uuid.UUID][]domain.Task, len(jobIDs))
|
||||
if len(jobIDs) == 0 {
|
||||
return out, nil
|
||||
}
|
||||
sql, args, err := psql.Select(taskColumns...).From("tasks").
|
||||
Where(sq.Eq{"job_id": jobIDs}).OrderBy("job_id ASC", "chunk_index ASC").ToSql()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
rows, err := conn(ctx, r.pool).Query(ctx, sql, args...)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("list tasks by jobs: %w", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
for rows.Next() {
|
||||
task, err := scanTask(rows)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out[task.JobID] = append(out[task.JobID], *task)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
func (r *UIReadRepo) ListWorkers(ctx context.Context, limit int) ([]domain.Worker, error) {
|
||||
if limit < 1 || limit > 100 {
|
||||
return nil, domain.ErrInvalidInput
|
||||
}
|
||||
sql, args, err := psql.Select(workerColumns...).From("workers").OrderBy("last_heartbeat_at DESC", "id DESC").Limit(uint64(limit)).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 workers: %w", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
workers := make([]domain.Worker, 0)
|
||||
for rows.Next() {
|
||||
worker, err := scanWorker(rows)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
workers = append(workers, *worker)
|
||||
}
|
||||
return workers, rows.Err()
|
||||
}
|
||||
|
||||
func (r *UIReadRepo) ListArtifactsByJob(ctx context.Context, jobID uuid.UUID) ([]domain.Artifact, error) {
|
||||
sql, args, err := psql.Select(artifactColumns...).From("artifacts").Where(sq.Eq{"job_id": jobID}).OrderBy("created_at ASC", "id 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 artifacts: %w", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
artifacts := make([]domain.Artifact, 0)
|
||||
for rows.Next() {
|
||||
var a domain.Artifact
|
||||
var kind string
|
||||
if err := rows.Scan(&a.ID, &a.JobID, &a.TaskID, &a.Attempt, &kind, &a.Filename, &a.StorageKey, &a.ContentType, &a.SizeBytes, &a.SHA256, &a.CreatedAt); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
a.Kind = domain.ArtifactKind(kind)
|
||||
artifacts = append(artifacts, a)
|
||||
}
|
||||
return artifacts, rows.Err()
|
||||
}
|
||||
@@ -111,13 +111,14 @@ type uploadJobResponse struct {
|
||||
}
|
||||
|
||||
type jobProgressResponse struct {
|
||||
ID uuid.UUID `json:"id"`
|
||||
Status string `json:"status"`
|
||||
Total int `json:"total"`
|
||||
Pending int `json:"pending"`
|
||||
Leased int `json:"leased"`
|
||||
Done int `json:"completed"`
|
||||
Failed int `json:"failed"`
|
||||
ID uuid.UUID `json:"id"`
|
||||
Status string `json:"status"`
|
||||
Total int `json:"total"`
|
||||
Pending int `json:"pending"`
|
||||
Leased int `json:"leased"`
|
||||
Done int `json:"completed"`
|
||||
Failed int `json:"failed"`
|
||||
Cancelled int `json:"cancelled"`
|
||||
}
|
||||
|
||||
type uploadArtifactResponse struct {
|
||||
@@ -153,12 +154,13 @@ func toClaimedTaskResponse(c domain.ClaimedTask) claimedTaskResponse {
|
||||
|
||||
func toJobProgressResponse(p domain.JobProgress) jobProgressResponse {
|
||||
return jobProgressResponse{
|
||||
ID: p.Job.ID,
|
||||
Status: string(p.DeriveStatus()),
|
||||
Total: p.Total,
|
||||
Pending: p.Pending,
|
||||
Leased: p.Leased,
|
||||
Done: p.Done,
|
||||
Failed: p.Failed,
|
||||
ID: p.Job.ID,
|
||||
Status: string(p.DeriveStatus()),
|
||||
Total: p.Total,
|
||||
Pending: p.Pending,
|
||||
Leased: p.Leased,
|
||||
Done: p.Done,
|
||||
Failed: p.Failed,
|
||||
Cancelled: p.Cancelled,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -50,7 +50,8 @@ func (s *Server) writeError(w http.ResponseWriter, r *http.Request, err error) {
|
||||
case errors.Is(err, domain.ErrLeaseConflict),
|
||||
errors.Is(err, domain.ErrStaleAttempt),
|
||||
errors.Is(err, domain.ErrResultConflict),
|
||||
errors.Is(err, domain.ErrTaskNotLeased):
|
||||
errors.Is(err, domain.ErrTaskNotLeased),
|
||||
errors.Is(err, domain.ErrJobNotCancellable):
|
||||
status = http.StatusConflict
|
||||
case errors.Is(err, usecase.ErrNotImplemented):
|
||||
status = http.StatusNotImplemented
|
||||
|
||||
@@ -79,6 +79,10 @@ func (s *Server) handleClaim(w http.ResponseWriter, r *http.Request) {
|
||||
s.writeError(w, r, domain.ErrInvalidInput)
|
||||
return
|
||||
}
|
||||
if _, err := uuid.Parse(req.WorkerID); err != nil {
|
||||
s.writeError(w, r, domain.ErrInvalidInput)
|
||||
return
|
||||
}
|
||||
|
||||
claimed, err := s.uc.ClaimTask.Execute(ctx, usecase.ClaimTaskInput{
|
||||
WorkerID: req.WorkerID,
|
||||
@@ -182,7 +186,7 @@ func (s *Server) handleFailure(w http.ResponseWriter, r *http.Request) {
|
||||
const defaultChunkRows = 1000
|
||||
|
||||
// handleUploadDataset accepts a multipart submission — the dataset file plus the
|
||||
// workload/parameters/chunk_rows fields — and hands the file, streamed, to the
|
||||
// workload/parameters/chunk_rows/max_rows fields — and hands the file, streamed, to the
|
||||
// chunker. The text fields MUST precede the file part: the file is streamed, not
|
||||
// buffered, so by the time it arrives the other fields are already parsed.
|
||||
func (s *Server) handleUploadDataset(w http.ResponseWriter, r *http.Request) {
|
||||
@@ -197,11 +201,13 @@ func (s *Server) handleUploadDataset(w http.ResponseWriter, r *http.Request) {
|
||||
workload string
|
||||
params map[string]any
|
||||
rows = defaultChunkRows
|
||||
maxRows int
|
||||
result usecase.SubmitDatasetResult
|
||||
gotDataset bool
|
||||
gotWorkload bool
|
||||
gotParams bool
|
||||
gotRows bool
|
||||
gotMaxRows bool
|
||||
)
|
||||
|
||||
for {
|
||||
@@ -249,6 +255,19 @@ func (s *Server) handleUploadDataset(w http.ResponseWriter, r *http.Request) {
|
||||
}
|
||||
rows = n
|
||||
gotRows = true
|
||||
case "max_rows":
|
||||
if gotDataset || gotMaxRows {
|
||||
s.writeError(w, r, domain.ErrInvalidInput)
|
||||
return
|
||||
}
|
||||
b, _ := io.ReadAll(io.LimitReader(part, 32))
|
||||
n, err := strconv.Atoi(strings.TrimSpace(string(b)))
|
||||
if err != nil || n < 1 {
|
||||
s.writeError(w, r, domain.ErrInvalidInput)
|
||||
return
|
||||
}
|
||||
maxRows = n
|
||||
gotMaxRows = true
|
||||
case "file", "dataset":
|
||||
if gotDataset || workload == "" {
|
||||
s.writeError(w, r, domain.ErrInvalidInput)
|
||||
@@ -262,6 +281,7 @@ func (s *Server) handleUploadDataset(w http.ResponseWriter, r *http.Request) {
|
||||
Workload: workload,
|
||||
Parameters: params,
|
||||
RowsPerShard: rows,
|
||||
MaxRows: maxRows,
|
||||
Filename: filename,
|
||||
ContentType: part.Header.Get("Content-Type"),
|
||||
Body: part,
|
||||
@@ -379,6 +399,27 @@ func (s *Server) handleGetJob(w http.ResponseWriter, r *http.Request) {
|
||||
writeJSON(w, http.StatusOK, toJobProgressResponse(progress))
|
||||
}
|
||||
|
||||
// handleCancelJob stops all non-terminal shards for an operator-requested job.
|
||||
// It is available to both the bearer API and the separately authenticated UI.
|
||||
func (s *Server) handleCancelJob(w http.ResponseWriter, r *http.Request) {
|
||||
ctx, cancel := s.reqCtx(r)
|
||||
defer cancel()
|
||||
jobID, ok := s.pathUUID(w, r, "job_id")
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
cancelled, err := s.uc.CancelJob.Execute(ctx, jobID)
|
||||
if err != nil {
|
||||
s.writeError(w, r, err)
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, map[string]any{
|
||||
"job_id": jobID,
|
||||
"status": domain.JobCancelled,
|
||||
"cancelled_tasks": cancelled,
|
||||
})
|
||||
}
|
||||
|
||||
// --- helpers ---
|
||||
|
||||
func (s *Server) reqCtx(r *http.Request) (context.Context, context.CancelFunc) {
|
||||
|
||||
@@ -64,6 +64,48 @@ func withAuth(token string) func(http.Handler) http.Handler {
|
||||
}
|
||||
}
|
||||
|
||||
// withBasicAuth protects the local operator UI with a credential distinct from
|
||||
// the worker bearer token. The username is intentionally ignored; the password
|
||||
// is the configured UI token. Basic Auth is suitable only for localhost or a
|
||||
// TLS-terminating trusted reverse proxy.
|
||||
func withBasicAuth(token string) func(http.Handler) http.Handler {
|
||||
return func(next http.Handler) http.Handler {
|
||||
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
_, password, ok := r.BasicAuth()
|
||||
if !ok || subtle.ConstantTimeCompare([]byte(password), []byte(token)) != 1 {
|
||||
w.Header().Set("WWW-Authenticate", `Basic realm="SciMesh UI", charset="UTF-8"`)
|
||||
writeJSON(w, http.StatusUnauthorized, errorResponse{Error: "unauthorized", RequestID: requestIDFrom(r.Context())})
|
||||
return
|
||||
}
|
||||
next.ServeHTTP(w, r)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// withSameOrigin rejects browser form/fetch writes initiated by another origin.
|
||||
// A missing Origin is allowed for direct local tools; authenticated UI pages use
|
||||
// the browser-supplied Origin header on state-changing requests.
|
||||
func withSameOrigin(next http.Handler) http.Handler {
|
||||
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method == http.MethodGet || r.Method == http.MethodHead || r.Method == http.MethodOptions {
|
||||
next.ServeHTTP(w, r)
|
||||
return
|
||||
}
|
||||
origin := r.Header.Get("Origin")
|
||||
if origin != "" {
|
||||
scheme := "http"
|
||||
if r.TLS != nil {
|
||||
scheme = "https"
|
||||
}
|
||||
if origin != scheme+"://"+r.Host {
|
||||
writeJSON(w, http.StatusForbidden, errorResponse{Error: "cross-origin request rejected", RequestID: requestIDFrom(r.Context())})
|
||||
return
|
||||
}
|
||||
}
|
||||
next.ServeHTTP(w, r)
|
||||
})
|
||||
}
|
||||
|
||||
// statusRecorder captures the status code for the access log.
|
||||
type statusRecorder struct {
|
||||
http.ResponseWriter
|
||||
|
||||
@@ -24,9 +24,11 @@ type UseCases struct {
|
||||
CompleteTask *usecase.CompleteTask
|
||||
FailTask *usecase.FailTask
|
||||
GetJobStatus *usecase.GetJobStatus
|
||||
CancelJob *usecase.CancelJob
|
||||
UploadArtifact *usecase.UploadArtifact
|
||||
DownloadArtifact *usecase.DownloadArtifact
|
||||
GetTaskInput *usecase.GetTaskInput
|
||||
Dashboard *usecase.Dashboard
|
||||
}
|
||||
|
||||
type Server struct {
|
||||
@@ -54,12 +56,13 @@ func NewServer(uc UseCases, log *slog.Logger, requestTimeout, heartbeatInterval
|
||||
|
||||
// Handler builds the router. Go 1.22's ServeMux matches on method and path
|
||||
// wildcards, so no third-party router is needed.
|
||||
func (s *Server) Handler(token string) http.Handler {
|
||||
func (s *Server) Handler(token string, uiToken ...string) http.Handler {
|
||||
protected := http.NewServeMux()
|
||||
protected.HandleFunc("POST /workers/register", s.handleRegister)
|
||||
protected.HandleFunc("POST /jobs", s.handleCreateJob)
|
||||
protected.HandleFunc("POST /jobs/upload", s.handleUploadDataset)
|
||||
protected.HandleFunc("GET /jobs/{job_id}", s.handleGetJob)
|
||||
protected.HandleFunc("POST /jobs/{job_id}/cancel", s.handleCancelJob)
|
||||
protected.HandleFunc("POST /tasks/claim", s.handleClaim)
|
||||
protected.HandleFunc("GET /tasks/{task_id}/input", s.handleGetTaskInput)
|
||||
protected.HandleFunc("POST /tasks/{task_id}/heartbeat", s.handleHeartbeat)
|
||||
@@ -70,6 +73,23 @@ func (s *Server) Handler(token string) http.Handler {
|
||||
|
||||
mux := http.NewServeMux()
|
||||
mux.HandleFunc("GET /health", s.handleHealth)
|
||||
if len(uiToken) > 0 && uiToken[0] != "" && s.uc.Dashboard != nil {
|
||||
ui := http.NewServeMux()
|
||||
ui.HandleFunc("GET /ui", s.handleUIHome)
|
||||
ui.HandleFunc("GET /ui/jobs/new", s.handleUINewJob)
|
||||
ui.HandleFunc("GET /ui/jobs/{job_id}", s.handleUIJob)
|
||||
ui.HandleFunc("GET /ui/api/jobs/{job_id}", s.handleUIJobJSON)
|
||||
ui.HandleFunc("POST /ui/api/jobs/{job_id}/cancel", s.handleCancelJob)
|
||||
ui.HandleFunc("POST /ui/api/jobs/upload", s.handleUploadDataset)
|
||||
ui.HandleFunc("GET /ui/jobs/{job_id}/artifacts/{artifact_id}", s.handleUIArtifactDownload)
|
||||
mux.Handle("/ui", chain(ui, withRequestID, withAccessLog(s.log), withBasicAuth(uiToken[0]), withSameOrigin))
|
||||
mux.Handle("/ui/", chain(ui, withRequestID, withAccessLog(s.log), withBasicAuth(uiToken[0]), withSameOrigin))
|
||||
} else {
|
||||
// More specific than the protected catch-all: UI absence is not an auth
|
||||
// failure and does not disclose that a UI feature is configured elsewhere.
|
||||
mux.HandleFunc("/ui", http.NotFound)
|
||||
mux.HandleFunc("/ui/", http.NotFound)
|
||||
}
|
||||
mux.Handle("/", chain(protected,
|
||||
withRequestID, // outermost: every response gets an ID,
|
||||
withAccessLog(s.log), // including the 401s below
|
||||
|
||||
@@ -20,13 +20,19 @@ import (
|
||||
)
|
||||
|
||||
const token = "secret"
|
||||
const uiToken = "ui-secret"
|
||||
|
||||
type env struct {
|
||||
ts *httptest.Server
|
||||
blobs *memstore.BlobStore
|
||||
ts *httptest.Server
|
||||
blobs *memstore.BlobStore
|
||||
workerID string
|
||||
}
|
||||
|
||||
func newEnv(t *testing.T, ready func(context.Context) error) *env {
|
||||
return newEnvWithUIToken(t, ready, uiToken)
|
||||
}
|
||||
|
||||
func newEnvWithUIToken(t *testing.T, ready func(context.Context) error, configuredUIToken string) *env {
|
||||
t.Helper()
|
||||
tasks := memstore.NewTaskRepo()
|
||||
jobs := memstore.NewJobRepo()
|
||||
@@ -40,20 +46,28 @@ func newEnv(t *testing.T, ready func(context.Context) error) *env {
|
||||
uc := coordhttp.UseCases{
|
||||
RegisterWorker: usecase.NewRegisterWorker(work, clk),
|
||||
CreateJob: usecase.NewCreateJob(jobs, tasks, tx, clk),
|
||||
SubmitDataset: usecase.NewSubmitDataset(blobs, arts, jobs, tasks, tx, clk),
|
||||
ClaimTask: usecase.NewClaimTask(tasks, clk, lease),
|
||||
SubmitDataset: usecase.NewSubmitDataset(blobs, arts, jobs, tasks, tx, clk, 3),
|
||||
ClaimTask: usecase.NewClaimTask(tasks, jobs, work, tx, clk, lease),
|
||||
RenewLease: usecase.NewRenewLease(tasks, work, tx, clk, lease),
|
||||
CompleteTask: usecase.NewCompleteTask(tasks, jobs, arts, tx, clk),
|
||||
FailTask: usecase.NewFailTask(tasks, jobs, tx, clk),
|
||||
GetJobStatus: usecase.NewGetJobStatus(jobs, tasks),
|
||||
UploadArtifact: usecase.NewUploadArtifact(tasks, arts, blobs, clk),
|
||||
CancelJob: usecase.NewCancelJob(jobs, tasks, tx, clk),
|
||||
UploadArtifact: usecase.NewUploadArtifact(tasks, arts, blobs, tx, clk),
|
||||
DownloadArtifact: usecase.NewDownloadArtifact(arts, blobs),
|
||||
GetTaskInput: usecase.NewGetTaskInput(tasks, arts, blobs),
|
||||
Dashboard: usecase.NewDashboard(memstore.NewUIReadRepo(jobs, tasks, work, arts)),
|
||||
}
|
||||
worker, err := uc.RegisterWorker.Execute(context.Background(), usecase.RegisterWorkerInput{
|
||||
Name: "test-worker", Capabilities: []string{"w", "similarity-search"},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("register test worker: %v", err)
|
||||
}
|
||||
srv := coordhttp.NewServer(uc, slog.New(slog.NewTextHandler(io.Discard, nil)), 5*time.Second, 15*time.Second, 1<<30, ready)
|
||||
ts := httptest.NewServer(srv.Handler(token))
|
||||
ts := httptest.NewServer(srv.Handler(token, configuredUIToken))
|
||||
t.Cleanup(ts.Close)
|
||||
return &env{ts: ts, blobs: blobs}
|
||||
return &env{ts: ts, blobs: blobs, workerID: worker.ID.String()}
|
||||
}
|
||||
|
||||
func healthy(context.Context) error { return nil }
|
||||
@@ -61,6 +75,7 @@ func healthy(context.Context) error { return nil }
|
||||
// do sends an authenticated JSON request and returns status + decoded body.
|
||||
func (e *env) do(t *testing.T, method, path, body string) (int, map[string]any) {
|
||||
t.Helper()
|
||||
body = strings.ReplaceAll(body, `"worker_id":"w1"`, `"worker_id":"`+e.workerID+`"`)
|
||||
req, _ := http.NewRequestWithContext(context.Background(), method, e.ts.URL+path, strings.NewReader(body))
|
||||
req.Header.Set("Authorization", "Bearer "+token)
|
||||
if body != "" {
|
||||
@@ -97,6 +112,181 @@ func TestHealthOK(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestUIRequiresDistinctCredentialAndRendersDashboard(t *testing.T) {
|
||||
e := newEnv(t, healthy)
|
||||
request := func() *http.Request {
|
||||
req, _ := http.NewRequestWithContext(context.Background(), "GET", e.ts.URL+"/ui", nil)
|
||||
return req
|
||||
}
|
||||
resp, err := http.DefaultClient.Do(request())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode != http.StatusUnauthorized {
|
||||
t.Fatalf("no UI auth: %d", resp.StatusCode)
|
||||
}
|
||||
|
||||
req := request()
|
||||
req.Header.Set("Authorization", "Bearer "+token)
|
||||
resp, err = http.DefaultClient.Do(req)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode != http.StatusUnauthorized {
|
||||
t.Fatalf("worker token authorized UI: %d", resp.StatusCode)
|
||||
}
|
||||
|
||||
req = request()
|
||||
req.SetBasicAuth("operator", uiToken)
|
||||
resp, err = http.DefaultClient.Do(req)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
t.Fatalf("UI status: %d", resp.StatusCode)
|
||||
}
|
||||
body, _ := io.ReadAll(resp.Body)
|
||||
if !strings.Contains(string(body), "SciMesh operator dashboard") {
|
||||
t.Errorf("dashboard body missing title")
|
||||
}
|
||||
}
|
||||
|
||||
func TestUIDisabledReturnsNotFound(t *testing.T) {
|
||||
e := newEnvWithUIToken(t, healthy, "")
|
||||
resp := e.get(t, "/ui")
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode != http.StatusNotFound {
|
||||
t.Fatalf("disabled UI = %d, want 404", resp.StatusCode)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUIRejectsCrossOriginUpload(t *testing.T) {
|
||||
e := newEnv(t, healthy)
|
||||
req, _ := http.NewRequestWithContext(context.Background(), "POST", e.ts.URL+"/ui/api/jobs/upload", strings.NewReader("dataset=x"))
|
||||
req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
|
||||
req.Header.Set("Origin", "https://attacker.example")
|
||||
req.SetBasicAuth("operator", uiToken)
|
||||
resp, err := http.DefaultClient.Do(req)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode != http.StatusForbidden {
|
||||
t.Fatalf("cross-origin upload = %d, want 403", resp.StatusCode)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUIUploadDatasetCreatesJob(t *testing.T) {
|
||||
e := newEnv(t, healthy)
|
||||
var body bytes.Buffer
|
||||
mw := multipart.NewWriter(&body)
|
||||
_ = mw.WriteField("workload", "similarity-search")
|
||||
_ = mw.WriteField("parameters", `{"query_smiles":"CCO","top_k":20,"progress_every":0}`)
|
||||
_ = mw.WriteField("chunk_rows", "1000")
|
||||
file, err := mw.CreateFormFile("file", "chembl.tsv")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
_, _ = io.WriteString(file, "chembl_id\tcanonical_smiles\nCHEMBL1\tCCO\n")
|
||||
if err := mw.Close(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
req, _ := http.NewRequestWithContext(context.Background(), "POST", e.ts.URL+"/ui/api/jobs/upload", &body)
|
||||
req.Header.Set("Content-Type", mw.FormDataContentType())
|
||||
req.SetBasicAuth("operator", uiToken)
|
||||
resp, err := http.DefaultClient.Do(req)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode != http.StatusCreated {
|
||||
result, _ := io.ReadAll(resp.Body)
|
||||
t.Fatalf("UI upload = %d: %s", resp.StatusCode, result)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCancelJobStopsUnfinishedTasks(t *testing.T) {
|
||||
e := newEnv(t, healthy)
|
||||
code, job := e.do(t, "POST", "/jobs", `{"workload":"w","input_uri":"s3://in","chunks":[{"chunk_index":0,"input_uri":"s3://c0","input_sha256":"sha"},{"chunk_index":1,"input_uri":"s3://c1","input_sha256":"sha"}]}`)
|
||||
if code != http.StatusCreated {
|
||||
t.Fatalf("create: %d", code)
|
||||
}
|
||||
jobID := job["id"].(string)
|
||||
if code, body := e.do(t, "POST", "/jobs/"+jobID+"/cancel", ""); code != http.StatusOK || body["cancelled_tasks"].(float64) != 2 {
|
||||
t.Fatalf("cancel = (%d, %v)", code, body)
|
||||
}
|
||||
if code, progress := e.do(t, "GET", "/jobs/"+jobID, ""); code != http.StatusOK || progress["status"] != "cancelled" || progress["cancelled"].(float64) != 2 {
|
||||
t.Fatalf("cancelled job progress = (%d, %v)", code, progress)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUICancelJobUsesOperatorCredential(t *testing.T) {
|
||||
e := newEnv(t, healthy)
|
||||
code, job := e.do(t, "POST", "/jobs", `{"workload":"w","input_uri":"s3://in","chunks":[{"chunk_index":0,"input_uri":"s3://c0","input_sha256":"sha"}]}`)
|
||||
if code != http.StatusCreated {
|
||||
t.Fatalf("create: %d", code)
|
||||
}
|
||||
req, _ := http.NewRequestWithContext(context.Background(), "POST", e.ts.URL+"/ui/api/jobs/"+job["id"].(string)+"/cancel", nil)
|
||||
req.SetBasicAuth("operator", uiToken)
|
||||
resp, err := http.DefaultClient.Do(req)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
t.Fatalf("UI cancel = %d", resp.StatusCode)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUIJobAndArtifactAreScopedToTheirJob(t *testing.T) {
|
||||
e := newEnv(t, healthy)
|
||||
code, job := e.do(t, "POST", "/jobs", `{"workload":"w","input_uri":"s3://in","chunks":[{"chunk_index":0,"input_uri":"s3://c","input_sha256":"sha"}]}`)
|
||||
if code != http.StatusCreated {
|
||||
t.Fatalf("create: %d", code)
|
||||
}
|
||||
req, _ := http.NewRequestWithContext(context.Background(), "GET", e.ts.URL+"/ui/jobs/"+job["id"].(string), nil)
|
||||
req.SetBasicAuth("operator", uiToken)
|
||||
resp, err := http.DefaultClient.Do(req)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
t.Fatalf("detail: %d", resp.StatusCode)
|
||||
}
|
||||
if got := resp.Header.Get("Content-Security-Policy"); got == "" {
|
||||
t.Error("missing UI CSP")
|
||||
}
|
||||
}
|
||||
|
||||
func TestUIArtifactDownloadRejectsAnotherJobsArtifact(t *testing.T) {
|
||||
e := newEnv(t, healthy)
|
||||
code, _ := e.do(t, "POST", "/jobs", `{"workload":"w","input_uri":"s3://in","chunks":[{"chunk_index":0,"input_uri":"s3://c","input_sha256":"sha"}]}`)
|
||||
if code != http.StatusCreated {
|
||||
t.Fatalf("first job: %d", code)
|
||||
}
|
||||
_, claim := e.do(t, "POST", "/tasks/claim", `{"worker_id":"w1","capabilities":["w"]}`)
|
||||
artifactID := e.putArtifact(t, claim["task_id"].(string), "w1", int(claim["attempt"].(float64)), "result")
|
||||
code, second := e.do(t, "POST", "/jobs", `{"workload":"w","input_uri":"s3://in","chunks":[{"chunk_index":0,"input_uri":"s3://c","input_sha256":"sha"}]}`)
|
||||
if code != http.StatusCreated {
|
||||
t.Fatalf("second job: %d", code)
|
||||
}
|
||||
req, _ := http.NewRequestWithContext(context.Background(), "GET", e.ts.URL+"/ui/jobs/"+second["id"].(string)+"/artifacts/"+artifactID, nil)
|
||||
req.SetBasicAuth("operator", uiToken)
|
||||
resp, err := http.DefaultClient.Do(req)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode != http.StatusNotFound {
|
||||
t.Errorf("cross-job artifact = %d, want 404", resp.StatusCode)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHealthUnavailableWhenDBDown(t *testing.T) {
|
||||
e := newEnv(t, func(context.Context) error { return context.DeadlineExceeded })
|
||||
resp := e.get(t, "/health")
|
||||
@@ -148,6 +338,26 @@ func TestRegisterRejectsNoCapabilities(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestClaimRequiresRegisteredWorkerAndUsesStoredCapabilities(t *testing.T) {
|
||||
e := newEnv(t, healthy)
|
||||
if code, _ := e.do(t, "POST", "/tasks/claim", `{"worker_id":"not-a-uuid"}`); code != http.StatusBadRequest {
|
||||
t.Fatalf("invalid worker id claim = %d, want 400", code)
|
||||
}
|
||||
if code, _ := e.do(t, "POST", "/tasks/claim", `{"worker_id":"11111111-1111-4111-8111-111111111111"}`); code != http.StatusNotFound {
|
||||
t.Fatalf("unregistered worker claim = %d, want 404", code)
|
||||
}
|
||||
if code, _ := e.do(t, "POST", "/jobs", `{"workload":"w","input_uri":"s3://in","chunks":[{"chunk_index":0,"input_uri":"s3://c","input_sha256":"sha"}]}`); code != http.StatusCreated {
|
||||
t.Fatalf("create job = %d", code)
|
||||
}
|
||||
code, worker := e.do(t, "POST", "/workers/register", `{"name":"search-only","capabilities":["similarity-search"]}`)
|
||||
if code != http.StatusCreated {
|
||||
t.Fatalf("register = %d", code)
|
||||
}
|
||||
if code, _ := e.do(t, "POST", "/tasks/claim", `{"worker_id":"`+worker["worker_id"].(string)+`","capabilities":["w"]}`); code != http.StatusNoContent {
|
||||
t.Fatalf("forged capability claim = %d, want 204", code)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFullLifecycle(t *testing.T) {
|
||||
e := newEnv(t, healthy)
|
||||
|
||||
@@ -211,9 +421,9 @@ func TestForeignArtifactResultConflict(t *testing.T) {
|
||||
|
||||
func TestUploadDatasetChunksAndServesInput(t *testing.T) {
|
||||
e := newEnv(t, healthy)
|
||||
tsv := "id\tsmiles\nA\tCC\nB\tCCC\nC\tCCCC\nD\tCCCCC\nE\tCCCCCC\n"
|
||||
tsv := "chembl_id\tcanonical_smiles\nA\tCC\nB\tCCC\nC\tCCCC\nD\tCCCCC\nE\tCCCCCC\n"
|
||||
|
||||
code, body := e.uploadDataset(t, "w", 2, tsv)
|
||||
code, body := e.uploadDataset(t, "similarity-search", 2, tsv)
|
||||
if code != 201 {
|
||||
t.Fatalf("upload: status = %d", code)
|
||||
}
|
||||
@@ -239,11 +449,60 @@ func TestUploadDatasetChunksAndServesInput(t *testing.T) {
|
||||
t.Fatalf("get input: status = %d", resp.StatusCode)
|
||||
}
|
||||
shard, _ := io.ReadAll(resp.Body)
|
||||
if !strings.HasPrefix(string(shard), "id\tsmiles\n") {
|
||||
if !strings.HasPrefix(string(shard), "chembl_id\tcanonical_smiles\n") {
|
||||
t.Errorf("shard missing header: %q", shard)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUploadDatasetLimitsRows(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")
|
||||
_ = mw.WriteField("max_rows", "3")
|
||||
fw, _ := mw.CreateFormFile("file", "chembl.tsv")
|
||||
_, _ = io.Copy(fw, strings.NewReader("chembl_id\tcanonical_smiles\nA\tCC\nB\tCCC\nC\tCCCC\nD\tCCCCC\n"))
|
||||
_ = mw.Close()
|
||||
req, _ := http.NewRequestWithContext(context.Background(), "POST", e.ts.URL+"/jobs/upload", &buf)
|
||||
req.Header.Set("Authorization", "Bearer "+token)
|
||||
req.Header.Set("Content-Type", mw.FormDataContentType())
|
||||
resp, err := http.DefaultClient.Do(req)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
var result map[string]any
|
||||
_ = json.NewDecoder(resp.Body).Decode(&result)
|
||||
if resp.StatusCode != http.StatusCreated || result["task_count"].(float64) != 2 {
|
||||
t.Fatalf("limited upload = (%d, %v)", resp.StatusCode, result)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUploadDatasetRejectsMissingChEMBLColumns(t *testing.T) {
|
||||
e := newEnv(t, healthy)
|
||||
var buf bytes.Buffer
|
||||
mw := multipart.NewWriter(&buf)
|
||||
_ = mw.WriteField("workload", "similarity-search")
|
||||
_ = mw.WriteField("parameters", `{"query_smiles":"CCO"}`)
|
||||
_ = mw.WriteField("chunk_rows", "2")
|
||||
fw, _ := mw.CreateFormFile("file", "not-chembl.tsv")
|
||||
_, _ = io.Copy(fw, strings.NewReader("id\tsmiles\nA\tCC\n"))
|
||||
_ = mw.Close()
|
||||
req, _ := http.NewRequestWithContext(context.Background(), "POST", e.ts.URL+"/jobs/upload", &buf)
|
||||
req.Header.Set("Authorization", "Bearer "+token)
|
||||
req.Header.Set("Content-Type", mw.FormDataContentType())
|
||||
resp, err := http.DefaultClient.Do(req)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode != http.StatusBadRequest {
|
||||
t.Errorf("missing ChEMBL columns = %d, want 400", resp.StatusCode)
|
||||
}
|
||||
}
|
||||
|
||||
func TestErrorMappings(t *testing.T) {
|
||||
e := newEnv(t, healthy)
|
||||
zero := "00000000-0000-0000-0000-000000000000"
|
||||
@@ -271,10 +530,11 @@ func TestUploadDatasetRejectsAmbiguousMultipartInput(t *testing.T) {
|
||||
e := newEnv(t, healthy)
|
||||
var buf bytes.Buffer
|
||||
mw := multipart.NewWriter(&buf)
|
||||
_ = mw.WriteField("workload", "w")
|
||||
_ = mw.WriteField("workload", "similarity-search")
|
||||
_ = mw.WriteField("parameters", `{"query_smiles":"CCO"}`)
|
||||
_ = mw.WriteField("chunk_rows", "not-a-number")
|
||||
fw, _ := mw.CreateFormFile("file", "chembl.tsv")
|
||||
_, _ = io.Copy(fw, strings.NewReader("id\tsmiles\nA\tCC\n"))
|
||||
_, _ = io.Copy(fw, strings.NewReader("chembl_id\tcanonical_smiles\nA\tCC\n"))
|
||||
_ = mw.Close()
|
||||
|
||||
req, _ := http.NewRequestWithContext(context.Background(), "POST", e.ts.URL+"/jobs/upload", &buf)
|
||||
@@ -298,6 +558,9 @@ func (e *env) putArtifact(t *testing.T, taskID, worker string, attempt int, data
|
||||
e.ts.URL+"/tasks/"+taskID+"/artifacts/r.csv", strings.NewReader(data))
|
||||
req.Header.Set("Authorization", "Bearer "+token)
|
||||
req.Header.Set("Content-Type", "text/csv")
|
||||
if worker == "w1" {
|
||||
worker = e.workerID
|
||||
}
|
||||
req.Header.Set("X-Worker-ID", worker)
|
||||
req.Header.Set("X-Task-Attempt", itoa(attempt))
|
||||
resp, err := http.DefaultClient.Do(req)
|
||||
@@ -319,6 +582,7 @@ func (e *env) uploadDataset(t *testing.T, workload string, rows int, tsv string)
|
||||
var buf bytes.Buffer
|
||||
mw := multipart.NewWriter(&buf)
|
||||
_ = mw.WriteField("workload", workload)
|
||||
_ = mw.WriteField("parameters", `{"query_smiles":"CCO"}`)
|
||||
_ = mw.WriteField("chunk_rows", itoa(rows))
|
||||
fw, _ := mw.CreateFormFile("file", "chembl.tsv")
|
||||
_, _ = io.Copy(fw, strings.NewReader(tsv))
|
||||
|
||||
@@ -0,0 +1,23 @@
|
||||
{{define "dashboard.html"}}
|
||||
<!doctype html>
|
||||
<html lang="en">
|
||||
<head>
|
||||
<meta charset="utf-8">
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1">
|
||||
<title>SciMesh operator dashboard</title>
|
||||
<style>
|
||||
:root{color:#172033;background:#f6f8fc;font:16px/1.5 system-ui,sans-serif}body{margin:0}.page{max-width:1180px;margin:auto;padding:32px 20px 56px}.top{display:flex;justify-content:space-between;gap:24px;align-items:start}.eyebrow{margin:0;color:#50617d;font-size:.86rem;font-weight:700;text-transform:uppercase;letter-spacing:.08em}h1{margin:.2rem 0;font-size:2rem}h2{margin:32px 0 12px;font-size:1.28rem}.lead{margin:0;color:#56657c}.button{display:inline-block;border:0;border-radius:8px;padding:11px 15px;background:#1f5eff;color:#fff;font-weight:700;text-decoration:none;white-space:nowrap}.notice{margin-top:24px;padding:16px 18px;border:1px solid #f2cb72;border-radius:10px;background:#fff8e6}.notice strong{display:block}.steps{display:grid;grid-template-columns:repeat(3,1fr);gap:12px;margin-top:14px}.step,.card{padding:16px;border:1px solid #dfe5f0;border-radius:10px;background:#fff}.step b{display:block;color:#1f5eff}.table-wrap{overflow-x:auto;background:#fff;border:1px solid #dfe5f0;border-radius:10px}table{width:100%;border-collapse:collapse}td,th{padding:13px 14px;border-bottom:1px solid #e8ecf4;text-align:left;vertical-align:top}th{color:#50617d;font-size:.78rem;text-transform:uppercase;letter-spacing:.06em}tr:last-child td{border:0}a{color:#174ecf}small,.muted{color:#68758b}.status{display:inline-block;border-radius:999px;padding:3px 9px;font-size:.84rem;font-weight:700}.status-success{background:#dff6e9;color:#126b3d}.status-danger{background:#ffe4e6;color:#a31135}.status-active{background:#e4edff;color:#174ecf}.status-waiting{background:#edf0f5;color:#50617d}.bar{height:7px;min-width:120px;margin-top:7px;overflow:hidden;border-radius:999px;background:#e6eaf1}.bar>span{display:block;height:100%;background:#1f5eff}.kicker{font-variant-numeric:tabular-nums}.empty{padding:28px;text-align:center;color:#68758b}.worker{display:grid;grid-template-columns:1.3fr .8fr 2fr 1fr;gap:12px;align-items:center}.worker+.worker{border-top:1px solid #e8ecf4;padding-top:12px;margin-top:12px}@media(max-width:760px){.top,.steps{display:block}.button{margin-top:12px}.step{margin-top:10px}.worker{grid-template-columns:1fr}.hide-mobile{display:none}}
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<main class="page">
|
||||
<header class="top"><div><p class="eyebrow">Local coordinator</p><h1>SciMesh operator dashboard</h1><p class="lead">See where a computation is and what should happen next.</p></div><a class="button" href="/ui/jobs/new">Start a check</a></header>
|
||||
<section class="notice" aria-label="Current pipeline limitation"><strong>This screen currently diagnoses shard jobs.</strong><span>Workers upload partial CSVs to the coordinator. Until a reducer is implemented, those files are not one final scientific result.</span><div class="steps"><div class="step"><b>1. Upload TSV</b>The coordinator splits the file into shard tasks.</div><div class="step"><b>2. Wait for a worker</b>A worker claims a shard, calculates similarity, and returns a CSV.</div><div class="step"><b>3. Inspect artifacts</b>Download a partial result from the job page.</div></div></section>
|
||||
<h2>Recent jobs</h2>
|
||||
<div class="table-wrap"><table><tr><th>Computation</th><th>State</th><th>Progress</th><th class="hide-mobile">Created</th></tr>{{range .Jobs}}<tr><td><a href="/ui/jobs/{{.ID}}"><strong>{{workloadLabel .Workload}}</strong></a><br><small>Open job details</small></td><td><span class="status status-{{statusClass .Status}}">{{statusLabel .Status}}</span><br><small>{{statusHint .Status}}</small></td><td class="kicker"><strong>{{.Completed}} / {{.Total}}</strong> complete{{if gt .Failed 0}} · <span style="color:#a31135">failed: {{.Failed}}</span>{{end}}{{if gt .Cancelled 0}} · <span>stopped: {{.Cancelled}}</span>{{end}}<div class="bar"><span style="width:{{progressPercent .Completed .Failed .Cancelled .Total}}%"></span></div></td><td class="hide-mobile"><small>{{time .CreatedAt}}</small></td></tr>{{else}}<tr><td colspan="4" class="empty"><strong>No jobs yet.</strong><br>Click “Start a check”, upload a small TSV, and leave a worker running.</td></tr>{{end}}</table></div>
|
||||
<h2>Workers</h2>
|
||||
<section class="card">{{range .Workers}}<div class="worker"><div><strong>{{.Name}}</strong><br><small>{{.ID}}</small></div><div><span class="status status-{{if eq .Status "online"}}success{{else}}waiting{{end}}">{{workerStatusLabel .Status}}</span></div><div>{{range .Capabilities}}<code>{{.}}</code> {{end}}</div><div class="muted">Last signal<br>{{time .LastHeartbeatAt}}</div></div>{{else}}<div class="empty"><strong>No worker is registered yet.</strong><br>Run <code>scimesh-worker</code> with the coordinator URL and worker token.</div>{{end}}</section>
|
||||
</main>
|
||||
</body>
|
||||
</html>
|
||||
{{end}}
|
||||
@@ -0,0 +1,29 @@
|
||||
{{define "job.html"}}
|
||||
<!doctype html>
|
||||
<html lang="en">
|
||||
<head>
|
||||
<meta charset="utf-8">
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1">
|
||||
<title>SciMesh job</title>
|
||||
<style>
|
||||
:root{color:#172033;background:#f6f8fc;font:16px/1.5 system-ui,sans-serif}body{margin:0}.page{max-width:1180px;margin:auto;padding:32px 20px 56px}a{color:#174ecf}.back{text-decoration:none}.eyebrow{margin:28px 0 4px;color:#50617d;font-size:.86rem;font-weight:700;text-transform:uppercase;letter-spacing:.08em}h1{margin:0;font-size:2rem}h2{margin:32px 0 12px;font-size:1.3rem}.summary,.card{padding:20px;border:1px solid #dfe5f0;border-radius:12px;background:#fff}.summary-head{display:flex;justify-content:space-between;gap:16px;align-items:start}.status{display:inline-block;border-radius:999px;padding:4px 10px;font-size:.9rem;font-weight:700}.status-success{background:#dff6e9;color:#126b3d}.status-danger{background:#ffe4e6;color:#a31135}.status-active{background:#e4edff;color:#174ecf}.status-waiting{background:#edf0f5;color:#50617d}.hint{margin:8px 0 0;color:#56657c}.bar{height:10px;margin:20px 0 8px;overflow:hidden;border-radius:999px;background:#e6eaf1}.bar>span{display:block;height:100%;background:#1f5eff;transition:width .3s}.numbers{display:grid;grid-template-columns:repeat(6,1fr);gap:10px}.number{padding:12px;border-radius:8px;background:#f6f8fc}.number b{display:block;font-size:1.35rem}.stop{display:block;margin-left:auto;border:1px solid #d43b51;border-radius:7px;padding:8px 11px;background:#fff;color:#b2223a;font:inherit;font-weight:700;cursor:pointer}.stop:disabled{opacity:.6}.notice{margin:20px 0;padding:15px 17px;border:1px solid #f2cb72;border-radius:10px;background:#fff8e6}.table-wrap{overflow-x:auto;border:1px solid #dfe5f0;border-radius:10px;background:#fff}table{width:100%;border-collapse:collapse}td,th{padding:12px 13px;border-bottom:1px solid #e8ecf4;text-align:left;vertical-align:top}th{color:#50617d;font-size:.78rem;text-transform:uppercase;letter-spacing:.06em}tr:last-child td{border:0}small,.muted{color:#68758b}.error{color:#a31135;max-width:360px;word-break:break-word}.download{display:inline-block;padding:7px 10px;border:1px solid #b9c9ee;border-radius:7px;text-decoration:none}.empty{padding:24px;text-align:center;color:#68758b}details{margin-top:18px;color:#56657c}code{word-break:break-all}@media(max-width:700px){.summary-head{display:block}.numbers{grid-template-columns:repeat(2,1fr)}}
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<main class="page">
|
||||
<a class="back" href="/ui">← Back to jobs</a><p class="eyebrow">{{workloadLabel .Workload}}</p><h1>Execution progress</h1>
|
||||
<section class="summary"><div class="summary-head"><div><span id="status" class="status status-{{statusClass .Status}}">{{statusLabel .Status}}</span><p id="hint" class="hint">{{statusHint .Status}}</p></div><div>{{if cancellable .Status}}<button id="stop-job" class="stop" type="button">Stop job</button><small>This cancels every shard that is not finished yet.</small>{{else}}<small>Summary refreshes automatically every two seconds.</small>{{end}}</div></div><div class="bar" aria-label="Progress"><span id="progress-bar" style="width:{{progressPercent .Completed .Failed .Cancelled .Total}}%"></span></div><p id="progress" class="muted">{{.Completed}} of {{.Total}} tasks complete</p><div class="numbers"><div class="number"><b id="total">{{.Total}}</b><small>total shards</small></div><div class="number"><b id="completed">{{.Completed}}</b><small>complete</small></div><div class="number"><b id="pending">{{.Pending}}</b><small>waiting</small></div><div class="number"><b id="active">{{add .Leased .Running}}</b><small>with workers</small></div><div class="number"><b id="failed">{{.Failed}}</b><small>failed</small></div><div class="number"><b id="cancelled">{{.Cancelled}}</b><small>stopped</small></div></div><details><summary>Technical details</summary><p>Job ID: <code>{{.ID}}</code><br>Workload: <code>{{.Workload}}</code><br>Created: {{time .CreatedAt}}</p></details></section>
|
||||
<section class="notice"><strong>What can be downloaded now?</strong><br><code>partial_result</code> files come from individual shards. They are useful for checking the pipeline, but are not a merged final CSV because the reducer is not implemented yet.</section>
|
||||
<h2>Shard tasks</h2><p class="muted">If a task fails, its code and message appear here. Refresh the page to update the detailed rows.</p>
|
||||
<div class="table-wrap"><table><tr><th>Shard</th><th>State</th><th>Attempt</th><th>Worker / lease</th><th>Error</th></tr>{{range .Tasks}}<tr><td>#{{.ChunkIndex}}</td><td><span class="status status-{{statusClass .Status}}">{{statusLabel .Status}}</span></td><td>{{.Attempt}} / {{.MaxAttempts}}</td><td>{{if .LeaseOwner}}<code>{{.LeaseOwner}}</code>{{if .LeaseExpiresAt}}<br><small>until {{time .LeaseExpiresAt}}</small>{{end}}{{else}}<span class="muted">—</span>{{end}}</td><td class="error">{{if .ErrorCode}}<strong>{{taskErrorLabel .ErrorCode}}</strong><br><small>{{taskErrorHint .ErrorCode}}</small>{{else}}<span class="muted">—</span>{{end}}</td></tr>{{else}}<tr><td colspan="5" class="empty">No tasks have appeared yet.</td></tr>{{end}}</table></div>
|
||||
<h2>Coordinator artifacts</h2>
|
||||
<div class="table-wrap"><table><tr><th>Type</th><th>File</th><th>Size</th><th>Integrity check</th><th></th></tr>{{range .Artifacts}}<tr><td>{{if .Diagnostic}}<strong>Partial result</strong><br><small>diagnostic</small>{{else}}{{.Kind}}{{end}}</td><td>{{.Filename}}</td><td>{{bytes .SizeBytes}}</td><td><code>{{.SHA256}}</code></td><td>{{if .Downloadable}}<a class="download" href="/ui/jobs/{{$.ID}}/artifacts/{{.ID}}">Download CSV</a>{{else}}<span class="muted">Unavailable</span>{{end}}</td></tr>{{else}}<tr><td colspan="5" class="empty">No artifacts yet. The worker uploads a CSV after it completes a shard.</td></tr>{{end}}</table></div>
|
||||
</main>
|
||||
<script>
|
||||
const id={{printf "%q" .ID}},state={pending:['Waiting for a worker','waiting','Waiting for an available worker with the required capability.'],leased:['Assigned to a worker','active','A worker has claimed the task and should begin processing shortly.'],running:['Running','active','A worker is reading a shard, calculating fingerprints, and uploading its result through the coordinator.'],completed:['Tasks complete','success','Every shard task is complete. Files below are still partial results.'],failed:['Needs attention','danger','One or more shard tasks failed. Open the task list below for details.'],cancelled:['Stopped','waiting','The operator stopped this job. No new shards can be claimed.']};
|
||||
const stop=document.querySelector('#stop-job');if(stop)stop.addEventListener('click',async()=>{if(!confirm('Stop this job? Unfinished shards will be cancelled.'))return;stop.disabled=true;const response=await fetch('/ui/api/jobs/'+id+'/cancel',{method:'POST'});if(!response.ok){stop.disabled=false;alert('Unable to stop this job.');return}location.reload()});
|
||||
const terminal=new Set(['completed','failed','cancelled']);let timer;const poll=async()=>{try{const response=await fetch('/ui/api/jobs/'+id);if(!response.ok)return;const job=await response.json(),info=state[job.status]||[job.status,'waiting','Status reported by the coordinator.'],done=job.completed+job.failed+job.cancelled,percent=job.total?Math.min(100,Math.floor(done*100/job.total)):0,badge=document.querySelector('#status');badge.textContent=info[0];badge.className='status status-'+info[1];document.querySelector('#hint').textContent=info[2];document.querySelector('#progress').textContent=job.completed+' of '+job.total+' tasks complete'+(job.failed?' · failed: '+job.failed:'')+(job.cancelled?' · stopped: '+job.cancelled:'');document.querySelector('#progress-bar').style.width=percent+'%';for(const key of ['total','completed','pending','failed','cancelled'])document.querySelector('#'+key).textContent=job[key];document.querySelector('#active').textContent=job.leased+job.running;if(terminal.has(job.status)&&timer){clearInterval(timer);timer=undefined}}catch(_){}};const start=()=>{if(!timer&&!document.hidden&&!terminal.has(document.querySelector('#status').textContent.toLowerCase()))timer=setInterval(poll,2000)};document.addEventListener('visibilitychange',()=>{if(document.hidden&&timer){clearInterval(timer);timer=undefined}else start()});start();
|
||||
</script>
|
||||
</body>
|
||||
</html>
|
||||
{{end}}
|
||||
@@ -0,0 +1,31 @@
|
||||
{{define "new-job.html"}}
|
||||
<!doctype html>
|
||||
<html lang="en">
|
||||
<head>
|
||||
<meta charset="utf-8">
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1">
|
||||
<title>Create a check — SciMesh</title>
|
||||
<style>
|
||||
:root{color:#172033;background:#f6f8fc;font:16px/1.5 system-ui,sans-serif}body{margin:0}.page{max-width:760px;margin:auto;padding:32px 20px 56px}a{color:#174ecf}.back{text-decoration:none}.eyebrow{margin:28px 0 4px;color:#50617d;font-size:.86rem;font-weight:700;text-transform:uppercase;letter-spacing:.08em}h1{margin:0;font-size:2rem}.lead{color:#56657c}.notice{margin:22px 0;padding:16px 18px;border:1px solid #f2cb72;border-radius:10px;background:#fff8e6}.notice strong{display:block}.card{padding:22px;border:1px solid #dfe5f0;border-radius:12px;background:#fff}label{display:block;margin:18px 0 4px;font-weight:700}input{box-sizing:border-box;width:100%;padding:10px;border:1px solid #bac5d8;border-radius:7px;font:inherit}input[type=file]{padding:8px;background:#f8faff}.hint{margin:4px 0;color:#68758b;font-size:.9rem}.button{margin-top:22px;border:0;border-radius:8px;padding:11px 16px;background:#1f5eff;color:#fff;font:inherit;font-weight:700;cursor:pointer}.button:disabled{opacity:.6;cursor:wait}.error{margin-top:16px;color:#a31135}.working{margin-top:16px;color:#174ecf}.checklist{margin:8px 0;padding-left:20px;color:#56657c}.checklist li{margin:5px 0}
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<main class="page">
|
||||
<a class="back" href="/ui">← Back to jobs</a><p class="eyebrow">Guided run</p><h1>Search for similar molecules</h1><p class="lead">Creates a diagnostic <code>similarity-search</code> job: a worker finds the top-k molecules most similar to a target SMILES.</p>
|
||||
<section class="notice"><strong>Before starting</strong><ul class="checklist"><li>Keep at least one <code>scimesh-worker</code> running.</li><li>Use a small TSV for a hands-on check.</li><li><b>“Rows per shard” does not limit the file size.</b> It splits the entire upload into tasks: a full ChEMBL TSV at 1,000 rows per shard creates thousands of tasks.</li></ul></section>
|
||||
<form id="run" class="card">
|
||||
<label for="file">ChEMBL TSV</label><input id="file" type="file" name="file" required accept=".tsv,.txt,text/tab-separated-values"><p class="hint">Expected columns: <code>chembl_id</code> and <code>canonical_smiles</code>.</p>
|
||||
<label for="query-smiles">Target molecule (SMILES)</label><input id="query-smiles" name="query_smiles" required maxlength="200" value="CCO" autocomplete="off"><p class="hint"><code>CCO</code> is ethanol. For gefitinib, use its SMILES here or the local CLI with <code>--query-id</code>.</p>
|
||||
<label for="top-k">Matches to return</label><input id="top-k" name="top_k" type="number" min="1" max="100000" value="20" required><p class="hint">This is the top-k within each shard, not a global top-k for the whole dataset yet.</p>
|
||||
<label for="chunk-rows">Rows per shard</label><input id="chunk-rows" name="chunk_rows" type="number" min="1" max="100000" value="1000" required><p class="hint">Fewer rows mean more tasks and more visible progress; more rows mean fewer, longer tasks.</p>
|
||||
<label for="max-rows">Maximum dataset rows to process <small>(optional)</small></label><input id="max-rows" name="max_rows" type="number" min="1" max="10000000" placeholder="For example: 500"><p class="hint">Useful for a quick check of a large TSV. The coordinator creates shards from only the first N data rows; it still stores the original upload.</p>
|
||||
<button class="button" id="submit" type="submit">Upload file and create job</button><p id="working" class="working" hidden aria-live="polite">Uploading the file and creating shard tasks… Keep this page open.</p><p id="error" class="error" role="alert"></p>
|
||||
</form>
|
||||
</main>
|
||||
<script>
|
||||
const form=document.querySelector('#run'),button=document.querySelector('#submit'),working=document.querySelector('#working'),error=document.querySelector('#error');
|
||||
form.addEventListener('submit',async event=>{event.preventDefault();error.textContent='';const fields=new FormData(form),file=fields.get('file'),maxRows=String(fields.get('max_rows')||'').trim();if(!(file instanceof File)||file.size===0){error.textContent='Choose a non-empty TSV file.';return}const parameters={query_smiles:fields.get('query_smiles'),top_k:Number(fields.get('top_k')),progress_every:0},upload=new FormData();upload.append('workload','similarity-search');upload.append('parameters',JSON.stringify(parameters));upload.append('chunk_rows',fields.get('chunk_rows'));if(maxRows)upload.append('max_rows',maxRows);upload.append('file',file,file.name);button.disabled=true;working.hidden=false;try{const response=await fetch('/ui/api/jobs/upload',{method:'POST',body:upload}),data=await response.json();if(!response.ok)throw Error(data.error||'Unable to create the job.');location.href='/ui/jobs/'+data.job_id}catch(err){error.textContent=err.message==='invalid input'?'Check the TSV and fields: the coordinator could not accept this request.':err.message;button.disabled=false;working.hidden=true}});
|
||||
</script>
|
||||
</body>
|
||||
</html>
|
||||
{{end}}
|
||||
@@ -0,0 +1,281 @@
|
||||
package http
|
||||
|
||||
import (
|
||||
"embed"
|
||||
"fmt"
|
||||
"html/template"
|
||||
"io"
|
||||
"mime"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
|
||||
"github.com/emil28092005/SciMesh/coordinator/internal/domain"
|
||||
)
|
||||
|
||||
//go:embed templates/*.html
|
||||
var uiFiles embed.FS
|
||||
|
||||
var uiTemplates = template.Must(template.New("ui").Funcs(template.FuncMap{
|
||||
"time": formatUITime,
|
||||
"statusLabel": uiStatusLabel,
|
||||
"statusHint": uiStatusHint,
|
||||
"statusClass": uiStatusClass,
|
||||
"taskErrorLabel": uiTaskErrorLabel,
|
||||
"taskErrorHint": uiTaskErrorHint,
|
||||
"workerStatusLabel": uiWorkerStatusLabel,
|
||||
"workloadLabel": uiWorkloadLabel,
|
||||
"progressPercent": uiProgressPercent,
|
||||
"cancellable": uiCancellable,
|
||||
"bytes": uiBytes,
|
||||
"add": func(a, b int) int { return a + b },
|
||||
}).ParseFS(uiFiles, "templates/*.html"))
|
||||
|
||||
func formatUITime(t time.Time) string {
|
||||
if t.IsZero() {
|
||||
return "—"
|
||||
}
|
||||
return t.UTC().Format("02.01.2006 15:04 UTC")
|
||||
}
|
||||
|
||||
func uiStatusLabel(status string) string {
|
||||
switch status {
|
||||
case "pending":
|
||||
return "Waiting for a worker"
|
||||
case "leased":
|
||||
return "Assigned to a worker"
|
||||
case "running":
|
||||
return "Running"
|
||||
case "completed":
|
||||
return "Tasks complete"
|
||||
case "failed":
|
||||
return "Needs attention"
|
||||
case "cancelled":
|
||||
return "Stopped"
|
||||
default:
|
||||
return status
|
||||
}
|
||||
}
|
||||
|
||||
func uiStatusHint(status string) string {
|
||||
switch status {
|
||||
case "pending":
|
||||
return "Waiting for an available worker with the required capability."
|
||||
case "leased":
|
||||
return "A worker has claimed the task and should begin processing shortly."
|
||||
case "running":
|
||||
return "A worker is reading a shard, calculating fingerprints, and uploading its result through the coordinator."
|
||||
case "completed":
|
||||
return "Every shard task is complete. Files below are still partial results."
|
||||
case "failed":
|
||||
return "One or more shard tasks failed. Open the task list below for details."
|
||||
case "cancelled":
|
||||
return "The operator stopped this job. No new shards can be claimed."
|
||||
default:
|
||||
return "Status reported by the coordinator."
|
||||
}
|
||||
}
|
||||
|
||||
func uiStatusClass(status string) string {
|
||||
switch status {
|
||||
case "completed":
|
||||
return "success"
|
||||
case "failed":
|
||||
return "danger"
|
||||
case "cancelled":
|
||||
return "waiting"
|
||||
case "running", "leased":
|
||||
return "active"
|
||||
default:
|
||||
return "waiting"
|
||||
}
|
||||
}
|
||||
|
||||
func uiWorkerStatusLabel(status string) string {
|
||||
switch status {
|
||||
case "online":
|
||||
return "Available"
|
||||
case "offline":
|
||||
return "Offline"
|
||||
default:
|
||||
return status
|
||||
}
|
||||
}
|
||||
|
||||
// uiTaskErrorLabel deliberately maps worker implementation errors to an
|
||||
// operator-facing diagnosis. Raw subprocess commands and local paths belong in
|
||||
// the worker terminal, not in the web UI.
|
||||
func uiTaskErrorLabel(errorCode string) string {
|
||||
switch errorCode {
|
||||
case "CalledProcessError":
|
||||
return "Local calculation failed"
|
||||
case "ValueError":
|
||||
return "Task input could not be processed"
|
||||
case "CoordinatorTransientError":
|
||||
return "Coordinator connection was interrupted"
|
||||
case "CoordinatorConflictError":
|
||||
return "Worker lease was no longer valid"
|
||||
case "FileNotFoundError":
|
||||
return "Local task file is missing"
|
||||
default:
|
||||
return errorCode
|
||||
}
|
||||
}
|
||||
|
||||
func uiTaskErrorHint(errorCode string) string {
|
||||
switch errorCode {
|
||||
case "CalledProcessError":
|
||||
return "The local SciMesh command stopped before it could upload a result. Check the worker terminal for the original error."
|
||||
case "ValueError":
|
||||
return "The coordinator task or its downloaded input did not meet the worker validation rules."
|
||||
case "CoordinatorTransientError":
|
||||
return "The worker will retry after the coordinator connection is available again."
|
||||
case "CoordinatorConflictError":
|
||||
return "Another worker or a lease timeout changed this task before completion."
|
||||
case "FileNotFoundError":
|
||||
return "The worker could not find one of its local task files. Restart it with an absolute --work-dir."
|
||||
default:
|
||||
return "Check the worker terminal for the original error details."
|
||||
}
|
||||
}
|
||||
|
||||
func uiWorkloadLabel(workload string) string {
|
||||
switch workload {
|
||||
case "similarity-search", "similarity_search":
|
||||
return "Molecule similarity search"
|
||||
case "similarity-graph", "similarity_graph":
|
||||
return "Molecular similarity graph"
|
||||
default:
|
||||
return workload
|
||||
}
|
||||
}
|
||||
|
||||
func uiCancellable(status string) bool {
|
||||
return status == "pending" || status == "running"
|
||||
}
|
||||
|
||||
func uiProgressPercent(completed, failed, cancelled, total int) int {
|
||||
if total <= 0 {
|
||||
return 0
|
||||
}
|
||||
percent := (completed + failed + cancelled) * 100 / total
|
||||
if percent > 100 {
|
||||
return 100
|
||||
}
|
||||
return percent
|
||||
}
|
||||
|
||||
func uiBytes(n int64) string {
|
||||
const kib = 1024
|
||||
if n < kib {
|
||||
return fmt.Sprintf("%d B", n)
|
||||
}
|
||||
if n < kib*kib {
|
||||
return fmt.Sprintf("%.1f KiB", float64(n)/kib)
|
||||
}
|
||||
if n < kib*kib*kib {
|
||||
return fmt.Sprintf("%.1f MiB", float64(n)/(kib*kib))
|
||||
}
|
||||
return fmt.Sprintf("%.1f GiB", float64(n)/(kib*kib*kib))
|
||||
}
|
||||
|
||||
func (s *Server) renderUI(w http.ResponseWriter, name string, data any) {
|
||||
w.Header().Set("Content-Type", "text/html; charset=utf-8")
|
||||
w.Header().Set("Cache-Control", "no-store")
|
||||
w.Header().Set("X-Content-Type-Options", "nosniff")
|
||||
w.Header().Set("Content-Security-Policy", "default-src 'self'; style-src 'self' 'unsafe-inline'; script-src 'self' 'unsafe-inline'; base-uri 'none'; frame-ancestors 'none'")
|
||||
if err := uiTemplates.ExecuteTemplate(w, name, data); err != nil {
|
||||
s.log.Error("render UI", "err", err)
|
||||
http.Error(w, "internal error", http.StatusInternalServerError)
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Server) handleUIHome(w http.ResponseWriter, r *http.Request) {
|
||||
ctx, cancel := s.reqCtx(r)
|
||||
defer cancel()
|
||||
view, err := s.uc.Dashboard.Overview(ctx, 20)
|
||||
if err != nil {
|
||||
s.writeError(w, r, err)
|
||||
return
|
||||
}
|
||||
s.renderUI(w, "dashboard.html", view)
|
||||
}
|
||||
|
||||
func (s *Server) handleUINewJob(w http.ResponseWriter, r *http.Request) {
|
||||
s.renderUI(w, "new-job.html", nil)
|
||||
}
|
||||
|
||||
func (s *Server) uiJobID(w http.ResponseWriter, r *http.Request) (uuid.UUID, bool) {
|
||||
return s.pathUUID(w, r, "job_id")
|
||||
}
|
||||
|
||||
func (s *Server) handleUIJob(w http.ResponseWriter, r *http.Request) {
|
||||
jobID, ok := s.uiJobID(w, r)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
ctx, cancel := s.reqCtx(r)
|
||||
defer cancel()
|
||||
view, err := s.uc.Dashboard.JobDetail(ctx, jobID)
|
||||
if err != nil {
|
||||
s.writeError(w, r, err)
|
||||
return
|
||||
}
|
||||
s.renderUI(w, "job.html", view)
|
||||
}
|
||||
|
||||
func (s *Server) handleUIJobJSON(w http.ResponseWriter, r *http.Request) {
|
||||
jobID, ok := s.uiJobID(w, r)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
ctx, cancel := s.reqCtx(r)
|
||||
defer cancel()
|
||||
view, err := s.uc.Dashboard.JobDetail(ctx, jobID)
|
||||
if err != nil {
|
||||
s.writeError(w, r, err)
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, view)
|
||||
}
|
||||
|
||||
func (s *Server) handleUIArtifactDownload(w http.ResponseWriter, r *http.Request) {
|
||||
jobID, ok := s.uiJobID(w, r)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
artifactID, err := uuid.Parse(r.PathValue("artifact_id"))
|
||||
if err != nil {
|
||||
s.writeError(w, r, domain.ErrInvalidInput)
|
||||
return
|
||||
}
|
||||
ctx, cancel := s.reqCtx(r)
|
||||
defer cancel()
|
||||
belongs, err := s.uc.Dashboard.ArtifactBelongsToJob(ctx, jobID, artifactID)
|
||||
if err != nil {
|
||||
s.writeError(w, r, err)
|
||||
return
|
||||
}
|
||||
if !belongs {
|
||||
s.writeError(w, r, domain.ErrArtifactNotFound)
|
||||
return
|
||||
}
|
||||
// Reuse the coordinator-owned blob stream after the job-scoped check above.
|
||||
art, body, err := s.uc.DownloadArtifact.Execute(ctx, artifactID)
|
||||
if err != nil {
|
||||
s.writeError(w, r, err)
|
||||
return
|
||||
}
|
||||
defer func() {
|
||||
if err := body.Close(); err != nil {
|
||||
s.log.Warn("close downloaded UI artifact", "artifact_id", artifactID, "err", err)
|
||||
}
|
||||
}()
|
||||
w.Header().Set("Content-Type", art.ContentType)
|
||||
w.Header().Set("Content-Disposition", mime.FormatMediaType("attachment", map[string]string{"filename": art.Filename}))
|
||||
w.Header().Set("Content-Length", strconv.FormatInt(art.SizeBytes, 10))
|
||||
w.Header().Set("X-Checksum-SHA256", art.SHA256)
|
||||
_, _ = io.Copy(w, body)
|
||||
}
|
||||
@@ -0,0 +1,44 @@
|
||||
package http
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestUIStatusPresentation(t *testing.T) {
|
||||
tests := []struct {
|
||||
status string
|
||||
label string
|
||||
class string
|
||||
}{
|
||||
{"pending", "Waiting for a worker", "waiting"},
|
||||
{"running", "Running", "active"},
|
||||
{"completed", "Tasks complete", "success"},
|
||||
{"failed", "Needs attention", "danger"},
|
||||
}
|
||||
for _, test := range tests {
|
||||
t.Run(test.status, func(t *testing.T) {
|
||||
if got := uiStatusLabel(test.status); got != test.label {
|
||||
t.Errorf("label = %q, want %q", got, test.label)
|
||||
}
|
||||
if got := uiStatusClass(test.status); got != test.class {
|
||||
t.Errorf("class = %q, want %q", got, test.class)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestUIProgressPercent(t *testing.T) {
|
||||
if got := uiProgressPercent(3, 1, 0, 8); got != 50 {
|
||||
t.Errorf("progress = %d, want 50", got)
|
||||
}
|
||||
if got := uiProgressPercent(1, 1, 0, 0); got != 0 {
|
||||
t.Errorf("empty progress = %d, want 0", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUITaskErrorPresentationDoesNotExposeCommand(t *testing.T) {
|
||||
if got := uiTaskErrorLabel("CalledProcessError"); got != "Local calculation failed" {
|
||||
t.Errorf("error label = %q", got)
|
||||
}
|
||||
if got := uiTaskErrorHint("CalledProcessError"); got == "" {
|
||||
t.Error("error hint must explain the failure")
|
||||
}
|
||||
}
|
||||
@@ -14,12 +14,13 @@ type UploadArtifact struct {
|
||||
tasks TaskRepository
|
||||
artifacts ArtifactRepository
|
||||
blobs BlobStore
|
||||
tx TxManager
|
||||
clk Clock
|
||||
}
|
||||
|
||||
func NewUploadArtifact(tasks TaskRepository, artifacts ArtifactRepository,
|
||||
blobs BlobStore, clk Clock) *UploadArtifact {
|
||||
return &UploadArtifact{tasks: tasks, artifacts: artifacts, blobs: blobs, clk: clk}
|
||||
blobs BlobStore, tx TxManager, clk Clock) *UploadArtifact {
|
||||
return &UploadArtifact{tasks: tasks, artifacts: artifacts, blobs: blobs, tx: tx, clk: clk}
|
||||
}
|
||||
|
||||
func (uc *UploadArtifact) Execute(ctx context.Context, in UploadArtifactInput) (*domain.Artifact, error) {
|
||||
@@ -32,6 +33,15 @@ func (uc *UploadArtifact) Execute(ctx context.Context, in UploadArtifactInput) (
|
||||
if !task.IsLeaseHeldBy(in.WorkerID, in.Attempt, uc.clk.Now()) {
|
||||
return nil, domain.ErrLeaseConflict
|
||||
}
|
||||
// A client can retry a PUT after losing the response. Return the one durable
|
||||
// result for this lease attempt instead of storing duplicate artifacts.
|
||||
existing, err := uc.artifacts.FindPartialResult(ctx, in.TaskID, in.Attempt)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if existing != nil {
|
||||
return existing, nil
|
||||
}
|
||||
|
||||
taskID := task.ID
|
||||
art, err := domain.NewArtifact(task.JobID, &taskID, domain.ArtifactPartialResult,
|
||||
@@ -50,25 +60,41 @@ func (uc *UploadArtifact) Execute(ctx context.Context, in UploadArtifactInput) (
|
||||
}
|
||||
art.SetContent(sum, size)
|
||||
|
||||
// The stream may take longer than the lease. Re-check after it finishes so
|
||||
// an expired worker cannot leave a durable result record behind. Completion
|
||||
// performs the same ownership check under its transaction.
|
||||
current, err := uc.tasks.Get(ctx, in.TaskID)
|
||||
// The stream may take longer than the lease. Lock the task while re-checking
|
||||
// ownership and inserting metadata: completion or another upload cannot race
|
||||
// this final decision. The database unique index is a second line of defence.
|
||||
var durable *domain.Artifact
|
||||
err = uc.tx.WithinTx(ctx, func(ctx context.Context) error {
|
||||
current, err := uc.tasks.GetForUpdate(ctx, in.TaskID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if !current.IsLeaseHeldBy(in.WorkerID, in.Attempt, uc.clk.Now()) {
|
||||
return domain.ErrLeaseConflict
|
||||
}
|
||||
existing, err := uc.artifacts.FindPartialResult(ctx, in.TaskID, in.Attempt)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if existing != nil {
|
||||
durable = existing
|
||||
return nil
|
||||
}
|
||||
if err := uc.artifacts.Insert(ctx, art); err != nil {
|
||||
return err
|
||||
}
|
||||
durable = art
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
_ = uc.blobs.Delete(ctx, art.StorageKey)
|
||||
return nil, err
|
||||
}
|
||||
if !current.IsLeaseHeldBy(in.WorkerID, in.Attempt, uc.clk.Now()) {
|
||||
if durable != art {
|
||||
// Another request won the race while this stream was being written.
|
||||
_ = uc.blobs.Delete(ctx, art.StorageKey)
|
||||
return nil, domain.ErrLeaseConflict
|
||||
}
|
||||
|
||||
// Persist the record. If that fails the blob would be an orphan, so remove it.
|
||||
if err := uc.artifacts.Insert(ctx, art); err != nil {
|
||||
_ = uc.blobs.Delete(ctx, art.StorageKey)
|
||||
return nil, err
|
||||
}
|
||||
return art, nil
|
||||
return durable, nil
|
||||
}
|
||||
|
||||
// DownloadArtifact returns an artifact's metadata together with a reader over
|
||||
|
||||
@@ -53,9 +53,12 @@ type SubmitDatasetInput struct {
|
||||
Workload string
|
||||
Parameters map[string]any
|
||||
RowsPerShard int
|
||||
Filename string
|
||||
ContentType string
|
||||
Body io.Reader
|
||||
// MaxRows limits how many data rows are turned into shards. Zero means the
|
||||
// whole uploaded dataset; the input artifact itself remains stored intact.
|
||||
MaxRows int
|
||||
Filename string
|
||||
ContentType string
|
||||
Body io.Reader
|
||||
}
|
||||
|
||||
type SubmitDatasetResult struct {
|
||||
|
||||
@@ -33,6 +33,17 @@ func NewCreateJob(jobs JobRepository, tasks TaskRepository, tx TxManager, clock
|
||||
// The all-or-none guarantee comes from TxManager: a half-created job would
|
||||
// leave chunks no worker could ever complete.
|
||||
func (uc *CreateJob) Execute(ctx context.Context, in CreateJobInput) (*domain.Job, error) {
|
||||
if in.Workload == "similarity-graph" || in.Workload == "similarity_graph" {
|
||||
// CTX-10 must plan triangular block pairs; ordinary independent input
|
||||
// chunks would silently omit every cross-chunk molecular pair.
|
||||
return nil, domain.ErrInvalidInput
|
||||
}
|
||||
if (in.Workload == "similarity-search" || in.Workload == "similarity_search") &&
|
||||
len(in.Chunks) > 1 && in.Parameters["query_id"] != nil {
|
||||
// Resolving once against the source dataset belongs to CTX-07. Letting
|
||||
// each shard resolve it would make most tasks fail or use inconsistent data.
|
||||
return nil, domain.ErrInvalidInput
|
||||
}
|
||||
chunks := make([]domain.ChunkSpec, 0, len(in.Chunks))
|
||||
for _, c := range in.Chunks {
|
||||
chunks = append(chunks, domain.ChunkSpec(c))
|
||||
@@ -62,6 +73,56 @@ type GetJobStatus struct {
|
||||
tasks TaskRepository
|
||||
}
|
||||
|
||||
// --- CancelJob -----------------------------------------------------------
|
||||
|
||||
type CancelJob struct {
|
||||
jobs JobRepository
|
||||
tasks TaskRepository
|
||||
tx TxManager
|
||||
clock Clock
|
||||
}
|
||||
|
||||
func NewCancelJob(jobs JobRepository, tasks TaskRepository, tx TxManager, clock Clock) *CancelJob {
|
||||
return &CancelJob{jobs: jobs, tasks: tasks, tx: tx, clock: clock}
|
||||
}
|
||||
|
||||
// Execute stops a job atomically. Completed and finally failed tasks are kept
|
||||
// as historical evidence; all other tasks are cancelled, including leased and
|
||||
// running ones. A repeated cancel of an already cancelled job is idempotent.
|
||||
func (uc *CancelJob) Execute(ctx context.Context, jobID uuid.UUID) (int64, error) {
|
||||
now := uc.clock.Now()
|
||||
var cancelled int64
|
||||
err := uc.tx.WithinTx(ctx, func(ctx context.Context) error {
|
||||
job, err := uc.jobs.Get(ctx, jobID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if job.Status == domain.JobCancelled {
|
||||
return nil
|
||||
}
|
||||
if job.Status == domain.JobCompleted || job.Status == domain.JobFailed {
|
||||
return domain.ErrJobNotCancellable
|
||||
}
|
||||
// The lease reaper can be the transition that exhausted the final task.
|
||||
// Check the authoritative task histogram as well as the cached job status,
|
||||
// so a stale status can never turn a failed/completed job into cancelled.
|
||||
counts, err := uc.tasks.CountByStatus(ctx, jobID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
derived := progressFrom(*job, counts).DeriveStatus()
|
||||
if derived == domain.JobCompleted || derived == domain.JobFailed {
|
||||
return domain.ErrJobNotCancellable
|
||||
}
|
||||
cancelled, err = uc.tasks.CancelByJob(ctx, jobID, now)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return uc.jobs.UpdateStatus(ctx, jobID, domain.JobCancelled, &now)
|
||||
})
|
||||
return cancelled, err
|
||||
}
|
||||
|
||||
func NewGetJobStatus(jobs JobRepository, tasks TaskRepository) *GetJobStatus {
|
||||
return &GetJobStatus{jobs: jobs, tasks: tasks}
|
||||
}
|
||||
@@ -144,9 +205,10 @@ func progressFrom(job domain.Job, counts map[domain.TaskStatus]int) domain.JobPr
|
||||
Job: job,
|
||||
Pending: counts[domain.TaskPending],
|
||||
// Leased and running are both "in flight" for progress purposes.
|
||||
Leased: counts[domain.TaskLeased] + counts[domain.TaskRunning],
|
||||
Done: counts[domain.TaskCompleted],
|
||||
Failed: counts[domain.TaskFailed],
|
||||
Leased: counts[domain.TaskLeased] + counts[domain.TaskRunning],
|
||||
Done: counts[domain.TaskCompleted],
|
||||
Failed: counts[domain.TaskFailed],
|
||||
Cancelled: counts[domain.TaskCancelled],
|
||||
}
|
||||
for _, n := range counts {
|
||||
p.Total += n
|
||||
@@ -172,3 +234,18 @@ func syncJobStatus(ctx context.Context, jobs JobRepository, tasks TaskRepository
|
||||
}
|
||||
return jobs.UpdateStatus(ctx, jobID, status, completedAt)
|
||||
}
|
||||
|
||||
func syncExpiredJobStatuses(ctx context.Context, jobs JobRepository, tasks TaskRepository,
|
||||
jobIDs []uuid.UUID, now time.Time) error {
|
||||
seen := make(map[uuid.UUID]struct{}, len(jobIDs))
|
||||
for _, jobID := range jobIDs {
|
||||
if _, duplicate := seen[jobID]; duplicate {
|
||||
continue
|
||||
}
|
||||
seen[jobID] = struct{}{}
|
||||
if err := syncJobStatus(ctx, jobs, tasks, jobID, now); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -56,9 +56,13 @@ type TaskRepository interface {
|
||||
// CountByStatus aggregates a job's tasks for progress reporting.
|
||||
CountByStatus(ctx context.Context, jobID uuid.UUID) (map[domain.TaskStatus]int, error)
|
||||
|
||||
// ExpireLeases applies the lease-expiry rule to every elapsed task and
|
||||
// reports how many were affected.
|
||||
ExpireLeases(ctx context.Context, now time.Time) (int64, error)
|
||||
// CancelByJob marks every non-terminal task as cancelled and invalidates its
|
||||
// lease. It returns how many tasks changed.
|
||||
CancelByJob(ctx context.Context, jobID uuid.UUID, now time.Time) (int64, error)
|
||||
|
||||
// ExpireLeases applies the lease-expiry rule to every elapsed task and returns
|
||||
// the distinct jobs whose aggregate status may have changed.
|
||||
ExpireLeases(ctx context.Context, now time.Time) ([]uuid.UUID, error)
|
||||
}
|
||||
|
||||
// JobRepository persists jobs.
|
||||
@@ -85,6 +89,9 @@ type WorkerRepository interface {
|
||||
type ArtifactRepository interface {
|
||||
Insert(ctx context.Context, a *domain.Artifact) error
|
||||
Get(ctx context.Context, id uuid.UUID) (*domain.Artifact, error)
|
||||
// FindPartialResult returns the durable result already uploaded for one task
|
||||
// attempt. A nil artifact means the attempt has not uploaded one yet.
|
||||
FindPartialResult(ctx context.Context, taskID uuid.UUID, attempt int) (*domain.Artifact, error)
|
||||
}
|
||||
|
||||
// BlobStore holds artifact bytes, addressed by an opaque storage key. It streams
|
||||
|
||||
@@ -21,12 +21,15 @@ import (
|
||||
|
||||
type ClaimTask struct {
|
||||
tasks TaskRepository
|
||||
jobs JobRepository
|
||||
workers WorkerRepository
|
||||
tx TxManager
|
||||
clock Clock
|
||||
leaseDuration time.Duration
|
||||
}
|
||||
|
||||
func NewClaimTask(tasks TaskRepository, clock Clock, leaseDuration time.Duration) *ClaimTask {
|
||||
return &ClaimTask{tasks: tasks, clock: clock, leaseDuration: leaseDuration}
|
||||
func NewClaimTask(tasks TaskRepository, jobs JobRepository, workers WorkerRepository, tx TxManager, clock Clock, leaseDuration time.Duration) *ClaimTask {
|
||||
return &ClaimTask{tasks: tasks, jobs: jobs, workers: workers, tx: tx, clock: clock, leaseDuration: leaseDuration}
|
||||
}
|
||||
|
||||
// Execute reclaims elapsed leases first, then hands out one task.
|
||||
@@ -42,27 +45,46 @@ func (uc *ClaimTask) Execute(ctx context.Context, in ClaimTaskInput) (*domain.Cl
|
||||
if in.WorkerID == "" {
|
||||
return nil, domain.ErrInvalidInput
|
||||
}
|
||||
now := uc.clock.Now()
|
||||
|
||||
if _, err := uc.tasks.ExpireLeases(ctx, now); err != nil {
|
||||
return nil, err
|
||||
workloads := in.Workloads
|
||||
if workerID, err := uuid.Parse(in.WorkerID); err == nil {
|
||||
worker, err := uc.workers.Get(ctx, workerID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
// Never trust caller-supplied capabilities: registration is the durable
|
||||
// worker identity and its allowlist.
|
||||
workloads = worker.Capabilities
|
||||
}
|
||||
var claimed *domain.ClaimedTask
|
||||
err := uc.tx.WithinTx(ctx, func(ctx context.Context) error {
|
||||
now := uc.clock.Now()
|
||||
affectedJobs, err := uc.tasks.ExpireLeases(ctx, now)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err := syncExpiredJobStatuses(ctx, uc.jobs, uc.tasks, affectedJobs, now); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
task, err := uc.tasks.ClaimNext(ctx, ClaimFilter{
|
||||
Workloads: in.Workloads,
|
||||
Owner: in.WorkerID,
|
||||
Now: now,
|
||||
LeaseUntil: now.Add(uc.leaseDuration),
|
||||
task, err := uc.tasks.ClaimNext(ctx, ClaimFilter{
|
||||
Workloads: workloads,
|
||||
Owner: in.WorkerID,
|
||||
Now: now,
|
||||
LeaseUntil: now.Add(uc.leaseDuration),
|
||||
})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if task != nil {
|
||||
value := task.AsClaimed()
|
||||
claimed = &value
|
||||
}
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if task == nil {
|
||||
return nil, nil // empty queue is a normal state, not an error
|
||||
}
|
||||
|
||||
claimed := task.AsClaimed()
|
||||
return &claimed, nil
|
||||
return claimed, nil // nil means an empty queue
|
||||
}
|
||||
|
||||
// --- RenewLease ----------------------------------------------------------
|
||||
@@ -230,18 +252,30 @@ func (uc *FailTask) Execute(ctx context.Context, in FailTaskInput) (*domain.Task
|
||||
|
||||
type ExpireLeases struct {
|
||||
tasks TaskRepository
|
||||
jobs JobRepository
|
||||
tx TxManager
|
||||
clock Clock
|
||||
}
|
||||
|
||||
func NewExpireLeases(tasks TaskRepository, clock Clock) *ExpireLeases {
|
||||
return &ExpireLeases{tasks: tasks, clock: clock}
|
||||
func NewExpireLeases(tasks TaskRepository, jobs JobRepository, tx TxManager, clock Clock) *ExpireLeases {
|
||||
return &ExpireLeases{tasks: tasks, jobs: jobs, tx: tx, clock: clock}
|
||||
}
|
||||
|
||||
// Execute reports how many tasks were reclaimed.
|
||||
// Execute reclaims elapsed tasks and persists the state of every affected job.
|
||||
//
|
||||
// The sweep is one set-based statement rather than a load-decide-save loop:
|
||||
// several coordinators run it concurrently, and a single atomic UPDATE makes
|
||||
// the duplicate work harmless — the loser simply updates 0 rows.
|
||||
func (uc *ExpireLeases) Execute(ctx context.Context) (int64, error) {
|
||||
return uc.tasks.ExpireLeases(ctx, uc.clock.Now())
|
||||
var affected []uuid.UUID
|
||||
err := uc.tx.WithinTx(ctx, func(ctx context.Context) error {
|
||||
now := uc.clock.Now()
|
||||
var err error
|
||||
affected, err = uc.tasks.ExpireLeases(ctx, now)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return syncExpiredJobStatuses(ctx, uc.jobs, uc.tasks, affected, now)
|
||||
})
|
||||
return int64(len(affected)), err
|
||||
}
|
||||
|
||||
@@ -0,0 +1,182 @@
|
||||
package usecase
|
||||
|
||||
import (
|
||||
"context"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
|
||||
"github.com/emil28092005/SciMesh/coordinator/internal/domain"
|
||||
)
|
||||
|
||||
// UIReadRepository is a read-only projection source for the local operator UI.
|
||||
// It intentionally exposes no storage paths or credentials.
|
||||
type UIReadRepository interface {
|
||||
GetJob(ctx context.Context, jobID uuid.UUID) (*domain.Job, error)
|
||||
ListJobs(ctx context.Context, limit int) ([]domain.Job, error)
|
||||
ListTasksByJob(ctx context.Context, jobID uuid.UUID) ([]domain.Task, error)
|
||||
ListTasksByJobs(ctx context.Context, jobIDs []uuid.UUID) (map[uuid.UUID][]domain.Task, error)
|
||||
ListWorkers(ctx context.Context, limit int) ([]domain.Worker, error)
|
||||
ListArtifactsByJob(ctx context.Context, jobID uuid.UUID) ([]domain.Artifact, error)
|
||||
}
|
||||
|
||||
type JobCard struct {
|
||||
ID string `json:"id"`
|
||||
Workload string `json:"workload"`
|
||||
Status string `json:"status"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
Total int `json:"total"`
|
||||
Pending int `json:"pending"`
|
||||
Leased int `json:"leased"`
|
||||
Running int `json:"running"`
|
||||
Completed int `json:"completed"`
|
||||
Failed int `json:"failed"`
|
||||
Cancelled int `json:"cancelled"`
|
||||
}
|
||||
|
||||
type TaskCard struct {
|
||||
ID string `json:"id"`
|
||||
ChunkIndex int `json:"chunk_index"`
|
||||
Status string `json:"status"`
|
||||
Attempt int `json:"attempt"`
|
||||
MaxAttempts int `json:"max_attempts"`
|
||||
LeaseOwner string `json:"lease_owner,omitempty"`
|
||||
LeaseExpiresAt *time.Time `json:"lease_expires_at,omitempty"`
|
||||
ErrorCode string `json:"error_code,omitempty"`
|
||||
ErrorMessage string `json:"error_message,omitempty"`
|
||||
}
|
||||
|
||||
type ArtifactCard struct {
|
||||
ID string `json:"id"`
|
||||
Kind string `json:"kind"`
|
||||
Filename string `json:"filename"`
|
||||
SizeBytes int64 `json:"size_bytes"`
|
||||
SHA256 string `json:"sha256"`
|
||||
Downloadable bool `json:"downloadable"`
|
||||
Diagnostic bool `json:"diagnostic"`
|
||||
}
|
||||
|
||||
type WorkerCard struct {
|
||||
ID string `json:"id"`
|
||||
Name string `json:"name"`
|
||||
Status string `json:"status"`
|
||||
Capabilities []string `json:"capabilities"`
|
||||
LastHeartbeatAt time.Time `json:"last_heartbeat_at"`
|
||||
}
|
||||
|
||||
type DashboardView struct {
|
||||
Jobs []JobCard
|
||||
Workers []WorkerCard
|
||||
}
|
||||
type JobDetailView struct {
|
||||
JobCard
|
||||
Tasks []TaskCard `json:"tasks"`
|
||||
Artifacts []ArtifactCard `json:"artifacts"`
|
||||
FinalResultAvailable bool `json:"final_result_available"`
|
||||
}
|
||||
|
||||
type Dashboard struct{ read UIReadRepository }
|
||||
|
||||
func NewDashboard(read UIReadRepository) *Dashboard { return &Dashboard{read: read} }
|
||||
|
||||
func (d *Dashboard) Overview(ctx context.Context, limit int) (DashboardView, error) {
|
||||
jobs, err := d.read.ListJobs(ctx, limit)
|
||||
if err != nil {
|
||||
return DashboardView{}, err
|
||||
}
|
||||
workers, err := d.read.ListWorkers(ctx, limit)
|
||||
if err != nil {
|
||||
return DashboardView{}, err
|
||||
}
|
||||
out := DashboardView{Jobs: make([]JobCard, 0, len(jobs)), Workers: make([]WorkerCard, 0, len(workers))}
|
||||
jobIDs := make([]uuid.UUID, 0, len(jobs))
|
||||
for _, job := range jobs {
|
||||
jobIDs = append(jobIDs, job.ID)
|
||||
}
|
||||
tasksByJob, err := d.read.ListTasksByJobs(ctx, jobIDs)
|
||||
if err != nil {
|
||||
return DashboardView{}, err
|
||||
}
|
||||
for _, job := range jobs {
|
||||
out.Jobs = append(out.Jobs, jobCard(job, tasksByJob[job.ID]))
|
||||
}
|
||||
for _, worker := range workers {
|
||||
out.Workers = append(out.Workers, WorkerCard{ID: worker.ID.String(), Name: worker.Name, Status: string(worker.Status), Capabilities: worker.Capabilities, LastHeartbeatAt: worker.LastHeartbeatAt})
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (d *Dashboard) JobDetail(ctx context.Context, jobID uuid.UUID) (JobDetailView, error) {
|
||||
job, err := d.read.GetJob(ctx, jobID)
|
||||
if err != nil {
|
||||
return JobDetailView{}, err
|
||||
}
|
||||
tasks, err := d.read.ListTasksByJob(ctx, jobID)
|
||||
if err != nil {
|
||||
return JobDetailView{}, err
|
||||
}
|
||||
artifacts, err := d.read.ListArtifactsByJob(ctx, jobID)
|
||||
if err != nil {
|
||||
return JobDetailView{}, err
|
||||
}
|
||||
out := JobDetailView{JobCard: jobCard(*job, tasks), Tasks: make([]TaskCard, 0, len(tasks)), Artifacts: make([]ArtifactCard, 0, len(artifacts))}
|
||||
for _, task := range tasks {
|
||||
card := TaskCard{ID: task.ID.String(), ChunkIndex: task.ChunkIndex, Status: string(task.Status), Attempt: task.Attempt, MaxAttempts: task.MaxAttempts, LeaseExpiresAt: task.LeaseExpiresAt}
|
||||
if task.LeaseOwner != nil {
|
||||
card.LeaseOwner = *task.LeaseOwner
|
||||
}
|
||||
if task.ErrorCode != nil {
|
||||
card.ErrorCode = *task.ErrorCode
|
||||
}
|
||||
if task.ErrorMessage != nil {
|
||||
card.ErrorMessage = *task.ErrorMessage
|
||||
}
|
||||
out.Tasks = append(out.Tasks, card)
|
||||
}
|
||||
for _, artifact := range artifacts {
|
||||
diagnostic := artifact.Kind == domain.ArtifactPartialResult
|
||||
downloadable := diagnostic || (artifact.Kind == domain.ArtifactFinalResult && out.Status == string(domain.JobCompleted))
|
||||
out.Artifacts = append(out.Artifacts, ArtifactCard{ID: artifact.ID.String(), Kind: string(artifact.Kind), Filename: artifact.Filename, SizeBytes: artifact.SizeBytes, SHA256: artifact.SHA256, Downloadable: downloadable, Diagnostic: diagnostic})
|
||||
if artifact.Kind == domain.ArtifactFinalResult && downloadable {
|
||||
out.FinalResultAvailable = true
|
||||
}
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (d *Dashboard) ArtifactBelongsToJob(ctx context.Context, jobID, artifactID uuid.UUID) (bool, error) {
|
||||
artifacts, err := d.read.ListArtifactsByJob(ctx, jobID)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
for _, a := range artifacts {
|
||||
if a.ID == artifactID {
|
||||
return true, nil
|
||||
}
|
||||
}
|
||||
return false, nil
|
||||
}
|
||||
|
||||
func jobCard(job domain.Job, tasks []domain.Task) JobCard {
|
||||
c := JobCard{ID: job.ID.String(), Workload: job.Workload, CreatedAt: job.CreatedAt}
|
||||
for _, task := range tasks {
|
||||
c.Total++
|
||||
switch task.Status {
|
||||
case domain.TaskPending:
|
||||
c.Pending++
|
||||
case domain.TaskLeased:
|
||||
c.Leased++
|
||||
case domain.TaskRunning:
|
||||
c.Running++
|
||||
case domain.TaskCompleted:
|
||||
c.Completed++
|
||||
case domain.TaskFailed:
|
||||
c.Failed++
|
||||
case domain.TaskCancelled:
|
||||
c.Cancelled++
|
||||
}
|
||||
}
|
||||
p := domain.JobProgress{Job: job, Total: c.Total, Pending: c.Pending, Leased: c.Leased + c.Running, Done: c.Completed, Failed: c.Failed, Cancelled: c.Cancelled}
|
||||
c.Status = string(p.DeriveStatus())
|
||||
return c
|
||||
}
|
||||
@@ -4,6 +4,7 @@ import (
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"math"
|
||||
|
||||
"github.com/google/uuid"
|
||||
|
||||
@@ -15,20 +16,27 @@ import (
|
||||
// creates the job with one task per shard — the coordinator-side counterpart of
|
||||
// a client submitting pre-chunked URIs.
|
||||
type SubmitDataset struct {
|
||||
blobs BlobStore
|
||||
artifacts ArtifactRepository
|
||||
jobs JobRepository
|
||||
tasks TaskRepository
|
||||
tx TxManager
|
||||
clk Clock
|
||||
blobs BlobStore
|
||||
artifacts ArtifactRepository
|
||||
jobs JobRepository
|
||||
tasks TaskRepository
|
||||
tx TxManager
|
||||
clk Clock
|
||||
maxAttempts int
|
||||
}
|
||||
|
||||
func NewSubmitDataset(blobs BlobStore, artifacts ArtifactRepository, jobs JobRepository,
|
||||
tasks TaskRepository, tx TxManager, clk Clock) *SubmitDataset {
|
||||
return &SubmitDataset{blobs: blobs, artifacts: artifacts, jobs: jobs, tasks: tasks, tx: tx, clk: clk}
|
||||
tasks TaskRepository, tx TxManager, clk Clock, maxAttempts int) *SubmitDataset {
|
||||
return &SubmitDataset{blobs: blobs, artifacts: artifacts, jobs: jobs, tasks: tasks, tx: tx, clk: clk, maxAttempts: maxAttempts}
|
||||
}
|
||||
|
||||
func (uc *SubmitDataset) Execute(ctx context.Context, in SubmitDatasetInput) (SubmitDatasetResult, error) {
|
||||
if err := validateUploadedWorkload(in.Workload, in.Parameters); err != nil {
|
||||
return SubmitDatasetResult{}, err
|
||||
}
|
||||
if uc.maxAttempts < 1 {
|
||||
return SubmitDatasetResult{}, domain.ErrInvalidInput
|
||||
}
|
||||
now := uc.clk.Now()
|
||||
|
||||
job, err := domain.NewUploadedJob(in.Workload, in.Parameters, now)
|
||||
@@ -64,7 +72,7 @@ func (uc *SubmitDataset) Execute(ctx context.Context, in SubmitDatasetInput) (Su
|
||||
cleanup()
|
||||
return SubmitDatasetResult{}, err
|
||||
}
|
||||
splitErr := chunk.SplitTSV(rc, in.RowsPerShard, func(index int, shard io.Reader) error {
|
||||
splitErr := chunk.SplitChEMBLTSVLimit(rc, in.RowsPerShard, in.MaxRows, func(index int, shard io.Reader) error {
|
||||
art, err := domain.NewArtifact(job.ID, nil, domain.ArtifactShard,
|
||||
fmt.Sprintf("shard-%d.tsv", index), in.ContentType, now)
|
||||
if err != nil {
|
||||
@@ -77,7 +85,7 @@ func (uc *SubmitDataset) Execute(ctx context.Context, in SubmitDatasetInput) (Su
|
||||
putKeys = append(putKeys, art.StorageKey)
|
||||
art.SetContent(ssum, ssize)
|
||||
|
||||
task, err := domain.NewShardTask(job.ID, index, in.Workload, art.ID, ssum, in.Parameters, 0, now)
|
||||
task, err := domain.NewShardTask(job.ID, index, in.Workload, art.ID, ssum, in.Parameters, uc.maxAttempts, now)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
@@ -88,7 +96,8 @@ func (uc *SubmitDataset) Execute(ctx context.Context, in SubmitDatasetInput) (Su
|
||||
_ = rc.Close()
|
||||
if splitErr != nil {
|
||||
cleanup()
|
||||
return SubmitDatasetResult{}, splitErr
|
||||
// Dataset shape is caller input, not an internal coordinator failure.
|
||||
return SubmitDatasetResult{}, domain.ErrInvalidInput
|
||||
}
|
||||
|
||||
// 3. Persist job + all artifacts + all tasks atomically.
|
||||
@@ -118,6 +127,76 @@ func (uc *SubmitDataset) Execute(ctx context.Context, in SubmitDatasetInput) (Su
|
||||
}, nil
|
||||
}
|
||||
|
||||
// validateUploadedWorkload is deliberately narrow until CTX-07/08/10 adds a
|
||||
// typed distributed-workload registry. In particular, running similarity-graph
|
||||
// independently per TSV shard is scientifically wrong: cross-shard pairs would
|
||||
// be absent from the apparent graph.
|
||||
func validateUploadedWorkload(workload string, parameters map[string]any) error {
|
||||
if workload != "similarity-search" {
|
||||
return domain.ErrInvalidInput
|
||||
}
|
||||
allowed := map[string]struct{}{
|
||||
"query_smiles": {}, "top_k": {}, "threshold": {},
|
||||
"threshold_direction": {}, "progress_every": {},
|
||||
}
|
||||
for key := range parameters {
|
||||
if _, ok := allowed[key]; !ok {
|
||||
return domain.ErrInvalidInput
|
||||
}
|
||||
}
|
||||
query, ok := parameters["query_smiles"].(string)
|
||||
if !ok || query == "" || len(query) > 200 {
|
||||
return domain.ErrInvalidInput
|
||||
}
|
||||
if value, ok := parameters["top_k"]; ok && !isPositiveJSONInteger(value) {
|
||||
return domain.ErrInvalidInput
|
||||
}
|
||||
if value, ok := parameters["progress_every"]; ok && !isNonNegativeJSONInteger(value) {
|
||||
return domain.ErrInvalidInput
|
||||
}
|
||||
if value, ok := parameters["threshold"]; ok && !isUnitIntervalNumber(value) {
|
||||
return domain.ErrInvalidInput
|
||||
}
|
||||
if value, ok := parameters["threshold_direction"]; ok && value != "greater" && value != "less" {
|
||||
return domain.ErrInvalidInput
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func isPositiveJSONInteger(value any) bool { return isJSONInteger(value, false) }
|
||||
func isNonNegativeJSONInteger(value any) bool { return isJSONInteger(value, true) }
|
||||
|
||||
func isJSONInteger(value any, allowZero bool) bool {
|
||||
var n int64
|
||||
switch v := value.(type) {
|
||||
case int:
|
||||
n = int64(v)
|
||||
case int64:
|
||||
n = v
|
||||
case float64:
|
||||
if math.Trunc(v) != v || v > math.MaxInt64 || v < math.MinInt64 {
|
||||
return false
|
||||
}
|
||||
n = int64(v)
|
||||
default:
|
||||
return false
|
||||
}
|
||||
return n >= 0 && (allowZero || n > 0)
|
||||
}
|
||||
|
||||
func isUnitIntervalNumber(value any) bool {
|
||||
switch v := value.(type) {
|
||||
case float64:
|
||||
return !math.IsNaN(v) && !math.IsInf(v, 0) && v >= 0 && v <= 1
|
||||
case int:
|
||||
return v >= 0 && v <= 1
|
||||
case int64:
|
||||
return v >= 0 && v <= 1
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
// GetTaskInput resolves a task's input shard and opens it for streaming. The
|
||||
// caller closes the reader.
|
||||
type GetTaskInput struct {
|
||||
|
||||
@@ -54,6 +54,7 @@ type harness struct {
|
||||
downloadArt *usecase.DownloadArtifact
|
||||
getInput *usecase.GetTaskInput
|
||||
expire *usecase.ExpireLeases
|
||||
cancel *usecase.CancelJob
|
||||
}
|
||||
|
||||
func newHarness() *harness {
|
||||
@@ -67,18 +68,19 @@ func newHarness() *harness {
|
||||
}
|
||||
tx := memstore.Tx{}
|
||||
h.createJob = usecase.NewCreateJob(h.jobs, h.tasks, tx, h.clk)
|
||||
h.submit = usecase.NewSubmitDataset(h.blobs, h.arts, h.jobs, h.tasks, tx, h.clk)
|
||||
h.claim = usecase.NewClaimTask(h.tasks, h.clk, lease)
|
||||
h.submit = usecase.NewSubmitDataset(h.blobs, h.arts, h.jobs, h.tasks, tx, h.clk, 3)
|
||||
h.claim = usecase.NewClaimTask(h.tasks, h.jobs, h.work, tx, h.clk, lease)
|
||||
h.renew = usecase.NewRenewLease(h.tasks, h.work, tx, h.clk, lease)
|
||||
h.complete = usecase.NewCompleteTask(h.tasks, h.jobs, h.arts, tx, h.clk)
|
||||
h.fail = usecase.NewFailTask(h.tasks, h.jobs, tx, h.clk)
|
||||
h.status = usecase.NewGetJobStatus(h.jobs, h.tasks)
|
||||
h.results = usecase.NewListResults(h.tasks)
|
||||
h.register = usecase.NewRegisterWorker(h.work, h.clk)
|
||||
h.uploadArt = usecase.NewUploadArtifact(h.tasks, h.arts, h.blobs, h.clk)
|
||||
h.uploadArt = usecase.NewUploadArtifact(h.tasks, h.arts, h.blobs, tx, h.clk)
|
||||
h.downloadArt = usecase.NewDownloadArtifact(h.arts, h.blobs)
|
||||
h.getInput = usecase.NewGetTaskInput(h.tasks, h.arts, h.blobs)
|
||||
h.expire = usecase.NewExpireLeases(h.tasks, h.clk)
|
||||
h.expire = usecase.NewExpireLeases(h.tasks, h.jobs, tx, h.clk)
|
||||
h.cancel = usecase.NewCancelJob(h.jobs, h.tasks, tx, h.clk)
|
||||
return h
|
||||
}
|
||||
|
||||
@@ -136,6 +138,44 @@ func TestClaimLeasesAndAdvancesAttempt(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestRegisteredWorkerCannotBroadenItsCapabilitiesAtClaim(t *testing.T) {
|
||||
h := newHarness()
|
||||
h.seedJob(t, "restricted", 1)
|
||||
worker, err := h.register.Execute(ctx, usecase.RegisterWorkerInput{
|
||||
Name: "search-only", Capabilities: []string{"similarity-search"},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
claimed, err := h.claim.Execute(ctx, usecase.ClaimTaskInput{
|
||||
WorkerID: worker.ID.String(), Workloads: []string{"restricted"},
|
||||
})
|
||||
if err != nil || claimed != nil {
|
||||
t.Fatalf("claim = (%v, %v), want no compatible task", claimed, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCreateJobRejectsUnsafeDistributedScientificPlans(t *testing.T) {
|
||||
h := newHarness()
|
||||
_, err := h.createJob.Execute(ctx, usecase.CreateJobInput{
|
||||
Workload: "similarity-graph", InputURI: "s3://input",
|
||||
Chunks: []usecase.ChunkInput{{ChunkIndex: 0, InputURI: "s3://chunk", InputSHA256: "sha"}},
|
||||
})
|
||||
if !errors.Is(err, domain.ErrInvalidInput) {
|
||||
t.Errorf("graph job err = %v, want ErrInvalidInput", err)
|
||||
}
|
||||
_, err = h.createJob.Execute(ctx, usecase.CreateJobInput{
|
||||
Workload: "similarity-search", InputURI: "s3://input", Parameters: map[string]any{"query_id": "CHEMBL1"},
|
||||
Chunks: []usecase.ChunkInput{
|
||||
{ChunkIndex: 0, InputURI: "s3://chunk0", InputSHA256: "sha"},
|
||||
{ChunkIndex: 1, InputURI: "s3://chunk1", InputSHA256: "sha"},
|
||||
},
|
||||
})
|
||||
if !errors.Is(err, domain.ErrInvalidInput) {
|
||||
t.Errorf("sharded query_id job err = %v, want ErrInvalidInput", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestClaimEmptyQueueReturnsNil(t *testing.T) {
|
||||
h := newHarness()
|
||||
c, err := h.claim.Execute(ctx, usecase.ClaimTaskInput{WorkerID: "w1", Workloads: []string{"w"}})
|
||||
@@ -279,7 +319,7 @@ func TestUploadRejectsLeaseThatExpiresDuringStreaming(t *testing.T) {
|
||||
h.seedJob(t, "w", 1)
|
||||
taskID, attempt := h.leaseOne(t, "w1", "w")
|
||||
h.uploadArt = usecase.NewUploadArtifact(
|
||||
h.tasks, h.arts, expiringBlobStore{BlobStore: h.blobs, clock: h.clk}, h.clk,
|
||||
h.tasks, h.arts, expiringBlobStore{BlobStore: h.blobs, clock: h.clk}, memstore.Tx{}, h.clk,
|
||||
)
|
||||
|
||||
_, err := h.uploadArt.Execute(ctx, usecase.UploadArtifactInput{
|
||||
@@ -412,6 +452,23 @@ func TestUploadArtifactRejectsForeignWorker(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestUploadArtifactIsIdempotentPerTaskAttempt(t *testing.T) {
|
||||
h := newHarness()
|
||||
h.seedJob(t, "w", 1)
|
||||
taskID, attempt := h.leaseOne(t, "w1", "w")
|
||||
first := h.uploadResult(t, taskID, "w1", attempt)
|
||||
second, err := h.uploadArt.Execute(ctx, usecase.UploadArtifactInput{
|
||||
TaskID: taskID, WorkerID: "w1", Attempt: attempt,
|
||||
Filename: "retry.csv", ContentType: "text/csv", Body: strings.NewReader("different bytes"),
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if second.ID != first {
|
||||
t.Errorf("retry artifact = %s, want existing %s", second.ID, first)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDownloadArtifactRoundTrips(t *testing.T) {
|
||||
h := newHarness()
|
||||
h.seedJob(t, "w", 1)
|
||||
@@ -432,10 +489,10 @@ func TestDownloadArtifactRoundTrips(t *testing.T) {
|
||||
|
||||
func TestSubmitDatasetChunksAndServesInput(t *testing.T) {
|
||||
h := newHarness()
|
||||
tsv := "id\tsmiles\nA\tCC\nB\tCCC\nC\tCCCC\nD\tCCCCC\nE\tCCCCCC\n"
|
||||
tsv := "chembl_id\tcanonical_smiles\nA\tCC\nB\tCCC\nC\tCCCC\nD\tCCCCC\nE\tCCCCCC\n"
|
||||
|
||||
res, err := h.submit.Execute(ctx, usecase.SubmitDatasetInput{
|
||||
Workload: "w", RowsPerShard: 2, Filename: "chembl.tsv",
|
||||
Workload: "similarity-search", Parameters: map[string]any{"query_smiles": "CCO"}, RowsPerShard: 2, Filename: "chembl.tsv",
|
||||
ContentType: "text/tab-separated-values", Body: strings.NewReader(tsv),
|
||||
})
|
||||
if err != nil {
|
||||
@@ -451,7 +508,7 @@ func TestSubmitDatasetChunksAndServesInput(t *testing.T) {
|
||||
t.Errorf("job total = %d, want 3", prog.Total)
|
||||
}
|
||||
|
||||
c, err := h.claim.Execute(ctx, usecase.ClaimTaskInput{WorkerID: "w1", Workloads: []string{"w"}})
|
||||
c, err := h.claim.Execute(ctx, usecase.ClaimTaskInput{WorkerID: "w1", Workloads: []string{"similarity-search"}})
|
||||
if err != nil || c == nil {
|
||||
t.Fatalf("claim shard: %v", err)
|
||||
}
|
||||
@@ -468,6 +525,64 @@ func TestSubmitDatasetChunksAndServesInput(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestSubmitDatasetLimitsRowsBeforeCreatingShards(t *testing.T) {
|
||||
h := newHarness()
|
||||
tsv := "chembl_id\tcanonical_smiles\nA\tCC\nB\tCCC\nC\tCCCC\nD\tCCCCC\nE\tCCCCCC\n"
|
||||
res, err := h.submit.Execute(ctx, usecase.SubmitDatasetInput{
|
||||
Workload: "similarity-search", Parameters: map[string]any{"query_smiles": "CCO"}, RowsPerShard: 2, MaxRows: 3, Filename: "chembl.tsv",
|
||||
ContentType: "text/tab-separated-values", Body: strings.NewReader(tsv),
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if res.TaskCount != 2 {
|
||||
t.Fatalf("task_count = %d, want 2", res.TaskCount)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSubmitDatasetRejectsUnsupportedDistributedWorkloads(t *testing.T) {
|
||||
h := newHarness()
|
||||
_, err := h.submit.Execute(ctx, usecase.SubmitDatasetInput{
|
||||
Workload: "similarity-graph", Parameters: map[string]any{"threshold": 0.7}, RowsPerShard: 2,
|
||||
Filename: "chembl.tsv", ContentType: "text/tab-separated-values",
|
||||
Body: strings.NewReader("chembl_id\tcanonical_smiles\nA\tCC\n"),
|
||||
})
|
||||
if !errors.Is(err, domain.ErrInvalidInput) {
|
||||
t.Errorf("graph submission err = %v, want ErrInvalidInput", err)
|
||||
}
|
||||
_, err = h.submit.Execute(ctx, usecase.SubmitDatasetInput{
|
||||
Workload: "similarity-search", Parameters: map[string]any{"query_id": "CHEMBL1"}, RowsPerShard: 2,
|
||||
Filename: "chembl.tsv", ContentType: "text/tab-separated-values",
|
||||
Body: strings.NewReader("chembl_id\tcanonical_smiles\nA\tCC\n"),
|
||||
})
|
||||
if !errors.Is(err, domain.ErrInvalidInput) {
|
||||
t.Errorf("query_id submission err = %v, want ErrInvalidInput", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCancelJobInvalidatesClaimedAndPendingTasks(t *testing.T) {
|
||||
h := newHarness()
|
||||
jobID := h.seedJob(t, "w", 3)
|
||||
claimed, err := h.claim.Execute(ctx, usecase.ClaimTaskInput{WorkerID: "w1", Workloads: []string{"w"}})
|
||||
if err != nil || claimed == nil {
|
||||
t.Fatalf("claim: %v", err)
|
||||
}
|
||||
cancelled, err := h.cancel.Execute(ctx, jobID)
|
||||
if err != nil || cancelled != 3 {
|
||||
t.Fatalf("cancel = (%d, %v), want (3, nil)", cancelled, err)
|
||||
}
|
||||
if _, err := h.renew.Execute(ctx, usecase.RenewLeaseInput{TaskID: claimed.TaskID, WorkerID: "w1", Attempt: claimed.Attempt}); !errors.Is(err, domain.ErrTaskNotLeased) {
|
||||
t.Errorf("cancelled lease heartbeat = %v, want ErrTaskNotLeased", err)
|
||||
}
|
||||
progress, err := h.status.Execute(ctx, jobID)
|
||||
if err != nil || progress.DeriveStatus() != domain.JobCancelled || progress.Cancelled != 3 {
|
||||
t.Errorf("cancelled progress = %+v, err = %v", progress, err)
|
||||
}
|
||||
if cancelled, err := h.cancel.Execute(ctx, jobID); err != nil || cancelled != 0 {
|
||||
t.Errorf("second cancel = (%d, %v), want (0, nil)", cancelled, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetTaskInputMissingForURITask(t *testing.T) {
|
||||
h := newHarness()
|
||||
h.seedJob(t, "w", 1) // URI-based task, no coordinator-stored input
|
||||
@@ -491,3 +606,22 @@ func TestExpireLeasesReclaims(t *testing.T) {
|
||||
t.Errorf("expire = (%d, %v), want (1, nil)", n, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFinalLeaseExpiryPersistsFailedJobAndCannotBeCancelled(t *testing.T) {
|
||||
h := newHarness()
|
||||
jobID := h.seedJob(t, "w", 1)
|
||||
for attempt := 1; attempt <= domain.DefaultMaxAttempts; attempt++ {
|
||||
h.leaseOne(t, "w1", "w")
|
||||
h.clk.Advance(lease + time.Second)
|
||||
if _, err := h.expire.Execute(ctx); err != nil {
|
||||
t.Fatalf("expire attempt %d: %v", attempt, err)
|
||||
}
|
||||
}
|
||||
progress, err := h.status.Execute(ctx, jobID)
|
||||
if err != nil || progress.Job.Status != domain.JobFailed || progress.DeriveStatus() != domain.JobFailed {
|
||||
t.Fatalf("progress = %+v, err = %v; want persisted failed job", progress, err)
|
||||
}
|
||||
if _, err := h.cancel.Execute(ctx, jobID); !errors.Is(err, domain.ErrJobNotCancellable) {
|
||||
t.Errorf("cancel terminal lease failure = %v, want ErrJobNotCancellable", err)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,5 @@
|
||||
BEGIN;
|
||||
|
||||
DROP INDEX IF EXISTS uq_partial_result_task_attempt;
|
||||
|
||||
COMMIT;
|
||||
@@ -0,0 +1,26 @@
|
||||
BEGIN;
|
||||
|
||||
-- Old deployments can contain more than one partial result because earlier
|
||||
-- versions accepted repeated PUTs. Preserve the one referenced by a completed
|
||||
-- task and discard stale rows; unfinished tasks must upload again after a
|
||||
-- deploy, just as they do after a lost lease.
|
||||
DELETE FROM artifacts AS a
|
||||
USING tasks AS t
|
||||
WHERE a.task_id = t.id
|
||||
AND a.kind = 'partial_result'::artifact_kind
|
||||
AND t.status <> 'completed'::task_status;
|
||||
|
||||
DELETE FROM artifacts AS a
|
||||
USING tasks AS t
|
||||
WHERE a.task_id = t.id
|
||||
AND a.kind = 'partial_result'::artifact_kind
|
||||
AND t.status = 'completed'::task_status
|
||||
AND a.id <> t.result_artifact_id;
|
||||
|
||||
-- One lease attempt has one durable partial result. This makes an upload retry
|
||||
-- idempotent and prevents repeated uploads from accumulating orphan artifacts.
|
||||
CREATE UNIQUE INDEX uq_partial_result_task_attempt
|
||||
ON artifacts (task_id, attempt)
|
||||
WHERE kind = 'partial_result'::artifact_kind;
|
||||
|
||||
COMMIT;
|
||||
@@ -51,6 +51,13 @@ check "register worker" 201 -X POST "${HOST}/workers/register" "
|
||||
-d '{"name":"smoke-worker","capabilities":["similarity_search"],"cpu_count":4,"memory_mb":8192}'
|
||||
check "register without capabilities → 400" 400 -X POST "${HOST}/workers/register" "${auth[@]}" \
|
||||
-d '{"name":"bad"}'
|
||||
registration=$(curl -sS "${auth[@]}" -X POST "${HOST}/workers/register" \
|
||||
-d '{"name":"smoke-worker-active","capabilities":["similarity_search","similarity-search"],"cpu_count":4}')
|
||||
worker_id=$(printf '%s' "$registration" | python3 -c 'import json,sys;print(json.load(sys.stdin)["worker_id"])' 2>/dev/null)
|
||||
if [[ -z "${worker_id:-}" ]]; then
|
||||
echo " ✗ could not register active worker: $registration"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
echo
|
||||
echo "job lifecycle"
|
||||
@@ -75,7 +82,7 @@ declare -A our_chunks
|
||||
task_id=""
|
||||
attempt=""
|
||||
for _ in $(seq 1 40); do
|
||||
claim=$(curl -sS "${auth[@]}" -X POST "${HOST}/tasks/claim" -d '{"worker_id":"w1"}')
|
||||
claim=$(curl -sS "${auth[@]}" -X POST "${HOST}/tasks/claim" -d "{\"worker_id\":\"${worker_id}\"}")
|
||||
[[ -z "$claim" ]] && break # 204: queue drained
|
||||
|
||||
read -r c_job c_task c_chunk c_attempt < <(printf '%s' "$claim" |
|
||||
@@ -100,7 +107,7 @@ else
|
||||
fi
|
||||
|
||||
check "heartbeat" 200 -X POST "${HOST}/tasks/${task_id}/heartbeat" "${auth[@]}" \
|
||||
-d "{\"worker_id\":\"w1\",\"attempt\":${attempt}}"
|
||||
-d "{\"worker_id\":\"${worker_id}\",\"attempt\":${attempt}}"
|
||||
|
||||
# --- artifacts + result (uploads happen while the task is still leased) ---
|
||||
bearer=(-H "Authorization: Bearer ${TOKEN}")
|
||||
@@ -108,43 +115,47 @@ bearer=(-H "Authorization: Bearer ${TOKEN}")
|
||||
# upload <filename> -> prints the artifact_id
|
||||
upload() {
|
||||
curl -sS -X PUT "${HOST}/tasks/${task_id}/artifacts/$1" "${bearer[@]}" \
|
||||
-H 'Content-Type: text/csv' -H 'X-Worker-ID: w1' -H "X-Task-Attempt: ${attempt}" \
|
||||
-H 'Content-Type: text/csv' -H "X-Worker-ID: ${worker_id}" -H "X-Task-Attempt: ${attempt}" \
|
||||
--data-binary $'query,match,score\nA,B,0.9\n' |
|
||||
python3 -c 'import json,sys;print(json.load(sys.stdin)["artifact_id"])' 2>/dev/null
|
||||
}
|
||||
|
||||
check "upload artifact" 200 -X PUT "${HOST}/tasks/${task_id}/artifacts/result.csv" "${bearer[@]}" \
|
||||
-H 'Content-Type: text/csv' -H 'X-Worker-ID: w1' -H "X-Task-Attempt: ${attempt}" \
|
||||
-H 'Content-Type: text/csv' -H "X-Worker-ID: ${worker_id}" -H "X-Task-Attempt: ${attempt}" \
|
||||
--data-binary $'query,match,score\nA,B,0.9\n'
|
||||
check "foreign worker upload → 409" 409 -X PUT "${HOST}/tasks/${task_id}/artifacts/x.csv" "${bearer[@]}" \
|
||||
-H 'Content-Type: text/csv' -H 'X-Worker-ID: impostor' -H "X-Task-Attempt: ${attempt}" \
|
||||
--data-binary 'x'
|
||||
|
||||
# Two result artifacts, uploaded now while the lease is held: one to complete
|
||||
# with, a second to prove a different manifest is rejected after completion.
|
||||
# A retry of a PUT returns the same durable artifact for the task attempt.
|
||||
art_id=$(upload primary.csv)
|
||||
art_id2=$(upload secondary.csv)
|
||||
if [[ "$art_id" == "$art_id2" ]]; then
|
||||
printf ' \033[32m✓\033[0m %-46s %s\n' "duplicate upload is idempotent" "$art_id"
|
||||
pass=$((pass + 1))
|
||||
else
|
||||
printf ' \033[31m✗\033[0m %-46s got %s and %s\n' "duplicate upload is idempotent" "$art_id" "$art_id2"
|
||||
fail=$((fail + 1))
|
||||
fi
|
||||
check "download artifact" 200 "${HOST}/artifacts/${art_id}/download" "${bearer[@]}"
|
||||
|
||||
check "foreign worker submits → 409" 409 -X POST "${HOST}/tasks/${task_id}/result" "${auth[@]}" \
|
||||
-d "{\"worker_id\":\"impostor\",\"attempt\":${attempt},\"result\":{\"artifact_id\":\"${art_id}\"}}"
|
||||
check "submit result" 200 -X POST "${HOST}/tasks/${task_id}/result" "${auth[@]}" \
|
||||
-d "{\"worker_id\":\"w1\",\"attempt\":${attempt},\"result\":{\"artifact_id\":\"${art_id}\"}}"
|
||||
-d "{\"worker_id\":\"${worker_id}\",\"attempt\":${attempt},\"result\":{\"artifact_id\":\"${art_id}\"}}"
|
||||
check "replay same result → idempotent" 200 -X POST "${HOST}/tasks/${task_id}/result" "${auth[@]}" \
|
||||
-d "{\"worker_id\":\"w1\",\"attempt\":${attempt},\"result\":{\"artifact_id\":\"${art_id}\"}}"
|
||||
check "different result → 409" 409 -X POST "${HOST}/tasks/${task_id}/result" "${auth[@]}" \
|
||||
-d "{\"worker_id\":\"w1\",\"attempt\":${attempt},\"result\":{\"artifact_id\":\"${art_id2}\"}}"
|
||||
-d "{\"worker_id\":\"${worker_id}\",\"attempt\":${attempt},\"result\":{\"artifact_id\":\"${art_id}\"}}"
|
||||
check "GET /jobs/{id}" 200 "${HOST}/jobs/${job_id}" "${auth[@]}"
|
||||
|
||||
echo
|
||||
echo "input validation"
|
||||
check "malformed uuid → 400" 400 -X POST "${HOST}/tasks/not-a-uuid/result" "${auth[@]}" \
|
||||
-d '{"worker_id":"w1","attempt":1,"result":{"artifact_id":"00000000-0000-0000-0000-000000000000"}}'
|
||||
-d "{\"worker_id\":\"${worker_id}\",\"attempt\":1,\"result\":{\"artifact_id\":\"00000000-0000-0000-0000-000000000000\"}}"
|
||||
# Note: Go's encoding/json matches field names case-insensitively, so
|
||||
# "worker_ID" would be accepted as "worker_id". Only a genuinely unknown key
|
||||
# trips DisallowUnknownFields.
|
||||
check "unknown json field → 400" 400 -X POST "${HOST}/tasks/claim" "${auth[@]}" \
|
||||
-d '{"worker_id":"w1","totally_unknown":1}'
|
||||
-d "{\"worker_id\":\"${worker_id}\",\"totally_unknown\":1}"
|
||||
check "unknown job → 404" 404 "${HOST}/jobs/00000000-0000-0000-0000-000000000000" "${auth[@]}"
|
||||
|
||||
echo
|
||||
@@ -152,11 +163,11 @@ echo "dataset upload → chunking"
|
||||
# Upload a 5-row TSV split at 2 rows/shard → expect 3 shard tasks. The text
|
||||
# fields precede the file part, which the coordinator streams.
|
||||
up=$(curl -sS "${bearer[@]}" -X POST "${HOST}/jobs/upload" \
|
||||
-F 'workload=similarity_search' \
|
||||
-F 'parameters={"top_k":10}' \
|
||||
-F 'workload=similarity-search' \
|
||||
-F 'parameters={"query_smiles":"CCO","top_k":10}' \
|
||||
-F 'chunk_rows=2' \
|
||||
-F 'file=@-;filename=chembl.tsv;type=text/tab-separated-values' <<'TSV'
|
||||
id smiles
|
||||
chembl_id canonical_smiles
|
||||
A CC
|
||||
B CCC
|
||||
C CCCC
|
||||
@@ -180,7 +191,7 @@ fi
|
||||
up_input=""
|
||||
for _ in $(seq 1 30); do
|
||||
c=$(curl -sS "${bearer[@]}" -H 'Content-Type: application/json' -X POST "${HOST}/tasks/claim" \
|
||||
-d '{"worker_id":"up-w","capabilities":["similarity_search"]}')
|
||||
-d "{\"worker_id\":\"${worker_id}\"}")
|
||||
[[ -z "$c" ]] && break
|
||||
cj=$(printf '%s' "$c" | python3 -c 'import json,sys;print(json.load(sys.stdin)["job_id"])' 2>/dev/null)
|
||||
[[ "$cj" != "$up_job" ]] && continue
|
||||
|
||||
+27
-5
@@ -52,9 +52,11 @@ Content-Type: multipart/form-data
|
||||
```
|
||||
|
||||
Fields, in order (text fields first, file last — the file is streamed):
|
||||
`workload`, `parameters` (JSON), `chunk_rows` (int, default 1000), and the file
|
||||
part `file`. The coordinator stores the input, splits the TSV into shard
|
||||
artifacts (header repeated per shard), and creates one task per shard.
|
||||
`workload`, `parameters` (JSON), `chunk_rows` (int, default 1000), optional
|
||||
`max_rows` (positive int), and the file part `file`. `max_rows` limits the
|
||||
leading data rows that become shards; it does not change the stored source
|
||||
artifact. The coordinator splits the selected TSV rows into shard artifacts
|
||||
(header repeated per shard) and creates one task per shard.
|
||||
|
||||
`201`:
|
||||
|
||||
@@ -65,6 +67,24 @@ artifacts (header repeated per shard), and creates one task per shard.
|
||||
Each resulting task's claim response carries `input.uri = /tasks/{id}/input`,
|
||||
served by §5.4.
|
||||
|
||||
## Stop a job
|
||||
|
||||
```http
|
||||
POST /jobs/{job_id}/cancel
|
||||
Authorization: Bearer <token>
|
||||
```
|
||||
|
||||
The coordinator transactionally marks every pending, leased, or running shard
|
||||
as `cancelled`, invalidates its lease, and marks the job `cancelled`. Completed
|
||||
and terminally failed shards remain as history. Repeating a cancellation of an
|
||||
already cancelled job is safe.
|
||||
|
||||
`200`:
|
||||
|
||||
```json
|
||||
{ "job_id": "uuid", "status": "cancelled", "cancelled_tasks": 12 }
|
||||
```
|
||||
|
||||
## Register worker
|
||||
|
||||
```http
|
||||
@@ -87,7 +107,8 @@ Content-Type: application/json
|
||||
```
|
||||
|
||||
`cpu_count`/`memory_mb` are accepted for forward compatibility and not yet
|
||||
persisted. `capabilities` must be non-empty (an allowlisted workload set).
|
||||
persisted. `capabilities` must be non-empty. A claim uses the capabilities
|
||||
stored at registration; the request cannot broaden them.
|
||||
|
||||
## Claim task
|
||||
|
||||
@@ -100,7 +121,8 @@ Content-Type: application/json
|
||||
```
|
||||
|
||||
- `204 No Content`: no compatible task.
|
||||
- `200 OK`: a task is leased atomically.
|
||||
- `200 OK`: a task is leased atomically. `worker_id` must be a registered UUID;
|
||||
its persisted capabilities, rather than this request field, decide eligibility.
|
||||
|
||||
```json
|
||||
{
|
||||
|
||||
@@ -0,0 +1,81 @@
|
||||
# Task: safely preview partial CSV artifacts in the UI
|
||||
|
||||
## Assignment
|
||||
|
||||
You are the junior developer implementing one contained UI feature: an
|
||||
authenticated operator can preview a small portion of a CSV artifact belonging
|
||||
to the job they are viewing. This is a diagnostic aid, not a final-results
|
||||
page.
|
||||
|
||||
## Read first
|
||||
|
||||
1. `AGENTS.md`
|
||||
2. `.agents/coordinator.md`
|
||||
3. `docs/web-interface-plan.md`
|
||||
4. `docs/api-contract.md`
|
||||
5. `coordinator/internal/transport/http/ui.go` and its tests
|
||||
|
||||
## Current baseline
|
||||
|
||||
The coordinator serves an authenticated local web UI. Job detail pages already
|
||||
list partial result artifacts and provide job-scoped downloads. The browser has
|
||||
no worker token and must not learn storage paths. A partial shard CSV is never
|
||||
a global or final molecular-search result.
|
||||
|
||||
## Scope
|
||||
|
||||
Add a **Preview** action next to eligible CSV artifacts on a job detail page.
|
||||
|
||||
- Preview only artifacts that belong to the requested job.
|
||||
- Show at most the first **30 rows** and read at most **64 KiB** from storage.
|
||||
- State clearly when content was truncated.
|
||||
- Preserve the existing download action.
|
||||
- For a non-CSV artifact, return a friendly, sanitized explanation rather than
|
||||
attempting to render bytes as text.
|
||||
- Use an existing UI route pattern or add a small UI-authenticated endpoint;
|
||||
keep it separate from worker API routes.
|
||||
|
||||
## Security rules
|
||||
|
||||
- Require UI Basic Auth for every preview request.
|
||||
- Verify job ownership in the coordinator before opening the artifact; an
|
||||
artifact ID from another job must not be previewable.
|
||||
- Never expose `storage_key`, filesystem paths, database errors, bearer tokens,
|
||||
or worker-local information.
|
||||
- Do not use `innerHTML` for CSV fields. Use `html/template` escaping or
|
||||
`textContent` so strings such as `<script>alert(1)</script>` are displayed as
|
||||
data, not executed.
|
||||
- Do not load the complete artifact into memory.
|
||||
|
||||
## Out of scope
|
||||
|
||||
- Charts, molecule imagery, RDKit rendering, client-side CSV libraries, React,
|
||||
and a new frontend service.
|
||||
- Changing job/task state, retrying tasks, or implementing reducer output.
|
||||
- Redesigning the broader dashboard or job-creation workflow; that belongs to
|
||||
`docs/user-space-task.md`.
|
||||
|
||||
## Suggested implementation shape
|
||||
|
||||
Keep UI transport, use case, and storage responsibilities separate. Return a
|
||||
small view model containing artifact metadata, column headers, rows, and a
|
||||
`truncated` flag. Reuse existing coordinator-owned artifact access rather than
|
||||
reading a path supplied by the browser. Keep the handler streaming/limited.
|
||||
|
||||
## Acceptance criteria
|
||||
|
||||
- A valid partial CSV can be previewed from its own job detail page.
|
||||
- The result shows no more than 30 data rows and marks 64 KiB/row truncation.
|
||||
- Empty and malformed CSV content fail safely with a clear message.
|
||||
- A non-CSV artifact is rejected safely.
|
||||
- Unauthenticated access is rejected; a cross-job artifact request is not
|
||||
disclosed or served.
|
||||
- HTML-like values are escaped in the rendered preview.
|
||||
- Existing artifact downloads still work.
|
||||
- Add Go tests for all cases above and run `go test ./...` and `go vet ./...`.
|
||||
|
||||
## Handoff
|
||||
|
||||
Work in one focused branch and one PR. Report files changed, any API impact,
|
||||
test commands/results, and known limitations. Do not commit datasets, generated
|
||||
CSV files, Docker volumes, `.venv`, or `worker-data/`.
|
||||
@@ -46,24 +46,25 @@ coordinator was started with. Never log it, never send it in an error body.
|
||||
|
||||
```http
|
||||
POST /workers/register
|
||||
{ "name": "lab-worker-01", "capabilities": ["similarity_search"] }
|
||||
{ "name": "lab-worker-01", "capabilities": ["similarity-search"] }
|
||||
```
|
||||
|
||||
Response: `{ "worker_id": "<uuid>", "heartbeat_interval_seconds": 15 }`.
|
||||
|
||||
- `capabilities` are the workload names you can run — the coordinator only hands
|
||||
you matching tasks.
|
||||
- `capabilities` are fixed at registration — the coordinator only hands you
|
||||
matching tasks and a later claim cannot broaden that set.
|
||||
- **Keep `worker_id`**. Use it as your identity in every later call. Using the
|
||||
registered UUID is what lets the coordinator track your liveness (it marks
|
||||
workers offline after they go silent).
|
||||
- Current coordinator jobs use `similarity_search` / `similarity_graph`; the
|
||||
reference Python worker also accepts the public CLI spellings with hyphens.
|
||||
- Current diagnostic uploads use `similarity-search` with `query_smiles`. The
|
||||
reference worker accepts the legacy `similarity_search` spelling too. Do not
|
||||
advertise `similarity-graph` until CTX-10 implements cross-shard pair planning.
|
||||
|
||||
## 2. Claim a task
|
||||
|
||||
```http
|
||||
POST /tasks/claim
|
||||
{ "worker_id": "<uuid>", "capabilities": ["similarity_search"] }
|
||||
{ "worker_id": "<uuid>", "capabilities": ["similarity-search"] }
|
||||
```
|
||||
|
||||
- `200` → a leased task (below).
|
||||
@@ -74,9 +75,9 @@ POST /tasks/claim
|
||||
"task_id": "<uuid>",
|
||||
"attempt": 1,
|
||||
"lease_expires_at": "2026-07-22T12:05:00Z",
|
||||
"workload": "similarity_search",
|
||||
"workload": "similarity-search",
|
||||
"input": { "uri": "/tasks/<uuid>/input", "sha256": "<hex>" },
|
||||
"parameters": { "query_id": "CHEMBL939", "top_k": 20 }
|
||||
"parameters": { "query_smiles": "CCO", "top_k": 20 }
|
||||
}
|
||||
```
|
||||
|
||||
@@ -193,6 +194,32 @@ Per the worker contract, at minimum:
|
||||
- poll interval and request timeout
|
||||
- a working directory for downloaded inputs and generated outputs
|
||||
|
||||
## Run the reference worker locally
|
||||
|
||||
Use one terminal per worker and a distinct work directory for each process:
|
||||
|
||||
```sh
|
||||
SCIMESH_COORDINATOR_URL=http://localhost:8080 \
|
||||
SCIMESH_BEARER_TOKEN=dev-token \
|
||||
SCIMESH_WORKER_NAME=worker-1 \
|
||||
scimesh-worker --work-dir "$PWD/worker-data-1"
|
||||
```
|
||||
|
||||
For a bounded manual check, use one of these lifecycle modes:
|
||||
|
||||
```sh
|
||||
# Make exactly one claim; exit immediately when no task is available.
|
||||
scimesh-worker --work-dir "$PWD/worker-data-check" --once
|
||||
|
||||
# Keep polling until two tasks complete successfully, then exit.
|
||||
scimesh-worker --work-dir "$PWD/worker-data-check" --max-tasks 2
|
||||
```
|
||||
|
||||
`SCIMESH_MAX_TASKS` provides the same limit through the environment. Pressing
|
||||
`Ctrl+C` stops the reference worker cleanly. If it interrupts an active task,
|
||||
the worker reports a sanitized retriable failure first, emits no traceback, and
|
||||
exits with status `130`.
|
||||
|
||||
## Generate a client from the spec
|
||||
|
||||
Instead of hand-writing request code, generate it:
|
||||
|
||||
@@ -0,0 +1,209 @@
|
||||
# CTX-07: distributed workload protocol and planner contract
|
||||
|
||||
## Status and scope
|
||||
|
||||
This document is the implementation contract for CTX-07. Its generic protocol,
|
||||
registry, strict JSON models, and deterministic reduction ordering are
|
||||
implemented in `scimesh/distributed/`. It does not implement a molecular
|
||||
planner, reducer, API endpoint, database migration, or final artifact. Until
|
||||
CTX-08 and CTX-09 are complete, shard CSVs remain diagnostic partial results.
|
||||
|
||||
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.
|
||||
+41
-9
@@ -95,8 +95,10 @@ paths:
|
||||
summary: Upload a dataset; the coordinator chunks it into shard tasks
|
||||
description: >
|
||||
multipart/form-data. The text fields (`workload`, `parameters`,
|
||||
`chunk_rows`) MUST precede the `file` part: the file is streamed, not
|
||||
buffered, so the fields have to be parsed before it arrives.
|
||||
`chunk_rows`, `max_rows`) MUST precede the `file` part: the file is streamed, not
|
||||
buffered, so the fields have to be parsed before it arrives. Currently
|
||||
only diagnostic `similarity-search` with `parameters.query_smiles` is
|
||||
accepted; distributed graph planning is not implemented.
|
||||
requestBody:
|
||||
required: true
|
||||
content:
|
||||
@@ -130,6 +132,23 @@ paths:
|
||||
"401": { $ref: "#/components/responses/Unauthorized" }
|
||||
"404": { $ref: "#/components/responses/NotFound" }
|
||||
|
||||
/jobs/{job_id}/cancel:
|
||||
post:
|
||||
tags: [jobs]
|
||||
summary: Cancel a job and invalidate all unfinished task leases
|
||||
parameters:
|
||||
- $ref: "#/components/parameters/JobID"
|
||||
responses:
|
||||
"200":
|
||||
description: The job is cancelled. Completed and terminally failed tasks remain unchanged.
|
||||
content:
|
||||
application/json:
|
||||
schema: { $ref: "#/components/schemas/CancelJobResponse" }
|
||||
"400": { $ref: "#/components/responses/BadRequest" }
|
||||
"401": { $ref: "#/components/responses/Unauthorized" }
|
||||
"404": { $ref: "#/components/responses/NotFound" }
|
||||
"409": { $ref: "#/components/responses/Conflict" }
|
||||
|
||||
/tasks/claim:
|
||||
post:
|
||||
tags: [tasks]
|
||||
@@ -371,7 +390,7 @@ components:
|
||||
type: array
|
||||
minItems: 1
|
||||
items: { type: string }
|
||||
example: [similarity_search, similarity_graph]
|
||||
example: [similarity-search]
|
||||
cpu_count:
|
||||
type: integer
|
||||
description: Accepted for forward compatibility; not yet persisted.
|
||||
@@ -402,7 +421,7 @@ components:
|
||||
type: object
|
||||
required: [workload, input_uri, chunks]
|
||||
properties:
|
||||
workload: { type: string, example: similarity_search }
|
||||
workload: { type: string, example: similarity-search }
|
||||
input_uri: { type: string }
|
||||
parameters: { type: object, additionalProperties: true }
|
||||
chunks:
|
||||
@@ -420,15 +439,20 @@ components:
|
||||
type: object
|
||||
required: [workload, file]
|
||||
properties:
|
||||
workload: { type: string, example: similarity_search }
|
||||
workload: { type: string, enum: [similarity-search], example: similarity-search }
|
||||
parameters:
|
||||
type: string
|
||||
description: JSON object, sent as a string form field.
|
||||
example: '{"top_k":10}'
|
||||
example: '{"query_smiles":"CCO","top_k":10}'
|
||||
chunk_rows:
|
||||
type: integer
|
||||
description: Data rows per shard. Default 1000.
|
||||
example: 1000
|
||||
max_rows:
|
||||
type: integer
|
||||
minimum: 1
|
||||
description: Optional leading data-row limit for a small pipeline check.
|
||||
example: 500
|
||||
file:
|
||||
type: string
|
||||
format: binary
|
||||
@@ -441,6 +465,13 @@ components:
|
||||
task_count: { type: integer, example: 3 }
|
||||
input_artifact_id: { type: string, format: uuid }
|
||||
|
||||
CancelJobResponse:
|
||||
type: object
|
||||
properties:
|
||||
job_id: { type: string, format: uuid }
|
||||
status: { type: string, enum: [cancelled] }
|
||||
cancelled_tasks: { type: integer }
|
||||
|
||||
JobProgress:
|
||||
type: object
|
||||
properties:
|
||||
@@ -451,16 +482,17 @@ components:
|
||||
leased: { type: integer }
|
||||
completed: { type: integer }
|
||||
failed: { type: integer }
|
||||
cancelled: { type: integer }
|
||||
|
||||
ClaimRequest:
|
||||
type: object
|
||||
required: [worker_id]
|
||||
properties:
|
||||
worker_id: { type: string }
|
||||
worker_id: { type: string, format: uuid, description: Registered worker identity. }
|
||||
capabilities:
|
||||
type: array
|
||||
items: { type: string }
|
||||
description: Workloads this worker can run. Empty means "any".
|
||||
description: Accepted for compatibility only; registration capabilities decide eligibility.
|
||||
max_concurrency:
|
||||
type: integer
|
||||
description: Accepted; the coordinator leases one task per call.
|
||||
@@ -549,4 +581,4 @@ components:
|
||||
|
||||
TaskStatus:
|
||||
type: string
|
||||
enum: [pending, leased, completed, failed, cancelled]
|
||||
enum: [pending, leased, running, completed, failed, cancelled]
|
||||
|
||||
@@ -0,0 +1,83 @@
|
||||
# Task: improve the operator user space
|
||||
|
||||
## Assignment
|
||||
|
||||
You are the senior developer responsible for the **user space**: the
|
||||
human-facing local operator interface served by the Go coordinator. In this
|
||||
task, “user space” means a clear UI and workflow for a trusted local operator;
|
||||
it does **not** mean public accounts, registration, roles, multi-tenancy, or
|
||||
remote deployment.
|
||||
|
||||
Create a small, coherent improvement to the existing UI so a person can
|
||||
understand and operate a SciMesh pipeline without reading API payloads or
|
||||
coordinator logs. Keep all interface copy in English.
|
||||
|
||||
## Read first
|
||||
|
||||
1. `AGENTS.md`
|
||||
2. `.agents/coordinator.md` and `.agents/integration.md`
|
||||
3. `docs/web-interface-plan.md`
|
||||
4. `docs/api-contract.md`
|
||||
5. `STATUS.md` and the current `coordinator/internal/transport/http/ui.go`
|
||||
|
||||
## Current baseline
|
||||
|
||||
`main` already provides local Basic Auth (`UI_AUTH_TOKEN`), a dashboard, job
|
||||
submission for diagnostic similarity-search runs, task progress, partial CSV
|
||||
downloads, a stop-job action, and an optional dataset row limit. A job is a
|
||||
**pipeline check** until CTX-09 adds a reducer; individual shard CSVs are not a
|
||||
final scientific result.
|
||||
|
||||
## Scope
|
||||
|
||||
Improve the end-to-end operator journey:
|
||||
|
||||
- make the dashboard explain service readiness, workers, jobs, and the next
|
||||
safe action in plain English;
|
||||
- make job creation validation and success/failure feedback understandable;
|
||||
- make job detail clearly distinguish queued, running, failed, stopped, and
|
||||
completed pipeline checks;
|
||||
- keep polling and all user-visible states reliable after a page refresh;
|
||||
- expose actionable, sanitized failure guidance without leaking paths, tokens,
|
||||
SQL errors, or tracebacks;
|
||||
- document the workflow in `coordinator/README.md` or `README.md`.
|
||||
|
||||
Use server-rendered Go templates, embedded assets, and small vanilla
|
||||
JavaScript only. Preserve the existing worker API and Basic Auth boundary.
|
||||
|
||||
## Explicitly out of scope
|
||||
|
||||
- user accounts, sign-up, roles, sessions, OAuth, or multi-tenancy;
|
||||
- executing or controlling workers from the browser;
|
||||
- direct browser access to PostgreSQL or worker endpoints;
|
||||
- final-result reduction, distributed workload planning, or graph execution;
|
||||
- artifact CSV preview/visualisation. That is assigned independently in
|
||||
`docs/artifact-preview-task.md`.
|
||||
|
||||
## Security and protocol rules
|
||||
|
||||
- `UI_AUTH_TOKEN` is never sent to HTML, JavaScript, URLs, logs, or storage.
|
||||
- Use `html/template`; JavaScript must use `textContent`, never `innerHTML` for
|
||||
received data.
|
||||
- UI artifact operations must be job-scoped and coordinator-owned.
|
||||
- Do not disclose raw worker commands, local paths, bearer tokens, or database
|
||||
errors.
|
||||
- Do not change worker/coordinator API contracts silently. Document any
|
||||
intentional API change in `docs/api-contract.md` and `docs/openapi.yaml`.
|
||||
|
||||
## Acceptance criteria
|
||||
|
||||
- A new operator can start the stack, authenticate, submit a small run, start
|
||||
workers, understand live progress, and safely stop a job from the UI.
|
||||
- The UI never calls or displays worker bearer-token endpoints.
|
||||
- All partial output is visibly labelled as diagnostic until a reducer exists.
|
||||
- Disabled UI remains `404`; unauthenticated UI requests remain rejected.
|
||||
- Go tests cover changed routes and states, including auth and a sanitized
|
||||
error case.
|
||||
- `go test ./...`, `go vet ./...`, and the relevant real-PostgreSQL tests pass.
|
||||
|
||||
## Handoff
|
||||
|
||||
Use one branch and one PR. In the PR description state the user journey that
|
||||
changed, screenshots if visual layout changed, API impact (`none` if none),
|
||||
and exact test commands/results. Do not stage local datasets or `worker-data/`.
|
||||
@@ -173,6 +173,9 @@ Rules:
|
||||
- TSV file, required, streamed; show expected columns
|
||||
`chembl_id` and `canonical_smiles`.
|
||||
- `chunk_rows`: integer 1--100000, default 1000.
|
||||
- `max_rows`: optional positive integer. The coordinator creates shards only
|
||||
from the first N data rows, so a user can test a large upload without
|
||||
creating thousands of tasks. It does not truncate the stored source blob.
|
||||
- optional human-readable run name is a later schema/API addition; v1 does not
|
||||
silently store it.
|
||||
- display file name and client-side size only as convenience; server limits and
|
||||
|
||||
@@ -59,7 +59,7 @@ Content-Type: application/json
|
||||
|
||||
{
|
||||
"name": "lab-worker-01",
|
||||
"capabilities": ["similarity-search", "similarity-graph"],
|
||||
"capabilities": ["similarity-search"],
|
||||
"cpu_count": 8,
|
||||
"memory_mb": 16384
|
||||
}
|
||||
@@ -74,8 +74,8 @@ POST /tasks/claim
|
||||
Content-Type: application/json
|
||||
|
||||
{
|
||||
"worker_id": "worker-01",
|
||||
"capabilities": ["similarity-search", "similarity-graph"],
|
||||
"worker_id": "<registered-uuid>",
|
||||
"capabilities": ["similarity-search"],
|
||||
"max_concurrency": 1
|
||||
}
|
||||
```
|
||||
@@ -95,7 +95,7 @@ When a task is available, it returns `200 OK`:
|
||||
"sha256": "..."
|
||||
},
|
||||
"parameters": {
|
||||
"query_id": "CHEMBL939",
|
||||
"query_smiles": "CCO",
|
||||
"top_k": 20
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,28 @@
|
||||
"""Coordinator-independent contracts for distributed SciMesh workloads.
|
||||
|
||||
This package defines the typed plan and reduction boundary shared by future
|
||||
planners, worker adapters, and coordinator bridges. It intentionally has no
|
||||
network, database, or coordinator imports.
|
||||
"""
|
||||
|
||||
from .models import (
|
||||
ArtifactReference,
|
||||
CompletedPartial,
|
||||
DistributedPlan,
|
||||
FinalResult,
|
||||
PlannedTask,
|
||||
)
|
||||
from .registry import DistributedWorkloadRegistry, PlanningService, WorkloadDescription
|
||||
from .workload import DistributedWorkload
|
||||
|
||||
__all__ = [
|
||||
"ArtifactReference",
|
||||
"CompletedPartial",
|
||||
"DistributedPlan",
|
||||
"DistributedWorkload",
|
||||
"DistributedWorkloadRegistry",
|
||||
"FinalResult",
|
||||
"PlannedTask",
|
||||
"PlanningService",
|
||||
"WorkloadDescription",
|
||||
]
|
||||
@@ -0,0 +1,262 @@
|
||||
"""Versioned, JSON-safe value objects for distributed workload contracts."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import math
|
||||
import re
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Mapping, Sequence
|
||||
from uuid import UUID
|
||||
|
||||
|
||||
SCHEMA_VERSION = 1
|
||||
_WORKLOAD_NAME = re.compile(r"^[a-z][a-z0-9]*(?:-[a-z0-9]+)*$")
|
||||
|
||||
|
||||
def _canonical_uuid(value: object, field: str) -> str:
|
||||
if not isinstance(value, str):
|
||||
raise ValueError(f"{field} must be a UUID string")
|
||||
try:
|
||||
return str(UUID(value))
|
||||
except ValueError as error:
|
||||
raise ValueError(f"{field} must be a UUID string") from error
|
||||
|
||||
|
||||
def _sha256(value: object, field: str) -> str:
|
||||
if not isinstance(value, str) or not re.fullmatch(r"[0-9a-f]{64}", value):
|
||||
raise ValueError(f"{field} must be a lowercase SHA-256 hex digest")
|
||||
return value
|
||||
|
||||
|
||||
def _content_type(value: object, field: str) -> str:
|
||||
if not isinstance(value, str) or not value or len(value) > 128:
|
||||
raise ValueError(f"{field} must be a non-empty content type")
|
||||
if any(character.isspace() or ord(character) < 32 for character in value):
|
||||
raise ValueError(f"{field} must be a non-empty content type")
|
||||
return value
|
||||
|
||||
|
||||
def _workload_name(value: object, field: str = "workload") -> str:
|
||||
if not isinstance(value, str) or not _WORKLOAD_NAME.fullmatch(value):
|
||||
raise ValueError(f"{field} must be a canonical hyphenated workload name")
|
||||
return value
|
||||
|
||||
|
||||
def _json_value(value: object, field: str) -> Any:
|
||||
"""Deep-copy a JSON value and reject non-finite or non-string-key data."""
|
||||
if value is None or isinstance(value, (bool, int)):
|
||||
return value
|
||||
if isinstance(value, str):
|
||||
# Coordinator-owned artifacts are represented exclusively by
|
||||
# ArtifactReference. A URI or a local path in a generic JSON payload
|
||||
# would let a planner accidentally leak a bridge/worker implementation
|
||||
# detail into durable task metadata.
|
||||
forbidden_prefixes = ("file://", "worker://", "http://", "https://", "s3://", "/")
|
||||
is_windows_path = len(value) >= 3 and value[0].isalpha() and value[1:3] in (":/", ":\\")
|
||||
if value.startswith(forbidden_prefixes) or is_windows_path:
|
||||
raise ValueError(f"{field} must not contain a URI or local path")
|
||||
return value
|
||||
if isinstance(value, float):
|
||||
if not math.isfinite(value):
|
||||
raise ValueError(f"{field} must not contain NaN or infinity")
|
||||
return value
|
||||
if isinstance(value, Mapping):
|
||||
copied: dict[str, Any] = {}
|
||||
for key, child in value.items():
|
||||
if not isinstance(key, str):
|
||||
raise ValueError(f"{field} must use string object keys")
|
||||
copied[key] = _json_value(child, f"{field}.{key}")
|
||||
return copied
|
||||
if isinstance(value, (list, tuple)):
|
||||
return [_json_value(child, f"{field}[]") for child in value]
|
||||
raise ValueError(f"{field} must contain only JSON-compatible values")
|
||||
|
||||
|
||||
def _json_mapping(value: object, field: str) -> dict[str, Any]:
|
||||
if not isinstance(value, Mapping):
|
||||
raise ValueError(f"{field} must be an object")
|
||||
return _json_value(value, field)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ArtifactReference:
|
||||
"""Immutable coordinator-owned artifact identity used in a plan."""
|
||||
|
||||
artifact_id: str
|
||||
sha256: str
|
||||
content_type: str
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
object.__setattr__(self, "artifact_id", _canonical_uuid(self.artifact_id, "artifact_id"))
|
||||
object.__setattr__(self, "sha256", _sha256(self.sha256, "sha256"))
|
||||
object.__setattr__(self, "content_type", _content_type(self.content_type, "content_type"))
|
||||
|
||||
def to_dict(self) -> dict[str, str]:
|
||||
return {
|
||||
"artifact_id": self.artifact_id,
|
||||
"sha256": self.sha256,
|
||||
"content_type": self.content_type,
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, value: object) -> "ArtifactReference":
|
||||
if not isinstance(value, Mapping):
|
||||
raise ValueError("artifact reference must be an object")
|
||||
_require_exact_keys(value, {"artifact_id", "sha256", "content_type"}, "artifact reference")
|
||||
return cls(
|
||||
artifact_id=value["artifact_id"],
|
||||
sha256=value["sha256"],
|
||||
content_type=value["content_type"],
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class PlannedTask:
|
||||
"""One deterministically indexed, artifact-backed worker task."""
|
||||
|
||||
chunk_index: int
|
||||
input_artifact: ArtifactReference
|
||||
parameters: Mapping[str, object]
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if isinstance(self.chunk_index, bool) or not isinstance(self.chunk_index, int) or self.chunk_index < 0:
|
||||
raise ValueError("chunk_index must be a non-negative integer")
|
||||
if not isinstance(self.input_artifact, ArtifactReference):
|
||||
raise ValueError("input_artifact must be an ArtifactReference")
|
||||
object.__setattr__(self, "parameters", _json_mapping(self.parameters, "task parameters"))
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
return {
|
||||
"chunk_index": self.chunk_index,
|
||||
"input_artifact": self.input_artifact.to_dict(),
|
||||
"parameters": _json_value(self.parameters, "task parameters"),
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, value: object) -> "PlannedTask":
|
||||
if not isinstance(value, Mapping):
|
||||
raise ValueError("planned task must be an object")
|
||||
_require_exact_keys(value, {"chunk_index", "input_artifact", "parameters"}, "planned task")
|
||||
return cls(
|
||||
chunk_index=value["chunk_index"],
|
||||
input_artifact=ArtifactReference.from_dict(value["input_artifact"]),
|
||||
parameters=value["parameters"],
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class DistributedPlan:
|
||||
"""The complete schema-versioned output of a distributed planner."""
|
||||
|
||||
workload: str
|
||||
resolved_parameters: Mapping[str, object]
|
||||
tasks: Sequence[PlannedTask]
|
||||
schema_version: int = SCHEMA_VERSION
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if self.schema_version != SCHEMA_VERSION:
|
||||
raise ValueError(f"schema_version must be {SCHEMA_VERSION}")
|
||||
object.__setattr__(self, "workload", _workload_name(self.workload))
|
||||
object.__setattr__(self, "resolved_parameters", _json_mapping(self.resolved_parameters, "resolved_parameters"))
|
||||
task_list = tuple(self.tasks)
|
||||
if not task_list:
|
||||
raise ValueError("plan must contain at least one task")
|
||||
if any(not isinstance(task, PlannedTask) for task in task_list):
|
||||
raise ValueError("tasks must contain PlannedTask values")
|
||||
indexes = [task.chunk_index for task in task_list]
|
||||
if indexes != sorted(indexes) or len(set(indexes)) != len(indexes):
|
||||
raise ValueError("tasks must have unique, ascending chunk_index values")
|
||||
object.__setattr__(self, "tasks", task_list)
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
return {
|
||||
"schema_version": self.schema_version,
|
||||
"workload": self.workload,
|
||||
"resolved_parameters": _json_value(self.resolved_parameters, "resolved_parameters"),
|
||||
"tasks": [task.to_dict() for task in self.tasks],
|
||||
}
|
||||
|
||||
def to_json(self) -> str:
|
||||
"""Return stable JSON suitable for hashing, tests, and durable payloads."""
|
||||
return json.dumps(self.to_dict(), sort_keys=True, separators=(",", ":"), allow_nan=False)
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, value: object) -> "DistributedPlan":
|
||||
if not isinstance(value, Mapping):
|
||||
raise ValueError("distributed plan must be an object")
|
||||
_require_exact_keys(
|
||||
value,
|
||||
{"schema_version", "workload", "resolved_parameters", "tasks"},
|
||||
"distributed plan",
|
||||
)
|
||||
raw_tasks = value["tasks"]
|
||||
if not isinstance(raw_tasks, list):
|
||||
raise ValueError("tasks must be an array")
|
||||
return cls(
|
||||
schema_version=value["schema_version"],
|
||||
workload=value["workload"],
|
||||
resolved_parameters=value["resolved_parameters"],
|
||||
tasks=tuple(PlannedTask.from_dict(task) for task in raw_tasks),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_json(cls, value: str) -> "DistributedPlan":
|
||||
try:
|
||||
decoded = json.loads(value)
|
||||
except (TypeError, json.JSONDecodeError) as error:
|
||||
raise ValueError("distributed plan must be valid JSON") from error
|
||||
return cls.from_dict(decoded)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class CompletedPartial:
|
||||
"""Coordinator-owned partial output supplied to a reducer."""
|
||||
|
||||
chunk_index: int
|
||||
artifact: ArtifactReference
|
||||
metrics: Mapping[str, int | float]
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if isinstance(self.chunk_index, bool) or not isinstance(self.chunk_index, int) or self.chunk_index < 0:
|
||||
raise ValueError("chunk_index must be a non-negative integer")
|
||||
if not isinstance(self.artifact, ArtifactReference):
|
||||
raise ValueError("artifact must be an ArtifactReference")
|
||||
if not isinstance(self.metrics, Mapping):
|
||||
raise ValueError("metrics must be an object")
|
||||
metrics: dict[str, int | float] = {}
|
||||
for name, value in self.metrics.items():
|
||||
if not isinstance(name, str) or not name:
|
||||
raise ValueError("metric names must be non-empty strings")
|
||||
if isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(value):
|
||||
raise ValueError("metric values must be finite JSON numbers")
|
||||
metrics[name] = value
|
||||
object.__setattr__(self, "metrics", metrics)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class FinalResult:
|
||||
"""A reducer's durable output, ready for coordinator persistence."""
|
||||
|
||||
artifact: ArtifactReference
|
||||
metrics: Mapping[str, int | float]
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if not isinstance(self.artifact, ArtifactReference):
|
||||
raise ValueError("artifact must be an ArtifactReference")
|
||||
# Reuse the CompletedPartial metric validation without inventing a fake
|
||||
# artifact lifecycle or widening the result contract.
|
||||
object.__setattr__(self, "metrics", CompletedPartial(0, self.artifact, self.metrics).metrics)
|
||||
|
||||
|
||||
def _require_exact_keys(value: Mapping[str, object], expected: set[str], label: str) -> None:
|
||||
actual = set(value)
|
||||
if actual != expected:
|
||||
missing = sorted(expected - actual)
|
||||
unknown = sorted(actual - expected)
|
||||
details: list[str] = []
|
||||
if missing:
|
||||
details.append(f"missing {', '.join(missing)}")
|
||||
if unknown:
|
||||
details.append(f"unknown {', '.join(unknown)}")
|
||||
raise ValueError(f"{label} has {'; '.join(details)} fields")
|
||||
@@ -0,0 +1,96 @@
|
||||
"""Registry and orchestration helpers for distributed workload contracts."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Mapping, Sequence
|
||||
|
||||
from .models import CompletedPartial, DistributedPlan, FinalResult, _workload_name
|
||||
from .workload import DistributedWorkload
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class WorkloadDescription:
|
||||
"""Safe metadata that a future coordinator or UI may display."""
|
||||
|
||||
name: str
|
||||
description: str
|
||||
|
||||
|
||||
class DistributedWorkloadRegistry:
|
||||
"""Collect distributed workloads without coupling them to the CLI registry."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._workloads: dict[str, DistributedWorkload] = {}
|
||||
|
||||
def register(self, workload: DistributedWorkload) -> None:
|
||||
name = _workload_name(workload.name)
|
||||
if name in self._workloads:
|
||||
raise ValueError(f"distributed workload already registered: {name}")
|
||||
if not isinstance(workload.description, str) or not workload.description.strip():
|
||||
raise ValueError("distributed workload description must be non-empty")
|
||||
self._workloads[name] = workload
|
||||
|
||||
def require(self, name: str) -> DistributedWorkload:
|
||||
try:
|
||||
return self._workloads[_workload_name(name)]
|
||||
except KeyError as error:
|
||||
raise ValueError(f"unknown distributed workload: {name}") from error
|
||||
|
||||
def descriptions(self) -> tuple[WorkloadDescription, ...]:
|
||||
return tuple(
|
||||
WorkloadDescription(name, workload.description)
|
||||
for name, workload in sorted(self._workloads.items())
|
||||
)
|
||||
|
||||
|
||||
class PlanningService:
|
||||
"""Small bridge-safe orchestration around a distributed workload registry.
|
||||
|
||||
It writes neither jobs nor artifacts. A Go coordinator bridge can therefore
|
||||
validate and produce a plan before opening its own all-or-nothing persistence
|
||||
transaction; CTX-08/09 will implement that concrete bridge and reducers.
|
||||
"""
|
||||
|
||||
def __init__(self, registry: DistributedWorkloadRegistry) -> None:
|
||||
self._registry = registry
|
||||
|
||||
def plan(
|
||||
self,
|
||||
workload_name: str,
|
||||
input_path: Path,
|
||||
input_artifact_id: str,
|
||||
parameters: Mapping[str, object],
|
||||
shard_rows: int,
|
||||
workspace: Path,
|
||||
) -> DistributedPlan:
|
||||
if isinstance(shard_rows, bool) or not isinstance(shard_rows, int) or shard_rows < 1:
|
||||
raise ValueError("shard_rows must be a positive integer")
|
||||
workload = self._registry.require(workload_name)
|
||||
workload.validate_job(parameters)
|
||||
plan = workload.plan(input_path, input_artifact_id, parameters, shard_rows, workspace)
|
||||
if not isinstance(plan, DistributedPlan):
|
||||
raise ValueError("distributed planner must return a DistributedPlan")
|
||||
if plan.workload != workload.name:
|
||||
raise ValueError("distributed planner returned a plan for another workload")
|
||||
# Round-trip through the strict wire schema now, before a future bridge
|
||||
# persists anything. This catches non-JSON values and undeclared fields.
|
||||
return DistributedPlan.from_json(plan.to_json())
|
||||
|
||||
def reduce(
|
||||
self,
|
||||
workload_name: str,
|
||||
partial_results: Sequence[CompletedPartial],
|
||||
parameters: Mapping[str, object],
|
||||
workspace: Path,
|
||||
) -> FinalResult:
|
||||
workload = self._registry.require(workload_name)
|
||||
indexes = [partial.chunk_index for partial in partial_results]
|
||||
if len(indexes) != len(set(indexes)):
|
||||
raise ValueError("partial results must have unique chunk_index values")
|
||||
ordered = tuple(sorted(partial_results, key=lambda partial: partial.chunk_index))
|
||||
result = workload.reduce(ordered, parameters, workspace)
|
||||
if not isinstance(result, FinalResult):
|
||||
raise ValueError("distributed reducer must return a FinalResult")
|
||||
return result
|
||||
@@ -0,0 +1,40 @@
|
||||
"""Protocol implemented by coordinator-independent distributed workloads."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
from typing import Mapping, Protocol, Sequence
|
||||
|
||||
from .models import CompletedPartial, DistributedPlan, FinalResult
|
||||
|
||||
|
||||
class DistributedWorkload(Protocol):
|
||||
"""Validate, plan, and reduce one explicit scientific workload.
|
||||
|
||||
``input_path`` and ``workspace`` are bridge-provided temporary local paths.
|
||||
They must never be included in returned plans or persisted task payloads.
|
||||
"""
|
||||
|
||||
name: str
|
||||
description: str
|
||||
|
||||
def validate_job(self, parameters: Mapping[str, object]) -> None:
|
||||
"""Reject invalid public parameters before the bridge writes metadata."""
|
||||
|
||||
def plan(
|
||||
self,
|
||||
input_path: Path,
|
||||
input_artifact_id: str,
|
||||
parameters: Mapping[str, object],
|
||||
shard_rows: int,
|
||||
workspace: Path,
|
||||
) -> DistributedPlan:
|
||||
"""Build a JSON-safe plan containing only coordinator artifact references."""
|
||||
|
||||
def reduce(
|
||||
self,
|
||||
partial_results: Sequence[CompletedPartial],
|
||||
parameters: Mapping[str, object],
|
||||
workspace: Path,
|
||||
) -> FinalResult:
|
||||
"""Reduce coordinator-owned partial artifacts in ascending chunk order."""
|
||||
+31
-6
@@ -13,15 +13,17 @@ from .daemon import WorkerDaemon
|
||||
from .runners import SciMeshRunner
|
||||
|
||||
|
||||
def main(argv: list[str] | None = None) -> int:
|
||||
def build_parser() -> argparse.ArgumentParser:
|
||||
"""Build the worker CLI parser for command-line use and focused tests."""
|
||||
parser = argparse.ArgumentParser(
|
||||
prog="scimesh-worker",
|
||||
epilog=(
|
||||
"Environment: SCIMESH_COORDINATOR_URL, SCIMESH_WORK_DIR, "
|
||||
"SCIMESH_WORKER_NAME, SCIMESH_CPU_COUNT, SCIMESH_MEMORY_MB, "
|
||||
"SCIMESH_POLL_INTERVAL, SCIMESH_REQUEST_TIMEOUT, "
|
||||
"SCIMESH_HEARTBEAT_INTERVAL, SCIMESH_CLEANUP_AFTER_SECONDS, and "
|
||||
"SCIMESH_BEARER_TOKEN. SCIMESH_WORKER_ID is a legacy/test override."
|
||||
"SCIMESH_HEARTBEAT_INTERVAL, SCIMESH_CLEANUP_AFTER_SECONDS, "
|
||||
"SCIMESH_MAX_TASKS, and SCIMESH_BEARER_TOKEN. "
|
||||
"SCIMESH_WORKER_ID is a legacy/test override."
|
||||
),
|
||||
)
|
||||
parser.add_argument("--coordinator-url")
|
||||
@@ -34,8 +36,31 @@ def main(argv: list[str] | None = None) -> int:
|
||||
parser.add_argument("--request-timeout", type=float)
|
||||
parser.add_argument("--heartbeat-interval", type=float)
|
||||
parser.add_argument("--cleanup-after-seconds", type=float)
|
||||
lifecycle = parser.add_mutually_exclusive_group()
|
||||
lifecycle.add_argument(
|
||||
"--once",
|
||||
action="store_true",
|
||||
help="Claim at most one task, then exit; exit immediately when the queue is empty",
|
||||
)
|
||||
lifecycle.add_argument(
|
||||
"--max-tasks",
|
||||
type=int,
|
||||
help="Process this many claimed tasks, then exit",
|
||||
)
|
||||
return parser
|
||||
|
||||
|
||||
def main(argv: list[str] | None = None) -> int:
|
||||
parser = build_parser()
|
||||
args = parser.parse_args(argv)
|
||||
overrides = {key: value for key, value in vars(args).items() if value is not None}
|
||||
overrides = {
|
||||
key: value
|
||||
for key, value in vars(args).items()
|
||||
if value is not None and key != "once"
|
||||
}
|
||||
if args.once:
|
||||
overrides["max_tasks"] = 1
|
||||
overrides["exit_when_idle"] = True
|
||||
if "work_dir" in overrides:
|
||||
overrides["work_dir"] = Path(overrides["work_dir"])
|
||||
try:
|
||||
@@ -44,13 +69,13 @@ def main(argv: list[str] | None = None) -> int:
|
||||
parser.error(str(error))
|
||||
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
|
||||
client = HttpCoordinatorClient(config.coordinator_url, config.request_timeout, config.bearer_token)
|
||||
WorkerDaemon(
|
||||
completed_without_interruption = WorkerDaemon(
|
||||
config,
|
||||
client,
|
||||
HttpArtifactClient(config.coordinator_url, config.request_timeout, config.bearer_token),
|
||||
SciMeshRunner(),
|
||||
).run_forever()
|
||||
return 0
|
||||
return 0 if completed_without_interruption else 130
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -36,13 +36,16 @@ class WorkerConfig:
|
||||
heartbeat_interval: float = 15.0
|
||||
bearer_token: str | None = None
|
||||
cleanup_after_seconds: float | None = None
|
||||
max_tasks: int | None = None
|
||||
exit_when_idle: bool = False
|
||||
# Distributed similarity-graph requires triangular block-pair planning and
|
||||
# is deliberately not advertised until CTX-10. A normal worker must never
|
||||
# make a multi-shard graph job appear scientifically complete.
|
||||
# The local CLI uses hyphens; the first coordinator contract used
|
||||
# underscores. Advertise both stable spellings while jobs are migrated.
|
||||
# underscores, so retain the search alias during migration.
|
||||
capabilities: tuple[str, ...] = (
|
||||
"similarity-search",
|
||||
"similarity-graph",
|
||||
"similarity_search",
|
||||
"similarity_graph",
|
||||
)
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
@@ -66,8 +69,21 @@ class WorkerConfig:
|
||||
_positive_number(self.heartbeat_interval, "heartbeat_interval")
|
||||
if self.cleanup_after_seconds is not None:
|
||||
_positive_number(self.cleanup_after_seconds, "cleanup_after_seconds", allow_zero=True)
|
||||
if self.max_tasks is not None:
|
||||
if (
|
||||
isinstance(self.max_tasks, bool)
|
||||
or not isinstance(self.max_tasks, int)
|
||||
or self.max_tasks < 1
|
||||
):
|
||||
raise ValueError("max_tasks must be positive when set")
|
||||
if not isinstance(self.exit_when_idle, bool):
|
||||
raise ValueError("exit_when_idle must be a boolean")
|
||||
if not self.capabilities:
|
||||
raise ValueError("capabilities cannot be empty")
|
||||
# Runner subprocesses use a task directory as their cwd. Keep the
|
||||
# configured root absolute so input/output paths remain valid there
|
||||
# even when the CLI received a convenient relative --work-dir value.
|
||||
object.__setattr__(self, "work_dir", self.work_dir.expanduser().resolve())
|
||||
|
||||
@classmethod
|
||||
def from_environment(
|
||||
@@ -86,6 +102,7 @@ class WorkerConfig:
|
||||
cleanup = value("cleanup_after_seconds", "SCIMESH_CLEANUP_AFTER_SECONDS")
|
||||
cpu_count = value("cpu_count", "SCIMESH_CPU_COUNT", os.cpu_count() or 1)
|
||||
memory_mb = value("memory_mb", "SCIMESH_MEMORY_MB")
|
||||
max_tasks = value("max_tasks", "SCIMESH_MAX_TASKS")
|
||||
return cls(
|
||||
coordinator_url=url.rstrip("/"),
|
||||
worker_id=value("worker_id", "SCIMESH_WORKER_ID"),
|
||||
@@ -98,4 +115,6 @@ class WorkerConfig:
|
||||
heartbeat_interval=float(value("heartbeat_interval", "SCIMESH_HEARTBEAT_INTERVAL", "15")),
|
||||
bearer_token=value("bearer_token", "SCIMESH_BEARER_TOKEN"),
|
||||
cleanup_after_seconds=float(cleanup) if cleanup else None,
|
||||
max_tasks=int(max_tasks) if max_tasks is not None else None,
|
||||
exit_when_idle=bool(values.get("exit_when_idle", False)),
|
||||
)
|
||||
|
||||
+94
-23
@@ -4,8 +4,10 @@ from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from dataclasses import replace
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
import random
|
||||
import re
|
||||
import shutil
|
||||
import threading
|
||||
import time
|
||||
@@ -72,6 +74,14 @@ class LeaseHeartbeat:
|
||||
return seconds
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class RunOnceOutcome:
|
||||
"""Whether a claim was made and whether that claimed task completed."""
|
||||
|
||||
claimed: bool
|
||||
completed: bool
|
||||
|
||||
|
||||
class WorkerDaemon:
|
||||
def __init__(self, config: WorkerConfig, coordinator: CoordinatorClient, artifacts: ArtifactClient, runner: Runner) -> None:
|
||||
self.config, self.coordinator, self.artifacts, self.runner = config, coordinator, artifacts, runner
|
||||
@@ -79,32 +89,68 @@ class WorkerDaemon:
|
||||
self._registered = False
|
||||
self.log = logging.getLogger("scimesh.worker")
|
||||
|
||||
def run_forever(self) -> None:
|
||||
def run_forever(self) -> bool:
|
||||
"""Run until stopped; return false only when interrupted by the operator."""
|
||||
failures = 0
|
||||
while True:
|
||||
try:
|
||||
if not self._registered:
|
||||
self._register_worker()
|
||||
self._cleanup_expired_directories()
|
||||
claimed = self.run_once()
|
||||
failures = 0
|
||||
if not claimed:
|
||||
self._sleep(self.config.poll_interval)
|
||||
except CoordinatorTransientError as error:
|
||||
failures += 1
|
||||
self._log("failed", error_type=type(error).__name__)
|
||||
self._sleep(min(self.config.poll_interval * 2 ** min(failures, 6), 60.0))
|
||||
completed_tasks = 0
|
||||
self._log(
|
||||
"started",
|
||||
max_tasks=self.config.max_tasks,
|
||||
exit_when_idle=self.config.exit_when_idle,
|
||||
)
|
||||
try:
|
||||
while True:
|
||||
try:
|
||||
if not self._registered:
|
||||
self._register_worker()
|
||||
self._cleanup_expired_directories()
|
||||
outcome = self.run_once()
|
||||
failures = 0
|
||||
if outcome.claimed:
|
||||
if outcome.completed:
|
||||
completed_tasks += 1
|
||||
if self.config.exit_when_idle:
|
||||
self._log(
|
||||
"stopped",
|
||||
reason="one_claim_processed",
|
||||
completed_tasks=completed_tasks,
|
||||
)
|
||||
return True
|
||||
if (
|
||||
outcome.completed
|
||||
and self.config.max_tasks is not None
|
||||
and completed_tasks >= self.config.max_tasks
|
||||
):
|
||||
self._log(
|
||||
"stopped",
|
||||
reason="max_tasks_reached",
|
||||
completed_tasks=completed_tasks,
|
||||
)
|
||||
return True
|
||||
elif self.config.exit_when_idle:
|
||||
self._log("stopped", reason="queue_empty", completed_tasks=completed_tasks)
|
||||
return True
|
||||
else:
|
||||
self._sleep(self.config.poll_interval)
|
||||
except CoordinatorTransientError as error:
|
||||
failures += 1
|
||||
self._log("failed", error_type=type(error).__name__)
|
||||
self._sleep(min(self.config.poll_interval * 2 ** min(failures, 6), 60.0))
|
||||
except KeyboardInterrupt:
|
||||
self._log("stopped", reason="interrupted", completed_tasks=completed_tasks)
|
||||
return False
|
||||
|
||||
def run_once(self) -> bool:
|
||||
def run_once(self) -> RunOnceOutcome:
|
||||
worker_id = self._worker_id()
|
||||
self._log("claiming")
|
||||
self._log("claiming", log_level=logging.DEBUG)
|
||||
task = self.coordinator.claim(worker_id, self.config.capabilities)
|
||||
if task is None:
|
||||
self._log("idle")
|
||||
return False
|
||||
self._log("idle", log_level=logging.DEBUG)
|
||||
return RunOnceOutcome(claimed=False, completed=False)
|
||||
started = time.monotonic()
|
||||
task_dir = self.config.work_dir / task.task_id / str(task.attempt)
|
||||
heartbeat = LeaseHeartbeat(task, self.coordinator, self.config)
|
||||
completed = False
|
||||
try:
|
||||
task_dir.mkdir(parents=True, exist_ok=False)
|
||||
heartbeat.start()
|
||||
@@ -135,7 +181,15 @@ class WorkerDaemon:
|
||||
},
|
||||
},
|
||||
)
|
||||
self._log("idle", task, elapsed_seconds=round(time.monotonic() - started, 3))
|
||||
completed = True
|
||||
self._log("completed", task, elapsed_seconds=round(time.monotonic() - started, 3))
|
||||
except KeyboardInterrupt:
|
||||
self._log("interrupted", task)
|
||||
try:
|
||||
self._report_failure(task, InterruptedError("worker interrupted by operator"))
|
||||
except CoordinatorTransientError:
|
||||
self._log("failed", task, error_type="FailureReportError")
|
||||
raise
|
||||
except CoordinatorConflictError as error:
|
||||
self._log("lease_lost", task, error_type=type(error).__name__)
|
||||
except Exception as error:
|
||||
@@ -143,10 +197,10 @@ class WorkerDaemon:
|
||||
self._report_failure(task, error)
|
||||
finally:
|
||||
heartbeat.stop()
|
||||
return True
|
||||
return RunOnceOutcome(claimed=True, completed=completed)
|
||||
|
||||
def _report_failure(self, task: ClaimedTask, error: Exception) -> None:
|
||||
message = str(error).replace(str(self.config.work_dir), "<worker-dir>")[:300]
|
||||
message = self._sanitize_error_message(error)
|
||||
try:
|
||||
self.coordinator.fail(task, {"worker_id": self._worker_id(), "attempt": task.attempt, "error_code": type(error).__name__, "error_message": message})
|
||||
except CoordinatorTransientError:
|
||||
@@ -154,6 +208,16 @@ class WorkerDaemon:
|
||||
except Exception:
|
||||
self._log("failed", task, error_type="FailureReportError")
|
||||
|
||||
def _sanitize_error_message(self, error: Exception) -> str:
|
||||
"""Keep coordinator-visible failures useful without exposing local paths."""
|
||||
message = str(error).replace(str(self.config.work_dir), "<worker-dir>")
|
||||
# CalledProcessError includes the complete argv, including sys.executable
|
||||
# outside work_dir. Replace POSIX and Windows absolute paths before the
|
||||
# message reaches the coordinator database or operator UI.
|
||||
message = re.sub(r"(?<![\w:])[A-Za-z]:\\[^\s'\"\],)]+", "<path>", message)
|
||||
message = re.sub(r"(?<![\w:])/(?:[^\s'\"\],)]+)", "<path>", message)
|
||||
return message[:300]
|
||||
|
||||
def _register_worker(self) -> None:
|
||||
registered = self.coordinator.register(
|
||||
self.config.worker_name,
|
||||
@@ -180,9 +244,16 @@ class WorkerDaemon:
|
||||
"""Keep completion payload exact: coordinator owns all artifact metadata."""
|
||||
return {"artifact_id": uploaded.artifact_id}
|
||||
|
||||
def _log(self, state: str, task: ClaimedTask | None = None, **extra: object) -> None:
|
||||
def _log(
|
||||
self,
|
||||
state: str,
|
||||
task: ClaimedTask | None = None,
|
||||
*,
|
||||
log_level: int = logging.INFO,
|
||||
**extra: object,
|
||||
) -> None:
|
||||
fields = {"worker_id": self.config.worker_id, "task_id": task.task_id if task else None, "attempt": task.attempt if task else None, "state": state, **extra}
|
||||
self.log.info("worker_event %s", fields)
|
||||
self.log.log(log_level, "worker_event %s", fields)
|
||||
|
||||
def _cleanup_expired_directories(self) -> None:
|
||||
"""Remove only old task attempt directories when retention was configured."""
|
||||
|
||||
@@ -18,6 +18,9 @@ class SciMeshRunner:
|
||||
"""Allowlisted adapter from coordinator workloads to the local SciMesh CLI."""
|
||||
|
||||
def run(self, task: ClaimedTask, task_dir: Path) -> RunResult:
|
||||
# The subprocess changes cwd to task_dir. Absolute paths keep a caller
|
||||
# supplied relative work directory from being resolved twice.
|
||||
task_dir = task_dir.resolve()
|
||||
input_path = task_dir / "input"
|
||||
output_path = task_dir / "result.csv"
|
||||
# The coordinator contract historically used underscores while the
|
||||
|
||||
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"
|
||||
@@ -0,0 +1,186 @@
|
||||
"""Contract tests for the coordinator-independent distributed workload boundary."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Mapping, Sequence
|
||||
from uuid import NAMESPACE_URL, uuid5
|
||||
|
||||
import pytest
|
||||
|
||||
from scimesh.distributed import (
|
||||
ArtifactReference,
|
||||
CompletedPartial,
|
||||
DistributedPlan,
|
||||
DistributedWorkloadRegistry,
|
||||
FinalResult,
|
||||
PlannedTask,
|
||||
PlanningService,
|
||||
)
|
||||
|
||||
|
||||
def artifact(seed: str, content_type: str = "text/tab-separated-values") -> ArtifactReference:
|
||||
return ArtifactReference(
|
||||
artifact_id=str(uuid5(NAMESPACE_URL, seed)),
|
||||
sha256=(seed.encode("utf-8").hex() * 64)[:64],
|
||||
content_type=content_type,
|
||||
)
|
||||
|
||||
|
||||
class DummyWorkload:
|
||||
"""A deterministic fake workload used to test the generic CTX-07 bridge."""
|
||||
|
||||
name = "dummy-workload"
|
||||
description = "A deterministic test workload."
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.plan_calls = 0
|
||||
self.received_partials: tuple[CompletedPartial, ...] = ()
|
||||
|
||||
def validate_job(self, parameters: Mapping[str, object]) -> None:
|
||||
if parameters != {"mode": "valid"}:
|
||||
raise ValueError("mode must be valid")
|
||||
|
||||
def plan(
|
||||
self,
|
||||
input_path: Path,
|
||||
input_artifact_id: str,
|
||||
parameters: Mapping[str, object],
|
||||
shard_rows: int,
|
||||
workspace: Path,
|
||||
) -> DistributedPlan:
|
||||
self.plan_calls += 1
|
||||
assert input_path.name == "input.tsv"
|
||||
assert workspace.name == "workspace"
|
||||
return DistributedPlan(
|
||||
workload=self.name,
|
||||
resolved_parameters={"mode": parameters["mode"], "source": input_artifact_id},
|
||||
tasks=(
|
||||
PlannedTask(0, artifact(f"{input_artifact_id}:0"), {"mode": "valid"}),
|
||||
PlannedTask(1, artifact(f"{input_artifact_id}:1"), {"mode": "valid"}),
|
||||
),
|
||||
)
|
||||
|
||||
def reduce(
|
||||
self,
|
||||
partial_results: Sequence[CompletedPartial],
|
||||
parameters: Mapping[str, object],
|
||||
workspace: Path,
|
||||
) -> FinalResult:
|
||||
self.received_partials = tuple(partial_results)
|
||||
return FinalResult(artifact("final", "text/csv"), {"partial_count": len(partial_results)})
|
||||
|
||||
|
||||
def service() -> tuple[PlanningService, DummyWorkload]:
|
||||
workload = DummyWorkload()
|
||||
registry = DistributedWorkloadRegistry()
|
||||
registry.register(workload)
|
||||
return PlanningService(registry), workload
|
||||
|
||||
|
||||
def test_unknown_workload_is_rejected_before_a_plan_is_written(tmp_path: Path) -> None:
|
||||
planner, workload = service()
|
||||
|
||||
with pytest.raises(ValueError, match="unknown distributed workload"):
|
||||
planner.plan(
|
||||
"unknown-workload", tmp_path / "input.tsv", artifact("input").artifact_id,
|
||||
{"mode": "valid"}, 10, tmp_path / "workspace",
|
||||
)
|
||||
|
||||
assert workload.plan_calls == 0
|
||||
|
||||
|
||||
def test_invalid_job_is_rejected_before_the_planner_runs(tmp_path: Path) -> None:
|
||||
planner, workload = service()
|
||||
|
||||
with pytest.raises(ValueError, match="mode must be valid"):
|
||||
planner.plan(
|
||||
"dummy-workload", tmp_path / "input.tsv", artifact("input").artifact_id,
|
||||
{"mode": "invalid"}, 10, tmp_path / "workspace",
|
||||
)
|
||||
|
||||
assert workload.plan_calls == 0
|
||||
|
||||
|
||||
def test_two_shard_plan_is_deterministic_and_json_serializable(tmp_path: Path) -> None:
|
||||
planner, _ = service()
|
||||
input_artifact_id = artifact("input").artifact_id
|
||||
first = planner.plan(
|
||||
"dummy-workload", tmp_path / "input.tsv", input_artifact_id,
|
||||
{"mode": "valid"}, 10, tmp_path / "workspace",
|
||||
)
|
||||
second = planner.plan(
|
||||
"dummy-workload", tmp_path / "input.tsv", input_artifact_id,
|
||||
{"mode": "valid"}, 10, tmp_path / "workspace",
|
||||
)
|
||||
|
||||
assert first.to_json() == second.to_json()
|
||||
payload = json.loads(first.to_json())
|
||||
assert [task["chunk_index"] for task in payload["tasks"]] == [0, 1]
|
||||
assert all(set(task) == {"chunk_index", "input_artifact", "parameters"} for task in payload["tasks"])
|
||||
assert DistributedPlan.from_json(first.to_json()) == first
|
||||
|
||||
|
||||
def test_plan_rejects_unsafe_or_non_deterministic_task_payloads() -> None:
|
||||
with pytest.raises(ValueError, match="unique, ascending"):
|
||||
DistributedPlan(
|
||||
workload="dummy-workload",
|
||||
resolved_parameters={},
|
||||
tasks=(
|
||||
PlannedTask(1, artifact("one"), {}),
|
||||
PlannedTask(0, artifact("zero"), {}),
|
||||
),
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="JSON-compatible"):
|
||||
PlannedTask(0, artifact("bad"), {"path": Path("not-serializable")})
|
||||
|
||||
with pytest.raises(ValueError, match="URI or local path"):
|
||||
PlannedTask(0, artifact("uri"), {"input": "file:///tmp/input.tsv"})
|
||||
|
||||
with pytest.raises(ValueError, match="canonical hyphenated"):
|
||||
DistributedPlan("dummy_workload", {}, (PlannedTask(0, artifact("one"), {}),))
|
||||
|
||||
|
||||
def test_reducer_receives_completed_partials_in_chunk_order(tmp_path: Path) -> None:
|
||||
planner, workload = service()
|
||||
result = planner.reduce(
|
||||
"dummy-workload",
|
||||
(
|
||||
CompletedPartial(3, artifact("three", "text/csv"), {"scanned_rows": 10}),
|
||||
CompletedPartial(1, artifact("one", "text/csv"), {"scanned_rows": 10}),
|
||||
),
|
||||
{"mode": "valid"},
|
||||
tmp_path / "workspace",
|
||||
)
|
||||
|
||||
assert [partial.chunk_index for partial in workload.received_partials] == [1, 3]
|
||||
assert result.metrics == {"partial_count": 2}
|
||||
|
||||
|
||||
def test_reducer_rejects_duplicate_chunk_indexes_before_invocation(tmp_path: Path) -> None:
|
||||
planner, workload = service()
|
||||
duplicate = CompletedPartial(0, artifact("partial", "text/csv"), {"scanned_rows": 1})
|
||||
|
||||
with pytest.raises(ValueError, match="unique chunk_index"):
|
||||
planner.reduce("dummy-workload", (duplicate, duplicate), {"mode": "valid"}, tmp_path)
|
||||
|
||||
assert workload.received_partials == ()
|
||||
|
||||
|
||||
def test_artifact_references_never_accept_paths_or_uris() -> None:
|
||||
with pytest.raises(ValueError, match="UUID"):
|
||||
ArtifactReference("file:///tmp/input.tsv", "a" * 64, "text/csv")
|
||||
with pytest.raises(ValueError, match="lowercase SHA-256"):
|
||||
ArtifactReference(str(uuid5(NAMESPACE_URL, "input")), "A" * 64, "text/csv")
|
||||
|
||||
|
||||
def test_registry_descriptions_are_stable_and_duplicate_names_are_rejected() -> None:
|
||||
registry = DistributedWorkloadRegistry()
|
||||
first, second = DummyWorkload(), DummyWorkload()
|
||||
registry.register(first)
|
||||
|
||||
assert registry.descriptions()[0].name == "dummy-workload"
|
||||
with pytest.raises(ValueError, match="already registered"):
|
||||
registry.register(second)
|
||||
+165
-5
@@ -1,7 +1,9 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import logging
|
||||
from pathlib import Path
|
||||
import subprocess
|
||||
import time
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from urllib.request import Request
|
||||
@@ -9,8 +11,10 @@ from urllib.request import Request
|
||||
import pytest
|
||||
|
||||
from scimesh.worker.config import WorkerConfig
|
||||
from scimesh.worker import cli as worker_cli
|
||||
from scimesh.worker.cli import build_parser
|
||||
from scimesh.worker.coordinator import CoordinatorTransientError
|
||||
from scimesh.worker.daemon import LeaseHeartbeat, WorkerDaemon
|
||||
from scimesh.worker.daemon import LeaseHeartbeat, RunOnceOutcome, WorkerDaemon
|
||||
from scimesh.worker.models import (
|
||||
ClaimedTask,
|
||||
InputArtifact,
|
||||
@@ -92,7 +96,7 @@ def daemon(tmp_path: Path, task: ClaimedTask | None, content: bytes):
|
||||
def test_claims_runs_uploads_and_submits_csv(tmp_path: Path) -> None:
|
||||
content = b"input fixture"
|
||||
worker, coordinator, artifacts, runner, _ = daemon(tmp_path, make_task(content), content)
|
||||
assert worker.run_once() is True
|
||||
assert worker.run_once() == RunOnceOutcome(claimed=True, completed=True)
|
||||
assert runner.calls == 1
|
||||
assert len(artifacts.uploaded) == 1
|
||||
assert coordinator.heartbeats == [("task-1", 1, "worker-1")]
|
||||
@@ -104,24 +108,153 @@ def test_claims_runs_uploads_and_submits_csv(tmp_path: Path) -> None:
|
||||
|
||||
def test_no_task_does_not_create_directory(tmp_path: Path) -> None:
|
||||
worker, _, _, runner, config = daemon(tmp_path, None, b"")
|
||||
assert worker.run_once() is False
|
||||
assert worker.run_once() == RunOnceOutcome(claimed=False, completed=False)
|
||||
assert runner.calls == 0
|
||||
assert not config.work_dir.exists()
|
||||
|
||||
|
||||
def test_once_worker_exits_after_an_empty_claim(tmp_path: Path, caplog: pytest.LogCaptureFixture) -> None:
|
||||
caplog.set_level(logging.INFO, logger="scimesh.worker")
|
||||
worker, _, _, runner, _ = daemon(tmp_path, None, b"")
|
||||
worker.config = WorkerConfig(**{**worker.config.__dict__, "exit_when_idle": True, "max_tasks": 1})
|
||||
assert worker.run_forever() is True
|
||||
assert runner.calls == 0
|
||||
assert "queue_empty" in caplog.text
|
||||
|
||||
|
||||
def test_worker_stops_after_the_configured_number_of_claims(tmp_path: Path, caplog: pytest.LogCaptureFixture) -> None:
|
||||
caplog.set_level(logging.INFO, logger="scimesh.worker")
|
||||
content = b"input fixture"
|
||||
worker, _, _, runner, _ = daemon(tmp_path, make_task(content), content)
|
||||
worker.config = WorkerConfig(**{**worker.config.__dict__, "max_tasks": 1})
|
||||
assert worker.run_forever() is True
|
||||
assert runner.calls == 1
|
||||
assert "max_tasks_reached" in caplog.text
|
||||
|
||||
|
||||
def test_keyboard_interrupt_stops_worker_without_propagating(tmp_path: Path, caplog: pytest.LogCaptureFixture) -> None:
|
||||
caplog.set_level(logging.INFO, logger="scimesh.worker")
|
||||
class InterruptingCoordinator(FakeCoordinator):
|
||||
def claim(self, worker_id: str, capabilities: tuple[str, ...]) -> ClaimedTask | None:
|
||||
raise KeyboardInterrupt
|
||||
|
||||
worker, _, _, _, _ = daemon(tmp_path, None, b"")
|
||||
worker.coordinator = InterruptingCoordinator(None)
|
||||
assert worker.run_forever() is False
|
||||
assert "interrupted" in caplog.text
|
||||
|
||||
|
||||
def test_interrupting_an_active_task_reports_a_sanitized_failure(tmp_path: Path) -> None:
|
||||
content = b"input fixture"
|
||||
worker, coordinator, _, _, _ = daemon(tmp_path, make_task(content), content)
|
||||
|
||||
class InterruptingRunner(FakeRunner):
|
||||
def run(self, task: ClaimedTask, task_dir: Path) -> RunResult:
|
||||
raise KeyboardInterrupt
|
||||
|
||||
worker.runner = InterruptingRunner()
|
||||
with pytest.raises(KeyboardInterrupt):
|
||||
worker.run_once()
|
||||
assert coordinator.failures == [
|
||||
{
|
||||
"worker_id": "worker-1",
|
||||
"attempt": 1,
|
||||
"error_code": "InterruptedError",
|
||||
"error_message": "worker interrupted by operator",
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
def test_max_tasks_counts_successes_not_failed_claims(tmp_path: Path) -> None:
|
||||
successful_content = b"successful input"
|
||||
|
||||
class SequencedCoordinator(FakeCoordinator):
|
||||
def __init__(self) -> None:
|
||||
super().__init__(None)
|
||||
self.tasks = [
|
||||
make_task(b"bad input", "wrong-checksum"),
|
||||
ClaimedTask(
|
||||
"task-2",
|
||||
1,
|
||||
(datetime.now(timezone.utc) + timedelta(seconds=60)).isoformat(),
|
||||
"similarity-search",
|
||||
InputArtifact(
|
||||
"https://example.test/input",
|
||||
hashlib.sha256(successful_content).hexdigest(),
|
||||
),
|
||||
{"query_id": "CHEMBL1"},
|
||||
),
|
||||
]
|
||||
|
||||
def claim(self, worker_id: str, capabilities: tuple[str, ...]) -> ClaimedTask | None:
|
||||
return self.tasks.pop(0) if self.tasks else None
|
||||
|
||||
coordinator = SequencedCoordinator()
|
||||
artifacts, runner = FakeArtifacts(successful_content), FakeRunner()
|
||||
config = WorkerConfig("https://example.test", "worker-1", tmp_path / "work", max_tasks=1)
|
||||
worker = WorkerDaemon(config, coordinator, artifacts, runner)
|
||||
assert worker.run_forever() is True
|
||||
assert len(coordinator.failures) == 1
|
||||
assert len(coordinator.submissions) == 1
|
||||
assert runner.calls == 1
|
||||
|
||||
|
||||
def test_worker_cli_lifecycle_options_are_explicit_and_exclusive() -> None:
|
||||
parser = build_parser()
|
||||
assert parser.parse_args(["--once"]).once is True
|
||||
assert parser.parse_args(["--max-tasks", "2"]).max_tasks == 2
|
||||
with pytest.raises(SystemExit):
|
||||
parser.parse_args(["--once", "--max-tasks", "2"])
|
||||
|
||||
|
||||
def test_worker_cli_uses_a_nonzero_exit_code_for_interruption(
|
||||
monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
||||
) -> None:
|
||||
class InterruptedDaemon:
|
||||
def __init__(self, *_: object) -> None:
|
||||
pass
|
||||
|
||||
def run_forever(self) -> bool:
|
||||
return False
|
||||
|
||||
monkeypatch.setattr(worker_cli, "WorkerDaemon", InterruptedDaemon)
|
||||
assert worker_cli.main(
|
||||
["--coordinator-url", "https://example.test", "--work-dir", str(tmp_path)]
|
||||
) == 130
|
||||
|
||||
|
||||
@pytest.mark.parametrize("value", [0, -1, True])
|
||||
def test_max_tasks_must_be_positive(value: object, tmp_path: Path) -> None:
|
||||
with pytest.raises(ValueError, match="max_tasks"):
|
||||
WorkerConfig("https://example.test", None, tmp_path, max_tasks=value) # type: ignore[arg-type]
|
||||
|
||||
|
||||
def test_bad_checksum_reports_failure_without_running(tmp_path: Path) -> None:
|
||||
worker, coordinator, _, runner, _ = daemon(tmp_path, make_task(b"actual", "not-the-hash"), b"actual")
|
||||
assert worker.run_once() is True
|
||||
assert worker.run_once() == RunOnceOutcome(claimed=True, completed=False)
|
||||
assert runner.calls == 0
|
||||
assert coordinator.failures[0]["error_code"] == "ValueError"
|
||||
assert not coordinator.submissions
|
||||
|
||||
|
||||
def test_failure_reporting_removes_paths_outside_the_worker_directory(tmp_path: Path) -> None:
|
||||
worker, coordinator, _, _, _ = daemon(tmp_path, make_task(b"input"), b"input")
|
||||
error = subprocess.CalledProcessError(
|
||||
1,
|
||||
["/home/alice/.venv/bin/python", "-m", "scimesh.cli", "/private/input.tsv"],
|
||||
)
|
||||
worker._report_failure(make_task(b"input"), error)
|
||||
message = coordinator.failures[0]["error_message"]
|
||||
assert "/home/alice" not in message
|
||||
assert "/private/input.tsv" not in message
|
||||
assert "<path>" in message
|
||||
|
||||
|
||||
def test_directory_creation_failure_is_reported(tmp_path: Path) -> None:
|
||||
content = b"input fixture"
|
||||
worker, coordinator, _, _, config = daemon(tmp_path, make_task(content), content)
|
||||
(config.work_dir / "task-1" / "1").mkdir(parents=True)
|
||||
assert worker.run_once() is True
|
||||
assert worker.run_once() == RunOnceOutcome(claimed=True, completed=False)
|
||||
assert coordinator.failures[0]["error_code"] == "FileExistsError"
|
||||
|
||||
|
||||
@@ -312,6 +445,33 @@ def test_environment_overrides_allow_cli_only_configuration(monkeypatch: pytest.
|
||||
assert config.worker_id is None
|
||||
assert "similarity-search" in config.capabilities
|
||||
assert "similarity_search" in config.capabilities
|
||||
assert "similarity-graph" not in config.capabilities
|
||||
|
||||
|
||||
def test_relative_work_dir_is_normalized_for_runner_subprocesses(
|
||||
monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
||||
) -> None:
|
||||
monkeypatch.chdir(tmp_path)
|
||||
config = WorkerConfig("https://coordinator.example", None, Path("./worker-data"))
|
||||
assert config.work_dir == tmp_path / "worker-data"
|
||||
|
||||
task_dir = config.work_dir / "task" / "1"
|
||||
task_dir.mkdir(parents=True)
|
||||
(task_dir / "input").write_text("fixture", encoding="utf-8")
|
||||
command: list[str] = []
|
||||
|
||||
def fake_run(args: list[str], **_: object) -> None:
|
||||
command.extend(args)
|
||||
output = Path(args[args.index("--output") + 1])
|
||||
output.write_text("id,score\n", encoding="utf-8")
|
||||
|
||||
monkeypatch.setattr("scimesh.worker.runners.subprocess.run", fake_run)
|
||||
task = ClaimedTask(
|
||||
"task", 1, "2026-07-30T00:00:00Z", "similarity-search",
|
||||
InputArtifact("https://example.test/input", "a" * 64), {"query_smiles": "CCO"},
|
||||
)
|
||||
SciMeshRunner().run(task, task_dir)
|
||||
assert command[4] == str(task_dir / "input")
|
||||
|
||||
|
||||
def test_worker_registration_sets_returned_identity(tmp_path: Path) -> None:
|
||||
|
||||
Reference in New Issue
Block a user