Compare commits

...
Author SHA1 Message Date
Nicolò Boschi 85d5b3bfd9 tests 2026-01-19 18:34:32 +01:00
Nicolò Boschi f15b76fba8 docs 2026-01-19 17:12:55 +01:00
Nicolò Boschi 61f457736e doc 2026-01-19 16:44:45 +01:00
Nicolò Boschi 5de1447ff5 feat: new 'worker' service 2026-01-19 16:32:18 +01:00
25 changed files with 1866 additions and 252 deletions
+16
View File
@@ -80,6 +80,22 @@ Control plane selector labels
app.kubernetes.io/component: control-plane
{{- end }}
{{/*
Worker labels
*/}}
{{- define "hindsight.worker.labels" -}}
{{ include "hindsight.labels" . }}
app.kubernetes.io/component: worker
{{- end }}
{{/*
Worker selector labels
*/}}
{{- define "hindsight.worker.selectorLabels" -}}
{{ include "hindsight.selectorLabels" . }}
app.kubernetes.io/component: worker
{{- end }}
{{/*
Create the name of the service account to use
*/}}
@@ -55,6 +55,11 @@ spec:
{{- end }}
- name: HINDSIGHT_API_DATABASE_URL
value: {{ include "hindsight.databaseUrl" . | quote }}
{{- /* Disable internal worker when dedicated workers are enabled */}}
{{- if .Values.worker.enabled }}
- name: HINDSIGHT_API_WORKER_ENABLED
value: "false"
{{- end }}
{{- range $key, $value := .Values.api.env }}
- name: {{ $key }}
value: {{ $value | quote }}
@@ -0,0 +1,25 @@
{{- if .Values.worker.enabled }}
apiVersion: v1
kind: Service
metadata:
name: {{ include "hindsight.fullname" . }}-worker
labels:
{{- include "hindsight.worker.labels" . | nindent 4 }}
{{- if .Values.podAnnotations }}
annotations:
{{- /* Common Prometheus annotations for metrics scraping */}}
prometheus.io/scrape: "true"
prometheus.io/port: {{ .Values.worker.service.port | quote }}
prometheus.io/path: "/metrics"
{{- end }}
spec:
# Headless service for StatefulSet (enables stable DNS names like worker-0.worker.namespace)
clusterIP: None
ports:
- port: {{ .Values.worker.service.port }}
targetPort: {{ .Values.worker.service.targetPort }}
protocol: TCP
name: http
selector:
{{- include "hindsight.worker.selectorLabels" . | nindent 4 }}
{{- end }}
@@ -0,0 +1,110 @@
{{- if .Values.worker.enabled }}
apiVersion: apps/v1
kind: StatefulSet
metadata:
name: {{ include "hindsight.fullname" . }}-worker
labels:
{{- include "hindsight.worker.labels" . | nindent 4 }}
spec:
serviceName: {{ include "hindsight.fullname" . }}-worker
replicas: {{ .Values.worker.replicaCount }}
selector:
matchLabels:
{{- include "hindsight.worker.selectorLabels" . | nindent 6 }}
template:
metadata:
annotations:
{{- if not .Values.existingSecret }}
checksum/secret: {{ include (print $.Template.BasePath "/secret.yaml") . | sha256sum }}
{{- end }}
{{- with .Values.podAnnotations }}
{{- toYaml . | nindent 8 }}
{{- end }}
labels:
{{- include "hindsight.worker.selectorLabels" . | nindent 8 }}
spec:
{{- if .Values.serviceAccount.create }}
serviceAccountName: {{ include "hindsight.serviceAccountName" . }}
{{- end }}
securityContext:
{{- toYaml .Values.podSecurityContext | nindent 8 }}
containers:
- name: worker
securityContext:
{{- toYaml .Values.securityContext | nindent 10 }}
image: "{{ .Values.worker.image.repository }}:{{ .Values.worker.image.tag | default .Values.version }}"
imagePullPolicy: {{ .Values.worker.image.pullPolicy }}
command: ["hindsight-worker"]
ports:
- name: http
containerPort: {{ .Values.worker.service.targetPort }}
protocol: TCP
{{- if .Values.existingSecret }}
envFrom:
- secretRef:
name: {{ .Values.existingSecret }}
{{- end }}
env:
{{- /* POSTGRES_PASSWORD must be defined before DATABASE_URL for $(VAR) interpolation */}}
{{- if not .Values.postgresql.enabled }}
- name: POSTGRES_PASSWORD
valueFrom:
secretKeyRef:
name: {{ include "hindsight.secretName" . }}
key: postgres-password
{{- end }}
- name: HINDSIGHT_API_DATABASE_URL
value: {{ include "hindsight.databaseUrl" . | quote }}
{{- /* Worker ID uses pod name (StatefulSet provides stable names like worker-0, worker-1) */}}
- name: HINDSIGHT_API_WORKER_ID
valueFrom:
fieldRef:
fieldPath: metadata.name
{{- /* Inherit LLM config from api.env */}}
{{- range $key, $value := .Values.api.env }}
- name: {{ $key }}
value: {{ $value | quote }}
{{- end }}
{{- /* Worker-specific env vars */}}
{{- range $key, $value := .Values.worker.env }}
- name: {{ $key }}
value: {{ $value | quote }}
{{- end }}
{{- /* Only use secrets when not using existingSecret */}}
{{- if not .Values.existingSecret }}
{{- /* Inherit secrets from api.secrets */}}
{{- range $key, $value := .Values.api.secrets }}
- name: {{ $key }}
valueFrom:
secretKeyRef:
name: {{ include "hindsight.secretName" $ }}
key: {{ $key }}
{{- end }}
{{- /* Worker-specific secrets (can override api.secrets) */}}
{{- range $key, $value := .Values.worker.secrets }}
- name: {{ $key }}
valueFrom:
secretKeyRef:
name: {{ include "hindsight.secretName" $ }}
key: {{ $key }}
{{- end }}
{{- end }}
livenessProbe:
{{- toYaml .Values.worker.livenessProbe | nindent 10 }}
readinessProbe:
{{- toYaml .Values.worker.readinessProbe | nindent 10 }}
resources:
{{- toYaml .Values.worker.resources | nindent 10 }}
{{- with .Values.nodeSelector }}
nodeSelector:
{{- toYaml . | nindent 8 }}
{{- end }}
{{- with .Values.affinity }}
affinity:
{{- toYaml . | nindent 8 }}
{{- end }}
{{- with .Values.tolerations }}
tolerations:
{{- toYaml . | nindent 8 }}
{{- end }}
{{- end }}
+57
View File
@@ -67,6 +67,63 @@ api:
# HINDSIGHT_API_LLM_API_KEY: "your-api-key"
# HINDSIGHT_API_LLM_BASE_URL: "https://api.groq.com/openai/v1"
# Worker settings (distributed task processing)
# When enabled, dedicated worker pods process tasks and the API's internal worker is disabled
worker:
enabled: false
replicaCount: 2
image:
repository: ghcr.io/vectorize-io/hindsight-api
pullPolicy: IfNotPresent
# tag defaults to .Values.version if not specified
service:
# Service for metrics scraping (headless for StatefulSet)
port: 8889
targetPort: 8889
# Resource limits and requests
resources:
limits:
cpu: 2000m
memory: 4Gi
requests:
cpu: 500m
memory: 1Gi
# Liveness and readiness probes
livenessProbe:
httpGet:
path: /health
port: 8889
initialDelaySeconds: 30
periodSeconds: 10
timeoutSeconds: 5
failureThreshold: 3
readinessProbe:
httpGet:
path: /health
port: 8889
initialDelaySeconds: 10
periodSeconds: 5
timeoutSeconds: 3
failureThreshold: 3
# Worker-specific environment variables
env:
# Poll interval in milliseconds (how often to check for new tasks)
HINDSIGHT_API_WORKER_POLL_INTERVAL_MS: "500"
# Number of tasks to claim per poll cycle
HINDSIGHT_API_WORKER_BATCH_SIZE: "10"
# Max retries before marking a task as failed
HINDSIGHT_API_WORKER_MAX_RETRIES: "3"
# HTTP port for metrics/health (matches service.targetPort)
HINDSIGHT_API_WORKER_HTTP_PORT: "8889"
# Secret environment variables (inherited from api.secrets if not specified)
secrets: {}
# Image settings for control plane
controlPlane:
enabled: true
+59
View File
@@ -244,6 +244,65 @@ def run_db_migration(
typer.echo("Database migrations completed successfully")
async def _decommission_worker(db_url: str, worker_id: str, schema: str = "public") -> int:
"""Release all tasks owned by a worker, setting them back to pending status."""
is_pg0, instance_name, _ = parse_pg0_url(db_url)
if is_pg0:
typer.echo(f"Starting embedded PostgreSQL (instance: {instance_name})...")
resolved_url = await resolve_database_url(db_url)
conn = await asyncpg.connect(resolved_url)
try:
table = _fq_table("async_operations", schema)
result = await conn.fetch(
f"""
UPDATE {table}
SET status = 'pending', worker_id = NULL, claimed_at = NULL, updated_at = now()
WHERE worker_id = $1 AND status = 'processing'
RETURNING operation_id
""",
worker_id,
)
return len(result)
finally:
await conn.close()
@app.command(name="decommission-worker")
def decommission_worker(
worker_id: str = typer.Argument(..., help="Worker ID to decommission"),
schema: str = typer.Option("public", "--schema", "-s", help="Database schema"),
yes: bool = typer.Option(False, "--yes", "-y", help="Skip confirmation prompt"),
):
"""Release all tasks owned by a worker (sets status back to pending).
Use this command when a worker has crashed or been removed without graceful shutdown.
All tasks that were being processed by the worker will be released back to the queue
so other workers can pick them up.
"""
config = HindsightConfig.from_env()
if not config.database_url:
typer.echo("Error: Database URL not configured.", err=True)
typer.echo("Set HINDSIGHT_API_DATABASE_URL environment variable.", err=True)
raise typer.Exit(1)
if not yes:
typer.confirm(
f"This will release all tasks owned by worker '{worker_id}' back to pending. Continue?",
abort=True,
)
typer.echo(f"Decommissioning worker '{worker_id}' (schema: {schema})...")
count = asyncio.run(_decommission_worker(config.database_url, worker_id, schema))
if count > 0:
typer.echo(f"Released {count} task(s) from worker '{worker_id}'")
else:
typer.echo(f"No tasks found for worker '{worker_id}'")
def main():
app()
@@ -0,0 +1,109 @@
"""add_worker_columns
Revision ID: l7g8h9i0j1k2
Revises: k6f7g8h9i0j1
Create Date: 2026-01-19 00:00:00.000000
This migration adds columns to async_operations for distributed worker support:
- worker_id: ID of the worker that claimed the task
- claimed_at: When the task was claimed
- retry_count: Number of retry attempts
- task_payload: The serialized task dictionary
"""
from collections.abc import Sequence
import sqlalchemy as sa
from alembic import context, op
from sqlalchemy.dialects import postgresql
# revision identifiers, used by Alembic.
revision: str = "l7g8h9i0j1k2"
down_revision: str | Sequence[str] | None = "k6f7g8h9i0j1"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
def _get_schema_prefix() -> str:
"""Get schema prefix for table names (required for multi-tenant support)."""
schema = context.config.get_main_option("target_schema")
return f'"{schema}".' if schema else ""
def upgrade() -> None:
"""Add worker columns to async_operations."""
schema = _get_schema_prefix()
# Add worker_id column (ID of worker that claimed the task)
op.add_column(
"async_operations",
sa.Column("worker_id", sa.Text(), nullable=True),
schema=context.config.get_main_option("target_schema") or None,
)
# Add claimed_at column (when task was claimed by worker)
op.add_column(
"async_operations",
sa.Column("claimed_at", postgresql.TIMESTAMP(timezone=True), nullable=True),
schema=context.config.get_main_option("target_schema") or None,
)
# Add retry_count column (number of retry attempts)
op.add_column(
"async_operations",
sa.Column("retry_count", sa.Integer(), server_default="0", nullable=False),
schema=context.config.get_main_option("target_schema") or None,
)
# Add task_payload column (serialized task dictionary)
op.add_column(
"async_operations",
sa.Column(
"task_payload",
postgresql.JSONB(astext_type=sa.Text()),
nullable=True,
),
schema=context.config.get_main_option("target_schema") or None,
)
# Add index for efficient worker polling (pending tasks ordered by creation time)
op.execute(
f"CREATE INDEX idx_async_operations_pending_claim ON {schema}async_operations (status, created_at) "
f"WHERE status = 'pending' AND task_payload IS NOT NULL"
)
# Add index for finding tasks by worker_id (for decommissioning)
op.execute(
f"CREATE INDEX idx_async_operations_worker_id ON {schema}async_operations (worker_id) WHERE worker_id IS NOT NULL"
)
def downgrade() -> None:
"""Remove worker columns from async_operations."""
schema = _get_schema_prefix()
# Drop indexes
op.execute(f"DROP INDEX IF EXISTS {schema}idx_async_operations_pending_claim")
op.execute(f"DROP INDEX IF EXISTS {schema}idx_async_operations_worker_id")
# Drop columns
op.drop_column(
"async_operations",
"task_payload",
schema=context.config.get_main_option("target_schema") or None,
)
op.drop_column(
"async_operations",
"retry_count",
schema=context.config.get_main_option("target_schema") or None,
)
op.drop_column(
"async_operations",
"claimed_at",
schema=context.config.get_main_option("target_schema") or None,
)
op.drop_column(
"async_operations",
"worker_id",
schema=context.config.get_main_option("target_schema") or None,
)
+35
View File
@@ -1408,6 +1408,16 @@ def create_app(
Lifespan context manager for startup and shutdown events.
Note: This only fires when running the app standalone, not when mounted.
"""
import asyncio
import socket
from hindsight_api.config import get_config
from hindsight_api.worker import WorkerPoller
config = get_config()
poller = None
poller_task = None
# Initialize OpenTelemetry metrics
try:
prometheus_reader = initialize_metrics(service_name="hindsight-api", service_version="1.0.0")
@@ -1430,6 +1440,20 @@ def create_app(
metrics_collector.set_db_pool(memory._pool)
logging.info("DB pool metrics configured")
# Start worker poller if enabled (standalone mode)
if config.worker_enabled and memory._pool is not None:
worker_id = config.worker_id or socket.gethostname()
poller = WorkerPoller(
pool=memory._pool,
worker_id=worker_id,
executor=memory.execute_task,
poll_interval_ms=config.worker_poll_interval_ms,
batch_size=config.worker_batch_size,
max_retries=config.worker_max_retries,
)
poller_task = asyncio.create_task(poller.run())
logging.info(f"Worker poller started (worker_id={worker_id})")
# Call HTTP extension startup hook
if http_extension:
await http_extension.on_startup()
@@ -1437,6 +1461,17 @@ def create_app(
yield
# Shutdown worker poller if running
if poller is not None:
await poller.shutdown_graceful(timeout=30.0)
if poller_task is not None:
poller_task.cancel()
try:
await poller_task
except asyncio.CancelledError:
pass
logging.info("Worker poller stopped")
# Call HTTP extension shutdown hook
if http_extension:
await http_extension.on_shutdown()
+28 -20
View File
@@ -106,10 +106,13 @@ ENV_DB_POOL_MAX_SIZE = "HINDSIGHT_API_DB_POOL_MAX_SIZE"
ENV_DB_COMMAND_TIMEOUT = "HINDSIGHT_API_DB_COMMAND_TIMEOUT"
ENV_DB_ACQUIRE_TIMEOUT = "HINDSIGHT_API_DB_ACQUIRE_TIMEOUT"
# Background task processing
ENV_TASK_BACKEND = "HINDSIGHT_API_TASK_BACKEND"
ENV_TASK_BACKEND_MEMORY_BATCH_SIZE = "HINDSIGHT_API_TASK_BACKEND_MEMORY_BATCH_SIZE"
ENV_TASK_BACKEND_MEMORY_BATCH_INTERVAL = "HINDSIGHT_API_TASK_BACKEND_MEMORY_BATCH_INTERVAL"
# Worker configuration (distributed task processing)
ENV_WORKER_ENABLED = "HINDSIGHT_API_WORKER_ENABLED"
ENV_WORKER_ID = "HINDSIGHT_API_WORKER_ID"
ENV_WORKER_POLL_INTERVAL_MS = "HINDSIGHT_API_WORKER_POLL_INTERVAL_MS"
ENV_WORKER_MAX_RETRIES = "HINDSIGHT_API_WORKER_MAX_RETRIES"
ENV_WORKER_BATCH_SIZE = "HINDSIGHT_API_WORKER_BATCH_SIZE"
ENV_WORKER_HTTP_PORT = "HINDSIGHT_API_WORKER_HTTP_PORT"
# Reflect agent settings
ENV_REFLECT_MAX_ITERATIONS = "HINDSIGHT_API_REFLECT_MAX_ITERATIONS"
@@ -177,10 +180,13 @@ DEFAULT_DB_POOL_MAX_SIZE = 100
DEFAULT_DB_COMMAND_TIMEOUT = 60 # seconds
DEFAULT_DB_ACQUIRE_TIMEOUT = 30 # seconds
# Background task processing
DEFAULT_TASK_BACKEND = "memory" # Options: "memory", "noop"
DEFAULT_TASK_BACKEND_MEMORY_BATCH_SIZE = 10
DEFAULT_TASK_BACKEND_MEMORY_BATCH_INTERVAL = 1.0 # seconds
# Worker configuration (distributed task processing)
DEFAULT_WORKER_ENABLED = True # API runs worker by default (standalone mode)
DEFAULT_WORKER_ID = None # Will use hostname if not specified
DEFAULT_WORKER_POLL_INTERVAL_MS = 500 # Poll database every 500ms
DEFAULT_WORKER_MAX_RETRIES = 3 # Max retries before marking task failed
DEFAULT_WORKER_BATCH_SIZE = 10 # Tasks to claim per poll cycle
DEFAULT_WORKER_HTTP_PORT = 8889 # HTTP port for worker metrics/health
# Reflect agent settings
DEFAULT_REFLECT_MAX_ITERATIONS = 10 # Max tool call iterations before forcing response
@@ -331,10 +337,13 @@ class HindsightConfig:
db_command_timeout: int
db_acquire_timeout: int
# Background task processing
task_backend: str
task_backend_memory_batch_size: int
task_backend_memory_batch_interval: float
# Worker configuration (distributed task processing)
worker_enabled: bool
worker_id: str | None
worker_poll_interval_ms: int
worker_max_retries: int
worker_batch_size: int
worker_http_port: int
# Reflect agent settings
reflect_max_iterations: int
@@ -424,14 +433,13 @@ class HindsightConfig:
db_pool_max_size=int(os.getenv(ENV_DB_POOL_MAX_SIZE, str(DEFAULT_DB_POOL_MAX_SIZE))),
db_command_timeout=int(os.getenv(ENV_DB_COMMAND_TIMEOUT, str(DEFAULT_DB_COMMAND_TIMEOUT))),
db_acquire_timeout=int(os.getenv(ENV_DB_ACQUIRE_TIMEOUT, str(DEFAULT_DB_ACQUIRE_TIMEOUT))),
# Background task processing
task_backend=os.getenv(ENV_TASK_BACKEND, DEFAULT_TASK_BACKEND),
task_backend_memory_batch_size=int(
os.getenv(ENV_TASK_BACKEND_MEMORY_BATCH_SIZE, str(DEFAULT_TASK_BACKEND_MEMORY_BATCH_SIZE))
),
task_backend_memory_batch_interval=float(
os.getenv(ENV_TASK_BACKEND_MEMORY_BATCH_INTERVAL, str(DEFAULT_TASK_BACKEND_MEMORY_BATCH_INTERVAL))
),
# Worker configuration
worker_enabled=os.getenv(ENV_WORKER_ENABLED, str(DEFAULT_WORKER_ENABLED)).lower() == "true",
worker_id=os.getenv(ENV_WORKER_ID) or DEFAULT_WORKER_ID,
worker_poll_interval_ms=int(os.getenv(ENV_WORKER_POLL_INTERVAL_MS, str(DEFAULT_WORKER_POLL_INTERVAL_MS))),
worker_max_retries=int(os.getenv(ENV_WORKER_MAX_RETRIES, str(DEFAULT_WORKER_MAX_RETRIES))),
worker_batch_size=int(os.getenv(ENV_WORKER_BATCH_SIZE, str(DEFAULT_WORKER_BATCH_SIZE))),
worker_http_port=int(os.getenv(ENV_WORKER_HTTP_PORT, str(DEFAULT_WORKER_HTTP_PORT))),
# Reflect agent settings
reflect_max_iterations=int(os.getenv(ENV_REFLECT_MAX_ITERATIONS, str(DEFAULT_REFLECT_MAX_ITERATIONS))),
)
@@ -159,7 +159,7 @@ from .retain.types import RetainContentDict
from .search import think_utils
from .search.reranking import CrossEncoderReranker
from .search.tags import TagsMatch
from .task_backend import AsyncIOQueueBackend, NoopTaskBackend, TaskBackend
from .task_backend import BrokerTaskBackend, SyncTaskBackend, TaskBackend
class Budget(str, Enum):
@@ -229,8 +229,6 @@ class MemoryEngine(MemoryEngineInterface):
db_command_timeout: int | None = None,
db_acquire_timeout: int | None = None,
task_backend: TaskBackend | None = None,
task_batch_size: int | None = None,
task_batch_interval: float | None = None,
run_migrations: bool = True,
operation_validator: "OperationValidatorExtension | None" = None,
tenant_extension: "TenantExtension | None" = None,
@@ -265,9 +263,7 @@ class MemoryEngine(MemoryEngineInterface):
pool_max_size: Maximum number of connections in the pool. Defaults to HINDSIGHT_API_DB_POOL_MAX_SIZE.
db_command_timeout: PostgreSQL command timeout in seconds. Defaults to HINDSIGHT_API_DB_COMMAND_TIMEOUT.
db_acquire_timeout: Connection acquisition timeout in seconds. Defaults to HINDSIGHT_API_DB_ACQUIRE_TIMEOUT.
task_backend: Custom task backend. If not provided, uses AsyncIOQueueBackend.
task_batch_size: Background task batch size. Defaults to HINDSIGHT_API_TASK_BATCH_SIZE.
task_batch_interval: Background task batch interval in seconds. Defaults to HINDSIGHT_API_TASK_BATCH_INTERVAL.
task_backend: Custom task backend. If not provided, uses BrokerTaskBackend for distributed processing.
run_migrations: Whether to run database migrations during initialize(). Default: True
operation_validator: Optional extension to validate operations before execution.
If provided, retain/recall/reflect operations will be validated.
@@ -405,13 +401,9 @@ class MemoryEngine(MemoryEngineInterface):
self._cross_encoder_reranker = CrossEncoderReranker(cross_encoder=cross_encoder)
# Initialize task backend
_task_batch_size = task_batch_size if task_batch_size is not None else config.task_backend_memory_batch_size
_task_batch_interval = (
task_batch_interval if task_batch_interval is not None else config.task_backend_memory_batch_interval
)
self._task_backend = task_backend or AsyncIOQueueBackend(
batch_size=_task_batch_size, batch_interval=_task_batch_interval
)
# If no custom backend provided, use BrokerTaskBackend which stores tasks in PostgreSQL
# The pool_getter lambda will return the pool once it's initialized
self._task_backend = task_backend or BrokerTaskBackend(pool_getter=lambda: self._pool)
# Backpressure mechanism: limit concurrent searches to prevent overwhelming the database
# Configurable via HINDSIGHT_API_RECALL_MAX_CONCURRENT (default: 50)
+106 -196
View File
@@ -1,31 +1,40 @@
"""
Abstract task backend for running async tasks.
Task backend for distributed task processing.
This provides an abstraction that can be adapted to different execution models:
- AsyncIO queue (default implementation)
- Pub/Sub architectures (future)
- Message brokers (future)
This provides an abstraction for task storage and execution:
- BrokerTaskBackend: Uses PostgreSQL as broker (production)
- SyncTaskBackend: Executes tasks immediately (testing/embedded)
"""
import asyncio
import json
import logging
from abc import ABC, abstractmethod
from collections.abc import Awaitable, Callable
from typing import Any
from typing import TYPE_CHECKING, Any
if TYPE_CHECKING:
import asyncpg
logger = logging.getLogger(__name__)
def fq_table(table: str, schema: str | None = None) -> str:
"""Get fully-qualified table name with optional schema prefix."""
if schema:
return f'"{schema}".{table}'
return table
class TaskBackend(ABC):
"""
Abstract base class for task execution backends.
Implementations must:
1. Store/publish task events (as serializable dicts)
2. Execute tasks through a provided executor callback
2. Execute tasks through a provided executor callback (optional)
The backend treats tasks as pure dictionaries that can be serialized
and sent over the network. The executor (typically MemoryEngine.execute_task)
and stored in the database. The executor (typically MemoryEngine.execute_task)
receives the dict and routes it to the appropriate handler.
"""
@@ -46,7 +55,7 @@ class TaskBackend(ABC):
@abstractmethod
async def initialize(self):
"""
Initialize the backend (e.g., start workers, connect to broker).
Initialize the backend (e.g., connect to database).
"""
pass
@@ -63,7 +72,7 @@ class TaskBackend(ABC):
@abstractmethod
async def shutdown(self):
"""
Shutdown the backend gracefully (e.g., stop workers, close connections).
Shutdown the backend gracefully.
"""
pass
@@ -93,9 +102,8 @@ class SyncTaskBackend(TaskBackend):
"""
Synchronous task backend that executes tasks immediately.
This is useful for embedded/CLI usage where we don't want background
workers that prevent clean exit. Tasks are executed inline rather than
being queued.
This is useful for tests and embedded/CLI usage where we don't want
background workers. Tasks are executed inline rather than being queued.
"""
async def initialize(self):
@@ -121,221 +129,123 @@ class SyncTaskBackend(TaskBackend):
logger.debug("SyncTaskBackend shutdown")
class NoopTaskBackend(TaskBackend):
class BrokerTaskBackend(TaskBackend):
"""
No-op task backend that discards all tasks.
Task backend using PostgreSQL as broker.
This is useful for tests where background task execution is not needed
and would only slow down the test suite.
submit_task() stores task_payload in async_operations table.
Actual polling and execution is handled separately by WorkerPoller.
This backend is used by the API to store tasks. Workers poll
the database separately to claim and execute tasks.
"""
async def initialize(self):
"""No-op."""
self._initialized = True
logger.debug("NoopTaskBackend initialized")
async def submit_task(self, task_dict: dict[str, Any]):
"""Discard the task (do nothing)."""
pass
async def shutdown(self):
"""No-op."""
self._initialized = False
logger.debug("NoopTaskBackend shutdown")
class AsyncIOQueueBackend(TaskBackend):
"""
Task backend implementation using asyncio queues.
This is the default implementation that uses in-process asyncio queues
and a periodic consumer worker.
"""
def __init__(self, batch_size: int = 10, batch_interval: float = 1.0):
def __init__(
self,
pool_getter: Callable[[], "asyncpg.Pool"],
schema: str | None = None,
):
"""
Initialize AsyncIO queue backend.
Initialize the broker task backend.
Args:
batch_size: Maximum number of tasks to process in one batch
batch_interval: Maximum time (seconds) to wait before processing batch
pool_getter: Callable that returns the asyncpg connection pool
schema: Database schema for multi-tenant support (optional)
"""
super().__init__()
self._queue: asyncio.Queue | None = None
self._worker_task: asyncio.Task | None = None
self._shutdown_event: asyncio.Event | None = None
self._batch_size = batch_size
self._batch_interval = batch_interval
self._in_flight_count = 0
self._in_flight_lock = asyncio.Lock()
self._pool_getter = pool_getter
self._schema = schema
async def initialize(self):
"""Initialize the queue and start the worker."""
if self._initialized:
return
self._queue = asyncio.Queue()
self._shutdown_event = asyncio.Event()
self._worker_task = asyncio.create_task(self._worker())
"""Initialize the backend."""
self._initialized = True
logger.info("AsyncIOQueueBackend initialized")
logger.info("BrokerTaskBackend initialized")
async def submit_task(self, task_dict: dict[str, Any]):
"""
Submit a task by putting it in the queue.
Store task payload in async_operations table.
The task_dict should contain an 'operation_id' if updating an existing
operation record, otherwise a new operation will be created.
Args:
task_dict: Task dictionary to execute
task_dict: Task dictionary to store (must be JSON serializable)
"""
if not self._initialized:
await self.initialize()
await self._queue.put(task_dict)
pool = self._pool_getter()
operation_id = task_dict.get("operation_id")
task_type = task_dict.get("type", "unknown")
bank_id = task_dict.get("bank_id")
payload_json = json.dumps(task_dict)
table = fq_table("async_operations", self._schema)
if operation_id:
# Update existing operation with task payload
await pool.execute(
f"""
UPDATE {table}
SET task_payload = $1::jsonb, updated_at = now()
WHERE operation_id = $2
""",
payload_json,
operation_id,
)
logger.debug(f"Updated task payload for operation {operation_id}")
else:
# Insert new operation (for tasks without pre-created records)
# e.g., access_count_update tasks
import uuid
new_id = uuid.uuid4()
await pool.execute(
f"""
INSERT INTO {table} (operation_id, bank_id, operation_type, status, task_payload)
VALUES ($1, $2, $3, 'pending', $4::jsonb)
""",
new_id,
bank_id,
task_type,
payload_json,
)
logger.debug(f"Created new operation {new_id} for task type {task_type}")
async def shutdown(self):
"""Shutdown the backend."""
self._initialized = False
logger.info("BrokerTaskBackend shutdown")
async def wait_for_pending_tasks(self, timeout: float = 120.0):
"""
Wait for all pending tasks in the queue and in-flight tasks to complete.
Wait for pending tasks to be processed.
This is useful in tests to ensure background tasks complete before assertions.
In the broker model, this polls the database to check if tasks
for this process have been completed. This is useful in tests
when worker_enabled=True (API processes its own tasks).
Args:
timeout: Maximum time to wait in seconds (default 120s for long-running tasks)
timeout: Maximum time to wait in seconds
"""
if not self._initialized or self._queue is None:
return
import asyncio
pool = self._pool_getter()
table = fq_table("async_operations", self._schema)
# Wait for queue to be empty AND no in-flight tasks
start_time = asyncio.get_event_loop().time()
while asyncio.get_event_loop().time() - start_time < timeout:
async with self._in_flight_lock:
in_flight = self._in_flight_count
# Check if there are any pending tasks with payloads
count = await pool.fetchval(
f"""
SELECT COUNT(*) FROM {table}
WHERE status = 'pending' AND task_payload IS NOT NULL
"""
)
if self._queue.empty() and in_flight == 0:
# Queue is empty and no tasks in flight, we're done
if count == 0:
return
# Wait a bit before checking again
await asyncio.sleep(0.5)
async def shutdown(self):
"""Shutdown the worker and drain the queue."""
if not self._initialized:
return
logger.info("Shutting down AsyncIOQueueBackend...")
# Signal shutdown
self._shutdown_event.set()
# Cancel worker
if self._worker_task is not None:
self._worker_task.cancel()
try:
await self._worker_task
except asyncio.CancelledError:
pass # Worker cancelled successfully
self._initialized = False
logger.info("AsyncIOQueueBackend shutdown complete")
async def _execute_task_with_tracking(self, task_dict: dict[str, Any]):
"""Execute a task and track its in-flight status."""
async with self._in_flight_lock:
self._in_flight_count += 1
try:
await self._execute_task(task_dict)
finally:
async with self._in_flight_lock:
self._in_flight_count -= 1
async def _execute_task_no_tracking(self, task_dict: dict[str, Any]):
"""Execute a task without in-flight tracking (tracking done at batch level)."""
await self._execute_task(task_dict)
def _get_queue_stats(self) -> tuple[int, dict[str, int]]:
"""Get current queue size and bank_id distribution."""
queue_size = self._queue.qsize() if self._queue else 0
bank_distribution: dict[str, int] = {}
if queue_size > 0 and self._queue:
# Peek at queue items without removing them
# Note: This is a snapshot and may not be perfectly accurate due to concurrency
try:
# Access internal deque for logging purposes only
items = list(self._queue._queue) # type: ignore[attr-defined]
for item in items:
bank_id = item.get("bank_id", "unknown")
bank_distribution[bank_id] = bank_distribution.get(bank_id, 0) + 1
except Exception:
pass # Queue access failed, return empty distribution
return queue_size, bank_distribution
async def _worker(self):
"""
Background worker that processes tasks in batches.
Collects tasks for up to batch_interval seconds or batch_size items,
then processes them.
"""
while not self._shutdown_event.is_set():
try:
# Collect tasks for batching
tasks = []
deadline = asyncio.get_event_loop().time() + self._batch_interval
while len(tasks) < self._batch_size and asyncio.get_event_loop().time() < deadline:
try:
remaining_time = max(0.1, deadline - asyncio.get_event_loop().time())
task_dict = await asyncio.wait_for(self._queue.get(), timeout=remaining_time)
# Track task as in-flight immediately when picked up from queue
# This prevents wait_for_pending_tasks from returning too early
async with self._in_flight_lock:
self._in_flight_count += 1
tasks.append(task_dict)
except TimeoutError:
break
# Process batch
if tasks:
# Log batch start with queue stats
queue_size, bank_distribution = self._get_queue_stats()
# Summarize batch by task type and bank
batch_summary: dict[str, dict[str, int]] = {}
for task_dict in tasks:
task_type = task_dict.get("type", "unknown")
bank_id = task_dict.get("bank_id", "unknown")
if task_type not in batch_summary:
batch_summary[task_type] = {}
batch_summary[task_type][bank_id] = batch_summary[task_type].get(bank_id, 0) + 1
# Build log message
batch_parts = []
for task_type, banks in sorted(batch_summary.items()):
bank_str = ", ".join(f"{b}:{c}" for b, c in sorted(banks.items()))
batch_parts.append(f"{task_type}[{bank_str}]")
batch_str = ", ".join(batch_parts)
if queue_size > 0:
pending_str = ", ".join(f"{k}:{v}" for k, v in sorted(bank_distribution.items()))
logger.info(
f"Processing {len(tasks)} tasks: {batch_str} (pending={queue_size} [{pending_str}])"
)
else:
logger.info(f"Processing {len(tasks)} tasks: {batch_str}")
# Execute tasks concurrently (in_flight already tracked when picked up)
await asyncio.gather(
*[self._execute_task_no_tracking(task_dict) for task_dict in tasks], return_exceptions=True
)
# Decrement in_flight count after all tasks complete
async with self._in_flight_lock:
self._in_flight_count -= len(tasks)
except asyncio.CancelledError:
break
except Exception as e:
logger.error(f"Worker error: {e}")
await asyncio.sleep(1) # Backoff on error
logger.warning(f"Timeout waiting for pending tasks after {timeout}s")
+6 -3
View File
@@ -219,9 +219,12 @@ def main():
db_pool_max_size=config.db_pool_max_size,
db_command_timeout=config.db_command_timeout,
db_acquire_timeout=config.db_acquire_timeout,
task_backend=config.task_backend,
task_backend_memory_batch_size=config.task_backend_memory_batch_size,
task_backend_memory_batch_interval=config.task_backend_memory_batch_interval,
worker_enabled=config.worker_enabled,
worker_id=config.worker_id,
worker_poll_interval_ms=config.worker_poll_interval_ms,
worker_max_retries=config.worker_max_retries,
worker_batch_size=config.worker_batch_size,
worker_http_port=config.worker_http_port,
reflect_max_iterations=config.reflect_max_iterations,
mental_model_refresh_concurrency=config.mental_model_refresh_concurrency,
)
@@ -0,0 +1,11 @@
"""
Worker package for distributed task processing.
This package provides:
- WorkerPoller: Polls PostgreSQL for pending tasks and executes them
- main: CLI entry point for hindsight-worker
"""
from .poller import WorkerPoller
__all__ = ["WorkerPoller"]
+285
View File
@@ -0,0 +1,285 @@
"""
Command-line interface for Hindsight Worker.
Run the worker with:
hindsight-worker
Stop with Ctrl+C (graceful shutdown).
"""
import argparse
import asyncio
import atexit
import logging
import os
import signal
import socket
import sys
import warnings
from ..config import get_config
from ..engine.task_backend import SyncTaskBackend
from .poller import WorkerPoller
# Filter deprecation warnings from third-party libraries
warnings.filterwarnings("ignore", message="websockets.legacy is deprecated")
warnings.filterwarnings("ignore", message="websockets.server.WebSocketServerProtocol is deprecated")
# Disable tokenizers parallelism to avoid warnings
os.environ["TOKENIZERS_PARALLELISM"] = "false"
logger = logging.getLogger(__name__)
def create_worker_app(poller: WorkerPoller, memory):
"""Create a minimal FastAPI app for worker metrics and health."""
from fastapi import FastAPI
from fastapi.responses import JSONResponse, Response
from prometheus_client import CONTENT_TYPE_LATEST, generate_latest
from ..metrics import create_metrics_collector, get_metrics_collector, initialize_metrics
app = FastAPI(
title="Hindsight Worker",
description="Worker process for distributed task execution",
)
# Initialize OpenTelemetry metrics
try:
prometheus_reader = initialize_metrics(service_name="hindsight-worker", service_version="1.0.0")
create_metrics_collector()
app.state.prometheus_reader = prometheus_reader
logger.info("Metrics initialized - available at /metrics endpoint")
except Exception as e:
logger.warning(f"Failed to initialize metrics: {e}. Metrics will be disabled.")
app.state.prometheus_reader = None
# Set up DB pool metrics if available
metrics_collector = get_metrics_collector()
if memory._pool is not None and hasattr(metrics_collector, "set_db_pool"):
metrics_collector.set_db_pool(memory._pool)
logger.info("DB pool metrics configured")
@app.get(
"/health",
summary="Health check endpoint",
description="Returns worker health status including database connectivity",
tags=["Monitoring"],
)
async def health_endpoint():
"""Health check endpoint."""
health = await memory.health_check()
health["worker_id"] = poller.worker_id
health["is_shutdown"] = poller.is_shutdown
status_code = 200 if health.get("status") == "healthy" else 503
return JSONResponse(content=health, status_code=status_code)
@app.get(
"/metrics",
summary="Prometheus metrics endpoint",
description="Exports metrics in Prometheus format for scraping",
tags=["Monitoring"],
)
async def metrics_endpoint():
"""Return Prometheus metrics."""
metrics_data = generate_latest()
return Response(content=metrics_data, media_type=CONTENT_TYPE_LATEST)
@app.get(
"/",
summary="Worker info",
description="Basic worker information",
tags=["Info"],
)
async def root():
"""Return basic worker info."""
return {
"service": "hindsight-worker",
"worker_id": poller.worker_id,
"is_shutdown": poller.is_shutdown,
}
return app
def main():
"""Main entry point for the hindsight-worker CLI."""
# Load configuration from environment
config = get_config()
parser = argparse.ArgumentParser(
prog="hindsight-worker",
description="Hindsight Worker - distributed task processor",
)
# Worker options
parser.add_argument(
"--worker-id",
default=config.worker_id or socket.gethostname(),
help="Worker identifier (default: hostname, env: HINDSIGHT_API_WORKER_ID)",
)
parser.add_argument(
"--poll-interval",
type=int,
default=config.worker_poll_interval_ms,
help=f"Poll interval in milliseconds (default: {config.worker_poll_interval_ms}, env: HINDSIGHT_API_WORKER_POLL_INTERVAL_MS)",
)
parser.add_argument(
"--batch-size",
type=int,
default=config.worker_batch_size,
help=f"Tasks to claim per poll (default: {config.worker_batch_size}, env: HINDSIGHT_API_WORKER_BATCH_SIZE)",
)
parser.add_argument(
"--max-retries",
type=int,
default=config.worker_max_retries,
help=f"Max retries before marking failed (default: {config.worker_max_retries}, env: HINDSIGHT_API_WORKER_MAX_RETRIES)",
)
# HTTP server options
parser.add_argument(
"--http-port",
type=int,
default=config.worker_http_port,
help=f"HTTP port for metrics/health endpoints (default: {config.worker_http_port}, env: HINDSIGHT_API_WORKER_HTTP_PORT)",
)
parser.add_argument(
"--http-host",
default="0.0.0.0",
help="HTTP host to bind (default: 0.0.0.0)",
)
# Logging options
parser.add_argument(
"--log-level",
default=config.log_level,
choices=["critical", "error", "warning", "info", "debug", "trace"],
help=f"Log level (default: {config.log_level}, env: HINDSIGHT_API_LOG_LEVEL)",
)
args = parser.parse_args()
# Configure logging
config.configure_logging()
# Import MemoryEngine here to avoid circular imports
from .. import MemoryEngine
print(f"Starting Hindsight Worker: {args.worker_id}")
print(f" Poll interval: {args.poll_interval}ms")
print(f" Batch size: {args.batch_size}")
print(f" Max retries: {args.max_retries}")
print(f" HTTP server: {args.http_host}:{args.http_port}")
print()
# Global references for cleanup
memory = None
poller = None
async def run():
nonlocal memory, poller
import uvicorn
# Initialize MemoryEngine
# Workers use SyncTaskBackend because they execute tasks directly,
# they don't need to store tasks (they poll from DB)
memory = MemoryEngine(
run_migrations=False, # Workers don't run migrations
task_backend=SyncTaskBackend(),
)
await memory.initialize()
print(f"Database connected: {config.database_url}")
# Create and start the poller
poller = WorkerPoller(
pool=memory._pool,
worker_id=args.worker_id,
executor=memory.execute_task,
poll_interval_ms=args.poll_interval,
batch_size=args.batch_size,
max_retries=args.max_retries,
)
# Create the HTTP app for metrics/health
app = create_worker_app(poller, memory)
# Setup signal handlers for graceful shutdown
shutdown_requested = asyncio.Event()
def signal_handler(signum, frame):
print(f"\nReceived signal {signum}, initiating graceful shutdown...")
shutdown_requested.set()
signal.signal(signal.SIGINT, signal_handler)
signal.signal(signal.SIGTERM, signal_handler)
# Create uvicorn config and server
uvicorn_config = uvicorn.Config(
app,
host=args.http_host,
port=args.http_port,
log_level="info", # Reduce uvicorn noise
access_log=False,
)
server = uvicorn.Server(uvicorn_config)
# Run the poller and HTTP server concurrently
poller_task = asyncio.create_task(poller.run())
http_task = asyncio.create_task(server.serve())
print(f"Worker started. Metrics available at http://{args.http_host}:{args.http_port}/metrics")
# Wait for shutdown signal
await shutdown_requested.wait()
# Graceful shutdown
print("Shutting down HTTP server...")
server.should_exit = True
print("Waiting for poller to finish...")
await poller.shutdown_graceful(timeout=30.0)
poller_task.cancel()
try:
await poller_task
except asyncio.CancelledError:
pass
# Wait for HTTP server to finish
try:
await asyncio.wait_for(http_task, timeout=5.0)
except asyncio.TimeoutError:
http_task.cancel()
try:
await http_task
except asyncio.CancelledError:
pass
# Close memory engine
await memory.close()
print("Worker shutdown complete")
def cleanup():
"""Synchronous cleanup for atexit."""
if memory is not None and memory._pg0 is not None:
try:
loop = asyncio.new_event_loop()
loop.run_until_complete(memory._pg0.stop())
loop.close()
print("\npg0 stopped.")
except Exception as e:
print(f"\nError stopping pg0: {e}")
atexit.register(cleanup)
try:
asyncio.run(run())
except KeyboardInterrupt:
print("\nWorker interrupted")
sys.exit(0)
if __name__ == "__main__":
main()
@@ -0,0 +1,281 @@
"""
Worker poller for distributed task execution.
Polls PostgreSQL for pending tasks and executes them using
FOR UPDATE SKIP LOCKED for safe concurrent claiming.
"""
import asyncio
import json
import logging
import traceback
from collections.abc import Awaitable, Callable
from typing import TYPE_CHECKING, Any
if TYPE_CHECKING:
import asyncpg
logger = logging.getLogger(__name__)
def fq_table(table: str, schema: str | None = None) -> str:
"""Get fully-qualified table name with optional schema prefix."""
if schema:
return f'"{schema}".{table}'
return table
class WorkerPoller:
"""
Polls PostgreSQL for pending tasks and executes them.
Uses FOR UPDATE SKIP LOCKED for safe distributed claiming,
allowing multiple workers to process tasks without conflicts.
"""
def __init__(
self,
pool: "asyncpg.Pool",
worker_id: str,
executor: Callable[[dict[str, Any]], Awaitable[None]],
poll_interval_ms: int = 500,
batch_size: int = 10,
max_retries: int = 3,
schema: str | None = None,
):
"""
Initialize the worker poller.
Args:
pool: asyncpg connection pool
worker_id: Unique identifier for this worker
executor: Async function to execute tasks (typically MemoryEngine.execute_task)
poll_interval_ms: Interval between polls when no tasks found (milliseconds)
batch_size: Maximum number of tasks to claim per poll cycle
max_retries: Maximum retry attempts before marking task as failed
schema: Database schema for multi-tenant support (optional)
"""
self._pool = pool
self._worker_id = worker_id
self._executor = executor
self._poll_interval_ms = poll_interval_ms
self._batch_size = batch_size
self._max_retries = max_retries
self._schema = schema
self._shutdown = asyncio.Event()
self._current_tasks: set[asyncio.Task] = set()
self._in_flight_count = 0
self._in_flight_lock = asyncio.Lock()
async def claim_batch(self) -> list[tuple[str, dict[str, Any]]]:
"""
Claim up to batch_size pending tasks atomically.
Uses FOR UPDATE SKIP LOCKED to ensure no conflicts with other workers.
Returns:
List of tuples (operation_id, task_dict)
"""
table = fq_table("async_operations", self._schema)
async with self._pool.acquire() as conn:
async with conn.transaction():
# Select and lock pending tasks
rows = await conn.fetch(
f"""
SELECT operation_id, task_payload
FROM {table}
WHERE status = 'pending' AND task_payload IS NOT NULL
ORDER BY created_at
LIMIT $1
FOR UPDATE SKIP LOCKED
""",
self._batch_size,
)
if not rows:
return []
# Claim the tasks by updating status and worker_id
operation_ids = [row["operation_id"] for row in rows]
await conn.execute(
f"""
UPDATE {table}
SET status = 'processing', worker_id = $1, claimed_at = now(), updated_at = now()
WHERE operation_id = ANY($2)
""",
self._worker_id,
operation_ids,
)
# Parse and return task payloads
return [(str(row["operation_id"]), json.loads(row["task_payload"])) for row in rows]
async def _mark_completed(self, operation_id: str):
"""Mark a task as completed."""
table = fq_table("async_operations", self._schema)
await self._pool.execute(
f"""
UPDATE {table}
SET status = 'completed', completed_at = now(), updated_at = now()
WHERE operation_id = $1
""",
operation_id,
)
async def _mark_failed(self, operation_id: str, error_message: str):
"""Mark a task as failed with error message."""
table = fq_table("async_operations", self._schema)
# Truncate error message if too long (max 5000 chars in schema)
error_message = error_message[:5000] if len(error_message) > 5000 else error_message
await self._pool.execute(
f"""
UPDATE {table}
SET status = 'failed', error_message = $2, completed_at = now(), updated_at = now()
WHERE operation_id = $1
""",
operation_id,
error_message,
)
async def _retry_or_fail(self, operation_id: str, error_message: str):
"""Increment retry count or mark as failed if max retries exceeded."""
table = fq_table("async_operations", self._schema)
# Get current retry count
row = await self._pool.fetchrow(
f"SELECT retry_count FROM {table} WHERE operation_id = $1",
operation_id,
)
if row is None:
logger.warning(f"Operation {operation_id} not found, cannot retry")
return
retry_count = row["retry_count"]
if retry_count >= self._max_retries:
# Max retries exceeded, mark as failed
await self._mark_failed(
operation_id, f"Max retries ({self._max_retries}) exceeded. Last error: {error_message}"
)
logger.error(f"Task {operation_id} failed after {retry_count} retries")
else:
# Increment retry and reset to pending
await self._pool.execute(
f"""
UPDATE {table}
SET status = 'pending', worker_id = NULL, claimed_at = NULL,
retry_count = retry_count + 1, updated_at = now()
WHERE operation_id = $1
""",
operation_id,
)
logger.warning(f"Task {operation_id} failed, will retry (attempt {retry_count + 1}/{self._max_retries})")
async def execute_task(self, operation_id: str, task_dict: dict[str, Any]):
"""Execute a single task and update its status."""
task_type = task_dict.get("type", "unknown")
bank_id = task_dict.get("bank_id", "unknown")
try:
logger.debug(f"Executing task {operation_id} (type={task_type}, bank={bank_id})")
await self._executor(task_dict)
await self._mark_completed(operation_id)
logger.debug(f"Task {operation_id} completed successfully")
except Exception as e:
error_msg = f"{type(e).__name__}: {e}\n{traceback.format_exc()}"
logger.error(f"Task {operation_id} failed: {e}")
await self._retry_or_fail(operation_id, error_msg)
async def run(self):
"""
Main polling loop.
Continuously polls for pending tasks, claims them, and executes them
until shutdown is signaled.
"""
logger.info(f"Worker {self._worker_id} starting polling loop")
while not self._shutdown.is_set():
try:
# Claim a batch of tasks
tasks = await self.claim_batch()
if tasks:
# Log batch info
task_types = {}
for _, task_dict in tasks:
t = task_dict.get("type", "unknown")
task_types[t] = task_types.get(t, 0) + 1
types_str = ", ".join(f"{k}:{v}" for k, v in task_types.items())
logger.info(f"Worker {self._worker_id} claimed {len(tasks)} tasks: {types_str}")
# Track in-flight tasks
async with self._in_flight_lock:
self._in_flight_count += len(tasks)
# Execute tasks concurrently
try:
await asyncio.gather(
*[self.execute_task(op_id, task_dict) for op_id, task_dict in tasks],
return_exceptions=True,
)
finally:
async with self._in_flight_lock:
self._in_flight_count -= len(tasks)
else:
# No tasks found, wait before polling again
try:
await asyncio.wait_for(
self._shutdown.wait(),
timeout=self._poll_interval_ms / 1000,
)
except asyncio.TimeoutError:
pass # Normal timeout, continue polling
except asyncio.CancelledError:
logger.info(f"Worker {self._worker_id} polling loop cancelled")
break
except Exception as e:
logger.error(f"Worker {self._worker_id} error in polling loop: {e}")
traceback.print_exc()
# Backoff on error
await asyncio.sleep(1)
logger.info(f"Worker {self._worker_id} polling loop stopped")
async def shutdown_graceful(self, timeout: float = 30.0):
"""
Signal shutdown and wait for current tasks to complete.
Args:
timeout: Maximum time to wait for in-flight tasks (seconds)
"""
logger.info(f"Worker {self._worker_id} initiating graceful shutdown")
self._shutdown.set()
# Wait for in-flight tasks to complete
start_time = asyncio.get_event_loop().time()
while asyncio.get_event_loop().time() - start_time < timeout:
async with self._in_flight_lock:
in_flight = self._in_flight_count
if in_flight == 0:
logger.info(f"Worker {self._worker_id} graceful shutdown complete")
return
logger.info(f"Worker {self._worker_id} waiting for {in_flight} in-flight tasks")
await asyncio.sleep(0.5)
logger.warning(f"Worker {self._worker_id} shutdown timeout after {timeout}s")
@property
def worker_id(self) -> str:
"""Get the worker ID."""
return self._worker_id
@property
def is_shutdown(self) -> bool:
"""Check if shutdown has been signaled."""
return self._shutdown.is_set()
+1
View File
@@ -63,6 +63,7 @@ test = [
[project.scripts]
hindsight-api = "hindsight_api.main:main"
hindsight-worker = "hindsight_api.worker.main:main"
hindsight-local-mcp = "hindsight_api.mcp_local:main"
hindsight-admin = "hindsight_api.admin.cli:main"
+3
View File
@@ -12,6 +12,7 @@ from hindsight_api import MemoryEngine, LLMConfig, LocalSTEmbeddings, RequestCon
from hindsight_api.engine.cross_encoder import LocalSTCrossEncoder
from hindsight_api.engine.query_analyzer import DateparserQueryAnalyzer
from hindsight_api.engine.task_backend import SyncTaskBackend
from hindsight_api.pg0 import EmbeddedPostgres
# Default pg0 instance configuration for tests
@@ -147,6 +148,7 @@ async def memory(pg0_db_url, embeddings, cross_encoder, query_analyzer):
Uses pg0_db_url (a postgresql:// URL) directly, so MemoryEngine won't try to
manage pg0 lifecycle - that's handled by the session-scoped pg0_db_url fixture.
Migrations are disabled here since they're run once at session scope in pg0_db_url.
Uses SyncTaskBackend so async tasks execute immediately (no worker needed).
"""
mem = MemoryEngine(
db_url=pg0_db_url, # Direct postgresql:// URL, not pg0://
@@ -160,6 +162,7 @@ async def memory(pg0_db_url, embeddings, cross_encoder, query_analyzer):
pool_min_size=1,
pool_max_size=5,
run_migrations=False, # Migrations already run at session scope
task_backend=SyncTaskBackend(), # Execute tasks immediately in tests
)
await mem.initialize()
yield mem
@@ -17,6 +17,7 @@ from hindsight_api import MemoryEngine, RequestContext
from hindsight_api.engine.embeddings import LocalSTEmbeddings, OpenAIEmbeddings, CohereEmbeddings
from hindsight_api.engine.cross_encoder import LocalSTCrossEncoder, CohereCrossEncoder
from hindsight_api.engine.query_analyzer import DateparserQueryAnalyzer
from hindsight_api.engine.task_backend import SyncTaskBackend
from hindsight_api.extensions import TenantExtension, TenantContext
from hindsight_api.migrations import run_migrations, ensure_embedding_dimension
@@ -323,6 +324,7 @@ class TestOpenAIEmbeddings:
pool_max_size=3,
run_migrations=False,
tenant_extension=SchemaTenantExtension(schema_name),
task_backend=SyncTaskBackend(),
)
try:
@@ -392,6 +394,7 @@ class TestOpenAIEmbeddings:
pool_max_size=3,
run_migrations=False,
tenant_extension=SchemaTenantExtension(schema_name),
task_backend=SyncTaskBackend(),
)
try:
@@ -559,6 +562,7 @@ class TestCohereIntegration:
pool_max_size=3,
run_migrations=False,
tenant_extension=SchemaTenantExtension(schema_name),
task_backend=SyncTaskBackend(),
)
try:
@@ -19,6 +19,7 @@ import pytest_asyncio
from hindsight_api import MemoryEngine, LLMConfig, LocalSTEmbeddings, RequestContext
from hindsight_api.engine.cross_encoder import LocalSTCrossEncoder
from hindsight_api.engine.query_analyzer import DateparserQueryAnalyzer
from hindsight_api.engine.task_backend import SyncTaskBackend
from hindsight_api.engine.retain.fact_extraction import FactExtractionResponse, ExtractedFact
from hindsight_api.engine.llm_wrapper import TokenUsage
@@ -106,6 +107,7 @@ class TestLargeBatchRetain:
pool_max_size=10,
run_migrations=False,
skip_llm_verification=True, # Skip LLM verification since we're mocking
task_backend=SyncTaskBackend(), # Execute tasks immediately in tests
)
await mem.initialize()
yield mem
+592
View File
@@ -0,0 +1,592 @@
"""
Tests for the distributed worker system.
Tests cover:
- BrokerTaskBackend task submission and storage
- WorkerPoller task claiming with FOR UPDATE SKIP LOCKED
- Concurrent workers claiming different tasks (no duplicates)
- Task completion and failure handling
- Retry mechanism
- Worker decommissioning
"""
import asyncio
import json
import uuid
import pytest
import pytest_asyncio
from hindsight_api.engine.task_backend import BrokerTaskBackend, SyncTaskBackend
# Use loadgroup to ensure these tests run in the same worker
# since they share database state
pytestmark = pytest.mark.xdist_group("worker_tests")
@pytest_asyncio.fixture
async def pool(pg0_db_url):
"""Create a dedicated connection pool for worker tests."""
import asyncpg
from hindsight_api.pg0 import resolve_database_url
# Resolve pg0:// URL to postgresql:// URL if needed
resolved_url = await resolve_database_url(pg0_db_url)
pool = await asyncpg.create_pool(
resolved_url,
min_size=2,
max_size=10,
command_timeout=30,
)
yield pool
await pool.close()
@pytest_asyncio.fixture
async def clean_operations(pool):
"""Clean up async_operations table before and after tests."""
# Clean before test
await pool.execute("DELETE FROM async_operations WHERE bank_id LIKE 'test-worker-%'")
yield
# Clean after test
await pool.execute("DELETE FROM async_operations WHERE bank_id LIKE 'test-worker-%'")
class TestBrokerTaskBackend:
"""Tests for BrokerTaskBackend task storage."""
@pytest.mark.asyncio
async def test_submit_task_updates_existing_operation(self, pool, clean_operations):
"""Test that submit_task updates task_payload for existing operations."""
# Create an operation record first
operation_id = uuid.uuid4()
bank_id = f"test-worker-{uuid.uuid4().hex[:8]}"
await pool.execute(
"""
INSERT INTO async_operations (operation_id, bank_id, operation_type, status)
VALUES ($1, $2, 'test_operation', 'pending')
""",
operation_id,
bank_id,
)
# Submit task with same operation_id
backend = BrokerTaskBackend(pool_getter=lambda: pool)
await backend.initialize()
task_dict = {
"operation_id": str(operation_id),
"type": "test_task",
"bank_id": bank_id,
"data": {"key": "value"},
}
await backend.submit_task(task_dict)
# Verify task_payload was stored
row = await pool.fetchrow(
"SELECT task_payload, status FROM async_operations WHERE operation_id = $1",
operation_id,
)
assert row is not None
assert row["status"] == "pending"
payload = json.loads(row["task_payload"])
assert payload["type"] == "test_task"
assert payload["data"] == {"key": "value"}
@pytest.mark.asyncio
async def test_submit_task_creates_new_operation(self, pool, clean_operations):
"""Test that submit_task creates new operation when no operation_id provided."""
backend = BrokerTaskBackend(pool_getter=lambda: pool)
await backend.initialize()
bank_id = f"test-worker-{uuid.uuid4().hex[:8]}"
task_dict = {
"type": "access_count_update",
"bank_id": bank_id,
"node_ids": ["node1", "node2"],
}
await backend.submit_task(task_dict)
# Verify new operation was created
row = await pool.fetchrow(
"SELECT operation_type, status, task_payload FROM async_operations WHERE bank_id = $1",
bank_id,
)
assert row is not None
assert row["operation_type"] == "access_count_update"
assert row["status"] == "pending"
payload = json.loads(row["task_payload"])
assert payload["node_ids"] == ["node1", "node2"]
class TestWorkerPoller:
"""Tests for WorkerPoller task claiming and execution."""
@pytest.mark.asyncio
async def test_claim_batch_claims_pending_tasks(self, pool, clean_operations):
"""Test that claim_batch claims pending tasks with task_payload."""
from hindsight_api.worker import WorkerPoller
# Create some pending tasks
bank_id = f"test-worker-{uuid.uuid4().hex[:8]}"
for i in range(3):
op_id = uuid.uuid4()
payload = json.dumps({"type": "test_task", "index": i, "bank_id": bank_id})
await pool.execute(
"""
INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload)
VALUES ($1, $2, 'test', 'pending', $3::jsonb)
""",
op_id,
bank_id,
payload,
)
# Create poller and claim tasks
executed_tasks = []
async def mock_executor(task_dict):
executed_tasks.append(task_dict)
poller = WorkerPoller(
pool=pool,
worker_id="test-worker-1",
executor=mock_executor,
batch_size=10,
)
claimed = await poller.claim_batch()
assert len(claimed) == 3
# Verify tasks are marked as processing with worker_id
rows = await pool.fetch(
"SELECT status, worker_id FROM async_operations WHERE bank_id = $1",
bank_id,
)
for row in rows:
assert row["status"] == "processing"
assert row["worker_id"] == "test-worker-1"
@pytest.mark.asyncio
async def test_claim_batch_respects_batch_size(self, pool, clean_operations):
"""Test that claim_batch respects the batch_size limit."""
from hindsight_api.worker import WorkerPoller
# Create 10 pending tasks
bank_id = f"test-worker-{uuid.uuid4().hex[:8]}"
for i in range(10):
op_id = uuid.uuid4()
payload = json.dumps({"type": "test_task", "index": i, "bank_id": bank_id})
await pool.execute(
"""
INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload)
VALUES ($1, $2, 'test', 'pending', $3::jsonb)
""",
op_id,
bank_id,
payload,
)
# Claim with batch_size=3
poller = WorkerPoller(
pool=pool,
worker_id="test-worker-1",
executor=lambda x: None,
batch_size=3,
)
claimed = await poller.claim_batch()
assert len(claimed) == 3
@pytest.mark.asyncio
async def test_execute_task_marks_completed(self, pool, clean_operations):
"""Test that successful task execution marks task as completed."""
from hindsight_api.worker import WorkerPoller
# Create a pending task
bank_id = f"test-worker-{uuid.uuid4().hex[:8]}"
op_id = uuid.uuid4()
payload = json.dumps({"type": "test_task", "operation_id": str(op_id), "bank_id": bank_id})
await pool.execute(
"""
INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload, worker_id)
VALUES ($1, $2, 'test', 'processing', $3::jsonb, 'test-worker-1')
""",
op_id,
bank_id,
payload,
)
executed = []
async def mock_executor(task_dict):
executed.append(task_dict)
poller = WorkerPoller(
pool=pool,
worker_id="test-worker-1",
executor=mock_executor,
)
# Execute the task
task_dict = json.loads(payload)
await poller.execute_task(str(op_id), task_dict)
assert len(executed) == 1
# Verify task is marked as completed
row = await pool.fetchrow(
"SELECT status, completed_at FROM async_operations WHERE operation_id = $1",
op_id,
)
assert row["status"] == "completed"
assert row["completed_at"] is not None
@pytest.mark.asyncio
async def test_execute_task_retries_on_failure(self, pool, clean_operations):
"""Test that failed task execution triggers retry mechanism."""
from hindsight_api.worker import WorkerPoller
# Create a pending task with retry_count=0
bank_id = f"test-worker-{uuid.uuid4().hex[:8]}"
op_id = uuid.uuid4()
payload = json.dumps({"type": "test_task", "operation_id": str(op_id), "bank_id": bank_id})
await pool.execute(
"""
INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload, worker_id, retry_count)
VALUES ($1, $2, 'test', 'processing', $3::jsonb, 'test-worker-1', 0)
""",
op_id,
bank_id,
payload,
)
async def failing_executor(task_dict):
raise ValueError("Simulated failure")
poller = WorkerPoller(
pool=pool,
worker_id="test-worker-1",
executor=failing_executor,
max_retries=3,
)
# Execute (should fail and retry)
task_dict = json.loads(payload)
await poller.execute_task(str(op_id), task_dict)
# Verify task is back to pending with incremented retry_count
row = await pool.fetchrow(
"SELECT status, retry_count, worker_id FROM async_operations WHERE operation_id = $1",
op_id,
)
assert row["status"] == "pending"
assert row["retry_count"] == 1
assert row["worker_id"] is None # Worker ID cleared for retry
@pytest.mark.asyncio
async def test_execute_task_fails_after_max_retries(self, pool, clean_operations):
"""Test that task is marked failed after exceeding max retries."""
from hindsight_api.worker import WorkerPoller
# Create a task that has already used all retries
bank_id = f"test-worker-{uuid.uuid4().hex[:8]}"
op_id = uuid.uuid4()
payload = json.dumps({"type": "test_task", "operation_id": str(op_id), "bank_id": bank_id})
await pool.execute(
"""
INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload, worker_id, retry_count)
VALUES ($1, $2, 'test', 'processing', $3::jsonb, 'test-worker-1', 3)
""",
op_id,
bank_id,
payload,
)
async def failing_executor(task_dict):
raise ValueError("Simulated failure")
poller = WorkerPoller(
pool=pool,
worker_id="test-worker-1",
executor=failing_executor,
max_retries=3,
)
# Execute (should fail permanently)
task_dict = json.loads(payload)
await poller.execute_task(str(op_id), task_dict)
# Verify task is marked as failed
row = await pool.fetchrow(
"SELECT status, error_message FROM async_operations WHERE operation_id = $1",
op_id,
)
assert row["status"] == "failed"
assert "Max retries" in row["error_message"]
class TestConcurrentWorkers:
"""Tests for concurrent worker task claiming (FOR UPDATE SKIP LOCKED)."""
@pytest.mark.asyncio
async def test_concurrent_workers_claim_different_tasks(self, pool, clean_operations):
"""Test that multiple workers claim different tasks (no duplicates)."""
from hindsight_api.worker import WorkerPoller
# Create 10 pending tasks
bank_id = f"test-worker-{uuid.uuid4().hex[:8]}"
task_ids = []
for i in range(10):
op_id = uuid.uuid4()
task_ids.append(op_id)
payload = json.dumps({"type": "test_task", "index": i, "bank_id": bank_id, "operation_id": str(op_id)})
await pool.execute(
"""
INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload)
VALUES ($1, $2, 'test', 'pending', $3::jsonb)
""",
op_id,
bank_id,
payload,
)
# Create 3 workers that will claim tasks concurrently
workers_claimed: dict[str, list[str]] = {"worker-1": [], "worker-2": [], "worker-3": []}
async def claim_for_worker(worker_id: str):
poller = WorkerPoller(
pool=pool,
worker_id=worker_id,
executor=lambda x: None,
batch_size=5, # Each worker tries to claim 5
)
claimed = await poller.claim_batch()
workers_claimed[worker_id] = [op_id for op_id, _ in claimed]
# Run all workers concurrently
await asyncio.gather(
claim_for_worker("worker-1"),
claim_for_worker("worker-2"),
claim_for_worker("worker-3"),
)
# Verify no duplicates - each task claimed by exactly one worker
all_claimed = workers_claimed["worker-1"] + workers_claimed["worker-2"] + workers_claimed["worker-3"]
assert len(all_claimed) == len(set(all_claimed)), "Duplicate task claimed by multiple workers!"
# Verify total claimed equals available tasks (10)
assert len(all_claimed) == 10, f"Expected 10 tasks claimed, got {len(all_claimed)}"
# Verify each task is assigned to exactly one worker in DB
rows = await pool.fetch(
"SELECT operation_id, worker_id FROM async_operations WHERE bank_id = $1",
bank_id,
)
worker_assignments = {str(row["operation_id"]): row["worker_id"] for row in rows}
# With FOR UPDATE SKIP LOCKED, it's a race condition which workers get tasks.
# The important invariant is no duplicates and all tasks claimed, which we verified above.
# Just verify that at least 1 worker got tasks and all tasks have a worker assigned.
assert len(set(worker_assignments.values())) >= 1, "At least one worker should have claimed tasks"
assert all(w is not None for w in worker_assignments.values()), "All tasks should have a worker assigned"
@pytest.mark.asyncio
async def test_workers_do_not_claim_already_processing_tasks(self, pool, clean_operations):
"""Test that workers skip tasks already being processed by another worker."""
from hindsight_api.worker import WorkerPoller
# Create tasks - some pending, some already processing
bank_id = f"test-worker-{uuid.uuid4().hex[:8]}"
# Create 3 pending tasks
for i in range(3):
op_id = uuid.uuid4()
payload = json.dumps({"type": "test_task", "index": i, "bank_id": bank_id})
await pool.execute(
"""
INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload)
VALUES ($1, $2, 'test', 'pending', $3::jsonb)
""",
op_id,
bank_id,
payload,
)
# Create 2 already-processing tasks owned by another worker
for i in range(2):
op_id = uuid.uuid4()
payload = json.dumps({"type": "test_task", "index": i + 10, "bank_id": bank_id})
await pool.execute(
"""
INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload, worker_id)
VALUES ($1, $2, 'test', 'processing', $3::jsonb, 'other-worker')
""",
op_id,
bank_id,
payload,
)
# New worker should only claim the 3 pending tasks
poller = WorkerPoller(
pool=pool,
worker_id="new-worker",
executor=lambda x: None,
batch_size=10,
)
claimed = await poller.claim_batch()
assert len(claimed) == 3, "Worker should only claim pending tasks"
# Verify other worker's tasks are still owned by them
row = await pool.fetchrow(
"SELECT COUNT(*) as count FROM async_operations WHERE bank_id = $1 AND worker_id = 'other-worker'",
bank_id,
)
assert row["count"] == 2
class TestWorkerDecommission:
"""Tests for worker decommissioning functionality."""
@pytest.mark.asyncio
async def test_decommission_releases_worker_tasks(self, pool, clean_operations):
"""Test that decommissioning a worker releases all its processing tasks."""
# Create tasks being processed by a worker
bank_id = f"test-worker-{uuid.uuid4().hex[:8]}"
worker_id = "worker-to-decommission"
for i in range(5):
op_id = uuid.uuid4()
payload = json.dumps({"type": "test_task", "index": i, "bank_id": bank_id})
await pool.execute(
"""
INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload, worker_id, claimed_at)
VALUES ($1, $2, 'test', 'processing', $3::jsonb, $4, now())
""",
op_id,
bank_id,
payload,
worker_id,
)
# Run decommission
result = await pool.fetch(
"""
UPDATE async_operations
SET status = 'pending', worker_id = NULL, claimed_at = NULL, updated_at = now()
WHERE worker_id = $1 AND status = 'processing'
RETURNING operation_id
""",
worker_id,
)
assert len(result) == 5
# Verify all tasks are back to pending
rows = await pool.fetch(
"SELECT status, worker_id, claimed_at FROM async_operations WHERE bank_id = $1",
bank_id,
)
for row in rows:
assert row["status"] == "pending"
assert row["worker_id"] is None
assert row["claimed_at"] is None
@pytest.mark.asyncio
async def test_decommission_does_not_affect_other_workers(self, pool, clean_operations):
"""Test that decommissioning one worker doesn't affect another worker's tasks."""
bank_id = f"test-worker-{uuid.uuid4().hex[:8]}"
# Create tasks for worker-1
for i in range(3):
op_id = uuid.uuid4()
payload = json.dumps({"type": "test_task", "index": i, "bank_id": bank_id})
await pool.execute(
"""
INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload, worker_id)
VALUES ($1, $2, 'test', 'processing', $3::jsonb, 'worker-1')
""",
op_id,
bank_id,
payload,
)
# Create tasks for worker-2
for i in range(3):
op_id = uuid.uuid4()
payload = json.dumps({"type": "test_task", "index": i + 10, "bank_id": bank_id})
await pool.execute(
"""
INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload, worker_id)
VALUES ($1, $2, 'test', 'processing', $3::jsonb, 'worker-2')
""",
op_id,
bank_id,
payload,
)
# Decommission worker-1 only
await pool.execute(
"""
UPDATE async_operations
SET status = 'pending', worker_id = NULL, claimed_at = NULL
WHERE worker_id = 'worker-1' AND status = 'processing'
"""
)
# Verify worker-1 tasks are released
worker1_rows = await pool.fetch(
"SELECT status, worker_id FROM async_operations WHERE bank_id = $1 AND worker_id IS NULL",
bank_id,
)
assert len(worker1_rows) == 3
# Verify worker-2 tasks are unaffected
worker2_rows = await pool.fetch(
"SELECT status, worker_id FROM async_operations WHERE bank_id = $1 AND worker_id = 'worker-2'",
bank_id,
)
assert len(worker2_rows) == 3
for row in worker2_rows:
assert row["status"] == "processing"
class TestSyncTaskBackend:
"""Tests for SyncTaskBackend (used in tests and embedded mode)."""
@pytest.mark.asyncio
async def test_sync_backend_executes_immediately(self):
"""Test that SyncTaskBackend executes tasks immediately."""
executed = []
async def mock_executor(task_dict):
executed.append(task_dict)
backend = SyncTaskBackend()
backend.set_executor(mock_executor)
await backend.initialize()
task_dict = {"type": "test", "data": "value"}
await backend.submit_task(task_dict)
assert len(executed) == 1
assert executed[0] == task_dict
@pytest.mark.asyncio
async def test_sync_backend_handles_errors(self):
"""Test that SyncTaskBackend handles executor errors gracefully."""
async def failing_executor(task_dict):
raise ValueError("Test error")
backend = SyncTaskBackend()
backend.set_executor(failing_executor)
await backend.initialize()
# Should not raise, error is logged
await backend.submit_task({"type": "test"})
@@ -127,6 +127,53 @@ Restore will **delete all existing data** in the target schema before importing
---
### decommission-worker
Release all tasks owned by a worker, resetting them from "processing" back to "pending" status so they can be picked up by other workers.
```bash
hindsight-admin decommission-worker WORKER_ID [OPTIONS]
```
**Arguments:**
| Argument | Description |
|----------|-------------|
| `WORKER_ID` | ID of the worker to decommission |
**Options:**
| Option | Description | Default |
|--------|-------------|---------|
| `--schema`, `-s` | Database schema | `public` |
**Examples:**
```bash
# Before scaling down - release tasks from workers being removed
hindsight-admin decommission-worker hindsight-worker-4
hindsight-admin decommission-worker hindsight-worker-3
# Release tasks from a crashed worker
hindsight-admin decommission-worker worker-2
# For a specific tenant schema
hindsight-admin decommission-worker worker-1 --schema tenant_acme
```
**When to Use:**
- **Scaling down**: Before removing worker replicas in Kubernetes
- **Graceful removal**: When taking a worker offline for maintenance
- **Crash recovery**: If a worker crashed while processing tasks
- **Stuck worker**: When a worker is unresponsive
:::tip Finding Worker IDs
Worker IDs default to the hostname. In Kubernetes StatefulSets, this is the pod name (e.g., `hindsight-worker-0`). You can also set a custom ID with `HINDSIGHT_API_WORKER_ID` or `--worker-id`.
:::
---
## Environment Variables
The admin CLI uses the same environment variables as the API service. The most important one is:
@@ -344,15 +344,18 @@ Configuration for the local MCP server (`hindsight-local-mcp` command).
export HINDSIGHT_API_MCP_INSTRUCTIONS="Also store every action you take, including tool calls and decisions made."
```
### Background Tasks
### Distributed Workers
Controls background task processing for async operations like opinion formation and entity observations.
Configuration for background task processing. By default, the API processes tasks internally. For high-throughput deployments, run dedicated workers. See [Services - Worker Service](./services#worker-service) for details.
| Variable | Description | Default |
|----------|-------------|---------|
| `HINDSIGHT_API_TASK_BACKEND` | Task backend implementation: `memory` (in-process queue) or `noop` (discard tasks, useful for tests) | `memory` |
| `HINDSIGHT_API_TASK_BACKEND_MEMORY_BATCH_SIZE` | Max tasks to process in one batch (memory backend only) | `10` |
| `HINDSIGHT_API_TASK_BACKEND_MEMORY_BATCH_INTERVAL` | Interval between batch processing in seconds (memory backend only) | `1.0` |
| `HINDSIGHT_API_WORKER_ENABLED` | Enable internal worker in API process | `true` |
| `HINDSIGHT_API_WORKER_ID` | Unique worker identifier | hostname |
| `HINDSIGHT_API_WORKER_POLL_INTERVAL_MS` | Database polling interval in milliseconds | `500` |
| `HINDSIGHT_API_WORKER_BATCH_SIZE` | Tasks to claim per poll cycle | `10` |
| `HINDSIGHT_API_WORKER_MAX_RETRIES` | Max retries before marking task failed | `3` |
| `HINDSIGHT_API_WORKER_HTTP_PORT` | HTTP port for worker metrics/health (worker CLI only) | `8889` |
### Performance Optimization
+13 -3
View File
@@ -36,8 +36,6 @@ See [Models](./models) for detailed comparison and configuration.
**Best for**: Quick start, development, small deployments
### Single Container (Quickest)
Run everything in one container with embedded PostgreSQL:
```bash
@@ -83,7 +81,19 @@ helm upgrade hindsight oci://ghcr.io/vectorize-io/charts/hindsight
- Kubernetes cluster (GKE, EKS, AKS, or self-hosted)
- Helm 3.8+
See the [Helm chart documentation](https://github.com/vectorize-io/hindsight/tree/main/helm) for advanced configuration.
### Distributed Workers
For high-throughput deployments, enable dedicated worker pods to scale task processing independently:
```bash
helm install hindsight oci://ghcr.io/vectorize-io/charts/hindsight \
--set worker.enabled=true \
--set worker.replicaCount=3
```
See [Services - Worker Service](./services#worker-service) for configuration details and architecture.
See the [Helm chart values.yaml](https://github.com/vectorize-io/hindsight/tree/main/helm/hindsight/values.yaml) for all chart options.
---
+38 -12
View File
@@ -1,6 +1,6 @@
# Services
Hindsight consists of two services that can run together or separately depending on your deployment needs.
Hindsight consists of three services that can run together or separately depending on your deployment needs.
## API Service
@@ -10,12 +10,48 @@ The core memory engine. Handles all memory operations:
- **Recall**: Semantic search across memories
- **Reflect**: Disposition-aware answer generation
```
```bash
hindsight-api # Default port: 8888
```
The API service is stateless and can be horizontally scaled behind a load balancer. All state is stored in PostgreSQL.
By default, the API also processes background tasks (opinion formation, entity observations) internally. For high-throughput deployments, you can disable this and run dedicated workers instead.
## Worker Service
Dedicated task processor for background operations. Uses the **same package and Docker image** as the API service, just with a different entry point.
```bash
hindsight-worker # Default metrics port: 8889
```
Workers use PostgreSQL as a task broker, polling for pending tasks. Multiple workers can run simultaneously without conflicts.
| Deployment | Internal Worker | Dedicated Workers |
|------------|-----------------|-------------------|
| **Development** | ✅ Simple, all-in-one | ❌ Overkill |
| **Small production** | ✅ Less infrastructure | ❌ Overkill |
| **High throughput** | ❌ API bottleneck | ✅ Scale independently |
| **Long-running tasks** | ❌ Blocks API resources | ✅ Isolated processing |
To use dedicated workers, disable the internal worker in the API and start worker processes:
```bash
# Disable internal worker in API
HINDSIGHT_API_WORKER_ENABLED=false hindsight-api
# Start dedicated workers (run multiple instances)
hindsight-worker --worker-id worker-1
hindsight-worker --worker-id worker-2
```
Each worker exposes `/health` and `/metrics` endpoints for monitoring.
Before scaling down or removing workers, release their tasks with `hindsight-admin decommission-worker <worker-id>`.
See [Configuration - Distributed Workers](./configuration#distributed-workers) for all worker settings and [Installation - Helm](./installation#distributed-workers) for Kubernetes deployment.
## Control Plane
Web UI for managing and exploring your memory banks:
@@ -28,13 +64,3 @@ Web UI for managing and exploring your memory banks:
The Control Plane connects to the API service and provides a visual interface for development and debugging.
For bare metal deployments, you can run the Control Plane standalone using npx. See [Installation - Bare Metal](./installation#control-plane) for details.
## Deployment Options
| Deployment | Services | Use Case |
|------------|----------|----------|
| **Docker (single container)** | Both bundled | Development, quick start |
| **Helm / Kubernetes** | Separate pods | Production, scaling |
| **Bare metal** | Run independently | Custom deployments |
In the Docker quickstart, both services run in a single container. For production Kubernetes deployments, they run as separate pods with independent scaling. For bare metal, you can run the API via pip and the Control Plane via npx.
+20
View File
@@ -0,0 +1,20 @@
#!/bin/bash
set -e
cd "$(dirname "$0")/../.."
ENV_FILE=".env"
if [ ! -f "$ENV_FILE" ]; then
echo "Error: Environment file $ENV_FILE not found at project root."
exit 1
fi
echo "Loading environment from $ENV_FILE"
echo ""
# Export all variables from env file
set -a
source "$ENV_FILE"
set +a
uv run hindsight-worker "$@"