Replace Python worker daemon with the Go worker agent
This commit is contained in:
Emil
@@ -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.
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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. |
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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}
|
||||
}
|
||||
@@ -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())
|
||||
}
|
||||
}
|
||||
@@ -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)}
|
||||
}
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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 == "" {
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
+1
-1
@@ -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.<module>` — 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
|
||||
|
||||
+2
-2
@@ -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.
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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.
|
||||
"""
|
||||
|
||||
@@ -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()
|
||||
@@ -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)
|
||||
@@ -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())
|
||||
@@ -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),
|
||||
)
|
||||
@@ -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 {}
|
||||
@@ -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), "<worker-dir>")
|
||||
# CalledProcessError includes the complete argv, including sys.executable
|
||||
# outside work_dir. Replace POSIX and Windows absolute paths before the
|
||||
# message reaches the coordinator database or operator UI.
|
||||
message = re.sub(r"(?<![\w:])[A-Za-z]:\\[^\s'\"\],)]+", "<path>", message)
|
||||
message = re.sub(r"(?<![\w:])/(?:[^\s'\"\],)]+)", "<path>", message)
|
||||
return message[:300]
|
||||
|
||||
def _register_worker(self) -> None:
|
||||
registered = self.coordinator.register(
|
||||
self.config.worker_name,
|
||||
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))
|
||||
+20
-61
@@ -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, ...]
|
||||
|
||||
+42
-38
@@ -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("_", "-")
|
||||
|
||||
@@ -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
|
||||
@@ -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=$!
|
||||
}
|
||||
|
||||
|
||||
@@ -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"))
|
||||
@@ -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 "<path>" in message
|
||||
|
||||
|
||||
def test_directory_creation_failure_is_reported(tmp_path: Path) -> None:
|
||||
content = b"input fixture"
|
||||
worker, coordinator, _, _, config = daemon(tmp_path, make_task(content), content)
|
||||
(config.work_dir / "task-1" / "1").mkdir(parents=True)
|
||||
assert worker.run_once() == 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}
|
||||
)
|
||||
Reference in New Issue
Block a user