From 706bc85e17d21f6161d973d00c567b672d34907b Mon Sep 17 00:00:00 2001 From: Emil Date: Sun, 2 Aug 2026 16:53:04 +0300 Subject: [PATCH] Replace Python worker daemon with the Go worker agent --- .opencode/session-goal.md | 16 +- AGENTS.md | 6 +- Makefile | 12 +- README.md | 6 +- STATUS.md | 2 +- coordinator/cmd/worker-agent/main.go | 8 +- coordinator/internal/agent/auth.go | 120 +++ coordinator/internal/agent/auth_test.go | 102 +++ coordinator/internal/agent/client.go | 54 +- coordinator/internal/agent/client_test.go | 2 +- coordinator/internal/agent/config.go | 17 +- coordinator/internal/agent/daemon.go | 36 + coordinator/internal/agent/daemon_test.go | 2 +- coordinator/scripts/demo-ui.sh | 39 +- docs/sdk-handoff.md | 2 +- mkdocs/sdk/cli.md | 4 +- mkdocs/sdk/worker-integration.md | 98 +-- pyproject.toml | 1 - scimesh/worker/__init__.py | 11 +- scimesh/worker/artifacts.py | 141 ---- scimesh/worker/auth.py | 128 ---- scimesh/worker/cli.py | 106 --- scimesh/worker/config.py | 221 ------ scimesh/worker/coordinator.py | 149 ---- scimesh/worker/daemon.py | 331 --------- scimesh/worker/models.py | 81 +-- scimesh/worker/runners.py | 80 +- scimesh/worker/transport.py | 53 -- scripts/two-worker-smoke.sh | 24 +- tests/test_worker_auth.py | 221 ------ tests/test_worker_daemon.py | 842 ---------------------- 31 files changed, 516 insertions(+), 2399 deletions(-) create mode 100644 coordinator/internal/agent/auth.go create mode 100644 coordinator/internal/agent/auth_test.go delete mode 100644 scimesh/worker/artifacts.py delete mode 100644 scimesh/worker/auth.py delete mode 100644 scimesh/worker/cli.py delete mode 100644 scimesh/worker/config.py delete mode 100644 scimesh/worker/coordinator.py delete mode 100644 scimesh/worker/daemon.py delete mode 100644 scimesh/worker/transport.py delete mode 100644 tests/test_worker_auth.py delete mode 100644 tests/test_worker_daemon.py diff --git a/.opencode/session-goal.md b/.opencode/session-goal.md index 4945f1e..6cf39db 100644 --- a/.opencode/session-goal.md +++ b/.opencode/session-goal.md @@ -12,12 +12,10 @@ COMPLETED ## Progress -- [x] Аудит: ruff F401/F811/F841 — 26 неиспользуемых импортов; vulture — кандидаты проверены grep'ом. -- [x] Удалено: - - 25 неиспользуемых импортов (ruff --fix) + 1 неиспользуемая локальная переменная в тесте; - - мёртвые функции descriptors/core.py `write_descriptor_shards`/`concatenate_descriptor_shards` (вытеснены дефолтами MapReduceWorkload; ссылки только в собственном `__init__`) + их экспорты; - - мёртвые атрибуты `MapReduceWorkload._resources/_execution` (записывались, нигде не читались); - - мёртвый `CancellationFlag.cancel` (0 использований); -- [x] `scripts/two-worker-smoke.sh` (рабочий E2E, но без точки входа) подключён как `make smoke-two-worker` — не мёртвый, а доступный. -- [x] Проверено: ruff clean; pytest 260 passed; pyright 0 ошибок (scimesh+tests); go test 11 пакетов + vet; mkdocs build без warnings. -- [x] Изменения не закоммичены. +Эта сессия (доп. задача): Go-агент доведён до паритета, Python-демон удалён. +- [x] Go-агент: token provider (static + worker-key exchange + 401 refresh на API/download/upload), CLEANUP_AFTER_SECONDS (очистка attempt-директорий), тесты auth (exchange/cache/reject/select/401-retry). +- [x] Python: удалены daemon.py, cli.py, config.py, coordinator.py, artifacts.py, auth.py, transport.py; `scimesh/worker/` = только task.py + runners.py (SDK-мост, allowlist из env) + models.py (ClaimedTask/RunResult); консольный скрипт scimesh-worker убран из pyproject. +- [x] Тесты: удалены test_worker_daemon.py, test_worker_auth.py; 208 pytest зелёные. +- [x] Демо/смок переведены на Go-агент (demo-ui.sh: build_agent + env; two-worker-smoke.sh: AGENT_BIN + TASK_RUNNER_JSON). `make smoke-two-worker` PASS: 4/4 шардов, worker-a=2, worker-b=2. +- [x] Документация: README, AGENTS.md, STATUS, handoff, mkdocs worker-integration/cli. +- [x] Верификация: ruff clean; 208 pytest; pyright 0 (scimesh+tests); go test 11 пакетов + vet; mkdocs 0 warnings. diff --git a/AGENTS.md b/AGENTS.md index 91c3530..4ea90ab 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -5,8 +5,10 @@ SciMesh is a Python package for molecular-similarity workloads. Source lives in `scimesh/`: `chemistry/` reads data and makes fingerprints, `workloads/` contains commands, and `core/` provides the workload protocol and registry. -The worker daemon in `scimesh/worker/` is a coordinator client, not a database -client. Tests are in `tests/`; specifications in `docs/`; roadmap: `PLAN.md`. +The Go worker agent (`coordinator/internal/agent/`) is a coordinator client, +not a database client; the Python side of a claimed task lives in +`scimesh/worker/` (the per-task SDK execution entry). Tests are in `tests/`; +specifications in `docs/`; roadmap: `PLAN.md`. For distributed work, read `.agents/`, `docs/api-contract.md`, and `STATUS.md`. Use one CTX task per pull request; local workloads are the diff --git a/Makefile b/Makefile index 7b791aa..134947b 100644 --- a/Makefile +++ b/Makefile @@ -1,15 +1,16 @@ .DEFAULT_GOAL := help -.PHONY: help demo-ui demo-down demo-logs smoke-two-worker docs docs-serve +.PHONY: help agent demo-ui demo-down demo-logs smoke-two-worker docs docs-serve help: @printf '%s\n' \ 'SciMesh developer commands:' \ - ' make demo-ui Start the local UI pipeline demo with 2 workers.' \ - ' make demo-ui WORKERS=3 Start the demo with 3 local workers.' \ + ' make agent Build the Go worker agent (coordinator/bin/worker-agent).' \ + ' make demo-ui Start the local UI pipeline demo with 2 Go worker agents.' \ + ' make demo-ui WORKERS=3 Start the demo with 3 workers.' \ ' make demo-logs Follow coordinator logs for the demo.' \ ' make demo-down Stop demo containers and workers.' \ - ' make smoke-two-worker E2E: two workers process 4 shards and the' \ + ' make smoke-two-worker E2E: two Go agents process 4 shards and the' \ ' result must match the local CLI reference.' \ ' make docs Build the MkDocs site into site/.' \ ' make docs-serve Serve the MkDocs site at http://localhost:8000.' \ @@ -18,6 +19,9 @@ help: # Convenient entry points from the repository root. Extra settings are passed # through, for example: make demo-ui WORKERS=3 +agent: + $(MAKE) -C coordinator agent + demo-ui: $(MAKE) -C coordinator demo-ui diff --git a/README.md b/README.md index b178e02..f4cd62d 100644 --- a/README.md +++ b/README.md @@ -3,7 +3,8 @@ SciMesh is a scientific-workload framework for molecular datasets. Its public CLI runs exact similarity search and sparse similarity-graph construction locally in one Python process; it creates no dense similarity matrix. The Go/PostgreSQL -coordinator and Python worker can run a shard-based `similarity-search` +coordinator and Go worker agents (which execute SDK workloads in a Python +subprocess) can run a shard-based `similarity-search` pipeline locally. After every shard succeeds, the coordinator deterministically merges its candidates into one final global top-k CSV. See [`STATUS.md`](STATUS.md). @@ -65,7 +66,8 @@ documentation rules the site is written by (`mkdocs/approach.md`). Open `http://localhost:18080/ui` and sign in with username `operator` and password `demo-ui-secret`. The command starts PostgreSQL, the coordinator, and -two local reference workers. Upload a small ChEMBL TSV, then use the job page +two Go worker agents (built by `make agent`; each executes the SDK workload +in a Python subprocess). Upload a small ChEMBL TSV, then use the job page to follow shard progress, inspect bounded **Preview CSV** results, and see a live processing-speed chart in shards per minute. The **Workloads** page shows the installed SDK workload library (descriptions, parameters, and artifact diff --git a/STATUS.md b/STATUS.md index 7e2af8b..133eb70 100644 --- a/STATUS.md +++ b/STATUS.md @@ -41,7 +41,7 @@ the complete result-artifact SHA-256 before a task is accepted. | CTX-03 Transactional queue | Implemented | Real-PostgreSQL integration tests cover atomic claims and concurrency. | | 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-06 Python Worker live-contract alignment | Superseded | The Python worker daemon was removed; the Go worker agent (`coordinator/internal/agent/` + `cmd/worker-agent`) now implements the lifecycle (register/claim/heartbeat/download/upload/submit/fail, token refresh, cleanup) and executes SDK workloads via the Python task entry `scimesh/worker/task.py`. E2E: `make smoke-two-worker` passes 4/4 shards with two agents. | | CTX-07 Distributed workload protocol | Implemented | Versioned Python contract models, registry, strict plan validation, and deterministic reduction ordering are in `scimesh/distributed/`. | | CTX-08 Distributed similarity-search | Implemented | Python planner resolves `query_id` once, creates deterministic shard plans, worker adapter emits exact partial top-k CSVs/metrics, and reducer matches the local reference. | | CTX-09 Reducer and final-result API | Implemented | Atomic `reducing` claim, deterministic coordinator-side top-k reducer, sanitized reducer failure, final artifact persistence, `result_uri`, and final CSV download. | diff --git a/coordinator/cmd/worker-agent/main.go b/coordinator/cmd/worker-agent/main.go index c31ed4f..225342d 100644 --- a/coordinator/cmd/worker-agent/main.go +++ b/coordinator/cmd/worker-agent/main.go @@ -16,7 +16,13 @@ func main() { os.Exit(2) } logger := slog.New(slog.NewTextHandler(os.Stderr, nil)) - client := agent.NewClient(config.CoordinatorURL, config.Token, config.RequestTimeout) + tokens := agent.NewTokenProvider( + config.WorkerKey, + config.UserserviceURL, + config.Token, + config.RequestTimeout, + ) + client := agent.NewClient(config.CoordinatorURL, tokens, config.RequestTimeout) runner := agent.NewTaskRunner(config.TaskRunner) daemon := agent.NewDaemon(config, client, runner, logger) if err := daemon.RunForever(); err != nil { diff --git a/coordinator/internal/agent/auth.go b/coordinator/internal/agent/auth.go new file mode 100644 index 0000000..029827b --- /dev/null +++ b/coordinator/internal/agent/auth.go @@ -0,0 +1,120 @@ +package agent + +import ( + "bytes" + "encoding/json" + "fmt" + "io" + "net/http" + "strings" + "sync" + "time" +) + +// TokenProvider supplies the current bearer token. A static token is served +// forever; a worker key is exchanged at the userservice for short-lived JWTs +// and refreshed before they expire (mirroring the former Python worker). +type TokenProvider interface { + Token() (string, error) + Refresh() error +} + +// StaticToken serves a fixed token forever; empty means no Authorization. +type StaticToken struct{ token string } + +func (s *StaticToken) Token() (string, error) { return s.token, nil } +func (s *StaticToken) Refresh() error { return nil } + +// WorkerKeyToken exchanges a long-lived worker key for short-lived JWTs. +type WorkerKeyToken struct { + userserviceURL string + workerKey string + timeout time.Duration + leeway float64 + mu sync.Mutex + token string + refreshAt time.Time +} + +func NewWorkerKeyToken(userserviceURL, workerKey string, timeout time.Duration) *WorkerKeyToken { + return &WorkerKeyToken{ + userserviceURL: strings.TrimRight(userserviceURL, "/"), + workerKey: workerKey, + timeout: timeout, + leeway: 0.2, + } +} + +// Token returns the current token, exchanging first when missing or stale. +func (p *WorkerKeyToken) Token() (string, error) { + p.mu.Lock() + defer p.mu.Unlock() + if p.token == "" || time.Now().After(p.refreshAt) { + if err := p.exchangeLocked(); err != nil { + return "", err + } + } + return p.token, nil +} + +// Refresh forces an immediate exchange. +func (p *WorkerKeyToken) Refresh() error { + p.mu.Lock() + defer p.mu.Unlock() + return p.exchangeLocked() +} + +func (p *WorkerKeyToken) exchangeLocked() error { + payload, err := json.Marshal(map[string]string{"key": p.workerKey}) + if err != nil { + return err + } + request, err := http.NewRequest(http.MethodPost, p.userserviceURL+"/worker-tokens/exchange", bytes.NewReader(payload)) + if err != nil { + return err + } + request.Header.Set("Content-Type", "application/json") + client := &http.Client{Timeout: p.timeout} + response, err := client.Do(request) + if err != nil { + return fmt.Errorf("worker key exchange request failed") + } + defer response.Body.Close() + if response.StatusCode != http.StatusOK { + return fmt.Errorf("worker key exchange rejected with status %d", response.StatusCode) + } + raw, err := io.ReadAll(io.LimitReader(response.Body, 1<<20)) + if err != nil { + return fmt.Errorf("worker key exchange request failed") + } + var data map[string]any + if err := json.Unmarshal(raw, &data); err != nil { + return fmt.Errorf("worker key exchange response is invalid") + } + token, _ := data["token"].(string) + if token == "" { + return fmt.Errorf("worker key exchange response is missing a token") + } + var ttl time.Duration + switch value := data["expires_in"].(type) { + case float64: + ttl = time.Duration(value * float64(time.Second)) + case int: + ttl = time.Duration(value) * time.Second + } + p.token = token + p.refreshAt = time.Time{} + if ttl > 0 { + p.refreshAt = time.Now().Add(time.Duration(float64(ttl) * (1.0 - p.leeway))) + } + return nil +} + +// NewTokenProvider picks the strategy: a worker key (with userservice) wins +// over a static bearer token. +func NewTokenProvider(workerKey, userserviceURL, bearerToken string, timeout time.Duration) TokenProvider { + if workerKey != "" && userserviceURL != "" { + return NewWorkerKeyToken(userserviceURL, workerKey, timeout) + } + return &StaticToken{token: bearerToken} +} diff --git a/coordinator/internal/agent/auth_test.go b/coordinator/internal/agent/auth_test.go new file mode 100644 index 0000000..845e942 --- /dev/null +++ b/coordinator/internal/agent/auth_test.go @@ -0,0 +1,102 @@ +package agent + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "sync/atomic" + "testing" + "time" +) + +func TestWorkerKeyTokenExchangesAndCaches(t *testing.T) { + var exchanges atomic.Int64 + var server *httptest.Server + server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/worker-tokens/exchange" { + http.NotFound(w, r) + return + } + var payload map[string]string + if err := json.NewDecoder(r.Body).Decode(&payload); err != nil || payload["key"] != "scimesh_wk_live_x" { + http.Error(w, "bad key", http.StatusUnauthorized) + return + } + exchanges.Add(1) + writeJSON(w, http.StatusOK, map[string]any{ + "token": "jwt-1", + "expires_in": 100, + }) + })) + defer server.Close() + + provider := NewWorkerKeyToken(server.URL, "scimesh_wk_live_x", 5*time.Second) + token, err := provider.Token() + if err != nil || token != "jwt-1" { + t.Fatalf("token = %q, err = %v", token, err) + } + // The second call within the TTL reuses the cache. + again, err := provider.Token() + if err != nil || again != "jwt-1" { + t.Fatalf("cached token = %q, err = %v", again, err) + } + if exchanges.Load() != 1 { + t.Errorf("exchanges = %d, want 1", exchanges.Load()) + } + // An explicit refresh re-exchanges. + if err := provider.Refresh(); err != nil { + t.Fatalf("Refresh: %v", err) + } + if exchanges.Load() != 2 { + t.Errorf("exchanges after refresh = %d, want 2", exchanges.Load()) + } +} + +func TestWorkerKeyTokenRejectsBadKey(t *testing.T) { + var server *httptest.Server + server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusUnauthorized) + })) + defer server.Close() + + provider := NewWorkerKeyToken(server.URL, "bad", 5*time.Second) + if _, err := provider.Token(); err == nil { + t.Error("expected exchange failure for a rejected key") + } +} + +func TestNewTokenProviderSelectsStrategy(t *testing.T) { + if _, ok := NewTokenProvider("", "", "static", time.Second).(*StaticToken); !ok { + t.Error("expected a static token provider") + } + if _, ok := NewTokenProvider("key", "http://users", "", time.Second).(*WorkerKeyToken); !ok { + t.Error("expected a worker-key provider") + } +} + +func TestClientRefreshesTokenOnceOn401(t *testing.T) { + var attempts atomic.Int64 + var server *httptest.Server + server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + attempts.Add(1) + if attempts.Load() == 1 { + w.WriteHeader(http.StatusUnauthorized) + return + } + writeJSON(w, http.StatusOK, map[string]any{}) + })) + defer server.Close() + + provider := &StaticToken{token: "t"} + client := NewClient(server.URL, provider, 5*time.Second) + status, _, err := client.requestJSON(http.MethodGet, "/ok", map[string]any{}) + if err != nil { + t.Fatalf("request: %v", err) + } + if status != http.StatusOK { + t.Errorf("status = %d", status) + } + if attempts.Load() != 2 { + t.Errorf("attempts = %d, want 2 (401 then retry)", attempts.Load()) + } +} diff --git a/coordinator/internal/agent/client.go b/coordinator/internal/agent/client.go index c171bbf..439c428 100644 --- a/coordinator/internal/agent/client.go +++ b/coordinator/internal/agent/client.go @@ -31,23 +31,25 @@ type ConflictError struct{ msg string } func (e *ConflictError) Error() string { return e.msg } -// Client speaks the v1 worker contract over HTTP with a static bearer token. +// Client speaks the v1 worker contract over HTTP with a token provider. // // API calls never follow redirects (a redirect is a contract violation); the // artifact download follows redirects but strips the Authorization header on // cross-origin hops, matching the Python worker's SameOriginAuthRedirectHandler. +// A 401 response refreshes the token exactly once and retries, so a lapsed JWT +// does not fail an in-flight task. type Client struct { baseURL string - token string + tokens TokenProvider timeout time.Duration apiClient *http.Client dlClient *http.Client } -func NewClient(baseURL, token string, timeout time.Duration) *Client { +func NewClient(baseURL string, tokens TokenProvider, timeout time.Duration) *Client { return &Client{ baseURL: strings.TrimRight(baseURL, "/"), - token: token, + tokens: tokens, timeout: timeout, apiClient: &http.Client{ Timeout: timeout, @@ -74,11 +76,19 @@ func origin(u *url.URL) string { return u.Scheme + "://" + u.Host } -func (c *Client) authHeaders() map[string]string { - if c.token == "" { - return map[string]string{} +func (c *Client) authHeaders() (map[string]string, error) { + token, err := c.tokens.Token() + if err != nil { + return nil, &CoordinatorError{msg: "token refresh failed: " + err.Error()} } - return map[string]string{"Authorization": "Bearer " + c.token} + if token == "" { + return map[string]string{}, nil + } + return map[string]string{"Authorization": "Bearer " + token}, nil +} + +func (c *Client) refreshAndRetry() bool { + return c.tokens.Refresh() == nil } // Register advertises the worker and returns its identity and heartbeat policy. @@ -209,7 +219,11 @@ func (c *Client) Download(uri, destination string) (string, error) { if err != nil { return "", err } - for name, value := range c.authHeaders() { + headers, err := c.authHeaders() + if err != nil { + return "", err + } + for name, value := range headers { request.Header.Set(name, value) } response, err := c.dlClient.Do(request) @@ -217,6 +231,9 @@ func (c *Client) Download(uri, destination string) (string, error) { return "", &TransientError{msg: "input download failed"} } defer response.Body.Close() + if response.StatusCode == http.StatusUnauthorized && c.refreshAndRetry() { + return c.Download(uri, destination) + } if response.StatusCode != http.StatusOK { return "", &CoordinatorError{msg: fmt.Sprintf("input download rejected with status %d", response.StatusCode)} } @@ -271,7 +288,12 @@ func (c *Client) Upload(task *Task, workerID string, path, contentType string) ( request.Header.Set("Content-Type", contentType) request.Header.Set("X-Worker-ID", workerID) request.Header.Set("X-Task-Attempt", strconv.Itoa(task.Attempt)) - for name, value := range c.authHeaders() { + headers, err := c.authHeaders() + if err != nil { + file.Close() + return nil, err + } + for name, value := range headers { request.Header.Set(name, value) } response, err := c.apiClient.Do(request) @@ -287,6 +309,9 @@ func (c *Client) Upload(task *Task, workerID string, path, contentType string) ( if response.StatusCode == http.StatusConflict { return nil, &ConflictError{msg: "artifact upload rejected because the task lease was lost"} } + if response.StatusCode == http.StatusUnauthorized && c.refreshAndRetry() { + return c.Upload(task, workerID, path, contentType) + } if response.StatusCode != http.StatusOK { return nil, &CoordinatorError{msg: fmt.Sprintf("artifact upload rejected with status %d", response.StatusCode)} } @@ -314,7 +339,11 @@ func (c *Client) requestJSON(method, path string, payload any) (int, map[string] return 0, nil, err } request.Header.Set("Content-Type", "application/json") - for name, value := range c.authHeaders() { + headers, err := c.authHeaders() + if err != nil { + return 0, nil, err + } + for name, value := range headers { request.Header.Set(name, value) } response, err := c.apiClient.Do(request) @@ -326,6 +355,9 @@ func (c *Client) requestJSON(method, path string, payload any) (int, map[string] if err != nil { return 0, nil, &TransientError{msg: "coordinator request interrupted"} } + if response.StatusCode == http.StatusUnauthorized && c.refreshAndRetry() { + return c.requestJSON(method, path, payload) + } if response.StatusCode >= 500 { return response.StatusCode, nil, &TransientError{msg: fmt.Sprintf("coordinator returned %d", response.StatusCode)} } diff --git a/coordinator/internal/agent/client_test.go b/coordinator/internal/agent/client_test.go index 61dab58..5f7fee3 100644 --- a/coordinator/internal/agent/client_test.go +++ b/coordinator/internal/agent/client_test.go @@ -15,7 +15,7 @@ import ( func newTestClient(t *testing.T, server *httptest.Server) *Client { t.Helper() - return NewClient(server.URL, "test-token", 5*time.Second) + return NewClient(server.URL, &StaticToken{token: "test-token"}, 5*time.Second) } func TestClientRegisterClaimHeartbeat(t *testing.T) { diff --git a/coordinator/internal/agent/config.go b/coordinator/internal/agent/config.go index c8782c8..a01a3a4 100644 --- a/coordinator/internal/agent/config.go +++ b/coordinator/internal/agent/config.go @@ -10,10 +10,13 @@ import ( "time" ) -// Config is read only from the environment, mirroring the Python worker. +// Config is read only from the environment, mirroring the former Python +// worker's configuration surface. type Config struct { CoordinatorURL string Token string + WorkerKey string + UserserviceURL string WorkerName string WorkerID string // set after registration; overridable for tests WorkDir string @@ -22,6 +25,7 @@ type Config struct { PollInterval time.Duration RequestTimeout time.Duration Heartbeat time.Duration + CleanupAfter time.Duration // 0 = keep attempt directories Capabilities []string TaskRunner []string // command + args; defaults to python -m scimesh.worker.task MaxTasks int // 0 = unlimited @@ -86,6 +90,10 @@ func LoadConfig() (*Config, error) { if err != nil { return nil, err } + cleanup, err := durationEnv("CLEANUP_AFTER_SECONDS", 0) + if err != nil { + return nil, err + } capabilities, err := envList("CAPABILITIES") if err != nil { return nil, err @@ -120,6 +128,8 @@ func LoadConfig() (*Config, error) { return &Config{ CoordinatorURL: strings.TrimRight(url, "/"), Token: os.Getenv("WORKER_AUTH_TOKEN"), + WorkerKey: os.Getenv("WORKER_KEY"), + UserserviceURL: strings.TrimRight(os.Getenv("USERSERVICE_URL"), "/"), WorkerName: name, WorkerID: os.Getenv("WORKER_ID"), WorkDir: absWorkDir, @@ -128,6 +138,7 @@ func LoadConfig() (*Config, error) { PollInterval: poll, RequestTimeout: timeout, Heartbeat: heartbeat, + CleanupAfter: cleanup, Capabilities: capabilities, TaskRunner: runner, MaxTasks: maxTasks, @@ -141,8 +152,8 @@ func durationEnv(name string, fallback time.Duration) (time.Duration, error) { return fallback, nil } parsed, err := time.ParseDuration(raw) - if err != nil || parsed <= 0 { - return 0, fmt.Errorf("%s must be a positive duration", name) + if err != nil || parsed < 0 { + return 0, fmt.Errorf("%s must be a non-negative duration", name) } return parsed, nil } diff --git a/coordinator/internal/agent/daemon.go b/coordinator/internal/agent/daemon.go index 5872a66..c63a40c 100644 --- a/coordinator/internal/agent/daemon.go +++ b/coordinator/internal/agent/daemon.go @@ -45,6 +45,7 @@ func (d *Daemon) RunForever() error { return err } } + d.cleanupExpiredDirectories() outcome, err := d.runOnce() if err != nil { failures++ @@ -109,6 +110,41 @@ func (d *Daemon) workerIDOrEmpty() string { return d.workerID } +// cleanupExpiredDirectories removes task attempt directories older than the +// configured retention, mirroring the former Python worker's cleanup. +func (d *Daemon) cleanupExpiredDirectories() { + if d.config.CleanupAfter <= 0 { + return + } + cutoff := time.Now().Add(-d.config.CleanupAfter) + tasks, err := os.ReadDir(d.config.WorkDir) + if err != nil { + return + } + for _, taskEntry := range tasks { + if !taskEntry.IsDir() { + continue + } + taskDir := filepath.Join(d.config.WorkDir, taskEntry.Name()) + attempts, err := os.ReadDir(taskDir) + if err != nil { + continue + } + for _, attemptEntry := range attempts { + if !attemptEntry.IsDir() { + continue + } + info, err := attemptEntry.Info() + if err == nil && info.ModTime().Before(cutoff) { + _ = os.RemoveAll(filepath.Join(taskDir, attemptEntry.Name())) + } + } + if entries, err := os.ReadDir(taskDir); err == nil && len(entries) == 0 { + _ = os.Remove(taskDir) + } + } +} + func (d *Daemon) runOnce() (Outcome, error) { workerID := d.workerIDOrEmpty() if workerID == "" { diff --git a/coordinator/internal/agent/daemon_test.go b/coordinator/internal/agent/daemon_test.go index 700c88d..ef0a252 100644 --- a/coordinator/internal/agent/daemon_test.go +++ b/coordinator/internal/agent/daemon_test.go @@ -139,7 +139,7 @@ func testDaemon(t *testing.T, fake *fakeCoordinator, script string) *Daemon { Capabilities: []string{"similarity-search"}, TaskRunner: []string{script}, } - client := NewClient(fake.server.URL, "test-token", 5*time.Second) + client := NewClient(fake.server.URL, &StaticToken{token: "test-token"}, 5*time.Second) runner := NewTaskRunner(config.TaskRunner) logger := slog.New(slog.NewTextHandler(io.Discard, nil)) daemon := NewDaemon(config, client, runner, logger) diff --git a/coordinator/scripts/demo-ui.sh b/coordinator/scripts/demo-ui.sh index 8705b01..cc03910 100755 --- a/coordinator/scripts/demo-ui.sh +++ b/coordinator/scripts/demo-ui.sh @@ -28,7 +28,7 @@ case "$demo_dir" in /*) ;; *) demo_dir="$coordinator_dir/$demo_dir" ;; esac -worker_bin=${SCIMESH_WORKER_BIN:-"$repo_dir/.venv/bin/scimesh-worker"} +agent_bin=${SCIMESH_AGENT_BIN:-"$coordinator_dir/bin/worker-agent"} pid_file="$demo_dir/workers.pids" logs_dir="$demo_dir/logs" @@ -79,13 +79,23 @@ stop_workers() { [[ "$pid" =~ ^[0-9]+$ ]] || continue command_line=$(ps -p "$pid" -o args= 2>/dev/null || true) # Never kill a recycled PID or a worker launched outside this demo. - if [[ "$command_line" == *"$demo_dir/worker-"* ]]; then + if [[ "$command_line" == *"worker-agent"* ]]; then kill "$pid" 2>/dev/null || true fi done < "$pid_file" rm -f "$pid_file" } +build_agent() { + if [[ ! -x "$agent_bin" ]]; then + echo "Building the Go worker agent..." >&2 + make -C "$coordinator_dir" agent >&2 || { + echo "Failed to build the Go worker agent." >&2 + exit 2 + } + fi +} + wait_for_coordinator() { local attempt=0 until curl --fail --silent --show-error "http://localhost:$coordinator_port/health" >/dev/null; do @@ -145,11 +155,7 @@ start() { echo "DEMO_WORKERS must be a positive integer (got $workers)." >&2 exit 2 fi - if [[ ! -x "$worker_bin" ]]; then - echo "Reference worker not found: $worker_bin" >&2 - echo "Create it first from the repository root: python3 -m venv .venv && .venv/bin/pip install -e '.[dev]'" >&2 - exit 2 - fi + build_agent command -v docker >/dev/null || { echo "Docker is required." >&2; exit 2; } command -v curl >/dev/null || { echo "curl is required." >&2; exit 2; } @@ -162,14 +168,21 @@ start() { wait_for_userservice : > "$pid_file" + task_runner="[\"$repo_dir/.venv/bin/python\",\"-m\",\"scimesh.worker.task\"]" for index in $(seq 1 "$workers"); do work_dir="$demo_dir/worker-$index" mkdir -p "$work_dir" - SCIMESH_COORDINATOR_URL="http://localhost:$coordinator_port" \ - SCIMESH_BEARER_TOKEN="$worker_token" \ - "$worker_bin" \ - --worker-name "demo-worker-$index" \ - --work-dir "$work_dir" \ + COORDINATOR_URL="http://localhost:$coordinator_port" \ + WORKER_AUTH_TOKEN="$worker_token" \ + WORKER_NAME="demo-worker-$index" \ + WORK_DIR="$work_dir" \ + CPU_COUNT=1 \ + MEMORY_MB=1024 \ + POLL_INTERVAL=0.5s \ + REQUEST_TIMEOUT=15s \ + HEARTBEAT_INTERVAL=15s \ + TASK_RUNNER="$task_runner" \ + "$agent_bin" \ >"$logs_dir/worker-$index.log" 2>&1 & echo "$!" >> "$pid_file" done @@ -184,7 +197,7 @@ SciMesh manual demo is ready. Userservice: http://localhost:$userservice_port Grafana: http://localhost:$grafana_port (anonymous view; admin/${GRAFANA_PASSWORD:-admin} to edit) Prometheus: http://localhost:$prometheus_port - Workers: $workers local reference workers + Workers: $workers Go worker agents (Python task execution) Sign in with the admin above, or register a new account from the login page. The admin sees every job; a plain user sees only their own. Upload a small diff --git a/docs/sdk-handoff.md b/docs/sdk-handoff.md index b639f20..b9d3c2a 100644 --- a/docs/sdk-handoff.md +++ b/docs/sdk-handoff.md @@ -28,7 +28,7 @@ resolves `query_id` per task and rejects plan-time `max_rows`. **MkDocs site (2026-08-02):** the standalone documentation site lives in `mkdocs/` (`docs_dir: mkdocs`) and does not use the project's `docs/` directory. It contains guides (`mkdocs/sdk/`: overview, authoring-workloads, cli, worker-integration), the full auto-generated API reference for all `scimesh.sdk` modules (`mkdocs/api/`, mkdocstrings `::: scimesh.sdk.` — set `show_if_no_docstring: true`), and the writing rules (`mkdocs/approach.md`). `make docs` builds it; the UI serves it at `/ui/docs/`. All public SDK members now carry Google-style docstrings. -**Go worker agent prototype (2026-08-02):** `coordinator/internal/agent/` (config, models, client, sanitize, taskrunner, daemon) + `coordinator/cmd/worker-agent`, built with `make agent`. It mirrors the Python worker's v1 lifecycle; per-task SDK execution happens in a Python subprocess (`scimesh/worker/task.py`: exits 0 on success, 3 permanent, 1 retryable). Default `TASK_RUNNER` is `python -m scimesh.worker.task` — set it to the venv python in source checkouts. Verified E2E against the demo coordinator. Open items: JWT refresh, resource slots/limits, attempt-dir cleanup, protocol-v2 features. +**Go worker agent (2026-08-02):** the Python worker daemon was removed. `coordinator/internal/agent/` + `cmd/worker-agent` (build: `make agent`) now owns the full lifecycle: register/claim/heartbeat/download(checksum)/spawn-task/upload/submit/fail, static bearer or worker-key JWT auth with 401 refresh, attempt-dir cleanup (`CLEANUP_AFTER_SECONDS`), backoff, idle/max-tasks exit. `scimesh/worker/` keeps only the per-task Python execution: `task.py` (exit 0/3/1), `runners.py` (SDK bridge, `SCIMESH_WORKLOAD_ALLOWLIST` discovery), `models.py` (claim payload). The `scimesh-worker` console script and the daemon/auth/transport modules are gone. Demo (`make demo-ui`) and smoke (`make smoke-two-worker`) run the Go agent; smoke passes 4/4 with two agents. `coordinator/internal/agent/` (config, models, client, sanitize, taskrunner, daemon) + `coordinator/cmd/worker-agent`, built with `make agent`. It mirrors the Python worker's v1 lifecycle; per-task SDK execution happens in a Python subprocess (`scimesh/worker/task.py`: exits 0 on success, 3 permanent, 1 retryable). Default `TASK_RUNNER` is `python -m scimesh.worker.task` — set it to the venv python in source checkouts. Verified E2E against the demo coordinator. Open items: JWT refresh, resource slots/limits, attempt-dir cleanup, protocol-v2 features. **Authoring scaffold (2026-08-01):** `scimesh/sdk/batch.py` adds `MapReduceWorkload` — the primary authoring surface for `core-batch-v1`. A diff --git a/mkdocs/sdk/cli.md b/mkdocs/sdk/cli.md index f8eb6e5..5083284 100644 --- a/mkdocs/sdk/cli.md +++ b/mkdocs/sdk/cli.md @@ -72,5 +72,5 @@ make workloads-export | `SCIMESH_WORKLOAD_ALLOWLIST` | JSON array of `{distribution, name, version, digest}` entries; discovery loads the matching installed `scimesh.workloads` entry points | | `SCIMESH_CAPABILITIES` | Comma-separated capabilities the worker advertises (default `similarity-search,similarity_search`) | -Both variables are read by the worker (`scimesh-worker`) and the workload -CLI. +Both variables are read by the Go worker agent's task subprocess and the +workload CLI. diff --git a/mkdocs/sdk/worker-integration.md b/mkdocs/sdk/worker-integration.md index 99de60d..b202159 100644 --- a/mkdocs/sdk/worker-integration.md +++ b/mkdocs/sdk/worker-integration.md @@ -1,73 +1,77 @@ # Worker integration -The Worker Agent (`scimesh-worker`) is a coordinator client, never a -database client. It polls the coordinator over HTTP, executes SDK-built -workloads, and uploads partial results through the coordinator — results -never carry `file://` or `worker://` URIs, and failures go to `/failure`. +The **Go worker agent** (`coordinator/cmd/worker-agent`, built with +`make agent`) is the worker: a coordinator client, never a database client. +It polls the coordinator over HTTP, executes SDK workloads, and uploads +partial results through the coordinator — results never carry `file://` or +`worker://` URIs, and failures go to `/failure`. The same scientific handlers run in three places: the local CLI cores, the -`LocalCoreBatchExecutor` conformance harness, and the worker — because the -worker executes the workload's own SDK runner. +`LocalCoreBatchExecutor` conformance harness, and the agent — because the +agent spawns the workload's own SDK runner in a Python subprocess per task. ## Claim lifecycle ```text register -> claim (one task) -> download input + verify sha256 - -> run via SDK bridge -> upload partial -> submit result + -> spawn python -m scimesh.worker.task -> upload partial -> submit ``` -- **Register**: the worker advertises its capabilities (`similarity-search` - by default; extend with `SCIMESH_CAPABILITIES`). +- **Register**: the agent advertises its capabilities + (`CAPABILITIES`, default `similarity-search,similarity_search`). - **Claim**: atomic lease of one task; `204` means idle. - **Download**: the input is streamed and its SHA-256 verified; the bearer token is stripped on cross-origin redirects. -- **Heartbeat**: a background thread renews the lease from the returned +- **Heartbeat**: a background goroutine renews the lease from the returned deadline at less than half the remaining TTL. -- **Upload**: the partial CSV is streamed to the coordinator with - `X-Worker-ID` / `X-Task-Attempt` headers, then the completion is submitted - referencing the coordinator-owned artifact id. +- **Task execution**: a Python subprocess (`TASK_RUNNER`, default + `python -m scimesh.worker.task`) runs the SDK workload; exit 0 writes the + result manifest, exit 3 means permanent failure, exit 1 retryable. +- **Upload**: the partial CSV is streamed with `X-Worker-ID` / + `X-Task-Attempt` headers, then completion references the + coordinator-owned artifact id. - **Failure**: sanitized `error_code` + message (≤300 chars, no local - paths, no tracebacks); transient transport errors are retried. + paths); transient errors are retried with backoff; lost leases stop + quietly. -## The SDK execution bridge +## Authentication -`scimesh/worker/runners.py` is workload-generic. For a claimed task it: +- `WORKER_AUTH_TOKEN` — a static bearer token (the shared service token). +- `WORKER_KEY` + `USERSERVICE_URL` — a long-lived worker key exchanged at + the userservice for short-lived JWTs; the agent refreshes them before + expiry and retries once after a 401. -1. normalizes the workload name (underscores → hyphens) and looks up the - loaded definition; -2. runs compatibility negotiation against a runtime derived from the loaded - definitions (capabilities + pinned environment digests) and the worker - inventory (CPU/memory from configuration); -3. verifies the workload's map stage fits the v1 contract — a single - `input` port and a single `partial` output — otherwise it fails closed - with a clear message; -4. imports the downloaded input into a content-addressed local store; -5. builds a digest-pinned `TaskSpec` (package/manifest/environment digests, - trust mode, negotiated features, stage resources and execution profile); -6. reserves resources through `ResourcePool` and runs the workload's own - `Runner` with a `LocalTaskContext` (scoped catalog/sink, provenance, - cancellation flag); -7. validates the returned `OutputManifest` (task key, provenance, sealed - vs. declared artifacts, byte budget) and returns the sealed partial for - upload. +## Configuration -Scientific policy lives in the workload: `query_id` resolution, parameter -validation, and `max_rows` rejection are all handled by the workload's own -hooks — the bridge passes task parameters through unchanged. +| Variable | Meaning | +| --- | --- | +| `COORDINATOR_URL` | Coordinator base URL (required) | +| `WORKER_AUTH_TOKEN` | Static bearer token (when no worker key) | +| `WORKER_KEY` / `USERSERVICE_URL` | Worker-key authentication | +| `WORK_DIR` | Attempt directory root (default `./scimesh-agent-data`) | +| `WORKER_NAME` | Registered name (default: hostname) | +| `WORKER_ID` | Fixed identity override (tests) | +| `CPU_COUNT` / `MEMORY_MB` | Advertised capacity | +| `POLL_INTERVAL` / `REQUEST_TIMEOUT` / `HEARTBEAT_INTERVAL` | Timings | +| `CLEANUP_AFTER_SECONDS` | Delete attempt dirs older than this | +| `CAPABILITIES` | JSON array of advertised capabilities | +| `TASK_RUNNER` | JSON command array for the task subprocess | +| `MAX_TASKS` / `EXIT_WHEN_IDLE` | Lifecycle limits | ## Loading workloads -The worker loads workloads from `SCIMESH_WORKLOAD_ALLOWLIST` (a JSON array -of `{distribution, name, version, digest}` entries matched against installed -`scimesh.workloads` entry points). Discovery measures the installed package -before and after importing and fails transactionally on any mismatch. When -no allowlist is configured, the worker falls back to the built-in -`similarity-search`. +The task subprocess loads workloads from `SCIMESH_WORKLOAD_ALLOWLIST` (a +JSON array of `{distribution, name, version, digest}` entries matched +against installed `scimesh.workloads` entry points) or falls back to the +built-in `similarity-search`. Discovery measures the installed package +before and after importing and fails transactionally on any mismatch. ```bash -SCIMESH_WORKLOAD_ALLOWLIST='[{"distribution": "scimesh", - "name": "descriptor-batch", "version": "1.0.0", - "digest": "sha256:..."}]' scimesh-worker --coordinator-url https://... +make agent +COORDINATOR_URL=https://coordinator.example \ +WORKER_AUTH_TOKEN=... \ +TASK_RUNNER='["/opt/scimesh/.venv/bin/python","-m","scimesh.worker.task"]' \ +./coordinator/bin/worker-agent ``` ## v1 contract limits @@ -76,7 +80,7 @@ The coordinator protocol v1 persists flat one-input/one-result tasks. Until a versioned protocol rollout: - map stages with more than one input port (for example - `similarity-graph`'s block pairs) are **rejected by the bridge** — the - coordinator does not create such tasks anyway; + `similarity-graph`'s block pairs) are **rejected by the task runner** — + the coordinator does not create such tasks anyway; - `max_rows` is a plan-time option and is rejected per task; - workloads beyond the allowlisted set are rejected as unsupported. diff --git a/pyproject.toml b/pyproject.toml index 6ebc843..330eaa4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -20,7 +20,6 @@ dev = [ [project.scripts] scimesh = "scimesh.cli:main" -scimesh-worker = "scimesh.worker.cli:main" [project.entry-points."scimesh.workloads"] "similarity-search@1.0.0" = "scimesh.workloads.search:workload_definition" diff --git a/scimesh/worker/__init__.py b/scimesh/worker/__init__.py index 6a351cb..732abb1 100644 --- a/scimesh/worker/__init__.py +++ b/scimesh/worker/__init__.py @@ -1,5 +1,8 @@ -"""Worker daemon for executing coordinator-assigned SciMesh workloads.""" +"""Per-task execution for the Go worker agent. -from .daemon import WorkerDaemon - -__all__ = ["WorkerDaemon"] +This package contains the Python side of a claimed task: the wire value +objects (``models``), the SDK execution bridge (``runners``), and the +command-line task entry (``task``). The agent lifecycle itself lives in the +Go worker agent (``coordinator/internal/agent``); this package is only ever +invoked by it, one process per task. +""" diff --git a/scimesh/worker/artifacts.py b/scimesh/worker/artifacts.py deleted file mode 100644 index 3f2c7f8..0000000 --- a/scimesh/worker/artifacts.py +++ /dev/null @@ -1,141 +0,0 @@ -"""Input/output artifact transport kept separate from the daemon state machine.""" - -from __future__ import annotations - -import hashlib -import http.client -import json -from pathlib import Path -from typing import Protocol -from urllib.error import HTTPError -from urllib.parse import quote, urljoin, urlsplit -from urllib.request import Request, build_opener - -from .auth import StaticTokenProvider, TokenProvider -from .coordinator import CoordinatorConflictError -from .models import ClaimedTask, ProducedArtifact, UploadedArtifact -from .transport import SameOriginAuthRedirectHandler, origin - -# Compatibility aliases for focused transport tests. -_SameOriginAuthRedirectHandler = SameOriginAuthRedirectHandler -_origin = origin - -class ArtifactClient(Protocol): - def download(self, uri: str, destination: Path) -> None: ... - - def upload( - self, task: ClaimedTask, worker_id: str, artifact: ProducedArtifact - ) -> UploadedArtifact: ... - - -class HttpArtifactClient: - """Transfers artifacts through the coordinator without leaking credentials.""" - - def __init__( - self, - coordinator_url: str, - timeout: float, - bearer_token: str | None = None, - *, - token_provider: TokenProvider | None = None, - ) -> None: - self.coordinator_url = coordinator_url.rstrip("/") - self.timeout = timeout - self._tokens: TokenProvider = token_provider or StaticTokenProvider(bearer_token) - self.coordinator_origin = origin(coordinator_url) - self._opener = build_opener(SameOriginAuthRedirectHandler(self.coordinator_origin)) - - @property - def bearer_token(self) -> str | None: - return self._tokens.token() - - def download(self, uri: str, destination: Path) -> None: - destination.parent.mkdir(parents=True, exist_ok=True) - resolved_uri = urljoin(f"{self.coordinator_url}/", uri) - self._download_once(resolved_uri, destination, allow_refresh=True) - - def _download_once(self, resolved_uri: str, destination: Path, *, allow_refresh: bool) -> None: - request = Request(resolved_uri, headers=self._auth_headers_for(resolved_uri)) - try: - with self._opener.open(request, timeout=self.timeout) as response, destination.open("wb") as target: - while chunk := response.read(1024 * 1024): - target.write(chunk) - except HTTPError as error: - # Refresh an expired token and retry once, mirroring the coordinator - # client, so a token that lapses mid-task does not fail the download. - if error.code == 401 and allow_refresh: - self._tokens.refresh() - self._download_once(resolved_uri, destination, allow_refresh=False) - return - raise - - def upload( - self, task: ClaimedTask, worker_id: str, artifact: ProducedArtifact - ) -> UploadedArtifact: - return self._upload_once(task, worker_id, artifact, allow_refresh=True) - - def _upload_once( - self, task: ClaimedTask, worker_id: str, artifact: ProducedArtifact, *, allow_refresh: bool - ) -> UploadedArtifact: - """Stream an artifact and require durable coordinator-owned metadata.""" - url = ( - f"{self.coordinator_url}/tasks/{quote(task.task_id, safe='')}/artifacts/" - f"{quote(artifact.path.name, safe='')}" - ) - parsed = urlsplit(url) - if parsed.scheme not in {"http", "https"} or not parsed.hostname: - raise ValueError("coordinator URL must be an absolute HTTP(S) URL") - connection_class = ( - http.client.HTTPSConnection if parsed.scheme == "https" else http.client.HTTPConnection - ) - connection = connection_class(parsed.hostname, parsed.port, timeout=self.timeout) - local_size = artifact.path.stat().st_size - local_sha256 = sha256_file(artifact.path) - try: - path = parsed.path + (f"?{parsed.query}" if parsed.query else "") - connection.putrequest("PUT", path) - connection.putheader("Content-Type", artifact.content_type) - connection.putheader("Content-Length", str(local_size)) - connection.putheader("X-Worker-ID", worker_id) - connection.putheader("X-Task-Attempt", str(task.attempt)) - for name, value in self._auth_headers_for(url).items(): - connection.putheader(name, value) - connection.endheaders() - with artifact.path.open("rb") as source: - while chunk := source.read(1024 * 1024): - connection.send(chunk) - response = connection.getresponse() - body = response.read() - if response.status == 401 and allow_refresh: - # Token lapsed mid-task: refresh and retry the upload once. - self._tokens.refresh() - connection.close() - return self._upload_once(task, worker_id, artifact, allow_refresh=False) - if response.status == 409: - raise CoordinatorConflictError("artifact upload rejected because the task lease was lost") - if response.status != 200: - raise RuntimeError(f"artifact upload rejected with status {response.status}") - try: - response_data = json.loads(body) - uploaded = UploadedArtifact.from_json(response_data) - except (ValueError, json.JSONDecodeError) as error: - raise RuntimeError("artifact upload returned invalid metadata") from error - if uploaded.sha256 != local_sha256 or uploaded.size_bytes != local_size: - raise RuntimeError("artifact upload metadata does not match local artifact") - return uploaded - finally: - connection.close() - - def _auth_headers_for(self, uri: str) -> dict[str, str]: - """Only coordinator-owned URLs receive the coordinator bearer token.""" - token = self._tokens.token() - if token and origin(uri) == self.coordinator_origin: - return {"Authorization": f"Bearer {token}"} - return {} - -def sha256_file(path: Path) -> str: - digest = hashlib.sha256() - with path.open("rb") as source: - for chunk in iter(lambda: source.read(1024 * 1024), b""): - digest.update(chunk) - return digest.hexdigest() diff --git a/scimesh/worker/auth.py b/scimesh/worker/auth.py deleted file mode 100644 index a447867..0000000 --- a/scimesh/worker/auth.py +++ /dev/null @@ -1,128 +0,0 @@ -"""Bearer-token strategies for the worker's coordinator calls. - -A worker authenticates in one of two ways: - -* a *static* token — the shared service token or a directly supplied JWT, fixed - for the life of the process; or -* a *worker key* — a long-lived per-user credential the worker trades for a - short-lived JWT at the userservice, refreshing before that JWT expires. - -Both are exposed through the small ``TokenProvider`` protocol so the HTTP -clients neither know nor care which one is in play. -""" - -from __future__ import annotations - -import json -import time -from typing import Callable, Protocol -from urllib.error import HTTPError, URLError -from urllib.request import Request, build_opener - -from .transport import NoRedirectHandler - - -class TokenExchangeError(RuntimeError): - """The userservice refused or failed to exchange a worker key.""" - - -class TokenProvider(Protocol): - def token(self) -> str | None: - """Return the current bearer token, refreshing it if necessary.""" - - def refresh(self) -> None: - """Force the next token to be re-fetched (e.g. after a 401).""" - - -class StaticTokenProvider: - """Serves a fixed token forever. ``None`` means "send no Authorization".""" - - def __init__(self, token: str | None) -> None: - self._token = token - - def token(self) -> str | None: - return self._token - - def refresh(self) -> None: # noqa: D401 - nothing to refresh - return None - - -class WorkerKeyTokenProvider: - """Exchanges a long-lived worker key for short-lived JWTs and refreshes them. - - The token is cached until roughly ``1 - refresh_leeway`` of its lifetime has - elapsed, so the worker renews ahead of expiry instead of waiting for a 401. - A monotonic clock is injectable to keep tests deterministic. - """ - - def __init__( - self, - userservice_url: str, - worker_key: str, - timeout: float, - *, - refresh_leeway: float = 0.2, - now: Callable[[], float] = time.monotonic, - ) -> None: - self._url = userservice_url.rstrip("/") - self._key = worker_key - self._timeout = timeout - self._leeway = refresh_leeway - self._now = now - self._token: str | None = None - self._refresh_at: float = 0.0 - self._opener = build_opener(NoRedirectHandler()) - - def token(self) -> str: - if self._token is None or self._now() >= self._refresh_at: - self._exchange() - assert self._token is not None # _exchange sets it or raises - return self._token - - def refresh(self) -> None: - self._exchange() - - def _exchange(self) -> None: - request = Request( - f"{self._url}/worker-tokens/exchange", - data=json.dumps({"key": self._key}).encode(), - method="POST", - headers={"Content-Type": "application/json"}, - ) - try: - with self._opener.open(request, timeout=self._timeout) as response: - raw = response.read() - data = json.loads(raw) if raw else {} - except HTTPError as error: - # A revoked or unknown key is a permanent 401; there is nothing the - # worker can do but stop, so surface it rather than retry forever. - raise TokenExchangeError( - f"worker key exchange rejected with status {error.code}" - ) from error - except (URLError, TimeoutError, json.JSONDecodeError) as error: - raise TokenExchangeError("worker key exchange request failed") from error - - token = data.get("token") - if not isinstance(token, str) or not token: - raise TokenExchangeError("worker key exchange response is missing a token") - - expires_in = data.get("expires_in") - ttl = float(expires_in) if isinstance(expires_in, (int, float)) and expires_in > 0 else 0.0 - self._token = token - # Renew once ~(1 - leeway) of the lifetime is gone. An unknown TTL falls - # back to re-exchanging on the next call — correct, just chattier. - self._refresh_at = self._now() + ttl * (1.0 - self._leeway) - - -def provider_from_config( - *, - worker_key: str | None, - userservice_url: str | None, - bearer_token: str | None, - request_timeout: float, -) -> TokenProvider: - """Pick the token strategy: a worker key (exchange mode) wins over a static - bearer token, which in turn wins over no credential at all.""" - if worker_key and userservice_url: - return WorkerKeyTokenProvider(userservice_url, worker_key, request_timeout) - return StaticTokenProvider(bearer_token) diff --git a/scimesh/worker/cli.py b/scimesh/worker/cli.py deleted file mode 100644 index 7237207..0000000 --- a/scimesh/worker/cli.py +++ /dev/null @@ -1,106 +0,0 @@ -"""Console entry point for ``scimesh-worker``.""" - -from __future__ import annotations - -import argparse -import logging -from pathlib import Path - -from .artifacts import HttpArtifactClient -from .auth import provider_from_config -from .config import WorkerConfig -from .coordinator import HttpCoordinatorClient -from .daemon import WorkerDaemon -from .runners import SciMeshRunner - - -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, " - "SCIMESH_MAX_TASKS, SCIMESH_BEARER_TOKEN, SCIMESH_WORKER_KEY, and " - "SCIMESH_USERSERVICE_URL. SCIMESH_WORKER_ID is a legacy/test override." - ), - ) - parser.add_argument("--coordinator-url") - parser.add_argument("--worker-id") - parser.add_argument("--work-dir") - parser.add_argument("--worker-name") - parser.add_argument( - "--worker-key", - help="Long-lived worker key from the web UI; the worker exchanges it for " - "short-lived tokens, binding it to your account. Requires --userservice-url.", - ) - parser.add_argument( - "--userservice-url", - help="Base URL of the userservice that issues tokens for --worker-key", - ) - parser.add_argument("--cpu-count", type=int) - parser.add_argument("--memory-mb", type=int) - parser.add_argument("--poll-interval", type=float) - 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 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: - config = WorkerConfig.from_environment(overrides) - except (TypeError, ValueError) as error: - parser.error(str(error)) - logging.basicConfig( - level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s" - ) - # One shared token strategy backs both clients: a worker key (exchanged and - # refreshed) or a static bearer token, decided by what the config carries. - tokens = provider_from_config( - worker_key=config.worker_key, - userservice_url=config.userservice_url, - bearer_token=config.bearer_token, - request_timeout=config.request_timeout, - ) - client = HttpCoordinatorClient( - config.coordinator_url, config.request_timeout, token_provider=tokens - ) - completed_without_interruption = WorkerDaemon( - config, - client, - HttpArtifactClient( - config.coordinator_url, config.request_timeout, token_provider=tokens - ), - SciMeshRunner.for_worker(config), - ).run_forever() - return 0 if completed_without_interruption else 130 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/scimesh/worker/config.py b/scimesh/worker/config.py deleted file mode 100644 index f92c653..0000000 --- a/scimesh/worker/config.py +++ /dev/null @@ -1,221 +0,0 @@ -"""Configuration parsing for the worker command.""" - -from __future__ import annotations - -from dataclasses import dataclass -from math import isfinite -from pathlib import Path -import os -import socket -from typing import Mapping -from urllib.parse import urlsplit - -from scimesh.sdk.registry import AllowedPackage, workload_allowlist_from_json - - -def _clean_url(value: object | None) -> str | None: - """Normalise an optional URL: drop a blank one, strip a trailing slash.""" - if value is None: - return None - text = str(value).strip() - return text.rstrip("/") or None - - -def _int_value(value: object | None, name: str) -> int | None: - if value is None: - return None - if isinstance(value, bool) or not isinstance(value, (int, float, str)): - raise ValueError(f"{name} must be a number") - return int(value) - - -def _float_value(value: object | None, name: str) -> float | None: - if value is None: - return None - if isinstance(value, bool) or not isinstance(value, (int, float, str)): - raise ValueError(f"{name} must be a number") - return float(value) - - -def _positive_number(value: object, name: str, *, allow_zero: bool = False) -> None: - if ( - isinstance(value, bool) - or not isinstance(value, (int, float)) - or not isfinite(value) - or value < 0 - or (not allow_zero and value == 0) - ): - qualifier = "non-negative" if allow_zero else "positive" - raise ValueError(f"{name} must be {qualifier}") - - -def _capabilities(value: object) -> tuple[str, ...]: - """Parse a comma-separated capability list into unique non-empty names.""" - if value is None: - return ("similarity-search", "similarity_search") - if not isinstance(value, str) or not value.strip(): - raise ValueError("capabilities must be a comma-separated list") - names = tuple( - dict.fromkeys(item.strip() for item in value.split(",") if item.strip()) - ) - if not names: - raise ValueError("capabilities cannot be empty") - return names - - -def _workload_allowlist(value: object) -> tuple[AllowedPackage, ...]: - return workload_allowlist_from_json(value) - - -@dataclass(frozen=True) -class WorkerConfig: - coordinator_url: str - worker_id: str | None - work_dir: Path - worker_name: str = "scimesh-worker" - cpu_count: int = 1 - memory_mb: int | None = None - poll_interval: float = 2.0 - request_timeout: float = 30.0 - heartbeat_interval: float = 15.0 - bearer_token: str | None = None - # A long-lived per-user credential. When set (with userservice_url), the - # worker exchanges it for short-lived JWTs instead of using bearer_token, - # binding the worker to that user's account. - worker_key: str | None = None - userservice_url: 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, so retain the search alias during migration. - capabilities: tuple[str, ...] = ( - "similarity-search", - "similarity_search", - ) - # Optional allowlist of installed SDK workload packages to execute. When - # empty, the worker runs the built-in similarity-search only. Entries are - # ``{distribution, name, version, digest}`` JSON objects matching the - # installed ``scimesh.workloads`` entry points. - workload_allowlist: tuple[AllowedPackage, ...] = () - - def __post_init__(self) -> None: - parsed = urlsplit(self.coordinator_url) - if parsed.scheme not in {"http", "https"} or not parsed.hostname: - raise ValueError("coordinator_url must be an absolute HTTP(S) URL") - if not isinstance(self.worker_name, str) or not self.worker_name.strip(): - raise ValueError("worker_name must be non-empty") - if self.userservice_url is not None: - us = urlsplit(self.userservice_url) - if us.scheme not in {"http", "https"} or not us.hostname: - raise ValueError("userservice_url must be an absolute HTTP(S) URL") - if self.worker_key is not None and not self.userservice_url: - raise ValueError( - "worker_key requires userservice_url (SCIMESH_USERSERVICE_URL)" - ) - if ( - isinstance(self.cpu_count, bool) - or not isinstance(self.cpu_count, int) - or self.cpu_count < 1 - ): - raise ValueError("cpu_count must be positive") - if self.worker_id is not None and not isinstance(self.worker_id, str): - raise ValueError("worker_id must be a string when set") - if self.memory_mb is not None and ( - isinstance(self.memory_mb, bool) - or not isinstance(self.memory_mb, int) - or self.memory_mb < 1 - ): - raise ValueError("memory_mb must be positive when set") - _positive_number(self.poll_interval, "poll_interval") - _positive_number(self.request_timeout, "request_timeout") - _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") - if any( - not isinstance(capability, str) or not capability.strip() - for capability in self.capabilities - ): - raise ValueError("capabilities must contain non-empty names") - if len(self.capabilities) != len(set(self.capabilities)): - raise ValueError("capabilities must be unique") - if any( - not isinstance(package, AllowedPackage) - for package in self.workload_allowlist - ): - raise ValueError("workload_allowlist must contain AllowedPackage values") - # 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( - cls, overrides: Mapping[str, object] | None = None - ) -> "WorkerConfig": - """Build config from environment, allowing typed CLI values to override it.""" - values = overrides or {} - - def value( - name: str, environment: str, default: object | None = None - ) -> object | None: - override = values.get(name) - return override if override is not None else os.getenv(environment, default) - - url = value("coordinator_url", "SCIMESH_COORDINATOR_URL") - if not isinstance(url, str) or not url: - raise ValueError("SCIMESH_COORDINATOR_URL or --coordinator-url is required") - 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") - capabilities = value("capabilities", "SCIMESH_CAPABILITIES") - allowlist = value("workload_allowlist", "SCIMESH_WORKLOAD_ALLOWLIST") - worker_id = value("worker_id", "SCIMESH_WORKER_ID") - work_dir = value("work_dir", "SCIMESH_WORK_DIR", "./scimesh-worker-data") - worker_name = value("worker_name", "SCIMESH_WORKER_NAME", socket.gethostname()) - poll_interval = value("poll_interval", "SCIMESH_POLL_INTERVAL", "2") - request_timeout = value("request_timeout", "SCIMESH_REQUEST_TIMEOUT", "30") - heartbeat_interval = value( - "heartbeat_interval", "SCIMESH_HEARTBEAT_INTERVAL", "15" - ) - bearer_token = value("bearer_token", "SCIMESH_BEARER_TOKEN") - worker_key = value("worker_key", "SCIMESH_WORKER_KEY") - userservice_url = _clean_url( - value("userservice_url", "SCIMESH_USERSERVICE_URL") - ) - return cls( - coordinator_url=url.rstrip("/"), - worker_id=str(worker_id) if worker_id is not None else None, - work_dir=Path(str(work_dir)), - worker_name=str(worker_name), - cpu_count=_int_value(cpu_count, "cpu_count") or 1, - memory_mb=_int_value(memory_mb, "memory_mb"), - poll_interval=_float_value(poll_interval, "poll_interval") or 2.0, - request_timeout=_float_value(request_timeout, "request_timeout") or 30.0, - heartbeat_interval=_float_value(heartbeat_interval, "heartbeat_interval") or 15.0, - bearer_token=str(bearer_token) if bearer_token is not None else None, - worker_key=str(worker_key) if worker_key is not None else None, - userservice_url=userservice_url, - cleanup_after_seconds=_float_value(cleanup, "cleanup_after_seconds"), - max_tasks=_int_value(max_tasks, "max_tasks"), - exit_when_idle=bool(values.get("exit_when_idle", False)), - capabilities=_capabilities(capabilities), - workload_allowlist=_workload_allowlist(allowlist), - ) diff --git a/scimesh/worker/coordinator.py b/scimesh/worker/coordinator.py deleted file mode 100644 index 4367232..0000000 --- a/scimesh/worker/coordinator.py +++ /dev/null @@ -1,149 +0,0 @@ -"""HTTP boundary for the coordinator; the daemon never accesses a database.""" - -from __future__ import annotations - -import json -from typing import Any, Protocol -from urllib.error import HTTPError, URLError -from urllib.request import Request, build_opener - -from .auth import StaticTokenProvider, TokenProvider -from .models import ClaimedTask, RegisteredWorker -from .transport import NoRedirectHandler - - -class CoordinatorError(RuntimeError): - """A non-retriable coordinator response.""" - - -class CoordinatorTransientError(CoordinatorError): - """A timeout, connection error, or 5xx coordinator response.""" - - -class CoordinatorConflictError(CoordinatorError): - """The worker no longer owns the task lease or attempted a conflicting mutation.""" - - -class CoordinatorClient(Protocol): - def register( - self, name: str, capabilities: tuple[str, ...], cpu_count: int, memory_mb: int | None - ) -> RegisteredWorker: ... - - def claim(self, worker_id: str, capabilities: tuple[str, ...]) -> ClaimedTask | None: ... - - def submit(self, task: ClaimedTask, payload: dict[str, Any]) -> None: ... - - def fail(self, task: ClaimedTask, payload: dict[str, Any]) -> None: ... - - def heartbeat(self, task: ClaimedTask, worker_id: str) -> str: ... - - -class HttpCoordinatorClient: - def __init__( - self, - base_url: str, - timeout: float, - bearer_token: str | None = None, - *, - token_provider: TokenProvider | None = None, - ) -> None: - self.base_url = base_url.rstrip("/") - self.timeout = timeout - # A bearer_token argument keeps older call sites working; internally - # everything goes through a provider so refresh is uniform. - self._tokens: TokenProvider = token_provider or StaticTokenProvider(bearer_token) - self._opener = build_opener(NoRedirectHandler()) - - @property - def bearer_token(self) -> str | None: - return self._tokens.token() - - def register( - self, name: str, capabilities: tuple[str, ...], cpu_count: int, memory_mb: int | None - ) -> RegisteredWorker: - payload: dict[str, Any] = { - "name": name, - "capabilities": list(capabilities), - "cpu_count": cpu_count, - } - if memory_mb is not None: - payload["memory_mb"] = memory_mb - status, body = self._request("POST", "/workers/register", payload) - if status != 201: - raise CoordinatorError(f"worker registration rejected with status {status}") - try: - return RegisteredWorker.from_json(body) - except ValueError as error: - raise CoordinatorError("invalid worker registration response") from error - - def claim(self, worker_id: str, capabilities: tuple[str, ...]) -> ClaimedTask | None: - status, body = self._request("POST", "/tasks/claim", { - "worker_id": worker_id, "capabilities": list(capabilities), "max_concurrency": 1, - }) - if status == 204: - return None - if status != 200: - raise CoordinatorError(f"unexpected claim status {status}") - return ClaimedTask.from_json(body) - - def submit(self, task: ClaimedTask, payload: dict[str, Any]) -> None: - status, _ = self._request("POST", f"/tasks/{task.task_id}/result", payload) - # 200/201/202 include a successful or idempotent duplicate result response. - if status not in (200, 201, 202): - if status == 409: - raise CoordinatorConflictError("result rejected because the task lease was lost") - raise CoordinatorError(f"result rejected with status {status}") - - def fail(self, task: ClaimedTask, payload: dict[str, Any]) -> None: - status, _ = self._request("POST", f"/tasks/{task.task_id}/failure", payload) - if status not in (200, 201, 202): - if status == 409: - raise CoordinatorConflictError("failure rejected because the task lease was lost") - raise CoordinatorError(f"failure report rejected with status {status}") - - def heartbeat(self, task: ClaimedTask, worker_id: str) -> str: - status, body = self._request( - "POST", f"/tasks/{task.task_id}/heartbeat", - {"worker_id": worker_id, "attempt": task.attempt}, - ) - if status != 200: - if status == 409: - raise CoordinatorConflictError("heartbeat rejected because the task lease was lost") - raise CoordinatorError(f"heartbeat rejected with status {status}") - lease_expires_at = body.get("lease_expires_at") - if not isinstance(lease_expires_at, str): - raise CoordinatorError("heartbeat response is missing lease_expires_at") - return lease_expires_at - - def _request(self, method: str, path: str, payload: dict[str, Any]) -> tuple[int, dict[str, Any]]: - return self._request_once(method, path, payload, allow_refresh=True) - - def _request_once( - self, method: str, path: str, payload: dict[str, Any], *, allow_refresh: bool - ) -> tuple[int, dict[str, Any]]: - request = Request( - f"{self.base_url}{path}", data=json.dumps(payload).encode(), method=method, - headers={"Content-Type": "application/json", **self._auth_header()}, - ) - try: - with self._opener.open(request, timeout=self.timeout) as response: - raw = response.read() - try: - return response.status, json.loads(raw) if raw else {} - except json.JSONDecodeError as error: - raise CoordinatorError("coordinator returned invalid JSON") from error - except HTTPError as error: - # A 401 usually means the short-lived JWT expired; mint a fresh one - # and retry exactly once so an in-flight worker rides over the gap. - if error.code == 401 and allow_refresh: - self._tokens.refresh() - return self._request_once(method, path, payload, allow_refresh=False) - if error.code >= 500: - raise CoordinatorTransientError(f"coordinator returned {error.code}") from error - return error.code, {} - except (URLError, TimeoutError) as error: - raise CoordinatorTransientError("coordinator request failed") from error - - def _auth_header(self) -> dict[str, str]: - token = self._tokens.token() - return {"Authorization": f"Bearer {token}"} if token else {} diff --git a/scimesh/worker/daemon.py b/scimesh/worker/daemon.py deleted file mode 100644 index 06b0b20..0000000 --- a/scimesh/worker/daemon.py +++ /dev/null @@ -1,331 +0,0 @@ -"""The worker state machine and its safe failure handling.""" - -from __future__ import annotations - -import logging -from dataclasses import replace -from dataclasses import dataclass -import random -import re -import shutil -import subprocess -import threading -import time -from datetime import datetime, timezone - -from .artifacts import ArtifactClient, sha256_file -from .config import WorkerConfig -from .coordinator import ( - CoordinatorClient, - CoordinatorConflictError, - CoordinatorTransientError, -) -from .models import ClaimedTask, UploadedArtifact -from .runners import Runner - - -class LeaseHeartbeat: - """Renews a claimed task lease while local work is in progress.""" - - def __init__( - self, task: ClaimedTask, coordinator: CoordinatorClient, config: WorkerConfig - ) -> None: - self.task, self.coordinator, self.config = task, coordinator, config - self._worker_id = config.worker_id or "" - self._stop = threading.Event() - self._error: Exception | None = None - self._thread: threading.Thread | None = None - self._lease_expires_at = task.lease_expires_at - - def start(self) -> None: - # Verify ownership before expensive download or calculation begins. - self._lease_expires_at = self.coordinator.heartbeat(self.task, self._worker_id) - self._next_delay() - self._thread = threading.Thread( - target=self._run, name=f"lease-{self.task.task_id}", daemon=True - ) - self._thread.start() - - def stop(self) -> None: - self._stop.set() - if self._thread: - self._thread.join() - - def raise_if_failed(self) -> None: - if self._error: - raise self._error - - def _run(self) -> None: - delay = self._next_delay() - while not self._stop.wait(max(delay, 0.01)): - try: - self._lease_expires_at = self.coordinator.heartbeat( - self.task, self._worker_id - ) - delay = self._next_delay() - except ( - Exception - ) as error: # Surface the lease loss in the main state machine. - self._error = error - return - - def _next_delay(self) -> float: - return min(self.config.heartbeat_interval, self._seconds_until_expiry() / 2) - - def _seconds_until_expiry(self) -> float: - try: - expiry = datetime.fromisoformat( - self._lease_expires_at.replace("Z", "+00:00") - ) - except ValueError as error: - raise ValueError("invalid lease_expires_at") from error - seconds = (expiry - datetime.now(timezone.utc)).total_seconds() - if seconds <= 0: - raise ValueError("claimed task lease has already expired") - 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, - ) - self.worker_id = config.worker_id - self._registered = False - self.log = logging.getLogger("scimesh.worker") - - def run_forever(self) -> bool: - """Run until stopped; return false only when interrupted by the operator.""" - failures = 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) -> RunOnceOutcome: - worker_id = self._worker_id() - self._log("claiming", log_level=logging.DEBUG) - task = self.coordinator.claim(worker_id, self.config.capabilities) - if task is None: - 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() - self._log("downloading", task) - input_path = task_dir / "input" - self.artifacts.download(task.input.uri, input_path) - if sha256_file(input_path).lower() != task.input.sha256.lower(): - raise ValueError("input checksum mismatch") - self._log("running", task) - result = self.runner.run(task, task_dir) - heartbeat.raise_if_failed() - if len(result.artifacts) != 1: - raise ValueError("runner must produce exactly one result artifact") - artifact = result.artifacts[0] - uploaded = self.artifacts.upload(task, worker_id, artifact) - manifest = self._result_manifest(uploaded) - self._log("submitting", task) - heartbeat.raise_if_failed() - self.coordinator.submit( - task, - { - "worker_id": worker_id, - "attempt": task.attempt, - "result": manifest, - "metrics": { - **result.metrics, - "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: - self._log("failed", task, error_type=type(error).__name__) - self._report_failure(task, error) - finally: - heartbeat.stop() - return RunOnceOutcome(claimed=True, completed=completed) - - def _report_failure(self, task: ClaimedTask, error: Exception) -> None: - message = self._sanitize_error_message(error) - try: - self.coordinator.fail( - task, - { - "worker_id": self._worker_id(), - "attempt": task.attempt, - "error_code": type(error).__name__, - "error_message": message, - "retryable": self._is_retryable(error), - }, - ) - except CoordinatorTransientError: - raise - except Exception: - self._log("failed", task, error_type="FailureReportError") - - @staticmethod - def _is_retryable(error: Exception) -> bool: - """Retry transient worker/transport failures, never invalid scientific input.""" - return not isinstance( - error, (ValueError, FileNotFoundError, subprocess.CalledProcessError) - ) - - def _sanitize_error_message(self, error: Exception) -> str: - """Keep coordinator-visible failures useful without exposing local paths.""" - message = str(error).replace(str(self.config.work_dir), "") - # 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"(?", message) - message = re.sub(r"(?", message) - return message[:300] - - def _register_worker(self) -> None: - registered = self.coordinator.register( - self.config.worker_name, - self.config.capabilities, - self.config.cpu_count, - self.config.memory_mb, - ) - self.worker_id = registered.worker_id - self.config = replace( - self.config, - worker_id=registered.worker_id, - heartbeat_interval=registered.heartbeat_interval_seconds, - ) - self._registered = True - self._log("registered") - - def _worker_id(self) -> str: - if not self.worker_id: - raise ValueError("worker is not registered") - return self.worker_id - - @staticmethod - def _result_manifest(uploaded: UploadedArtifact) -> dict[str, object]: - """Keep completion payload exact: coordinator owns all artifact metadata.""" - return {"artifact_id": uploaded.artifact_id} - - 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.log(log_level, "worker_event %s", fields) - - def _cleanup_expired_directories(self) -> None: - """Remove only old task attempt directories when retention was configured.""" - if ( - self.config.cleanup_after_seconds is None - or not self.config.work_dir.exists() - ): - return - cutoff = time.time() - self.config.cleanup_after_seconds - for task_dir in self.config.work_dir.iterdir(): - if not task_dir.is_dir(): - continue - for attempt_dir in task_dir.iterdir(): - if attempt_dir.is_dir() and attempt_dir.stat().st_mtime < cutoff: - shutil.rmtree(attempt_dir) - if not any(task_dir.iterdir()): - task_dir.rmdir() - - @staticmethod - def _sleep(delay: float) -> None: - time.sleep(delay * random.uniform(0.75, 1.25)) diff --git a/scimesh/worker/models.py b/scimesh/worker/models.py index 9cf4f97..b25eae1 100644 --- a/scimesh/worker/models.py +++ b/scimesh/worker/models.py @@ -1,10 +1,13 @@ -"""Value objects shared by the worker daemon components.""" +"""Value objects for one claimed task and its run result. + +These are the wire types the Go worker agent hands to the per-task Python +entry point; the daemon lifecycle itself lives in the Go agent. +""" from __future__ import annotations from dataclasses import dataclass from datetime import datetime -from math import isfinite from pathlib import Path from typing import Any from urllib.parse import urlsplit @@ -34,7 +37,9 @@ def _coordinator_uri(value: object, field: str) -> str: def _sha256(value: object, field: str) -> str: digest = _required_string(value, field).lower() - if len(digest) != 64 or any(character not in "0123456789abcdef" for character in digest): + if len(digest) != 64 or any( + character not in "0123456789abcdef" for character in digest + ): raise ValueError(f"{field} must be a SHA-256 hex digest") return digest @@ -61,11 +66,20 @@ class ClaimedTask: if not isinstance(input_data, dict): raise ValueError("input must be an object") raw_attempt = data["attempt"] - if isinstance(raw_attempt, bool) or not isinstance(raw_attempt, int) or raw_attempt < 1: + if ( + isinstance(raw_attempt, bool) + or not isinstance(raw_attempt, int) + or raw_attempt < 1 + ): raise ValueError("attempt must be a positive integer") task_id = str(UUID(_required_string(data["task_id"], "task_id"))) - lease_expires_at = _required_string(data["lease_expires_at"], "lease_expires_at") - if datetime.fromisoformat(lease_expires_at.replace("Z", "+00:00")).tzinfo is None: + lease_expires_at = _required_string( + data["lease_expires_at"], "lease_expires_at" + ) + if ( + datetime.fromisoformat(lease_expires_at.replace("Z", "+00:00")).tzinfo + is None + ): raise ValueError("lease_expires_at must include a timezone") parameters = data.get("parameters", {}) if not isinstance(parameters, dict): @@ -91,61 +105,6 @@ class ProducedArtifact: content_type: str -@dataclass(frozen=True) -class UploadedArtifact: - """Coordinator-owned artifact metadata returned after a successful upload.""" - - artifact_id: str - uri: str - sha256: str - size_bytes: int - - @classmethod - def from_json(cls, data: object) -> "UploadedArtifact": - if not isinstance(data, dict): - raise ValueError("artifact upload response must be an object") - raw_size = data.get("size_bytes") - if isinstance(raw_size, bool) or not isinstance(raw_size, int) or raw_size < 0: - raise ValueError("artifact size_bytes must be a non-negative integer") - try: - return cls( - artifact_id=str(UUID(_required_string(data.get("artifact_id"), "artifact_id"))), - uri=_coordinator_uri(data.get("uri"), "uri"), - sha256=_sha256(data.get("sha256"), "sha256"), - size_bytes=raw_size, - ) - except ValueError as error: - raise ValueError("invalid artifact upload response") from error - - -@dataclass(frozen=True) -class RegisteredWorker: - """Identity and heartbeat policy returned by worker registration.""" - - worker_id: str - heartbeat_interval_seconds: float - - @classmethod - def from_json(cls, data: object) -> "RegisteredWorker": - if not isinstance(data, dict): - raise ValueError("worker registration response must be an object") - raw_interval = data.get("heartbeat_interval_seconds") - if ( - isinstance(raw_interval, bool) - or not isinstance(raw_interval, (int, float)) - or not isfinite(raw_interval) - or raw_interval <= 0 - ): - raise ValueError("heartbeat_interval_seconds must be positive") - try: - return cls( - worker_id=str(UUID(_required_string(data.get("worker_id"), "worker_id"))), - heartbeat_interval_seconds=float(raw_interval), - ) - except ValueError as error: - raise ValueError("invalid worker registration response") from error - - @dataclass(frozen=True) class RunResult: artifacts: tuple[ProducedArtifact, ...] diff --git a/scimesh/worker/runners.py b/scimesh/worker/runners.py index 8bb4946..f0bd0b5 100644 --- a/scimesh/worker/runners.py +++ b/scimesh/worker/runners.py @@ -46,9 +46,9 @@ from scimesh.sdk.runtime import ( ) from scimesh.sdk.workflow import StageKind -from .config import WorkerConfig from .models import ClaimedTask, ProducedArtifact, RunResult + def _utc_now() -> str: return datetime.now(timezone.utc).isoformat().replace("+00:00", "Z") @@ -57,6 +57,38 @@ class Runner(Protocol): def run(self, task: ClaimedTask, task_dir: Path) -> RunResult: ... +def _definitions_from_environment() -> dict[str, WorkloadDefinition]: + """Load definitions from ``SCIMESH_WORKLOAD_ALLOWLIST`` or the built-ins. + + The environment-driven discovery mirrors the former worker configuration; + the Go agent passes the allowlist through unchanged. + """ + import os + + from scimesh.sdk.registry import workload_allowlist_from_json + + allowlist = workload_allowlist_from_json(os.getenv("SCIMESH_WORKLOAD_ALLOWLIST")) + if allowlist: + registry = WorkloadRegistry() + registry.discover_installed(allowlist) + definitions: dict[str, WorkloadDefinition] = {} + for description in registry.descriptions(): + definition, _ = registry.require( + description.workload.name, + description.workload.version, + description.package_digest, + ) + definitions[description.workload.name] = definition + if not definitions: + raise ValueError("workload_allowlist discovered no workloads") + return definitions + from scimesh.workloads.search import similarity_search_sdk_definition + + return { + "similarity-search": similarity_search_sdk_definition().definition(), + } + + def _inventory_for( definitions: Mapping[str, WorkloadDefinition], *, @@ -91,7 +123,12 @@ def _runtime_for( class SciMeshRunner: - """Execute claimed coordinator tasks through SDK-built workloads.""" + """Execute claimed coordinator tasks through SDK-built workloads. + + Definitions come from ``SCIMESH_WORKLOAD_ALLOWLIST`` when set, otherwise + the built-in ``similarity-search``; callers may supply an explicit mapping + for tests. + """ def __init__( self, @@ -100,13 +137,9 @@ class SciMeshRunner: inventory: ResourceInventory | None = None, runtime: RuntimeCapabilities | None = None, ) -> None: - self._definitions = dict(definitions or {}) - if "similarity-search" not in self._definitions: - from scimesh.workloads.search import similarity_search_sdk_definition - - self._definitions["similarity-search"] = ( - similarity_search_sdk_definition().definition() - ) + if definitions is None: + definitions = _definitions_from_environment() + self._definitions = dict(definitions) self._inventory = inventory or _inventory_for( self._definitions, cpu_cores=1, @@ -115,35 +148,6 @@ class SciMeshRunner: self._runtime = runtime or _runtime_for(self._definitions, self._inventory) self._pool = ResourcePool(self._runtime.inventory, max_concurrency=1) - @classmethod - def for_worker(cls, config: WorkerConfig) -> "SciMeshRunner": - """Build a runner for one worker: discover allowlisted workloads or use built-ins.""" - definitions: dict[str, WorkloadDefinition] = {} - if config.workload_allowlist: - registry = WorkloadRegistry() - registry.discover_installed(config.workload_allowlist) - for description in registry.descriptions(): - definition, _ = registry.require( - description.workload.name, - description.workload.version, - description.package_digest, - ) - definitions[description.workload.name] = definition - if not definitions: - raise ValueError("workload_allowlist discovered no workloads") - else: - from scimesh.workloads.search import similarity_search_sdk_definition - - definitions["similarity-search"] = ( - similarity_search_sdk_definition().definition() - ) - inventory = _inventory_for( - definitions, - cpu_cores=config.cpu_count, - memory_mb=config.memory_mb or 1024, - ) - return cls(definitions=definitions, inventory=inventory) - def run(self, task: ClaimedTask, task_dir: Path) -> RunResult: task_dir = task_dir.resolve() workload = task.workload.replace("_", "-") diff --git a/scimesh/worker/transport.py b/scimesh/worker/transport.py deleted file mode 100644 index 55cf3f2..0000000 --- a/scimesh/worker/transport.py +++ /dev/null @@ -1,53 +0,0 @@ -"""Small HTTP transport helpers shared by coordinator and artifact clients.""" - -from __future__ import annotations - -from http.client import HTTPMessage -from typing import IO -from urllib.request import HTTPRedirectHandler, Request -from urllib.parse import urlsplit - - -def origin(uri: str) -> tuple[str, str, int | None]: - """Return a normalized HTTP origin for authorization decisions.""" - parsed = urlsplit(uri) - scheme = parsed.scheme.lower() - default_port = {"http": 80, "https": 443}.get(scheme) - return scheme, (parsed.hostname or "").lower(), parsed.port or default_port - - -class SameOriginAuthRedirectHandler(HTTPRedirectHandler): - """Strip coordinator authorization when an artifact redirect changes origin.""" - - def __init__(self, coordinator_origin: tuple[str, str, int | None]) -> None: - super().__init__() - self.coordinator_origin = coordinator_origin - - def redirect_request( - self, - req: Request, - fp: IO[bytes], - code: int, - msg: str, - headers: HTTPMessage, - newurl: str, - ) -> Request | None: - redirected = super().redirect_request(req, fp, code, msg, headers, newurl) - if redirected and origin(newurl) != self.coordinator_origin: - redirected.remove_header("Authorization") - return redirected - - -class NoRedirectHandler(HTTPRedirectHandler): - """Reject redirects for mutating coordinator API calls.""" - - def redirect_request( - self, - req: Request, - fp: object, - code: int, - msg: str, - headers: object, - newurl: str, - ) -> Request | None: - return None diff --git a/scripts/two-worker-smoke.sh b/scripts/two-worker-smoke.sh index c8d96a5..0da0bc5 100755 --- a/scripts/two-worker-smoke.sh +++ b/scripts/two-worker-smoke.sh @@ -19,6 +19,7 @@ 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"} +AGENT_BIN=${SCIMESH_AGENT_BIN:-"$ROOT_DIR/coordinator/bin/worker-agent"} cleanup() { local exit_code=$? @@ -59,6 +60,12 @@ for command in docker curl python3; do require "$command"; done printf 'Set SCIMESH_WORKER_PYTHON to a Python environment with SciMesh and RDKit.\n' >&2 exit 2 } +[[ -x "$AGENT_BIN" ]] || { + printf 'Go worker agent is not built: %s\n' "$AGENT_BIN" >&2 + printf 'Run: make -C coordinator agent\n' >&2 + exit 2 +} +TASK_RUNNER_JSON="[\"$WORKER_PYTHON\",\"-m\",\"scimesh.worker.task\"]" printf 'Starting isolated coordinator on %s (project %s)\n' "$HOST" "$COMPOSE_PROJECT" ( @@ -80,11 +87,18 @@ curl -fsS "$HOST/health" >/dev/null || { 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 & + COORDINATOR_URL="$HOST" \ + WORKER_AUTH_TOKEN="$TOKEN" \ + WORKER_NAME="$worker_name" \ + WORK_DIR="$worker_dir" \ + CPU_COUNT=1 \ + MEMORY_MB=1024 \ + POLL_INTERVAL=0.2s \ + REQUEST_TIMEOUT=15s \ + HEARTBEAT_INTERVAL=15s \ + MAX_TASKS=2 \ + TASK_RUNNER="$TASK_RUNNER_JSON" \ + "$AGENT_BIN" >"$worker_dir.log" 2>&1 & STARTED_WORKER_PID=$! } diff --git a/tests/test_worker_auth.py b/tests/test_worker_auth.py deleted file mode 100644 index bfb0f56..0000000 --- a/tests/test_worker_auth.py +++ /dev/null @@ -1,221 +0,0 @@ -"""Tests for worker token strategies and the client's 401 refresh.""" - -from __future__ import annotations - -import json -from pathlib import Path -from urllib.error import HTTPError - -import pytest - -from scimesh.worker.auth import ( - StaticTokenProvider, - TokenExchangeError, - WorkerKeyTokenProvider, - provider_from_config, -) -from scimesh.worker.config import WorkerConfig -from scimesh.worker.coordinator import HttpCoordinatorClient - - -class FakeResponse: - def __init__(self, status: int, body: bytes) -> None: - self.status = status - self._body = body - - def read(self) -> bytes: - return self._body - - def __enter__(self) -> "FakeResponse": - return self - - def __exit__(self, *exc) -> bool: - return False - - -class SeqOpener: - """Returns/raises a scripted sequence of responses, recording each request.""" - - def __init__(self, actions: list) -> None: - self.actions = list(actions) - self.requests: list = [] - - def open(self, request, timeout=None): - self.requests.append(request) - action = self.actions.pop(0) - if isinstance(action, Exception): - raise action - return action - - -def _exchange_response(token: str, expires_in: int) -> FakeResponse: - return FakeResponse( - 200, json.dumps({"token": token, "expires_in": expires_in}).encode() - ) - - -def test_static_provider_returns_fixed_token_and_never_refreshes(): - provider = StaticTokenProvider("tok") - assert provider.token() == "tok" - provider.refresh() - assert provider.token() == "tok" - - -def test_static_provider_none_means_no_auth(): - assert StaticTokenProvider(None).token() is None - - -def test_worker_key_provider_exchanges_once_then_caches(): - clock = {"t": 1000.0} - provider = WorkerKeyTokenProvider( - "http://users", "scimesh_wk_live_x", timeout=5, now=lambda: clock["t"] - ) - provider._opener = SeqOpener([_exchange_response("jwt-1", 100)]) # type: ignore[reportAttributeAccessIssue] - - # First call exchanges; a second call well within the TTL reuses the cache. - assert provider.token() == "jwt-1" - clock["t"] = 1050.0 # 50s later, TTL 100s with 0.2 leeway → refresh at +80s - assert provider.token() == "jwt-1" - assert len(provider._opener.requests) == 1 # type: ignore[reportAttributeAccessIssue] - - -def test_worker_key_provider_refreshes_after_leeway(): - clock = {"t": 0.0} - provider = WorkerKeyTokenProvider( - "http://users", "k", timeout=5, now=lambda: clock["t"] - ) - provider._opener = SeqOpener( # type: ignore[reportAttributeAccessIssue] - [ - _exchange_response("jwt-1", 100), - _exchange_response("jwt-2", 100), - ] - ) - assert provider.token() == "jwt-1" - clock["t"] = 85.0 # past the 80s refresh point - assert provider.token() == "jwt-2" - assert len(provider._opener.requests) == 2 # type: ignore[reportAttributeAccessIssue] - - -def test_worker_key_provider_force_refresh(): - provider = WorkerKeyTokenProvider("http://users", "k", timeout=5, now=lambda: 0.0) - provider._opener = SeqOpener( # type: ignore[reportAttributeAccessIssue] - [ - _exchange_response("jwt-1", 100), - _exchange_response("jwt-2", 100), - ] - ) - assert provider.token() == "jwt-1" - provider.refresh() - assert provider.token() == "jwt-2" - - -def test_worker_key_provider_raises_on_rejected_key(): - provider = WorkerKeyTokenProvider("http://users", "bad", timeout=5, now=lambda: 0.0) - provider._opener = SeqOpener( # type: ignore[reportAttributeAccessIssue] - [HTTPError("http://users", 401, "unauthorized", {}, None)] # type: ignore[reportArgumentType] - ) - with pytest.raises(TokenExchangeError): - provider.token() - - -def test_worker_key_provider_raises_when_token_missing(): - provider = WorkerKeyTokenProvider("http://users", "k", timeout=5, now=lambda: 0.0) - provider._opener = SeqOpener( # type: ignore[reportAttributeAccessIssue] - [FakeResponse(200, json.dumps({"expires_in": 100}).encode())] - ) - with pytest.raises(TokenExchangeError): - provider.token() - - -def test_provider_from_config_selects_worker_key_mode(): - provider = provider_from_config( - worker_key="scimesh_wk_live_x", - userservice_url="http://users", - bearer_token="ignored", - request_timeout=5, - ) - assert isinstance(provider, WorkerKeyTokenProvider) - - -def test_provider_from_config_falls_back_to_static(): - provider = provider_from_config( - worker_key=None, userservice_url=None, bearer_token="tok", request_timeout=5 - ) - assert isinstance(provider, StaticTokenProvider) - assert provider.token() == "tok" - - -class RefreshCountingProvider: - def __init__(self) -> None: - self.tokens = ["stale", "fresh"] - self.index = 0 - self.refreshes = 0 - - def token(self) -> str: - return self.tokens[min(self.index, len(self.tokens) - 1)] - - def refresh(self) -> None: - self.refreshes += 1 - self.index += 1 - - -def test_coordinator_client_refreshes_and_retries_once_on_401(): - provider = RefreshCountingProvider() - client = HttpCoordinatorClient("http://coord", timeout=5, token_provider=provider) - client._opener = SeqOpener( # type: ignore[reportAttributeAccessIssue] - [ - HTTPError("http://coord/tasks/claim", 401, "unauthorized", {}, None), # type: ignore[reportArgumentType] - FakeResponse(204, b""), - ] - ) - - status, _ = client._request("POST", "/tasks/claim", {"worker_id": "w"}) - - assert status == 204 - assert provider.refreshes == 1 - # The retry carried the refreshed token. - assert provider.index == 1 - - -def test_coordinator_client_does_not_loop_on_persistent_401(): - provider = RefreshCountingProvider() - client = HttpCoordinatorClient("http://coord", timeout=5, token_provider=provider) - client._opener = SeqOpener( # type: ignore[reportAttributeAccessIssue] - [ - HTTPError("http://coord/x", 401, "unauthorized", {}, None), # type: ignore[reportArgumentType] - HTTPError("http://coord/x", 401, "unauthorized", {}, None), # type: ignore[reportArgumentType] - ] - ) - - status, _ = client._request("POST", "/x", {}) - - # One refresh, one retry, then the second 401 is surfaced rather than retried. - assert status == 401 - assert provider.refreshes == 1 - - -def _base_config(**extra) -> dict: - return { - "coordinator_url": "http://coord", - "worker_id": None, - "work_dir": Path("."), - **extra, - } - - -def test_worker_key_requires_userservice_url(): - with pytest.raises(ValueError, match="userservice_url"): - WorkerConfig(**_base_config(worker_key="scimesh_wk_live_x")) - - -def test_worker_key_with_userservice_url_is_valid(): - cfg = WorkerConfig( - **_base_config(worker_key="scimesh_wk_live_x", userservice_url="http://users") - ) - assert cfg.worker_key == "scimesh_wk_live_x" - assert cfg.userservice_url == "http://users" - - -def test_userservice_url_must_be_absolute(): - with pytest.raises(ValueError, match="userservice_url"): - WorkerConfig(**_base_config(userservice_url="not-a-url")) diff --git a/tests/test_worker_daemon.py b/tests/test_worker_daemon.py deleted file mode 100644 index 2e7c180..0000000 --- a/tests/test_worker_daemon.py +++ /dev/null @@ -1,842 +0,0 @@ -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 - -import pytest - -from scimesh.worker.config import WorkerConfig, _workload_allowlist -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, RunOnceOutcome, WorkerDaemon -from scimesh.worker.models import ( - ClaimedTask, - InputArtifact, - ProducedArtifact, - RegisteredWorker, - RunResult, - UploadedArtifact, -) -from scimesh.worker.artifacts import ( - HttpArtifactClient, - _SameOriginAuthRedirectHandler, - _origin, -) -from scimesh.worker.runners import SciMeshRunner -from scimesh.worker.transport import NoRedirectHandler - - -class FakeCoordinator: - def __init__(self, task: ClaimedTask | None) -> None: - self.task, self.submissions, self.failures, self.heartbeats = task, [], [], [] - - def claim( - self, worker_id: str, capabilities: tuple[str, ...] - ) -> ClaimedTask | None: - task, self.task = self.task, None - return task - - def register( - self, - name: str, - capabilities: tuple[str, ...], - cpu_count: int, - memory_mb: int | None, - ) -> RegisteredWorker: - return RegisteredWorker("11111111-1111-4111-8111-111111111111", 15) - - def submit(self, task: ClaimedTask, payload: dict) -> None: - self.submissions.append(payload) - - def fail(self, task: ClaimedTask, payload: dict) -> None: - self.failures.append(payload) - - def heartbeat(self, task: ClaimedTask, worker_id: str) -> str: - self.heartbeats.append((task.task_id, task.attempt, worker_id)) - return (datetime.now(timezone.utc) + timedelta(seconds=1)).isoformat() - - -class FakeArtifacts: - def __init__(self, content: bytes) -> None: - self.content, self.uploaded = content, [] - - def download(self, uri: str, destination: Path) -> None: - destination.write_bytes(self.content) - - def upload( - self, task: ClaimedTask, worker_id: str, artifact: ProducedArtifact - ) -> UploadedArtifact: - self.uploaded.append((task.task_id, worker_id, artifact.path)) - content = artifact.path.read_bytes() - return UploadedArtifact( - "22222222-2222-4222-8222-222222222222", - f"https://example.test/tasks/{task.task_id}/artifacts/{artifact.path.name}", - hashlib.sha256(content).hexdigest(), - len(content), - ) - - -class FakeRunner: - def __init__(self) -> None: - self.calls = 0 - - def run(self, task: ClaimedTask, task_dir: Path) -> RunResult: - self.calls += 1 - output = task_dir / "result.csv" - output.write_text("id,score\na,1\n", encoding="utf-8") - return RunResult((ProducedArtifact(output, "text/csv"),), {"processed_rows": 1}) - - -def make_task(content: bytes, checksum: str | None = None) -> ClaimedTask: - lease = (datetime.now(timezone.utc) + timedelta(seconds=60)).isoformat() - return ClaimedTask( - "task-1", - 1, - lease, - "similarity-search", - InputArtifact( - "https://example.test/input", - checksum or hashlib.sha256(content).hexdigest(), - ), - {"query_id": "CHEMBL1"}, - ) - - -def daemon(tmp_path: Path, task: ClaimedTask | None, content: bytes): - coordinator, artifacts, runner = ( - FakeCoordinator(task), - FakeArtifacts(content), - FakeRunner(), - ) - config = WorkerConfig("https://example.test", "worker-1", tmp_path / "work") - return ( - WorkerDaemon(config, coordinator, artifacts, runner), - coordinator, - artifacts, - runner, - config, - ) - - -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() == RunOnceOutcome(claimed=True, completed=True) - assert runner.calls == 1 - assert len(artifacts.uploaded) == 1 - assert coordinator.heartbeats == [("task-1", 1, "worker-1")] - assert "status" not in coordinator.submissions[0] - assert coordinator.submissions[0]["result"] == { - "artifact_id": "22222222-2222-4222-8222-222222222222" - } - - -def test_worker_executes_a_resolved_similarity_search_shard(tmp_path: Path) -> None: - content = ( - b"chembl_id\tcanonical_smiles\nQUERY\tCCO\nMATCH\tCCCO\nINVALID\tnot-a-smiles\n" - ) - task = make_task(content) - task = ClaimedTask( - task.task_id, - task.attempt, - task.lease_expires_at, - task.workload, - task.input, - {"query_smiles": "CCO", "top_k": 5, "progress_every": 0}, - ) - worker, coordinator, artifacts, _, _ = daemon(tmp_path, task, content) - worker.runner = SciMeshRunner() - - assert worker.run_once() == RunOnceOutcome(claimed=True, completed=True) - output = artifacts.uploaded[0][2].read_text(encoding="utf-8") - assert output.startswith("rank,chembl_id,canonical_smiles,similarity\n") - metrics = coordinator.submissions[0]["metrics"] - assert metrics["scanned_rows"] == 3 - assert metrics["valid_molecules"] == 2 - assert metrics["invalid_smiles"] == 1 - assert metrics["matches_emitted"] == 1 - assert isinstance(metrics["elapsed_seconds"], float) - - -def test_two_workers_complete_resolved_shards_after_one_retry(tmp_path: Path) -> None: - content = b"chembl_id\tcanonical_smiles\nQUERY\tCCO\nMATCH\tCCCO\n" - first = make_task(content) - first = ClaimedTask( - "retry-task", - 1, - first.lease_expires_at, - "similarity-search", - first.input, - {"query_smiles": "CCO", "top_k": 5}, - ) - second = ClaimedTask( - "other-task", - 1, - first.lease_expires_at, - "similarity-search", - first.input, - {"query_smiles": "CCO", "top_k": 5}, - ) - - class RetryCoordinator(FakeCoordinator): - def __init__(self) -> None: - super().__init__(None) - self.queue = [first, second] - self.claimants: list[str] = [] - - def claim( - self, worker_id: str, capabilities: tuple[str, ...] - ) -> ClaimedTask | None: - self.claimants.append(worker_id) - return self.queue.pop(0) if self.queue else None - - def fail(self, task: ClaimedTask, payload: dict) -> None: - self.failures.append(payload) - if ( - task.task_id == "retry-task" - and task.attempt == 1 - and payload["retryable"] - ): - self.queue.append( - ClaimedTask( - task.task_id, - 2, - task.lease_expires_at, - task.workload, - task.input, - task.parameters, - ) - ) - - class FailFirstAttempt: - def __init__(self) -> None: - self.calls = 0 - self.delegate = SciMeshRunner() - - def run(self, task: ClaimedTask, task_dir: Path) -> RunResult: - self.calls += 1 - if self.calls == 1: - raise RuntimeError("simulated retryable shard failure") - return self.delegate.run(task, task_dir) - - coordinator = RetryCoordinator() - artifacts = FakeArtifacts(content) - worker_a = WorkerDaemon( - WorkerConfig("https://example.test", "worker-a", tmp_path / "worker-a"), - coordinator, - artifacts, - FailFirstAttempt(), - ) - worker_b = WorkerDaemon( - WorkerConfig("https://example.test", "worker-b", tmp_path / "worker-b"), - coordinator, - artifacts, - SciMeshRunner(), - ) - - assert worker_a.run_once() == RunOnceOutcome(claimed=True, completed=False) - assert worker_b.run_once() == RunOnceOutcome(claimed=True, completed=True) - assert worker_a.run_once() == RunOnceOutcome(claimed=True, completed=True) - assert coordinator.claimants == ["worker-a", "worker-b", "worker-a"] - assert len(coordinator.failures) == 1 - assert coordinator.failures[0]["retryable"] is True - assert len(coordinator.submissions) == 2 - - -def test_no_task_does_not_create_directory(tmp_path: Path) -> None: - worker, _, _, runner, config = daemon(tmp_path, None, b"") - assert worker.run_once() == RunOnceOutcome(claimed=False, completed=False) - 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", - "retryable": True, - } - ] - - -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() == RunOnceOutcome(claimed=True, completed=False) - assert runner.calls == 0 - assert coordinator.failures[0]["error_code"] == "ValueError" - assert coordinator.failures[0]["retryable"] is False - assert not coordinator.submissions - - -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 "" 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() == RunOnceOutcome(claimed=True, completed=False) - assert coordinator.failures[0]["error_code"] == "FileExistsError" - - -def test_transient_claim_error_is_propagated_for_bounded_backoff( - tmp_path: Path, -) -> None: - class UnavailableCoordinator(FakeCoordinator): - def claim( - self, worker_id: str, capabilities: tuple[str, ...] - ) -> ClaimedTask | None: - raise CoordinatorTransientError("temporary outage") - - worker, _, _, _, _ = daemon(tmp_path, None, b"") - worker.coordinator = UnavailableCoordinator(None) - with pytest.raises(CoordinatorTransientError): - worker.run_once() - - -def test_task_directories_are_retained_until_cleanup_is_enabled(tmp_path: Path) -> None: - content = b"input fixture" - worker, _, _, _, config = daemon(tmp_path, make_task(content), content) - worker.run_once() - task_dir = config.work_dir / "task-1" / "1" - assert task_dir.is_dir() - worker.config = WorkerConfig(**{**config.__dict__, "cleanup_after_seconds": 0}) - worker._cleanup_expired_directories() - assert not task_dir.exists() - - -def test_input_token_is_sent_only_to_the_coordinator_origin() -> None: - client = HttpArtifactClient("https://coordinator.example/api", 10, "secret") - assert client._auth_headers_for("https://coordinator.example/tasks/1/input") == { - "Authorization": "Bearer secret" - } - assert client._auth_headers_for("https://bucket.example/presigned") == {} - - -def test_relative_input_uri_is_resolved_against_the_coordinator() -> None: - client = HttpArtifactClient("https://coordinator.example/api", 10, "secret") - assert client._auth_headers_for("https://coordinator.example/tasks/1/input") == { - "Authorization": "Bearer secret" - } - # The coordinator's contract returns root-relative artifact paths. - assert client.coordinator_url == "https://coordinator.example/api" - - -def test_redirect_to_external_storage_strips_authorization() -> None: - handler = _SameOriginAuthRedirectHandler(_origin("https://coordinator.example")) - source = Request( - "https://coordinator.example/tasks/1/input", - headers={"Authorization": "Bearer secret"}, - ) - redirected = handler.redirect_request( - source, None, 302, "Found", {}, "https://bucket.example/presigned" # type: ignore[arg-type] - ) - assert redirected is not None - assert redirected.get_header("Authorization") is None - - -def test_api_requests_never_follow_redirects() -> None: - handler = NoRedirectHandler() - request = Request( - "https://coordinator.example/tasks/claim", - headers={"Authorization": "Bearer secret"}, - ) - assert ( - handler.redirect_request( - request, None, 302, "Found", {}, "https://other.example" - ) - is None - ) - - -def test_lease_is_renewed_while_a_runner_is_still_working(tmp_path: Path) -> None: - content = b"input fixture" - worker, coordinator, _, _, config = daemon(tmp_path, make_task(content), content) - - class SlowRunner(FakeRunner): - def run(self, task: ClaimedTask, task_dir: Path) -> RunResult: - time.sleep(0.05) - return super().run(task, task_dir) - - worker.runner = SlowRunner() - worker.config = WorkerConfig(**{**config.__dict__, "heartbeat_interval": 0.01}) - worker.run_once() - assert len(coordinator.heartbeats) >= 2 - - -def test_heartbeat_reschedules_from_the_renewed_lease(tmp_path: Path) -> None: - class ShortLeaseCoordinator(FakeCoordinator): - def heartbeat(self, task: ClaimedTask, worker_id: str) -> str: - self.heartbeats.append((task.task_id, task.attempt, worker_id)) - return (datetime.now(timezone.utc) + timedelta(seconds=0.02)).isoformat() - - config = WorkerConfig( - "https://example.test", "worker-1", tmp_path / "work", heartbeat_interval=1 - ) - coordinator = ShortLeaseCoordinator(None) - heartbeat = LeaseHeartbeat(make_task(b"fixture"), coordinator, config) - heartbeat.start() - time.sleep(0.06) - heartbeat.stop() - assert len(coordinator.heartbeats) >= 3 - - -def test_runner_executes_search_through_the_sdk_and_rejects_graph( - tmp_path: Path, -) -> None: - runner = SciMeshRunner() - graph = ClaimedTask( - "graph", - 1, - "2026-07-30T00:00:00Z", - "similarity-graph", - InputArtifact("https://example/input", "x"), - { - "threshold": 0.2, - "threshold_direction": "less", - "block_size": 42, - "max_rows": 7, - "progress_every": 0, - }, - ) - search = ClaimedTask( - "search", - 1, - "2026-07-30T00:00:00Z", - "similarity-search", - InputArtifact("https://example/input", "x"), - {"query_smiles": "CCO", "top_k": 3}, - ) - search_dir = tmp_path / "search" - search_dir.mkdir() - (search_dir / "input").write_text( - "chembl_id\tcanonical_smiles\nA\tCCO\nB\tCCCO\n", encoding="utf-8" - ) - with pytest.raises(ValueError, match="unsupported workload"): - runner.run(graph, tmp_path / "graph") - result = runner.run(search, search_dir) - assert result.metrics == { - "scanned_rows": 2, - "valid_molecules": 2, - "invalid_smiles": 0, - "matches_emitted": 1, - } - assert result.artifacts[0].content_type == "text/csv" - assert ( - result.artifacts[0] - .path.read_text(encoding="utf-8") - .startswith("rank,chembl_id,canonical_smiles,similarity\n") - ) - - -def test_runner_resolves_query_id_from_the_shard_and_rejects_plan_time_parameters( - tmp_path: Path, -) -> None: - task_dir = tmp_path / "search" - task_dir.mkdir() - (task_dir / "input").write_text( - "chembl_id\tcanonical_smiles\nQUERY\tCCO\nMATCH\tCCCO\n", encoding="utf-8" - ) - task = ClaimedTask( - "search", - 1, - "2026-07-30T00:00:00Z", - "similarity-search", - InputArtifact("https://example/input", "a" * 64), - {"query_id": "QUERY", "top_k": 5}, - ) - result = SciMeshRunner().run(task, task_dir) - assert result.metrics["matches_emitted"] == 1 - assert (task_dir / "result.csv").is_file() - - with_max_rows = ClaimedTask( - "search", - 1, - "2026-07-30T00:00:00Z", - "similarity-search", - InputArtifact("https://example/input", "a" * 64), - {"query_smiles": "CCO", "max_rows": 1}, - ) - with pytest.raises(ValueError, match="outside the stage projection"): - SciMeshRunner().run(with_max_rows, task_dir) - - -def test_runner_accepts_coordinator_workload_names( - tmp_path: Path, monkeypatch: pytest.MonkeyPatch -) -> None: - task_dir = tmp_path / "search" - task_dir.mkdir() - (task_dir / "input").write_text( - "chembl_id\tcanonical_smiles\nA\tCCO\nB\tCCCO\n", encoding="utf-8" - ) - task = ClaimedTask( - "search", - 1, - "2026-07-30T00:00:00Z", - "similarity_search", - InputArtifact("https://example/input", "a" * 64), - {"query_smiles": "CCO"}, - ) - result = SciMeshRunner().run(task, task_dir) - assert result.metrics["matches_emitted"] == 1 - assert (task_dir / "result.csv").is_file() - - -def test_claimed_task_rejects_path_traversal_and_invalid_metadata() -> None: - payload = { - "task_id": "../outside", - "attempt": 1, - "lease_expires_at": "2026-07-30T00:00:00Z", - "workload": "similarity-search", - "input": {"uri": "https://example.test/input", "sha256": "a" * 64}, - "parameters": {}, - } - with pytest.raises(ValueError, match="invalid claimed-task response"): - ClaimedTask.from_json(payload) - - payload["task_id"] = "11111111-1111-4111-8111-111111111111" - payload["input"] = {"uri": "//outside.example/input", "sha256": "a" * 64} - with pytest.raises(ValueError, match="invalid claimed-task response"): - ClaimedTask.from_json(payload) - - payload["input"] = {"uri": "/tasks/../outside/input", "sha256": "a" * 64} - with pytest.raises(ValueError, match="invalid claimed-task response"): - ClaimedTask.from_json(payload) - - -def test_claimed_task_accepts_a_coordinator_relative_input_path() -> None: - task = ClaimedTask.from_json( - { - "task_id": "11111111-1111-4111-8111-111111111111", - "attempt": 1, - "lease_expires_at": "2026-07-30T00:00:00Z", - "workload": "similarity_search", - "input": { - "uri": "/tasks/11111111-1111-4111-8111-111111111111/input", - "sha256": "a" * 64, - }, - "parameters": {}, - } - ) - assert task.input.uri.startswith("/tasks/") - - -def test_uploaded_artifact_requires_complete_durable_metadata() -> None: - artifact = UploadedArtifact.from_json( - { - "artifact_id": "22222222-2222-4222-8222-222222222222", - "uri": "https://coordinator.example/artifacts/222/download", - "sha256": "a" * 64, - "size_bytes": 12, - } - ) - assert artifact.size_bytes == 12 - with pytest.raises(ValueError, match="artifact size_bytes"): - UploadedArtifact.from_json({"artifact_id": "missing"}) - - -def test_environment_overrides_allow_cli_only_configuration( - monkeypatch: pytest.MonkeyPatch, tmp_path: Path -) -> None: - monkeypatch.delenv("SCIMESH_COORDINATOR_URL", raising=False) - config = WorkerConfig.from_environment( - { - "coordinator_url": "https://coordinator.example", - "work_dir": tmp_path, - "worker_name": "test-worker", - } - ) - assert config.coordinator_url == "https://coordinator.example" - 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( - "chembl_id\tcanonical_smiles\nA\tCCO\nB\tCCCO\n", encoding="utf-8" - ) - task = ClaimedTask( - "task", - 1, - "2026-07-30T00:00:00Z", - "similarity-search", - InputArtifact("https://example.test/input", "a" * 64), - {"query_smiles": "CCO"}, - ) - SciMeshRunner().run(task, task_dir) - assert (task_dir / "result.csv").is_file() - - -def test_worker_registration_sets_returned_identity(tmp_path: Path) -> None: - worker, _, _, _, _ = daemon(tmp_path, None, b"") - worker._register_worker() - assert worker.worker_id == "11111111-1111-4111-8111-111111111111" - assert worker.config.heartbeat_interval == 15 - - -def test_runner_executes_an_arbitrary_sdk_workload(tmp_path: Path) -> None: - from scimesh.workloads.descriptors import descriptor_batch_sdk_definition - - content = b"chembl_id\tcanonical_smiles\nA\tCCO\nB\tCCCC\n" - task = ClaimedTask( - "task-1", - 1, - (datetime.now(timezone.utc) + timedelta(seconds=60)).isoformat(), - "descriptor-batch", - InputArtifact("https://example.test/input", hashlib.sha256(content).hexdigest()), - {"skip_invalid": True}, - ) - task_dir = tmp_path / "task-1" / "1" - task_dir.mkdir(parents=True) - (task_dir / "input").write_bytes(content) - runner = SciMeshRunner( - definitions={ - "descriptor-batch": descriptor_batch_sdk_definition().definition() - } - ) - - result = runner.run(task, task_dir) - header = result.artifacts[0].path.read_text(encoding="utf-8").splitlines()[0] - assert header.startswith("chembl_id,canonical_smiles,ExactMolWt") - assert result.metrics["rows_emitted"] == 2 - - -def test_runner_rejects_workloads_outside_the_v1_single_input_contract( - tmp_path: Path, -) -> None: - from scimesh.workloads.graph import similarity_graph_sdk_definition - - content = b"chembl_id\tcanonical_smiles\nA\tCCO\nB\tCCCC\n" - task = ClaimedTask( - "graph-task", - 1, - (datetime.now(timezone.utc) + timedelta(seconds=60)).isoformat(), - "similarity-graph", - InputArtifact("https://example.test/input", hashlib.sha256(content).hexdigest()), - {"threshold": 0.5}, - ) - task_dir = tmp_path / "graph" - task_dir.mkdir(parents=True) - (task_dir / "input").write_bytes(content) - runner = SciMeshRunner( - definitions={ - "similarity-graph": similarity_graph_sdk_definition().definition() - } - ) - - with pytest.raises(ValueError, match="v1 single-input contract"): - runner.run(task, task_dir) - - -def test_runner_for_worker_discovers_allowlisted_installed_workloads( - monkeypatch: pytest.MonkeyPatch, tmp_path: Path -) -> None: - from scimesh.sdk.registry import WorkloadRegistry - from scimesh.workloads.search import similarity_search_sdk_definition - - definition = similarity_search_sdk_definition().definition() - - def fake_discover(self: WorkloadRegistry, allowlist) -> None: - assert len(allowlist) == 1 - self.register(definition, enabled=True) - - monkeypatch.setattr(WorkloadRegistry, "discover_installed", fake_discover) - allowlist = _workload_allowlist( - '[{"distribution": "scimesh", "name": "similarity-search", ' - '"version": "1.0.0", "digest": "sha256:' + "a" * 64 + '"}]' - ) - config = WorkerConfig( - "https://example.test", "worker-1", tmp_path / "work", - capabilities=("similarity-search",), - workload_allowlist=allowlist, - ) - runner = SciMeshRunner.for_worker(config) - assert set(runner._definitions) == {"similarity-search"} - - -def test_worker_config_parses_capabilities_and_workload_allowlist( - monkeypatch: pytest.MonkeyPatch, tmp_path: Path -) -> None: - monkeypatch.setenv("SCIMESH_CAPABILITIES", "similarity-search,descriptor-batch") - config = WorkerConfig.from_environment( - {"coordinator_url": "https://example.test", "work_dir": tmp_path} - ) - assert config.capabilities == ("similarity-search", "descriptor-batch") - assert config.workload_allowlist == () - - monkeypatch.setenv( - "SCIMESH_WORKLOAD_ALLOWLIST", - '[{"distribution": "scimesh", "name": "descriptor-batch", ' - '"version": "1.0.0", "digest": "sha256:' + "b" * 64 + '"}]', - ) - config = WorkerConfig.from_environment( - {"coordinator_url": "https://example.test", "work_dir": tmp_path} - ) - assert len(config.workload_allowlist) == 1 - assert config.workload_allowlist[0].workload.name == "descriptor-batch" - - monkeypatch.setenv("SCIMESH_WORKLOAD_ALLOWLIST", "not-json") - with pytest.raises(ValueError, match="valid JSON"): - WorkerConfig.from_environment( - {"coordinator_url": "https://example.test", "work_dir": tmp_path} - )