Harden distributed pipeline
This commit is contained in:
@@ -7,12 +7,19 @@ import (
|
||||
"bytes"
|
||||
"fmt"
|
||||
"io"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// ErrNoRows is returned when the input has a header but no data rows: a job with
|
||||
// zero tasks could never complete, so it is rejected at the source.
|
||||
var ErrNoRows = fmt.Errorf("input has no data rows")
|
||||
|
||||
// maxShardBytes bounds the coordinator memory used by one in-progress shard.
|
||||
// The uploaded file may be much larger: it is first stored on disk, then split
|
||||
// in small bounded pieces. Operators can lower rowsPerShard when this limit is
|
||||
// reached rather than exhausting the coordinator process.
|
||||
const maxShardBytes = 64 << 20 // 64 MiB
|
||||
|
||||
// SplitTSV reads a header-plus-rows text stream and cuts it into shards of at
|
||||
// most rowsPerShard data rows. Every shard repeats the header, so a worker can
|
||||
// parse its shard in isolation. emit is called once per shard, in order, with a
|
||||
@@ -26,13 +33,25 @@ var ErrNoRows = fmt.Errorf("input has no data rows")
|
||||
// Only one shard is buffered at a time, so memory is bounded by shard size (a
|
||||
// worker-sized slice of the data), not by the size of the whole dataset.
|
||||
func SplitTSV(r io.Reader, rowsPerShard int, emit func(index int, shard io.Reader) error) error {
|
||||
return SplitTSVLimit(r, rowsPerShard, 0, emit)
|
||||
return splitTSVLimit(r, rowsPerShard, 0, nil, emit)
|
||||
}
|
||||
|
||||
// SplitTSVLimit behaves like SplitTSV but emits no more than maxRows data rows.
|
||||
// A maxRows value of zero means unlimited. This lets an operator make a small,
|
||||
// representative pipeline check without materialising a second dataset file.
|
||||
func SplitTSVLimit(r io.Reader, rowsPerShard, maxRows int, emit func(index int, shard io.Reader) error) error {
|
||||
return splitTSVLimit(r, rowsPerShard, maxRows, nil, emit)
|
||||
}
|
||||
|
||||
// SplitChEMBLTSVLimit is the coordinator's scientific-upload splitter. It
|
||||
// validates the two columns every local SciMesh workload requires before any
|
||||
// shard task is persisted, while generic SplitTSV remains reusable for future
|
||||
// non-chemistry workloads.
|
||||
func SplitChEMBLTSVLimit(r io.Reader, rowsPerShard, maxRows int, emit func(index int, shard io.Reader) error) error {
|
||||
return splitTSVLimit(r, rowsPerShard, maxRows, validateChEMBLHeader, emit)
|
||||
}
|
||||
|
||||
func splitTSVLimit(r io.Reader, rowsPerShard, maxRows int, validateHeader func([]byte) error, emit func(index int, shard io.Reader) error) error {
|
||||
if rowsPerShard <= 0 {
|
||||
return fmt.Errorf("rowsPerShard must be positive, got %d", rowsPerShard)
|
||||
}
|
||||
@@ -51,6 +70,11 @@ func SplitTSVLimit(r io.Reader, rowsPerShard, maxRows int, emit func(index int,
|
||||
return ErrNoRows // completely empty input
|
||||
}
|
||||
header := append([]byte(nil), sc.Bytes()...)
|
||||
if validateHeader != nil {
|
||||
if err := validateHeader(header); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
var (
|
||||
buf bytes.Buffer
|
||||
@@ -71,9 +95,15 @@ func SplitTSVLimit(r io.Reader, rowsPerShard, maxRows int, emit func(index int,
|
||||
|
||||
for sc.Scan() {
|
||||
if rows == 0 {
|
||||
if len(header)+1 > maxShardBytes {
|
||||
return fmt.Errorf("TSV header exceeds maximum shard size of %d bytes", maxShardBytes)
|
||||
}
|
||||
buf.Write(header)
|
||||
buf.WriteByte('\n')
|
||||
}
|
||||
if buf.Len()+len(sc.Bytes())+1 > maxShardBytes {
|
||||
return fmt.Errorf("shard exceeds maximum size of %d bytes; lower rowsPerShard", maxShardBytes)
|
||||
}
|
||||
buf.Write(sc.Bytes())
|
||||
buf.WriteByte('\n')
|
||||
rows++
|
||||
@@ -103,3 +133,17 @@ func SplitTSVLimit(r io.Reader, rowsPerShard, maxRows int, emit func(index int,
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func validateChEMBLHeader(header []byte) error {
|
||||
seen := make(map[string]struct{})
|
||||
for _, field := range strings.Split(strings.TrimPrefix(string(header), "\ufeff"), "\t") {
|
||||
seen[field] = struct{}{}
|
||||
}
|
||||
if _, ok := seen["chembl_id"]; !ok {
|
||||
return fmt.Errorf("TSV is missing required column chembl_id")
|
||||
}
|
||||
if _, ok := seen["canonical_smiles"]; !ok {
|
||||
return fmt.Errorf("TSV is missing required column canonical_smiles")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -119,6 +119,14 @@ func TestSplitLimitUsesOnlyLeadingDataRows(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestChEMBLSplitRejectsMissingRequiredColumns(t *testing.T) {
|
||||
err := SplitChEMBLTSVLimit(strings.NewReader("id\tsmiles\nA\tCC\n"), 1, 0,
|
||||
func(int, io.Reader) error { return nil })
|
||||
if err == nil || !strings.Contains(err.Error(), "chembl_id") {
|
||||
t.Errorf("err = %v, want missing-column error", err)
|
||||
}
|
||||
}
|
||||
|
||||
// The scanned bytes are reused by bufio; the shard buffer must copy them, or a
|
||||
// later row would corrupt an earlier one. This guards that copy.
|
||||
func TestSplitDoesNotAliasScannerBuffer(t *testing.T) {
|
||||
|
||||
@@ -129,6 +129,9 @@ func LoadConfig() (Config, error) {
|
||||
if cfg.DefaultMaxAttempts, err = getEnvInt("DEFAULT_MAX_ATTEMPTS", cfg.DefaultMaxAttempts); err != nil {
|
||||
return Config{}, err
|
||||
}
|
||||
if cfg.DefaultMaxAttempts < 1 {
|
||||
return Config{}, fmt.Errorf("DEFAULT_MAX_ATTEMPTS must be positive")
|
||||
}
|
||||
|
||||
return cfg, nil
|
||||
}
|
||||
|
||||
@@ -32,3 +32,14 @@ func TestLoadConfigAllowsDistinctUIAndWorkerTokens(t *testing.T) {
|
||||
t.Fatalf("unexpected tokens: %+v", cfg)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoadConfigRejectsNonPositiveDefaultMaxAttempts(t *testing.T) {
|
||||
t.Setenv("ENV_FILE", filepath.Join(t.TempDir(), "missing.env"))
|
||||
t.Setenv("DATABASE_URL", "postgres://test")
|
||||
t.Setenv("DEFAULT_MAX_ATTEMPTS", "0")
|
||||
|
||||
_, err := LoadConfig()
|
||||
if err == nil || !strings.Contains(err.Error(), "DEFAULT_MAX_ATTEMPTS") {
|
||||
t.Fatalf("LoadConfig error = %v, want default-attempt validation", err)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -160,17 +160,18 @@ func (r *TaskRepo) CancelByJob(_ context.Context, jobID uuid.UUID, now time.Time
|
||||
return cancelled, nil
|
||||
}
|
||||
|
||||
func (r *TaskRepo) ExpireLeases(ctx context.Context, now time.Time) (int64, error) {
|
||||
func (r *TaskRepo) ExpireLeases(ctx context.Context, now time.Time) ([]uuid.UUID, error) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
var n int64
|
||||
affected := make([]uuid.UUID, 0)
|
||||
for _, t := range r.tasks {
|
||||
if t.Status == domain.TaskLeased && t.LeaseExpiresAt != nil && t.LeaseExpiresAt.Before(now) {
|
||||
if (t.Status == domain.TaskLeased || t.Status == domain.TaskRunning) &&
|
||||
t.LeaseExpiresAt != nil && t.LeaseExpiresAt.Before(now) {
|
||||
t.ExpireLease(now)
|
||||
n++
|
||||
affected = append(affected, t.JobID)
|
||||
}
|
||||
}
|
||||
return n, nil
|
||||
return affected, nil
|
||||
}
|
||||
|
||||
// --- JobRepo -------------------------------------------------------------
|
||||
@@ -298,6 +299,23 @@ func (r *ArtifactRepo) Get(ctx context.Context, id uuid.UUID) (*domain.Artifact,
|
||||
return &cp, nil
|
||||
}
|
||||
|
||||
func (r *ArtifactRepo) FindPartialResult(_ context.Context, taskID uuid.UUID, attempt int) (*domain.Artifact, error) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
for _, a := range r.arts {
|
||||
if a.TaskID != nil && *a.TaskID == taskID && a.Kind == domain.ArtifactPartialResult &&
|
||||
a.Attempt != nil && *a.Attempt == attempt {
|
||||
return cloneArtifact(a), nil
|
||||
}
|
||||
}
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
func cloneArtifact(a *domain.Artifact) *domain.Artifact {
|
||||
cp := *a
|
||||
return &cp
|
||||
}
|
||||
|
||||
// --- BlobStore -----------------------------------------------------------
|
||||
|
||||
type BlobStore struct {
|
||||
|
||||
@@ -60,6 +60,18 @@ func (r *UIReadRepo) ListTasksByJob(_ context.Context, jobID uuid.UUID) ([]domai
|
||||
sort.Slice(out, func(i, j int) bool { return out[i].ChunkIndex < out[j].ChunkIndex })
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (r *UIReadRepo) ListTasksByJobs(ctx context.Context, jobIDs []uuid.UUID) (map[uuid.UUID][]domain.Task, error) {
|
||||
out := make(map[uuid.UUID][]domain.Task, len(jobIDs))
|
||||
for _, id := range jobIDs {
|
||||
tasks, err := r.ListTasksByJob(ctx, id)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out[id] = tasks
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
func (r *UIReadRepo) ListWorkers(_ context.Context, limit int) ([]domain.Worker, error) {
|
||||
if limit < 1 || limit > 100 {
|
||||
return nil, domain.ErrInvalidInput
|
||||
|
||||
@@ -70,3 +70,33 @@ func (r *ArtifactRepo) Get(ctx context.Context, id uuid.UUID) (*domain.Artifact,
|
||||
a.Kind = domain.ArtifactKind(kind)
|
||||
return &a, nil
|
||||
}
|
||||
|
||||
func (r *ArtifactRepo) FindPartialResult(ctx context.Context, taskID uuid.UUID, attempt int) (*domain.Artifact, error) {
|
||||
sql, args, err := psql.Select(artifactColumns...).
|
||||
From("artifacts").
|
||||
Where(sq.Eq{
|
||||
"task_id": taskID,
|
||||
"attempt": attempt,
|
||||
"kind": string(domain.ArtifactPartialResult),
|
||||
}).
|
||||
ToSql()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
var (
|
||||
a domain.Artifact
|
||||
kind string
|
||||
)
|
||||
err = conn(ctx, r.pool).QueryRow(ctx, sql, args...).Scan(
|
||||
&a.ID, &a.JobID, &a.TaskID, &a.Attempt, &kind, &a.Filename, &a.StorageKey,
|
||||
&a.ContentType, &a.SizeBytes, &a.SHA256, &a.CreatedAt)
|
||||
if errors.Is(err, pgx.ErrNoRows) {
|
||||
return nil, nil
|
||||
}
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("find partial result: %w", err)
|
||||
}
|
||||
a.Kind = domain.ArtifactKind(kind)
|
||||
return &a, nil
|
||||
}
|
||||
|
||||
@@ -360,12 +360,23 @@ func TestCompleteTaskReplayIsIdempotent(t *testing.T) {
|
||||
t.Errorf("replay must be idempotent, got %v", err)
|
||||
}
|
||||
|
||||
// A different result artifact for the same task is a genuine conflict.
|
||||
art2 := seedArtifact(t, pool, job.ID, &claimed.ID, domain.ArtifactPartialResult)
|
||||
other := in
|
||||
other.ResultArtifactID = art2.ID
|
||||
if _, err := uc.Execute(ctx, other); !errors.Is(err, domain.ErrResultConflict) {
|
||||
t.Errorf("err = %v, want ErrResultConflict", err)
|
||||
}
|
||||
|
||||
func TestPartialResultIsUniquePerTaskAttempt(t *testing.T) {
|
||||
pool := testPool(t)
|
||||
ctx := context.Background()
|
||||
job, tasks := seedJob(t, pool, 1)
|
||||
taskID := tasks[0].ID
|
||||
first := seedArtifact(t, pool, job.ID, &taskID, domain.ArtifactPartialResult)
|
||||
second, err := domain.NewArtifact(job.ID, &taskID, domain.ArtifactPartialResult, "retry.csv", "text/csv", time.Now().UTC())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
attempt := 1
|
||||
second.Attempt = &attempt
|
||||
second.SetContent("other-sha", 5)
|
||||
if err := NewArtifactRepo(pool).Insert(ctx, second); err == nil {
|
||||
t.Fatalf("second partial artifact for %s/%d was accepted after %s", taskID, attempt, first.ID)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -337,17 +337,26 @@ SET status = CASE WHEN attempt < max_attempts THEN 'pending'::task_sta
|
||||
ELSE error_message END,
|
||||
completed_at = CASE WHEN attempt >= max_attempts THEN $1 ELSE completed_at END,
|
||||
version = version + 1
|
||||
WHERE status IN ('leased','running') AND lease_expires_at < $1`
|
||||
WHERE status IN ('leased','running') AND lease_expires_at < $1
|
||||
RETURNING job_id`
|
||||
|
||||
func (r *TaskRepo) ExpireLeases(ctx context.Context, now time.Time) (int64, error) {
|
||||
var affected int64
|
||||
func (r *TaskRepo) ExpireLeases(ctx context.Context, now time.Time) ([]uuid.UUID, error) {
|
||||
var affected []uuid.UUID
|
||||
err := withRetry(ctx, func(ctx context.Context) error {
|
||||
tag, err := conn(ctx, r.pool).Exec(ctx, expireLeasesSQL, now, domain.ErrCodeLeaseExpired)
|
||||
rows, err := conn(ctx, r.pool).Query(ctx, expireLeasesSQL, now, domain.ErrCodeLeaseExpired)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
affected = tag.RowsAffected()
|
||||
return nil
|
||||
defer rows.Close()
|
||||
affected = affected[:0]
|
||||
for rows.Next() {
|
||||
var jobID uuid.UUID
|
||||
if err := rows.Scan(&jobID); err != nil {
|
||||
return err
|
||||
}
|
||||
affected = append(affected, jobID)
|
||||
}
|
||||
return rows.Err()
|
||||
})
|
||||
return affected, err
|
||||
}
|
||||
|
||||
@@ -75,6 +75,31 @@ func (r *UIReadRepo) ListTasksByJob(ctx context.Context, jobID uuid.UUID) ([]dom
|
||||
return tasks, rows.Err()
|
||||
}
|
||||
|
||||
func (r *UIReadRepo) ListTasksByJobs(ctx context.Context, jobIDs []uuid.UUID) (map[uuid.UUID][]domain.Task, error) {
|
||||
out := make(map[uuid.UUID][]domain.Task, len(jobIDs))
|
||||
if len(jobIDs) == 0 {
|
||||
return out, nil
|
||||
}
|
||||
sql, args, err := psql.Select(taskColumns...).From("tasks").
|
||||
Where(sq.Eq{"job_id": jobIDs}).OrderBy("job_id ASC", "chunk_index ASC").ToSql()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
rows, err := conn(ctx, r.pool).Query(ctx, sql, args...)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("list tasks by jobs: %w", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
for rows.Next() {
|
||||
task, err := scanTask(rows)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out[task.JobID] = append(out[task.JobID], *task)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
func (r *UIReadRepo) ListWorkers(ctx context.Context, limit int) ([]domain.Worker, error) {
|
||||
if limit < 1 || limit > 100 {
|
||||
return nil, domain.ErrInvalidInput
|
||||
|
||||
@@ -79,6 +79,10 @@ func (s *Server) handleClaim(w http.ResponseWriter, r *http.Request) {
|
||||
s.writeError(w, r, domain.ErrInvalidInput)
|
||||
return
|
||||
}
|
||||
if _, err := uuid.Parse(req.WorkerID); err != nil {
|
||||
s.writeError(w, r, domain.ErrInvalidInput)
|
||||
return
|
||||
}
|
||||
|
||||
claimed, err := s.uc.ClaimTask.Execute(ctx, usecase.ClaimTaskInput{
|
||||
WorkerID: req.WorkerID,
|
||||
|
||||
@@ -23,8 +23,9 @@ const token = "secret"
|
||||
const uiToken = "ui-secret"
|
||||
|
||||
type env struct {
|
||||
ts *httptest.Server
|
||||
blobs *memstore.BlobStore
|
||||
ts *httptest.Server
|
||||
blobs *memstore.BlobStore
|
||||
workerID string
|
||||
}
|
||||
|
||||
func newEnv(t *testing.T, ready func(context.Context) error) *env {
|
||||
@@ -45,22 +46,28 @@ func newEnvWithUIToken(t *testing.T, ready func(context.Context) error, configur
|
||||
uc := coordhttp.UseCases{
|
||||
RegisterWorker: usecase.NewRegisterWorker(work, clk),
|
||||
CreateJob: usecase.NewCreateJob(jobs, tasks, tx, clk),
|
||||
SubmitDataset: usecase.NewSubmitDataset(blobs, arts, jobs, tasks, tx, clk),
|
||||
ClaimTask: usecase.NewClaimTask(tasks, clk, lease),
|
||||
SubmitDataset: usecase.NewSubmitDataset(blobs, arts, jobs, tasks, tx, clk, 3),
|
||||
ClaimTask: usecase.NewClaimTask(tasks, jobs, work, tx, clk, lease),
|
||||
RenewLease: usecase.NewRenewLease(tasks, work, tx, clk, lease),
|
||||
CompleteTask: usecase.NewCompleteTask(tasks, jobs, arts, tx, clk),
|
||||
FailTask: usecase.NewFailTask(tasks, jobs, tx, clk),
|
||||
GetJobStatus: usecase.NewGetJobStatus(jobs, tasks),
|
||||
CancelJob: usecase.NewCancelJob(jobs, tasks, tx, clk),
|
||||
UploadArtifact: usecase.NewUploadArtifact(tasks, arts, blobs, clk),
|
||||
UploadArtifact: usecase.NewUploadArtifact(tasks, arts, blobs, tx, clk),
|
||||
DownloadArtifact: usecase.NewDownloadArtifact(arts, blobs),
|
||||
GetTaskInput: usecase.NewGetTaskInput(tasks, arts, blobs),
|
||||
Dashboard: usecase.NewDashboard(memstore.NewUIReadRepo(jobs, tasks, work, arts)),
|
||||
}
|
||||
worker, err := uc.RegisterWorker.Execute(context.Background(), usecase.RegisterWorkerInput{
|
||||
Name: "test-worker", Capabilities: []string{"w", "similarity-search"},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("register test worker: %v", err)
|
||||
}
|
||||
srv := coordhttp.NewServer(uc, slog.New(slog.NewTextHandler(io.Discard, nil)), 5*time.Second, 15*time.Second, 1<<30, ready)
|
||||
ts := httptest.NewServer(srv.Handler(token, configuredUIToken))
|
||||
t.Cleanup(ts.Close)
|
||||
return &env{ts: ts, blobs: blobs}
|
||||
return &env{ts: ts, blobs: blobs, workerID: worker.ID.String()}
|
||||
}
|
||||
|
||||
func healthy(context.Context) error { return nil }
|
||||
@@ -68,6 +75,7 @@ func healthy(context.Context) error { return nil }
|
||||
// do sends an authenticated JSON request and returns status + decoded body.
|
||||
func (e *env) do(t *testing.T, method, path, body string) (int, map[string]any) {
|
||||
t.Helper()
|
||||
body = strings.ReplaceAll(body, `"worker_id":"w1"`, `"worker_id":"`+e.workerID+`"`)
|
||||
req, _ := http.NewRequestWithContext(context.Background(), method, e.ts.URL+path, strings.NewReader(body))
|
||||
req.Header.Set("Authorization", "Bearer "+token)
|
||||
if body != "" {
|
||||
@@ -330,6 +338,26 @@ func TestRegisterRejectsNoCapabilities(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestClaimRequiresRegisteredWorkerAndUsesStoredCapabilities(t *testing.T) {
|
||||
e := newEnv(t, healthy)
|
||||
if code, _ := e.do(t, "POST", "/tasks/claim", `{"worker_id":"not-a-uuid"}`); code != http.StatusBadRequest {
|
||||
t.Fatalf("invalid worker id claim = %d, want 400", code)
|
||||
}
|
||||
if code, _ := e.do(t, "POST", "/tasks/claim", `{"worker_id":"11111111-1111-4111-8111-111111111111"}`); code != http.StatusNotFound {
|
||||
t.Fatalf("unregistered worker claim = %d, want 404", code)
|
||||
}
|
||||
if code, _ := e.do(t, "POST", "/jobs", `{"workload":"w","input_uri":"s3://in","chunks":[{"chunk_index":0,"input_uri":"s3://c","input_sha256":"sha"}]}`); code != http.StatusCreated {
|
||||
t.Fatalf("create job = %d", code)
|
||||
}
|
||||
code, worker := e.do(t, "POST", "/workers/register", `{"name":"search-only","capabilities":["similarity-search"]}`)
|
||||
if code != http.StatusCreated {
|
||||
t.Fatalf("register = %d", code)
|
||||
}
|
||||
if code, _ := e.do(t, "POST", "/tasks/claim", `{"worker_id":"`+worker["worker_id"].(string)+`","capabilities":["w"]}`); code != http.StatusNoContent {
|
||||
t.Fatalf("forged capability claim = %d, want 204", code)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFullLifecycle(t *testing.T) {
|
||||
e := newEnv(t, healthy)
|
||||
|
||||
@@ -393,9 +421,9 @@ func TestForeignArtifactResultConflict(t *testing.T) {
|
||||
|
||||
func TestUploadDatasetChunksAndServesInput(t *testing.T) {
|
||||
e := newEnv(t, healthy)
|
||||
tsv := "id\tsmiles\nA\tCC\nB\tCCC\nC\tCCCC\nD\tCCCCC\nE\tCCCCCC\n"
|
||||
tsv := "chembl_id\tcanonical_smiles\nA\tCC\nB\tCCC\nC\tCCCC\nD\tCCCCC\nE\tCCCCCC\n"
|
||||
|
||||
code, body := e.uploadDataset(t, "w", 2, tsv)
|
||||
code, body := e.uploadDataset(t, "similarity-search", 2, tsv)
|
||||
if code != 201 {
|
||||
t.Fatalf("upload: status = %d", code)
|
||||
}
|
||||
@@ -421,7 +449,7 @@ func TestUploadDatasetChunksAndServesInput(t *testing.T) {
|
||||
t.Fatalf("get input: status = %d", resp.StatusCode)
|
||||
}
|
||||
shard, _ := io.ReadAll(resp.Body)
|
||||
if !strings.HasPrefix(string(shard), "id\tsmiles\n") {
|
||||
if !strings.HasPrefix(string(shard), "chembl_id\tcanonical_smiles\n") {
|
||||
t.Errorf("shard missing header: %q", shard)
|
||||
}
|
||||
}
|
||||
@@ -430,11 +458,12 @@ func TestUploadDatasetLimitsRows(t *testing.T) {
|
||||
e := newEnv(t, healthy)
|
||||
var buf bytes.Buffer
|
||||
mw := multipart.NewWriter(&buf)
|
||||
_ = mw.WriteField("workload", "w")
|
||||
_ = mw.WriteField("workload", "similarity-search")
|
||||
_ = mw.WriteField("parameters", `{"query_smiles":"CCO"}`)
|
||||
_ = mw.WriteField("chunk_rows", "2")
|
||||
_ = mw.WriteField("max_rows", "3")
|
||||
fw, _ := mw.CreateFormFile("file", "chembl.tsv")
|
||||
_, _ = io.Copy(fw, strings.NewReader("id\tsmiles\nA\tCC\nB\tCCC\nC\tCCCC\nD\tCCCCC\n"))
|
||||
_, _ = io.Copy(fw, strings.NewReader("chembl_id\tcanonical_smiles\nA\tCC\nB\tCCC\nC\tCCCC\nD\tCCCCC\n"))
|
||||
_ = mw.Close()
|
||||
req, _ := http.NewRequestWithContext(context.Background(), "POST", e.ts.URL+"/jobs/upload", &buf)
|
||||
req.Header.Set("Authorization", "Bearer "+token)
|
||||
@@ -451,6 +480,29 @@ func TestUploadDatasetLimitsRows(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestUploadDatasetRejectsMissingChEMBLColumns(t *testing.T) {
|
||||
e := newEnv(t, healthy)
|
||||
var buf bytes.Buffer
|
||||
mw := multipart.NewWriter(&buf)
|
||||
_ = mw.WriteField("workload", "similarity-search")
|
||||
_ = mw.WriteField("parameters", `{"query_smiles":"CCO"}`)
|
||||
_ = mw.WriteField("chunk_rows", "2")
|
||||
fw, _ := mw.CreateFormFile("file", "not-chembl.tsv")
|
||||
_, _ = io.Copy(fw, strings.NewReader("id\tsmiles\nA\tCC\n"))
|
||||
_ = mw.Close()
|
||||
req, _ := http.NewRequestWithContext(context.Background(), "POST", e.ts.URL+"/jobs/upload", &buf)
|
||||
req.Header.Set("Authorization", "Bearer "+token)
|
||||
req.Header.Set("Content-Type", mw.FormDataContentType())
|
||||
resp, err := http.DefaultClient.Do(req)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode != http.StatusBadRequest {
|
||||
t.Errorf("missing ChEMBL columns = %d, want 400", resp.StatusCode)
|
||||
}
|
||||
}
|
||||
|
||||
func TestErrorMappings(t *testing.T) {
|
||||
e := newEnv(t, healthy)
|
||||
zero := "00000000-0000-0000-0000-000000000000"
|
||||
@@ -478,10 +530,11 @@ func TestUploadDatasetRejectsAmbiguousMultipartInput(t *testing.T) {
|
||||
e := newEnv(t, healthy)
|
||||
var buf bytes.Buffer
|
||||
mw := multipart.NewWriter(&buf)
|
||||
_ = mw.WriteField("workload", "w")
|
||||
_ = mw.WriteField("workload", "similarity-search")
|
||||
_ = mw.WriteField("parameters", `{"query_smiles":"CCO"}`)
|
||||
_ = mw.WriteField("chunk_rows", "not-a-number")
|
||||
fw, _ := mw.CreateFormFile("file", "chembl.tsv")
|
||||
_, _ = io.Copy(fw, strings.NewReader("id\tsmiles\nA\tCC\n"))
|
||||
_, _ = io.Copy(fw, strings.NewReader("chembl_id\tcanonical_smiles\nA\tCC\n"))
|
||||
_ = mw.Close()
|
||||
|
||||
req, _ := http.NewRequestWithContext(context.Background(), "POST", e.ts.URL+"/jobs/upload", &buf)
|
||||
@@ -505,6 +558,9 @@ func (e *env) putArtifact(t *testing.T, taskID, worker string, attempt int, data
|
||||
e.ts.URL+"/tasks/"+taskID+"/artifacts/r.csv", strings.NewReader(data))
|
||||
req.Header.Set("Authorization", "Bearer "+token)
|
||||
req.Header.Set("Content-Type", "text/csv")
|
||||
if worker == "w1" {
|
||||
worker = e.workerID
|
||||
}
|
||||
req.Header.Set("X-Worker-ID", worker)
|
||||
req.Header.Set("X-Task-Attempt", itoa(attempt))
|
||||
resp, err := http.DefaultClient.Do(req)
|
||||
@@ -526,6 +582,7 @@ func (e *env) uploadDataset(t *testing.T, workload string, rows int, tsv string)
|
||||
var buf bytes.Buffer
|
||||
mw := multipart.NewWriter(&buf)
|
||||
_ = mw.WriteField("workload", workload)
|
||||
_ = mw.WriteField("parameters", `{"query_smiles":"CCO"}`)
|
||||
_ = mw.WriteField("chunk_rows", itoa(rows))
|
||||
fw, _ := mw.CreateFormFile("file", "chembl.tsv")
|
||||
_, _ = io.Copy(fw, strings.NewReader(tsv))
|
||||
|
||||
@@ -22,7 +22,7 @@
|
||||
<script>
|
||||
const id={{printf "%q" .ID}},state={pending:['Waiting for a worker','waiting','Waiting for an available worker with the required capability.'],leased:['Assigned to a worker','active','A worker has claimed the task and should begin processing shortly.'],running:['Running','active','A worker is reading a shard, calculating fingerprints, and uploading its result through the coordinator.'],completed:['Tasks complete','success','Every shard task is complete. Files below are still partial results.'],failed:['Needs attention','danger','One or more shard tasks failed. Open the task list below for details.'],cancelled:['Stopped','waiting','The operator stopped this job. No new shards can be claimed.']};
|
||||
const stop=document.querySelector('#stop-job');if(stop)stop.addEventListener('click',async()=>{if(!confirm('Stop this job? Unfinished shards will be cancelled.'))return;stop.disabled=true;const response=await fetch('/ui/api/jobs/'+id+'/cancel',{method:'POST'});if(!response.ok){stop.disabled=false;alert('Unable to stop this job.');return}location.reload()});
|
||||
setInterval(async()=>{try{const response=await fetch('/ui/api/jobs/'+id);if(!response.ok)return;const job=await response.json(),info=state[job.status]||[job.status,'waiting','Status reported by the coordinator.'],done=job.completed+job.failed+job.cancelled,percent=job.total?Math.min(100,Math.floor(done*100/job.total)):0,badge=document.querySelector('#status');badge.textContent=info[0];badge.className='status status-'+info[1];document.querySelector('#hint').textContent=info[2];document.querySelector('#progress').textContent=job.completed+' of '+job.total+' tasks complete'+(job.failed?' · failed: '+job.failed:'')+(job.cancelled?' · stopped: '+job.cancelled:'');document.querySelector('#progress-bar').style.width=percent+'%';for(const key of ['total','completed','pending','failed','cancelled'])document.querySelector('#'+key).textContent=job[key];document.querySelector('#active').textContent=job.leased+job.running}catch(_){}} ,2000);
|
||||
const terminal=new Set(['completed','failed','cancelled']);let timer;const poll=async()=>{try{const response=await fetch('/ui/api/jobs/'+id);if(!response.ok)return;const job=await response.json(),info=state[job.status]||[job.status,'waiting','Status reported by the coordinator.'],done=job.completed+job.failed+job.cancelled,percent=job.total?Math.min(100,Math.floor(done*100/job.total)):0,badge=document.querySelector('#status');badge.textContent=info[0];badge.className='status status-'+info[1];document.querySelector('#hint').textContent=info[2];document.querySelector('#progress').textContent=job.completed+' of '+job.total+' tasks complete'+(job.failed?' · failed: '+job.failed:'')+(job.cancelled?' · stopped: '+job.cancelled:'');document.querySelector('#progress-bar').style.width=percent+'%';for(const key of ['total','completed','pending','failed','cancelled'])document.querySelector('#'+key).textContent=job[key];document.querySelector('#active').textContent=job.leased+job.running;if(terminal.has(job.status)&&timer){clearInterval(timer);timer=undefined}}catch(_){}};const start=()=>{if(!timer&&!document.hidden&&!terminal.has(document.querySelector('#status').textContent.toLowerCase()))timer=setInterval(poll,2000)};document.addEventListener('visibilitychange',()=>{if(document.hidden&&timer){clearInterval(timer);timer=undefined}else start()});start();
|
||||
</script>
|
||||
</body>
|
||||
</html>
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -33,6 +33,17 @@ func NewCreateJob(jobs JobRepository, tasks TaskRepository, tx TxManager, clock
|
||||
// The all-or-none guarantee comes from TxManager: a half-created job would
|
||||
// leave chunks no worker could ever complete.
|
||||
func (uc *CreateJob) Execute(ctx context.Context, in CreateJobInput) (*domain.Job, error) {
|
||||
if in.Workload == "similarity-graph" || in.Workload == "similarity_graph" {
|
||||
// CTX-10 must plan triangular block pairs; ordinary independent input
|
||||
// chunks would silently omit every cross-chunk molecular pair.
|
||||
return nil, domain.ErrInvalidInput
|
||||
}
|
||||
if (in.Workload == "similarity-search" || in.Workload == "similarity_search") &&
|
||||
len(in.Chunks) > 1 && in.Parameters["query_id"] != nil {
|
||||
// Resolving once against the source dataset belongs to CTX-07. Letting
|
||||
// each shard resolve it would make most tasks fail or use inconsistent data.
|
||||
return nil, domain.ErrInvalidInput
|
||||
}
|
||||
chunks := make([]domain.ChunkSpec, 0, len(in.Chunks))
|
||||
for _, c := range in.Chunks {
|
||||
chunks = append(chunks, domain.ChunkSpec(c))
|
||||
@@ -92,6 +103,17 @@ func (uc *CancelJob) Execute(ctx context.Context, jobID uuid.UUID) (int64, error
|
||||
if job.Status == domain.JobCompleted || job.Status == domain.JobFailed {
|
||||
return domain.ErrJobNotCancellable
|
||||
}
|
||||
// The lease reaper can be the transition that exhausted the final task.
|
||||
// Check the authoritative task histogram as well as the cached job status,
|
||||
// so a stale status can never turn a failed/completed job into cancelled.
|
||||
counts, err := uc.tasks.CountByStatus(ctx, jobID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
derived := progressFrom(*job, counts).DeriveStatus()
|
||||
if derived == domain.JobCompleted || derived == domain.JobFailed {
|
||||
return domain.ErrJobNotCancellable
|
||||
}
|
||||
cancelled, err = uc.tasks.CancelByJob(ctx, jobID, now)
|
||||
if err != nil {
|
||||
return err
|
||||
@@ -212,3 +234,18 @@ func syncJobStatus(ctx context.Context, jobs JobRepository, tasks TaskRepository
|
||||
}
|
||||
return jobs.UpdateStatus(ctx, jobID, status, completedAt)
|
||||
}
|
||||
|
||||
func syncExpiredJobStatuses(ctx context.Context, jobs JobRepository, tasks TaskRepository,
|
||||
jobIDs []uuid.UUID, now time.Time) error {
|
||||
seen := make(map[uuid.UUID]struct{}, len(jobIDs))
|
||||
for _, jobID := range jobIDs {
|
||||
if _, duplicate := seen[jobID]; duplicate {
|
||||
continue
|
||||
}
|
||||
seen[jobID] = struct{}{}
|
||||
if err := syncJobStatus(ctx, jobs, tasks, jobID, now); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -60,9 +60,9 @@ type TaskRepository interface {
|
||||
// lease. It returns how many tasks changed.
|
||||
CancelByJob(ctx context.Context, jobID uuid.UUID, now time.Time) (int64, error)
|
||||
|
||||
// ExpireLeases applies the lease-expiry rule to every elapsed task and
|
||||
// reports how many were affected.
|
||||
ExpireLeases(ctx context.Context, now time.Time) (int64, error)
|
||||
// ExpireLeases applies the lease-expiry rule to every elapsed task and returns
|
||||
// the distinct jobs whose aggregate status may have changed.
|
||||
ExpireLeases(ctx context.Context, now time.Time) ([]uuid.UUID, error)
|
||||
}
|
||||
|
||||
// JobRepository persists jobs.
|
||||
@@ -89,6 +89,9 @@ type WorkerRepository interface {
|
||||
type ArtifactRepository interface {
|
||||
Insert(ctx context.Context, a *domain.Artifact) error
|
||||
Get(ctx context.Context, id uuid.UUID) (*domain.Artifact, error)
|
||||
// FindPartialResult returns the durable result already uploaded for one task
|
||||
// attempt. A nil artifact means the attempt has not uploaded one yet.
|
||||
FindPartialResult(ctx context.Context, taskID uuid.UUID, attempt int) (*domain.Artifact, error)
|
||||
}
|
||||
|
||||
// BlobStore holds artifact bytes, addressed by an opaque storage key. It streams
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -15,6 +15,7 @@ type UIReadRepository interface {
|
||||
GetJob(ctx context.Context, jobID uuid.UUID) (*domain.Job, error)
|
||||
ListJobs(ctx context.Context, limit int) ([]domain.Job, error)
|
||||
ListTasksByJob(ctx context.Context, jobID uuid.UUID) ([]domain.Task, error)
|
||||
ListTasksByJobs(ctx context.Context, jobIDs []uuid.UUID) (map[uuid.UUID][]domain.Task, error)
|
||||
ListWorkers(ctx context.Context, limit int) ([]domain.Worker, error)
|
||||
ListArtifactsByJob(ctx context.Context, jobID uuid.UUID) ([]domain.Artifact, error)
|
||||
}
|
||||
@@ -88,12 +89,16 @@ func (d *Dashboard) Overview(ctx context.Context, limit int) (DashboardView, err
|
||||
return DashboardView{}, err
|
||||
}
|
||||
out := DashboardView{Jobs: make([]JobCard, 0, len(jobs)), Workers: make([]WorkerCard, 0, len(workers))}
|
||||
jobIDs := make([]uuid.UUID, 0, len(jobs))
|
||||
for _, job := range jobs {
|
||||
tasks, err := d.read.ListTasksByJob(ctx, job.ID)
|
||||
if err != nil {
|
||||
return DashboardView{}, err
|
||||
}
|
||||
out.Jobs = append(out.Jobs, jobCard(job, tasks))
|
||||
jobIDs = append(jobIDs, job.ID)
|
||||
}
|
||||
tasksByJob, err := d.read.ListTasksByJobs(ctx, jobIDs)
|
||||
if err != nil {
|
||||
return DashboardView{}, err
|
||||
}
|
||||
for _, job := range jobs {
|
||||
out.Jobs = append(out.Jobs, jobCard(job, tasksByJob[job.ID]))
|
||||
}
|
||||
for _, worker := range workers {
|
||||
out.Workers = append(out.Workers, WorkerCard{ID: worker.ID.String(), Name: worker.Name, Status: string(worker.Status), Capabilities: worker.Capabilities, LastHeartbeatAt: worker.LastHeartbeatAt})
|
||||
|
||||
@@ -4,6 +4,7 @@ import (
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"math"
|
||||
|
||||
"github.com/google/uuid"
|
||||
|
||||
@@ -15,20 +16,27 @@ import (
|
||||
// creates the job with one task per shard — the coordinator-side counterpart of
|
||||
// a client submitting pre-chunked URIs.
|
||||
type SubmitDataset struct {
|
||||
blobs BlobStore
|
||||
artifacts ArtifactRepository
|
||||
jobs JobRepository
|
||||
tasks TaskRepository
|
||||
tx TxManager
|
||||
clk Clock
|
||||
blobs BlobStore
|
||||
artifacts ArtifactRepository
|
||||
jobs JobRepository
|
||||
tasks TaskRepository
|
||||
tx TxManager
|
||||
clk Clock
|
||||
maxAttempts int
|
||||
}
|
||||
|
||||
func NewSubmitDataset(blobs BlobStore, artifacts ArtifactRepository, jobs JobRepository,
|
||||
tasks TaskRepository, tx TxManager, clk Clock) *SubmitDataset {
|
||||
return &SubmitDataset{blobs: blobs, artifacts: artifacts, jobs: jobs, tasks: tasks, tx: tx, clk: clk}
|
||||
tasks TaskRepository, tx TxManager, clk Clock, maxAttempts int) *SubmitDataset {
|
||||
return &SubmitDataset{blobs: blobs, artifacts: artifacts, jobs: jobs, tasks: tasks, tx: tx, clk: clk, maxAttempts: maxAttempts}
|
||||
}
|
||||
|
||||
func (uc *SubmitDataset) Execute(ctx context.Context, in SubmitDatasetInput) (SubmitDatasetResult, error) {
|
||||
if err := validateUploadedWorkload(in.Workload, in.Parameters); err != nil {
|
||||
return SubmitDatasetResult{}, err
|
||||
}
|
||||
if uc.maxAttempts < 1 {
|
||||
return SubmitDatasetResult{}, domain.ErrInvalidInput
|
||||
}
|
||||
now := uc.clk.Now()
|
||||
|
||||
job, err := domain.NewUploadedJob(in.Workload, in.Parameters, now)
|
||||
@@ -64,7 +72,7 @@ func (uc *SubmitDataset) Execute(ctx context.Context, in SubmitDatasetInput) (Su
|
||||
cleanup()
|
||||
return SubmitDatasetResult{}, err
|
||||
}
|
||||
splitErr := chunk.SplitTSVLimit(rc, in.RowsPerShard, in.MaxRows, func(index int, shard io.Reader) error {
|
||||
splitErr := chunk.SplitChEMBLTSVLimit(rc, in.RowsPerShard, in.MaxRows, func(index int, shard io.Reader) error {
|
||||
art, err := domain.NewArtifact(job.ID, nil, domain.ArtifactShard,
|
||||
fmt.Sprintf("shard-%d.tsv", index), in.ContentType, now)
|
||||
if err != nil {
|
||||
@@ -77,7 +85,7 @@ func (uc *SubmitDataset) Execute(ctx context.Context, in SubmitDatasetInput) (Su
|
||||
putKeys = append(putKeys, art.StorageKey)
|
||||
art.SetContent(ssum, ssize)
|
||||
|
||||
task, err := domain.NewShardTask(job.ID, index, in.Workload, art.ID, ssum, in.Parameters, 0, now)
|
||||
task, err := domain.NewShardTask(job.ID, index, in.Workload, art.ID, ssum, in.Parameters, uc.maxAttempts, now)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
@@ -88,7 +96,8 @@ func (uc *SubmitDataset) Execute(ctx context.Context, in SubmitDatasetInput) (Su
|
||||
_ = rc.Close()
|
||||
if splitErr != nil {
|
||||
cleanup()
|
||||
return SubmitDatasetResult{}, splitErr
|
||||
// Dataset shape is caller input, not an internal coordinator failure.
|
||||
return SubmitDatasetResult{}, domain.ErrInvalidInput
|
||||
}
|
||||
|
||||
// 3. Persist job + all artifacts + all tasks atomically.
|
||||
@@ -118,6 +127,76 @@ func (uc *SubmitDataset) Execute(ctx context.Context, in SubmitDatasetInput) (Su
|
||||
}, nil
|
||||
}
|
||||
|
||||
// validateUploadedWorkload is deliberately narrow until CTX-07/08/10 adds a
|
||||
// typed distributed-workload registry. In particular, running similarity-graph
|
||||
// independently per TSV shard is scientifically wrong: cross-shard pairs would
|
||||
// be absent from the apparent graph.
|
||||
func validateUploadedWorkload(workload string, parameters map[string]any) error {
|
||||
if workload != "similarity-search" {
|
||||
return domain.ErrInvalidInput
|
||||
}
|
||||
allowed := map[string]struct{}{
|
||||
"query_smiles": {}, "top_k": {}, "threshold": {},
|
||||
"threshold_direction": {}, "progress_every": {},
|
||||
}
|
||||
for key := range parameters {
|
||||
if _, ok := allowed[key]; !ok {
|
||||
return domain.ErrInvalidInput
|
||||
}
|
||||
}
|
||||
query, ok := parameters["query_smiles"].(string)
|
||||
if !ok || query == "" || len(query) > 200 {
|
||||
return domain.ErrInvalidInput
|
||||
}
|
||||
if value, ok := parameters["top_k"]; ok && !isPositiveJSONInteger(value) {
|
||||
return domain.ErrInvalidInput
|
||||
}
|
||||
if value, ok := parameters["progress_every"]; ok && !isNonNegativeJSONInteger(value) {
|
||||
return domain.ErrInvalidInput
|
||||
}
|
||||
if value, ok := parameters["threshold"]; ok && !isUnitIntervalNumber(value) {
|
||||
return domain.ErrInvalidInput
|
||||
}
|
||||
if value, ok := parameters["threshold_direction"]; ok && value != "greater" && value != "less" {
|
||||
return domain.ErrInvalidInput
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func isPositiveJSONInteger(value any) bool { return isJSONInteger(value, false) }
|
||||
func isNonNegativeJSONInteger(value any) bool { return isJSONInteger(value, true) }
|
||||
|
||||
func isJSONInteger(value any, allowZero bool) bool {
|
||||
var n int64
|
||||
switch v := value.(type) {
|
||||
case int:
|
||||
n = int64(v)
|
||||
case int64:
|
||||
n = v
|
||||
case float64:
|
||||
if math.Trunc(v) != v || v > math.MaxInt64 || v < math.MinInt64 {
|
||||
return false
|
||||
}
|
||||
n = int64(v)
|
||||
default:
|
||||
return false
|
||||
}
|
||||
return n >= 0 && (allowZero || n > 0)
|
||||
}
|
||||
|
||||
func isUnitIntervalNumber(value any) bool {
|
||||
switch v := value.(type) {
|
||||
case float64:
|
||||
return !math.IsNaN(v) && !math.IsInf(v, 0) && v >= 0 && v <= 1
|
||||
case int:
|
||||
return v >= 0 && v <= 1
|
||||
case int64:
|
||||
return v >= 0 && v <= 1
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
// GetTaskInput resolves a task's input shard and opens it for streaming. The
|
||||
// caller closes the reader.
|
||||
type GetTaskInput struct {
|
||||
|
||||
@@ -68,18 +68,18 @@ func newHarness() *harness {
|
||||
}
|
||||
tx := memstore.Tx{}
|
||||
h.createJob = usecase.NewCreateJob(h.jobs, h.tasks, tx, h.clk)
|
||||
h.submit = usecase.NewSubmitDataset(h.blobs, h.arts, h.jobs, h.tasks, tx, h.clk)
|
||||
h.claim = usecase.NewClaimTask(h.tasks, h.clk, lease)
|
||||
h.submit = usecase.NewSubmitDataset(h.blobs, h.arts, h.jobs, h.tasks, tx, h.clk, 3)
|
||||
h.claim = usecase.NewClaimTask(h.tasks, h.jobs, h.work, tx, h.clk, lease)
|
||||
h.renew = usecase.NewRenewLease(h.tasks, h.work, tx, h.clk, lease)
|
||||
h.complete = usecase.NewCompleteTask(h.tasks, h.jobs, h.arts, tx, h.clk)
|
||||
h.fail = usecase.NewFailTask(h.tasks, h.jobs, tx, h.clk)
|
||||
h.status = usecase.NewGetJobStatus(h.jobs, h.tasks)
|
||||
h.results = usecase.NewListResults(h.tasks)
|
||||
h.register = usecase.NewRegisterWorker(h.work, h.clk)
|
||||
h.uploadArt = usecase.NewUploadArtifact(h.tasks, h.arts, h.blobs, h.clk)
|
||||
h.uploadArt = usecase.NewUploadArtifact(h.tasks, h.arts, h.blobs, tx, h.clk)
|
||||
h.downloadArt = usecase.NewDownloadArtifact(h.arts, h.blobs)
|
||||
h.getInput = usecase.NewGetTaskInput(h.tasks, h.arts, h.blobs)
|
||||
h.expire = usecase.NewExpireLeases(h.tasks, h.clk)
|
||||
h.expire = usecase.NewExpireLeases(h.tasks, h.jobs, tx, h.clk)
|
||||
h.cancel = usecase.NewCancelJob(h.jobs, h.tasks, tx, h.clk)
|
||||
return h
|
||||
}
|
||||
@@ -138,6 +138,44 @@ func TestClaimLeasesAndAdvancesAttempt(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestRegisteredWorkerCannotBroadenItsCapabilitiesAtClaim(t *testing.T) {
|
||||
h := newHarness()
|
||||
h.seedJob(t, "restricted", 1)
|
||||
worker, err := h.register.Execute(ctx, usecase.RegisterWorkerInput{
|
||||
Name: "search-only", Capabilities: []string{"similarity-search"},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
claimed, err := h.claim.Execute(ctx, usecase.ClaimTaskInput{
|
||||
WorkerID: worker.ID.String(), Workloads: []string{"restricted"},
|
||||
})
|
||||
if err != nil || claimed != nil {
|
||||
t.Fatalf("claim = (%v, %v), want no compatible task", claimed, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCreateJobRejectsUnsafeDistributedScientificPlans(t *testing.T) {
|
||||
h := newHarness()
|
||||
_, err := h.createJob.Execute(ctx, usecase.CreateJobInput{
|
||||
Workload: "similarity-graph", InputURI: "s3://input",
|
||||
Chunks: []usecase.ChunkInput{{ChunkIndex: 0, InputURI: "s3://chunk", InputSHA256: "sha"}},
|
||||
})
|
||||
if !errors.Is(err, domain.ErrInvalidInput) {
|
||||
t.Errorf("graph job err = %v, want ErrInvalidInput", err)
|
||||
}
|
||||
_, err = h.createJob.Execute(ctx, usecase.CreateJobInput{
|
||||
Workload: "similarity-search", InputURI: "s3://input", Parameters: map[string]any{"query_id": "CHEMBL1"},
|
||||
Chunks: []usecase.ChunkInput{
|
||||
{ChunkIndex: 0, InputURI: "s3://chunk0", InputSHA256: "sha"},
|
||||
{ChunkIndex: 1, InputURI: "s3://chunk1", InputSHA256: "sha"},
|
||||
},
|
||||
})
|
||||
if !errors.Is(err, domain.ErrInvalidInput) {
|
||||
t.Errorf("sharded query_id job err = %v, want ErrInvalidInput", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestClaimEmptyQueueReturnsNil(t *testing.T) {
|
||||
h := newHarness()
|
||||
c, err := h.claim.Execute(ctx, usecase.ClaimTaskInput{WorkerID: "w1", Workloads: []string{"w"}})
|
||||
@@ -281,7 +319,7 @@ func TestUploadRejectsLeaseThatExpiresDuringStreaming(t *testing.T) {
|
||||
h.seedJob(t, "w", 1)
|
||||
taskID, attempt := h.leaseOne(t, "w1", "w")
|
||||
h.uploadArt = usecase.NewUploadArtifact(
|
||||
h.tasks, h.arts, expiringBlobStore{BlobStore: h.blobs, clock: h.clk}, h.clk,
|
||||
h.tasks, h.arts, expiringBlobStore{BlobStore: h.blobs, clock: h.clk}, memstore.Tx{}, h.clk,
|
||||
)
|
||||
|
||||
_, err := h.uploadArt.Execute(ctx, usecase.UploadArtifactInput{
|
||||
@@ -414,6 +452,23 @@ func TestUploadArtifactRejectsForeignWorker(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestUploadArtifactIsIdempotentPerTaskAttempt(t *testing.T) {
|
||||
h := newHarness()
|
||||
h.seedJob(t, "w", 1)
|
||||
taskID, attempt := h.leaseOne(t, "w1", "w")
|
||||
first := h.uploadResult(t, taskID, "w1", attempt)
|
||||
second, err := h.uploadArt.Execute(ctx, usecase.UploadArtifactInput{
|
||||
TaskID: taskID, WorkerID: "w1", Attempt: attempt,
|
||||
Filename: "retry.csv", ContentType: "text/csv", Body: strings.NewReader("different bytes"),
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if second.ID != first {
|
||||
t.Errorf("retry artifact = %s, want existing %s", second.ID, first)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDownloadArtifactRoundTrips(t *testing.T) {
|
||||
h := newHarness()
|
||||
h.seedJob(t, "w", 1)
|
||||
@@ -434,10 +489,10 @@ func TestDownloadArtifactRoundTrips(t *testing.T) {
|
||||
|
||||
func TestSubmitDatasetChunksAndServesInput(t *testing.T) {
|
||||
h := newHarness()
|
||||
tsv := "id\tsmiles\nA\tCC\nB\tCCC\nC\tCCCC\nD\tCCCCC\nE\tCCCCCC\n"
|
||||
tsv := "chembl_id\tcanonical_smiles\nA\tCC\nB\tCCC\nC\tCCCC\nD\tCCCCC\nE\tCCCCCC\n"
|
||||
|
||||
res, err := h.submit.Execute(ctx, usecase.SubmitDatasetInput{
|
||||
Workload: "w", RowsPerShard: 2, Filename: "chembl.tsv",
|
||||
Workload: "similarity-search", Parameters: map[string]any{"query_smiles": "CCO"}, RowsPerShard: 2, Filename: "chembl.tsv",
|
||||
ContentType: "text/tab-separated-values", Body: strings.NewReader(tsv),
|
||||
})
|
||||
if err != nil {
|
||||
@@ -453,7 +508,7 @@ func TestSubmitDatasetChunksAndServesInput(t *testing.T) {
|
||||
t.Errorf("job total = %d, want 3", prog.Total)
|
||||
}
|
||||
|
||||
c, err := h.claim.Execute(ctx, usecase.ClaimTaskInput{WorkerID: "w1", Workloads: []string{"w"}})
|
||||
c, err := h.claim.Execute(ctx, usecase.ClaimTaskInput{WorkerID: "w1", Workloads: []string{"similarity-search"}})
|
||||
if err != nil || c == nil {
|
||||
t.Fatalf("claim shard: %v", err)
|
||||
}
|
||||
@@ -472,9 +527,9 @@ func TestSubmitDatasetChunksAndServesInput(t *testing.T) {
|
||||
|
||||
func TestSubmitDatasetLimitsRowsBeforeCreatingShards(t *testing.T) {
|
||||
h := newHarness()
|
||||
tsv := "id\tsmiles\nA\tCC\nB\tCCC\nC\tCCCC\nD\tCCCCC\nE\tCCCCCC\n"
|
||||
tsv := "chembl_id\tcanonical_smiles\nA\tCC\nB\tCCC\nC\tCCCC\nD\tCCCCC\nE\tCCCCCC\n"
|
||||
res, err := h.submit.Execute(ctx, usecase.SubmitDatasetInput{
|
||||
Workload: "w", RowsPerShard: 2, MaxRows: 3, Filename: "chembl.tsv",
|
||||
Workload: "similarity-search", Parameters: map[string]any{"query_smiles": "CCO"}, RowsPerShard: 2, MaxRows: 3, Filename: "chembl.tsv",
|
||||
ContentType: "text/tab-separated-values", Body: strings.NewReader(tsv),
|
||||
})
|
||||
if err != nil {
|
||||
@@ -485,6 +540,26 @@ func TestSubmitDatasetLimitsRowsBeforeCreatingShards(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestSubmitDatasetRejectsUnsupportedDistributedWorkloads(t *testing.T) {
|
||||
h := newHarness()
|
||||
_, err := h.submit.Execute(ctx, usecase.SubmitDatasetInput{
|
||||
Workload: "similarity-graph", Parameters: map[string]any{"threshold": 0.7}, RowsPerShard: 2,
|
||||
Filename: "chembl.tsv", ContentType: "text/tab-separated-values",
|
||||
Body: strings.NewReader("chembl_id\tcanonical_smiles\nA\tCC\n"),
|
||||
})
|
||||
if !errors.Is(err, domain.ErrInvalidInput) {
|
||||
t.Errorf("graph submission err = %v, want ErrInvalidInput", err)
|
||||
}
|
||||
_, err = h.submit.Execute(ctx, usecase.SubmitDatasetInput{
|
||||
Workload: "similarity-search", Parameters: map[string]any{"query_id": "CHEMBL1"}, RowsPerShard: 2,
|
||||
Filename: "chembl.tsv", ContentType: "text/tab-separated-values",
|
||||
Body: strings.NewReader("chembl_id\tcanonical_smiles\nA\tCC\n"),
|
||||
})
|
||||
if !errors.Is(err, domain.ErrInvalidInput) {
|
||||
t.Errorf("query_id submission err = %v, want ErrInvalidInput", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCancelJobInvalidatesClaimedAndPendingTasks(t *testing.T) {
|
||||
h := newHarness()
|
||||
jobID := h.seedJob(t, "w", 3)
|
||||
@@ -531,3 +606,22 @@ func TestExpireLeasesReclaims(t *testing.T) {
|
||||
t.Errorf("expire = (%d, %v), want (1, nil)", n, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFinalLeaseExpiryPersistsFailedJobAndCannotBeCancelled(t *testing.T) {
|
||||
h := newHarness()
|
||||
jobID := h.seedJob(t, "w", 1)
|
||||
for attempt := 1; attempt <= domain.DefaultMaxAttempts; attempt++ {
|
||||
h.leaseOne(t, "w1", "w")
|
||||
h.clk.Advance(lease + time.Second)
|
||||
if _, err := h.expire.Execute(ctx); err != nil {
|
||||
t.Fatalf("expire attempt %d: %v", attempt, err)
|
||||
}
|
||||
}
|
||||
progress, err := h.status.Execute(ctx, jobID)
|
||||
if err != nil || progress.Job.Status != domain.JobFailed || progress.DeriveStatus() != domain.JobFailed {
|
||||
t.Fatalf("progress = %+v, err = %v; want persisted failed job", progress, err)
|
||||
}
|
||||
if _, err := h.cancel.Execute(ctx, jobID); !errors.Is(err, domain.ErrJobNotCancellable) {
|
||||
t.Errorf("cancel terminal lease failure = %v, want ErrJobNotCancellable", err)
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user