labtasker-server 2.1.0__tar.gz → 2.2.0__tar.gz

{labtasker_server-2.1.0 → labtasker_server-2.2.0}/.gitignore

@@ -1,6 +1,7 @@
 .venv/
 .pytest_cache/
 .coverage
+coverage/
 coverage.xml
 htmlcov/
 .mypy_cache/
@@ -12,3 +13,6 @@ build/
 dist/
 site/
 .labtasker/
+.DS_Store
+.idea/
+tests/skill/runs/

{labtasker_server-2.1.0 → labtasker_server-2.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: labtasker-server
-Version: 2.1.0
+Version: 2.2.0
 Summary: SQLite/FastAPI server for parallel model inference and evaluation
 Project-URL: Homepage, https://github.com/luocfprime/labtasker
 Project-URL: Repository, https://github.com/luocfprime/labtasker.git
@@ -10,7 +10,7 @@ License-File: LICENSE
 Requires-Python: >=3.11
 Requires-Dist: alembic<2,>=1.14
 Requires-Dist: fastapi<1,>=0.115
-Requires-Dist: pydantic<3,>=2.10
+Requires-Dist: pydantic<3,>=2.13.5
 Requires-Dist: sqlalchemy<3,>=2.0
 Requires-Dist: typer<1,>=0.16
 Requires-Dist: uvicorn<1,>=0.34

{labtasker_server-2.1.0 → labtasker_server-2.2.0}/pyproject.toml

@@ -1,10 +1,10 @@
 [build-system]
-requires = ["hatchling>=1.27"]
+requires = ["hatchling>=1.32.0"]
 build-backend = "hatchling.build"
 
 [project]
 name = "labtasker-server"
-version = "2.1.0"
+version = "2.2.0"
 description = "SQLite/FastAPI server for parallel model inference and evaluation"
 requires-python = ">=3.11"
 license = "Apache-2.0"
@@ -13,7 +13,7 @@ authors = [{ name = "lcf", email = "luocfprime@gmail.com" }]
 dependencies = [
   "alembic>=1.14,<2",
   "fastapi>=0.115,<1",
-  "pydantic>=2.10,<3",
+  "pydantic>=2.13.5,<3",
   "sqlalchemy>=2.0,<3",
   "typer>=0.16,<1",
   "uvicorn>=0.34,<1",

{labtasker_server-2.1.0 → labtasker_server-2.2.0}/src/labtasker_server/__init__.py

@@ -1,3 +1,3 @@
 """Labtasker v2 HTTP server package."""
 
-__version__ = "2.1.0"
+__version__ = "2.2.0"

{labtasker_server-2.1.0 → labtasker_server-2.2.0}/src/labtasker_server/app.py

@@ -1,22 +1,30 @@
 from __future__ import annotations
 
 import asyncio
-import hmac
 import logging
 from collections.abc import AsyncIterator, Callable
 from contextlib import asynccontextmanager, suppress
 from typing import Annotated, Any
 
 from fastapi import Depends, FastAPI, Query, Request, Response
+from fastapi.exception_handlers import http_exception_handler
 from fastapi.exceptions import RequestValidationError
 from fastapi.responses import JSONResponse
+from fastapi.routing import APIRoute
 from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
 from sqlalchemy import text
+from starlette.exceptions import HTTPException
 
+from labtasker_server import __version__
 from labtasker_server.config import ServerSettings
 from labtasker_server.database import Database
 from labtasker_server.errors import DomainError
-from labtasker_server.middleware import RequestBodyLimitMiddleware
+from labtasker_server.grouping import request_error
+from labtasker_server.middleware import (
+    RequestBodyLimitMiddleware,
+    ServerVersionMiddleware,
+    is_authenticated,
+)
 from labtasker_server.schemas import (
     BulkUpdateRequest,
     BulkUpdateResult,
@@ -26,6 +34,7 @@ from labtasker_server.schemas import (
     CountResponse,
     ErrorEnvelope,
     FailRequest,
+    GroupCountPage,
     HealthyResponse,
     HeartbeatResponse,
     Queue,
@@ -37,10 +46,13 @@ from labtasker_server.schemas import (
     TaskStatus,
     TaskUpdate,
     UnhealthyResponse,
+    WorkerPage,
+    WorkerReport,
 )
 from labtasker_server.services.queues import QueueService
 from labtasker_server.services.tasks import TaskService, system_now_us
-from labtasker_server.validation import MAX_TASK_DATA_BYTES
+from labtasker_server.services.workers import WorkerService
+from labtasker_server.validation import MAX_JSON_DEPTH, MAX_TASK_DATA_BYTES
 
 EXPIRY_SCAN_INTERVAL_SECONDS = 60
 logger = logging.getLogger(__name__)
@@ -58,16 +70,18 @@ def create_app(
     database = Database(settings.database, ownership_fd=settings.database_fd)
     try:
         database.initialize()
+        queue_service = QueueService(database)
+        task_service = TaskService(database, now_us=now_us)
+        task_service.expire_leases()
+        worker_service = WorkerService(database, now_us=now_us)
+        worker_service.expire()
     except BaseException:
         database.dispose()
         raise
-    queue_service = QueueService(database)
-    task_service = TaskService(database, now_us=now_us)
-    task_service.expire_leases()
 
     @asynccontextmanager
     async def lifespan(_: FastAPI) -> AsyncIterator[None]:
-        scanner = asyncio.create_task(_expiry_scanner(task_service))
+        scanner = asyncio.create_task(_expiry_scanner(task_service, worker_service))
         try:
             yield
         finally:
@@ -78,9 +92,11 @@ def create_app(
 
     app = FastAPI(docs_url=None, redoc_url=None, lifespan=lifespan)
     app.add_middleware(RequestBodyLimitMiddleware, max_bytes=MAX_TASK_DATA_BYTES)
+    app.add_middleware(ServerVersionMiddleware, version=__version__, token=settings.token)
     app.state.database = database
     app.state.settings = settings
     app.state.task_service = task_service
+    app.state.worker_service = worker_service
 
     @app.exception_handler(DomainError)
     async def handle_domain_error(_: Request, exc: DomainError) -> JSONResponse:
@@ -134,15 +150,37 @@ def create_app(
             },
         )
 
+    @app.exception_handler(HTTPException)
+    async def handle_http_error(request: Request, exc: HTTPException) -> Response:
+        # FastAPI wraps decoder limits and invalid byte encodings in HTTP 400
+        # before Pydantic sees the body. Keep those in the validation contract.
+        if (
+            exc.status_code == 400
+            and exc.detail == "There was an error parsing the body"
+            and isinstance(exc.__cause__, (ValueError, RecursionError))
+        ):
+            if isinstance(exc.__cause__, RecursionError):
+                error = DomainError(
+                    422,
+                    "json_too_deep",
+                    "JSON value is too deeply nested.",
+                    {"max_depth": MAX_JSON_DEPTH},
+                )
+            else:
+                error = DomainError(
+                    422,
+                    "invalid_request",
+                    "Request validation failed.",
+                    {"errors": [{"location": ["body"], "message": "Malformed JSON body."}]},
+                )
+            return await handle_domain_error(request, error)
+        return await http_exception_handler(request, exc)
+
     def require_auth(
+        request: Request,
         credentials: Annotated[HTTPAuthorizationCredentials | None, Depends(BEARER)],
     ) -> None:
-        token = settings.token
-        if token is None:
-            return
-        if credentials is None or credentials.scheme.lower() != "bearer":
-            raise _unauthorized()
-        if not hmac.compare_digest(credentials.credentials.encode(), token.encode("ascii")):
+        if not is_authenticated(request.headers.get("authorization"), settings.token):
             raise _unauthorized()
 
     authenticated = [Depends(require_auth)]
@@ -217,6 +255,7 @@ def create_app(
         queue: str,
         status: TaskStatus | None = None,
         name: str | None = None,
+        name_fuzzy: str | None = None,
         filter_expression: Annotated[str | None, Query(alias="filter")] = None,
         order_by: TaskOrderField = "created_at",
         descending: bool = True,
@@ -227,6 +266,7 @@ def create_app(
             queue,
             status=status,
             name=name,
+            name_fuzzy=name_fuzzy,
             filter_expression=filter_expression,
             order_by=order_by,
             descending=descending,
@@ -236,24 +276,93 @@ def create_app(
 
     @app.get(
         "/api/v2/queues/{queue}/tasks/count",
-        response_model=CountResponse,
+        response_model=CountResponse | GroupCountPage,
         dependencies=authenticated,
         responses=API_ERROR_RESPONSES,
     )
     def count_tasks(
+        request: Request,
         queue: str,
         status: TaskStatus | None = None,
         name: str | None = None,
+        name_fuzzy: str | None = None,
         filter_expression: Annotated[str | None, Query(alias="filter")] = None,
-    ) -> CountResponse:
-        return CountResponse(
-            count=task_service.count_tasks(
-                queue,
-                status=status,
-                name=name,
-                filter_expression=filter_expression,
-            )
+        group_by: str | None = None,
+        limit: Annotated[int | None, Query(ge=1, le=1000)] = None,
+        cursor: str | None = None,
+    ) -> CountResponse | GroupCountPage:
+        _single_grouping(request)
+        result = task_service.count_tasks(
+            queue,
+            status=status,
+            name=name,
+            name_fuzzy=name_fuzzy,
+            filter_expression=filter_expression,
+            group_by=group_by,
+            limit=limit,
+            cursor=cursor,
+        )
+        return CountResponse(count=result) if isinstance(result, int) else result
+
+    @app.get(
+        "/api/v2/queues/{queue}/workers",
+        response_model=WorkerPage,
+        dependencies=authenticated,
+        responses=API_ERROR_RESPONSES,
+    )
+    def list_workers(
+        queue: str,
+        filter_expression: Annotated[str | None, Query(alias="filter")] = None,
+        limit: Annotated[int, Query(ge=1, le=1000)] = 100,
+        cursor: str | None = None,
+    ) -> WorkerPage:
+        return worker_service.list(
+            queue, filter_expression=filter_expression, limit=limit, cursor=cursor
+        )
+
+    @app.get(
+        "/api/v2/queues/{queue}/workers/count",
+        response_model=CountResponse | GroupCountPage,
+        dependencies=authenticated,
+        responses=API_ERROR_RESPONSES,
+    )
+    def count_workers(
+        request: Request,
+        queue: str,
+        filter_expression: Annotated[str | None, Query(alias="filter")] = None,
+        group_by: str | None = None,
+        limit: Annotated[int | None, Query(ge=1, le=1000)] = None,
+        cursor: str | None = None,
+    ) -> CountResponse | GroupCountPage:
+        _single_grouping(request)
+        result = worker_service.count(
+            queue,
+            filter_expression=filter_expression,
+            group_by=group_by,
+            limit=limit,
+            cursor=cursor,
         )
+        return CountResponse(count=result) if isinstance(result, int) else result
+
+    @app.put(
+        "/api/v2/queues/{queue}/workers/{id}",
+        status_code=204,
+        dependencies=authenticated,
+        responses=API_ERROR_RESPONSES,
+    )
+    def report_worker(queue: str, id: str, report: WorkerReport) -> Response:
+        worker_service.report(queue, id, report)
+        return Response(status_code=204)
+
+    @app.delete(
+        "/api/v2/queues/{queue}/workers/{id}",
+        status_code=204,
+        dependencies=authenticated,
+        responses=API_ERROR_RESPONSES,
+    )
+    def withdraw_worker(queue: str, id: str) -> Response:
+        worker_service.withdraw(queue, id)
+        return Response(status_code=204)
 
     @app.get(
         "/api/v2/queues/{queue}/tasks/{task_id}",
@@ -363,18 +472,42 @@ def create_app(
         task_service.delete(queue, task_id)
         return Response(status_code=204)
 
+    for route in app.routes:
+        if isinstance(route, APIRoute) and route.path.startswith("/api/"):
+            for status in {route.status_code or 200, *route.responses}:
+                if str(status) == "401":
+                    continue
+                route.responses.setdefault(status, {}).setdefault("headers", {})[
+                    "Labtasker-Server-Version"
+                ] = {
+                    "description": "Server package version; present only with valid credentials "
+                    "when authentication is enabled.",
+                    "schema": {"type": "string"},
+                }
     return app
 
 
-async def _expiry_scanner(task_service: TaskService) -> None:
+async def _expiry_scanner(task_service: TaskService, worker_service: WorkerService) -> None:
     while True:
         await asyncio.sleep(EXPIRY_SCAN_INTERVAL_SECONDS)
+        scan = asyncio.create_task(asyncio.to_thread(_expire_records, task_service, worker_service))
         try:
-            await asyncio.to_thread(task_service.expire_leases)
+            await asyncio.shield(scan)
+        except asyncio.CancelledError:
+            # Cancelling to_thread cannot stop its database command. Keep the
+            # ownership descriptor until that command has actually finished.
+            with suppress(Exception):
+                await scan
+            raise
         except Exception:
             logger.exception("Heartbeat expiry scan failed; it will retry in 60 seconds.")
 
 
+def _expire_records(task_service: TaskService, worker_service: WorkerService) -> None:
+    task_service.expire_leases()
+    worker_service.expire()
+
+
 def _unauthorized() -> DomainError:
     return DomainError(401, "unauthorized", "Authentication is required.", {})
 
@@ -393,6 +526,10 @@ def _validation_error(
 ) -> tuple[str, dict[str, object] | None]:
     for error in exc.errors():
         error_type = str(error.get("type", ""))
+        if error_type == "recursion_loop":
+            # A JSON body cannot contain reference cycles. This is Pydantic's
+            # own nesting limit, reached before the domain depth validator.
+            return "json_too_deep", {"max_depth": MAX_JSON_DEPTH}
         if error_type == "json_invalid":
             return "invalid_request", {
                 "errors": [{"location": ["body"], "message": "Malformed JSON body."}]
@@ -414,3 +551,8 @@ def _specific_validation_message(code: str) -> str:
     if code == "invalid_task_name":
         return "Task name is invalid."
     raise AssertionError(f"Unknown specific validation code: {code}")
+
+
+def _single_grouping(request: Request) -> None:
+    if len(request.query_params.getlist("group_by")) > 1:
+        request_error("group_by", "Specify group_by only once.")

{labtasker_server-2.1.0 → labtasker_server-2.2.0}/src/labtasker_server/cli.py

@@ -27,9 +27,7 @@ from labtasker_server.local import (
     metadata_matches_database,
     metadata_owner_is_verified,
     read_metadata,
-    remove_generation_artifacts,
-    remove_generation_socket,
-    remove_stale_artifacts,
+    remove_stopped_artifacts,
     require_local_capabilities,
     socket_health,
     startup_age,
@@ -209,7 +207,7 @@ def stop(
     if database_is_free(paths):
         if has_runtime_artifacts(paths):
             try:
-                remove_stale_artifacts(paths)
+                remove_stopped_artifacts(paths)
             except RuntimeError as error:
                 typer.echo(f"[labtasker-server] Local Server stop error: {error}", err=True)
                 raise typer.Exit(1) from error
@@ -231,7 +229,7 @@ def stop(
         os.kill(metadata.pid, signal.SIGTERM)
     typer.echo(f"[labtasker-server] stopping local daemon pid={metadata.pid}", err=True)
     if _wait_for_exit(paths.directory, timeout=30.0):
-        remove_generation_artifacts(paths, metadata.generation)
+        remove_stopped_artifacts(paths, generation=metadata.generation)
         typer.echo(f"[labtasker-server] stopped local daemon pid={metadata.pid}", err=True)
         return
     if not force:
@@ -262,7 +260,7 @@ def stop(
             err=True,
         )
         raise typer.Exit(1)
-    remove_generation_artifacts(paths, current.generation)
+    remove_stopped_artifacts(paths, generation=current.generation)
     typer.echo(f"[labtasker-server] stopped local daemon pid={current.pid}", err=True)
 
 
@@ -331,7 +329,7 @@ def daemon(
             listener.close()
         # Preserve the attempt metadata so an unexpected exit remains throttled.
         # An explicit successful stop removes the full generation itself.
-        remove_generation_socket(paths, generation)
+        remove_stopped_artifacts(paths, generation=generation, preserve_metadata=True)
 
 
 def _local_status(directory: Path) -> dict[str, object]:

{labtasker_server-2.1.0 → labtasker_server-2.2.0}/src/labtasker_server/database.py

@@ -5,19 +5,16 @@ from collections.abc import Iterator
 from contextlib import contextmanager
 from pathlib import Path
 
-try:
-    import fcntl
-except ImportError:  # pragma: no cover - explicit HTTP remains best effort off POSIX
-    fcntl = None  # type: ignore[assignment]
-
 from alembic import command
 from alembic.config import Config
-from sqlalchemy import Engine, create_engine, event, inspect, text
+from sqlalchemy import URL, Connection, Engine, create_engine, event, insert, inspect, text
 from sqlalchemy.exc import OperationalError
 from sqlalchemy.orm import Session, sessionmaker
 
 from labtasker_server.errors import DomainError
 from labtasker_server.models import QueueRow
+from labtasker_server.name_search import name_matches_fuzzy
+from labtasker_server.ownership import lock_database
 
 LOCAL_GITIGNORE = "*\n!.gitignore\n"
 
@@ -37,32 +34,43 @@ class Database:
         if labtasker_dir is not None:
             _ensure_local_gitignore(labtasker_dir)
         self._ownership_fd: int | None = _acquire_database_ownership(self.path, ownership_fd)
-        self.engine = _create_sqlite_engine(self.path)
-        self._session_factory = sessionmaker(self.engine, expire_on_commit=False)
+        try:
+            self.engine = _create_sqlite_engine(self.path)
+            self._session_factory = sessionmaker(self.engine, expire_on_commit=False)
+        except BaseException:
+            os.close(self._ownership_fd)
+            self._ownership_fd = None
+            raise
 
     def initialize(self) -> None:
-        existing_tables = set(inspect(self.engine).get_table_names())
-        is_fresh = not existing_tables
-        if existing_tables and "alembic_version" not in existing_tables:
-            raise RuntimeError("Database has tables but is not a recognized Labtasker v2 schema.")
-
         alembic_config = Config()
         alembic_config.set_main_option(
             "script_location",
-            str(Path(__file__).resolve().parent / "migrations"),
+            # Alembic's ConfigParser interpolates percent signs in option values.
+            str(Path(__file__).resolve().parent / "migrations").replace("%", "%%"),
         )
         with self.engine.begin() as connection:
+            # sqlite3's legacy mode does not begin a transaction for DDL.
+            # Commit schema changes and initial Queue together, including when
+            # startup fails or is interrupted between migration operations.
+            connection.exec_driver_sql("BEGIN IMMEDIATE")
+            existing_tables = set(inspect(connection).get_table_names())
+            if existing_tables and "alembic_version" not in existing_tables:
+                raise RuntimeError(
+                    "Database has tables but is not a recognized Labtasker v2 schema."
+                )
             alembic_config.attributes["connection"] = connection
             command.upgrade(alembic_config, "head")
-        _verify_sqlite_settings(self.engine)
-
-        if is_fresh:
-            with self.write_session() as session:
-                session.add(QueueRow(name="default"))
+            _verify_sqlite_settings(connection)
+            if not existing_tables:
+                connection.execute(insert(QueueRow).values(name="default"))
 
     @contextmanager
     def read_session(self) -> Iterator[Session]:
         with self._session_factory() as session:
+            # sqlite3's legacy transaction mode does not begin on SELECT. Keep
+            # all reads (including relationship loaders) in one SQLite snapshot.
+            session.execute(text("BEGIN"))
             yield session
 
     @contextmanager
@@ -87,23 +95,27 @@ class Database:
                 raise
 
     def dispose(self) -> None:
-        self.engine.dispose()
-        if self._ownership_fd is not None:
-            os.close(self._ownership_fd)
-            self._ownership_fd = None
+        try:
+            self.engine.dispose()
+        finally:
+            if self._ownership_fd is not None:
+                os.close(self._ownership_fd)
+                self._ownership_fd = None
 
 
 def _acquire_database_ownership(path: Path, inherited_fd: int | None) -> int:
     if inherited_fd is None:
         fd = os.open(path, os.O_RDWR | os.O_CREAT, 0o600)
-        if fcntl is not None:
-            try:
-                fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
-            except BlockingIOError as error:
-                os.close(fd)
-                raise DatabaseOwnershipError(
-                    f"Another Server process already owns database {path}."
-                ) from error
+        try:
+            lock_database(fd)
+        except BlockingIOError as error:
+            os.close(fd)
+            raise DatabaseOwnershipError(
+                f"Another Server process already owns database {path}."
+            ) from error
+        except BaseException:
+            os.close(fd)
+            raise
     else:
         fd = os.dup(inherited_fd)
         os.set_inheritable(fd, False)
@@ -127,12 +139,15 @@ def _acquire_database_ownership(path: Path, inherited_fd: int | None) -> int:
 
 def _create_sqlite_engine(path: Path) -> Engine:
     engine = create_engine(
-        f"sqlite+pysqlite:///{path}",
+        URL.create("sqlite+pysqlite", database=str(path)),
         connect_args={"check_same_thread": False, "timeout": 5.0},
     )
 
     @event.listens_for(engine, "connect")
     def configure_connection(dbapi_connection: object, _: object) -> None:
+        dbapi_connection.create_function(  # type: ignore[attr-defined]
+            "labtasker_name_fuzzy", 2, name_matches_fuzzy, deterministic=True
+        )
         cursor = dbapi_connection.cursor()  # type: ignore[attr-defined]
         try:
             cursor.execute("PRAGMA foreign_keys=ON")
@@ -158,14 +173,13 @@ def _is_sqlite_busy(error: OperationalError) -> bool:
     return code in {5, 6} or "database is locked" in str(error.orig).lower()
 
 
-def _verify_sqlite_settings(engine: Engine) -> None:
-    with engine.connect() as connection:
-        actual = {
-            "journal_mode": connection.scalar(text("PRAGMA journal_mode")),
-            "foreign_keys": connection.scalar(text("PRAGMA foreign_keys")),
-            "busy_timeout": connection.scalar(text("PRAGMA busy_timeout")),
-            "synchronous": connection.scalar(text("PRAGMA synchronous")),
-        }
+def _verify_sqlite_settings(connection: Connection) -> None:
+    actual = {
+        "journal_mode": connection.scalar(text("PRAGMA journal_mode")),
+        "foreign_keys": connection.scalar(text("PRAGMA foreign_keys")),
+        "busy_timeout": connection.scalar(text("PRAGMA busy_timeout")),
+        "synchronous": connection.scalar(text("PRAGMA synchronous")),
+    }
     expected = {
         "journal_mode": "wal",
         "foreign_keys": 1,