pyfilament 0.0.0__tar.gz

pyfilament-0.0.0/PKG-INFO

@@ -0,0 +1,26 @@
+Metadata-Version: 2.4
+Name: pyfilament
+Version: 0.0.0
+Summary: 
+Author: James Wu
+Author-email: james@centauri-ai.tech
+Requires-Python: >=3.14,<4.0
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.14
+Provides-Extra: dev
+Requires-Dist: aiosqlite (>=0.22.1,<0.23.0) ; extra == "dev"
+Requires-Dist: alembic (>=1.18.4,<2.0.0)
+Requires-Dist: anyio (>=4.13.0,<5.0.0)
+Requires-Dist: beartype (>=0.22.9,<0.23.0)
+Requires-Dist: dotenv (>=0.9.9,<0.10.0)
+Requires-Dist: fastapi (>=0.136.3,<0.137.0)
+Requires-Dist: inflection (>=0.5.1,<0.6.0)
+Requires-Dist: pandas (>=3.0.3,<4.0.0)
+Requires-Dist: polars (>=1.41.2,<2.0.0)
+Requires-Dist: pydantic (>=2.13.4,<3.0.0)
+Requires-Dist: redis (>=8.0.0,<9.0.0)
+Requires-Dist: sqlalchemy[asyncio] (>=2.0.50,<3.0.0)
+Requires-Dist: starlette (>=1.2.1,<2.0.0)
+Requires-Dist: strawberry-graphql (>=0.316.0,<0.317.0)
+Requires-Dist: uvicorn (>=0.49.0,<0.50.0)
+Requires-Dist: werkzeug (>=3.1.8,<4.0.0)

pyfilament-0.0.0/pyproject.toml

@@ -0,0 +1,64 @@
+[tool.ruff]
+line-length = 120
+target-version = 'py314'
+
+[tool.ruff.format]
+quote-style = 'single'
+
+[tool.mypy]
+plugins = ['pydantic.mypy']
+disable_error_code = ['import-untyped', 'empty-body', 'prop-decorator']
+
+[tool.poetry]
+name = "pyfilament"
+version = "0.0.0"
+description = ""
+authors = ["James Wu <james@centauri-ai.tech>"]
+packages = [{ include = "filament", from = "src" }]
+
+[tool.poetry.dependencies]
+python = "^3.14"
+sqlalchemy = {extras = ["asyncio"], version = "^2.0.50"}
+alembic = "^1.18.4"
+uvicorn = "^0.49.0"
+anyio = "^4.13.0"
+inflection = "^0.5.1"
+pydantic = "^2.13.4"
+beartype = "^0.22.9"
+pandas = "^3.0.3"
+polars = "^1.41.2"
+dotenv = "^0.9.9"
+fastapi = "^0.136.3"
+strawberry-graphql = "^0.316.0"
+starlette = "^1.2.1"
+werkzeug = "^3.1.8"
+redis = "^8.0.0"
+aiosqlite = { version = "^0.22.1", optional = true }
+
+[tool.poetry.extras]
+dev = ["aiosqlite"]
+
+[tool.poetry.group.dev.dependencies]
+pytest = "^9.0.3"
+hupper = "^1.12.1"
+watchdog = "^6.0.0"
+ruff = "^0.15.16"
+mypy = "^1.20.2"
+pytest-asyncio = "^1.4.0"
+pytest-cov = "^7.1.0"
+diff-cover = "^10.3.0"
+coverage = "^7.14.1"
+mock = "^5.2.0"
+
+[build-system]
+requires = ["poetry-core"]
+build-backend = "poetry.core.masonry.api"
+
+[tool.pytest.ini_options]
+asyncio_mode = "auto"
+asyncio_default_test_loop_scope = "session"
+asyncio_default_fixture_loop_scope = "session"
+log_cli = true
+log_cli_level = "INFO"
+log_cli_format = "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
+log_cli_date_format = "%Y-%m-%d %H:%M:%S"

pyfilament-0.0.0/src/filament/.gitignore

@@ -0,0 +1,2 @@
+__pycache__
+.venv/

pyfilament-0.0.0/src/filament/api/app.py

@@ -0,0 +1,3 @@
+from fastapi import FastAPI
+
+app = FastAPI()

pyfilament-0.0.0/src/filament/api/graphql.py

@@ -0,0 +1,63 @@
+import strawberry
+from fastapi import Depends, Request
+from strawberry.extensions import SchemaExtension
+from strawberry.fastapi import GraphQLRouter
+
+import filament.api.resolvers.task as task_resolver
+from filament.api.types.task import TaskRun, TaskType
+from filament.db.models import Base
+
+
+def get_session_from_request(request: Request):
+    return request.state.session
+
+
+async def get_context(
+    session=Depends(get_session_from_request),
+):
+    return {
+        'session': session,
+    }
+
+
+class SessionFlusher(SchemaExtension):
+    def resolve(self, _next, root, info, *args, **kwargs):
+        if (
+            info.path is not None
+            and info.path.key is not None
+            and (info.path.key == 'id' or info.path.key.endswith('_id'))
+        ):
+            if isinstance(root, Base) and getattr(root, info.path.key, None) is None:
+                if info.context is not None and 'session' in info.context:
+                    info.context['session'].flush()
+        return _next(root, info, *args, **kwargs)
+
+
+@strawberry.type
+class Query:
+    get_task_run: TaskRun = strawberry.field(resolver=task_resolver.get_task_run)
+    get_task_type: TaskType = strawberry.field(resolver=task_resolver.get_task_type)
+    get_task_types: list[TaskType] = strawberry.field(resolver=task_resolver.get_task_types)
+    get_task_runs: list[TaskRun] = strawberry.field(resolver=task_resolver.get_task_runs)
+    get_task_runs_by_ids: list[TaskRun] = strawberry.field(resolver=task_resolver.get_task_runs_by_ids)
+    get_task_types_by_ids: list[TaskType] = strawberry.field(resolver=task_resolver.get_task_types_by_ids)
+    get_task_type_stack_runs: list[TaskRun] = strawberry.field(resolver=task_resolver.get_task_type_stack_runs)
+
+
+@strawberry.type
+class Mutation:
+    cancel_task_run: TaskRun = strawberry.field(resolver=task_resolver.cancel_task_run)
+    run_task: TaskRun = strawberry.field(resolver=task_resolver.run_task)
+
+
+schema = strawberry.Schema(
+    query=Query,
+    mutation=Mutation,
+    extensions=[SessionFlusher],
+)
+
+
+graphql_app = GraphQLRouter(
+    schema,
+    context_getter=get_context,
+)

pyfilament-0.0.0/src/filament/api/logic/task_run_dict.py

@@ -0,0 +1,27 @@
+from beartype import beartype
+
+from filament.db.models import TaskRun as TaskRunModel
+from filament.logic.utils import avoid_nans, get_json_dict
+
+
+@beartype
+async def deep_get_task_run_dict(task_run: TaskRunModel, max_child_tasks: int = 100, child_depth: int = 0) -> dict:
+    task_run_dict = get_json_dict(task_run)
+    task_run_dict['task_type'] = get_json_dict(await task_run.awaitable_attrs.task_type)
+    if child_depth > 0:
+        sorted_child_tasks = sorted(await task_run.awaitable_attrs.child_tasks, key=lambda x: x.id)
+        task_run_dict['child_tasks'] = [
+            await deep_get_task_run_dict(child_task_run, max_child_tasks, child_depth - 1)
+            for child_task_run in sorted_child_tasks[:max_child_tasks]
+        ]
+    else:
+        task_run_dict['child_tasks'] = []
+    sorted_state_transitions = sorted(await task_run.awaitable_attrs.state_transitions, key=lambda x: x.id)
+    if task_run.parameters_json is not None:
+        task_run_dict['parameters_json'] = avoid_nans(task_run.parameters_json)
+    if task_run.result_json is not None:
+        task_run_dict['result_json'] = avoid_nans(task_run.result_json)
+    task_run_dict['state_transitions'] = [
+        get_json_dict(state_transition) for state_transition in sorted_state_transitions
+    ]
+    return task_run_dict

pyfilament-0.0.0/src/filament/api/main.py

@@ -0,0 +1,11 @@
+from filament.api.setup_logging import setup_logging
+
+setup_logging()
+
+import filament.api.routes  # noqa: F401
+from filament.api.app import app
+from filament.api.graphql import graphql_app
+from filament.api.middleware import SessionMiddleware
+
+app.add_middleware(SessionMiddleware)
+app.include_router(graphql_app, prefix='/graphql')

pyfilament-0.0.0/src/filament/api/middleware.py

@@ -0,0 +1,11 @@
+from fastapi import Request
+from starlette.middleware.base import BaseHTTPMiddleware
+
+from filament.db.session import async_session_scope
+
+
+class SessionMiddleware(BaseHTTPMiddleware):
+    async def dispatch(self, request: Request, call_next):
+        async with async_session_scope() as session:
+            request.state.session = session
+            return await call_next(request)

pyfilament-0.0.0/src/filament/api/resolvers/task.py

@@ -0,0 +1,168 @@
+import datetime
+import json
+
+from sqlalchemy import select
+from sqlalchemy.orm import aliased
+from sqlalchemy.sql import func
+from strawberry import ID, Info
+from werkzeug.exceptions import BadRequest, NotFound
+
+from filament.api.types.task import TaskRun, TaskType
+from filament.db.models import TaskRun as TaskRunModel
+from filament.db.models import TaskType as TaskTypeModel
+from filament.state.task_run_state import cancel_task_run as logic_cancel_task_run
+from filament.task.registry.task_type_registry import lookup
+
+DEFAULT_MAX_DAYS = 3
+
+
+async def get_task_run(self, info: Info, id: ID | None = None, task_uuid: str | None = None) -> TaskRun:
+    session = info.context['session']
+    if task_uuid is not None:
+        statement = select(TaskRunModel).where(TaskRunModel.task_uuid == task_uuid)
+    elif id is not None:
+        statement = select(TaskRunModel).where(TaskRunModel.id == int(id))
+    else:
+        raise BadRequest('Either id or task_uuid must be provided')
+    task_run = (await session.execute(statement)).scalars().one_or_none()
+    if not task_run:
+        raise NotFound(f'TaskRun with ID {id} or UUID {task_uuid} not found')
+    return task_run
+
+
+async def get_task_type(self, info: Info, id: ID | None = None, func_address: str | None = None) -> TaskType:
+    session = info.context['session']
+    if func_address is not None:
+        statement = select(TaskTypeModel).where(TaskTypeModel.func_address == func_address)
+    elif id is not None:
+        statement = select(TaskTypeModel).where(TaskTypeModel.id == int(id))
+    else:
+        raise BadRequest('Either id or func_address must be provided')
+    task_type = (await session.execute(statement)).scalars().one_or_none()
+    if not task_type:
+        raise NotFound(f'TaskType with id {id} or func_address {func_address} not found')
+    return task_type
+
+
+async def get_task_types_by_ids(
+    self, info: Info, ids: list[int] | None = None, uuids: list[str] | None = None
+) -> list[TaskType]:
+    session = info.context['session']
+    if ids:
+        statement = select(TaskTypeModel).where(TaskTypeModel.id.in_(ids))
+        task_types = (await session.execute(statement)).scalars().all()
+        ids_to_task_types = {task_type.id: task_type for task_type in task_types}
+        task_types = [ids_to_task_types[id] for id in ids if id in ids_to_task_types]
+    elif uuids:
+        statement = select(TaskTypeModel).where(TaskTypeModel.uuid.in_(uuids))
+        task_types = (await session.execute(statement)).scalars().all()
+        uuids_to_task_types = {task_type.uuid: task_type for task_type in task_types}
+        task_types = [uuids_to_task_types[uuid] for uuid in uuids if uuid in uuids_to_task_types]
+    else:
+        raise BadRequest('Either ids or uuids must be provided')
+    return task_types
+
+
+async def get_task_type_stack_runs(
+    self, info: Info, task_type_ids: list[int], states: list[str] | None = None
+) -> list[TaskRun]:
+    session = info.context['session']
+    if len(task_type_ids) == 0:
+        raise BadRequest('task_type_ids must be provided')
+    MAX_RESULTS = 100
+    final_task_type_id = task_type_ids[-1]
+    statement = select(TaskRunModel).filter(TaskRunModel.task_type_id == final_task_type_id)
+    if states:
+        statement = statement.where(TaskRunModel.state.in_(states))
+    last_model = TaskRunModel
+    for task_type_id in reversed(task_type_ids[:-1]):
+        current_model = aliased(TaskRunModel)
+        statement = statement.join(current_model, last_model.parent_task_uuid == current_model.task_uuid)
+        statement = statement.where(current_model.task_type_id == task_type_id)
+        last_model = current_model
+    statement = (
+        statement.where(last_model.parent_task_uuid.is_(None))
+        .order_by(TaskRunModel.created_at.desc())
+        .limit(MAX_RESULTS)
+    )
+    task_runs = (await session.execute(statement)).scalars().all()
+    return task_runs
+
+
+async def get_task_runs_by_ids(self, info: Info, ids: list[int]) -> list[TaskRun]:
+    session = info.context['session']
+    if len(ids) == 0:
+        raise BadRequest('ids must be provided')
+    statement = select(TaskRunModel).where(TaskRunModel.id.in_(ids))
+    task_runs = (await session.execute(statement)).scalars().all()
+    ids_to_task_runs = {task_run.id: task_run for task_run in task_runs}
+    for id in ids:
+        if id not in ids_to_task_runs:
+            raise NotFound(f'TaskRun with ID {id} not found')
+    task_runs = [ids_to_task_runs[id] for id in ids if id in ids_to_task_runs]
+    return task_runs
+
+
+async def get_task_types(self, info: Info, days: int = DEFAULT_MAX_DAYS):
+    session = info.context['session']
+    today = datetime.datetime.now()
+    before = today - datetime.timedelta(days=days)
+    subquery = (
+        select(TaskRunModel.task_type_id, func.max(TaskRunModel.id).label('task_run_id'))
+        .filter(TaskRunModel.created_at > before)
+        .group_by(TaskRunModel.task_type_id)
+        .subquery()
+    )
+    task_types_statement = select(TaskTypeModel).join(
+        subquery,
+        TaskTypeModel.id == subquery.c.task_type_id,
+    )
+    task_types = (await session.execute(task_types_statement)).scalars().all()
+    return task_types
+
+
+async def cancel_task_run(self, info: Info, id: ID | None = None, task_uuid: str | None = None) -> TaskRun:
+    session = info.context['session']
+    if task_uuid is not None:
+        statement = select(TaskRunModel).where(TaskRunModel.task_uuid == task_uuid)
+    elif id is not None:
+        statement = select(TaskRunModel).where(TaskRunModel.id == int(id))
+    else:
+        raise BadRequest('Either id or task_uuid must be provided')
+    task_run = (await session.execute(statement)).scalars().one_or_none()
+    if not task_run:
+        raise NotFound(f'TaskRun with UUID {task_uuid} not found')
+    await logic_cancel_task_run(session, task_run)
+    return task_run
+
+
+async def get_task_runs(
+    self, info: Info, task_type_id: ID, states: list[str] | None = None, days: int = DEFAULT_MAX_DAYS
+):
+    session = info.context['session']
+    today = datetime.datetime.now()
+    before = today - datetime.timedelta(days=days)
+    statement = (
+        select(TaskRunModel)
+        .where(TaskRunModel.task_type_id == int(task_type_id))
+        .where(TaskRunModel.created_at > before)
+    )
+    if states:
+        statement = statement.where(TaskRunModel.state.in_(states))
+    statement = statement.order_by(TaskRunModel.created_at.desc()).limit(99)
+    task_runs = (await session.execute(statement)).scalars().all()
+    return task_runs
+
+
+async def run_task(self, info: Info, task_type_id: ID, parameters_json: str) -> TaskRun:
+    session = info.context['session']
+    statement = select(TaskTypeModel).where(TaskTypeModel.id == int(task_type_id))
+    task_type = (await session.execute(statement)).scalars().one()
+    func_address = task_type.func_address
+    filament_task_type = lookup(func_address)
+    parameters = json.loads(parameters_json)
+    parameters.update({'start_immediately': True})
+    filament_task_run = await filament_task_type._request(task_args=[], task_kwargs=parameters)
+    statement = select(TaskRunModel).where(TaskRunModel.task_uuid == filament_task_run.uuid)
+    task_run = (await session.execute(statement)).scalars().one()
+    return task_run

pyfilament-0.0.0/src/filament/api/routes.py

@@ -0,0 +1,56 @@
+import logging
+
+from fastapi import Request, Response
+from werkzeug.exceptions import NotFound
+
+from filament.api.app import app
+from filament.api.logic.task_run_dict import deep_get_task_run_dict
+from filament.db.models import TaskRun as TaskRunModel
+from filament.logic.utils import rename_keys_to_camel_case, safe_json_dumps
+
+logger = logging.getLogger(__name__)
+
+
+@app.get('/tasks')
+async def root():
+    return {'message': 'Hello World'}
+
+
+@app.get('/api/task-run/{task_run_id}')
+async def get_task_run(request: Request, task_run_id: int, max_child_tasks: int = 100, child_depth: int = 3):
+    session = request.state.session
+    task_run = await session.get(TaskRunModel, task_run_id)
+    if task_run is None:
+        raise NotFound(f'TaskRun with ID {task_run_id} not found')
+    task_run_dict = await deep_get_task_run_dict(task_run, max_child_tasks, child_depth)
+    return rename_keys_to_camel_case(task_run_dict)
+
+
+@app.get('/api/task-runs/{task_run_ids_str}')
+async def get_task_runs(request: Request, task_run_ids_str: str, max_child_tasks: int = 100, child_depth: int = 3):
+    task_run_ids = [int(id) for id in task_run_ids_str.split(',')]
+    session = request.state.session
+    task_runs = []
+    for task_run_id in task_run_ids:
+        task_run = await session.get(TaskRunModel, task_run_id)
+        if task_run is None:
+            raise NotFound(f'TaskRun with ID {task_run_id} not found')
+        task_runs.append(await deep_get_task_run_dict(task_run, max_child_tasks, child_depth))
+    return rename_keys_to_camel_case(task_runs)
+
+
+@app.get('/api/task-run/{task_run_id}/download')
+async def download_task_run(request: Request, task_run_id: int, max_child_tasks: int = 100, child_depth: int = 3):
+    session = request.state.session
+    task_run = await session.get(TaskRunModel, task_run_id)
+    if task_run is None:
+        raise NotFound(f'TaskRun with ID {task_run_id} not found')
+    task_run_dict = await deep_get_task_run_dict(task_run, max_child_tasks, child_depth)
+    file_content = safe_json_dumps(rename_keys_to_camel_case(task_run_dict), indent=2, sort_keys=True).encode('utf-8')
+    filename = f'task_run_{task_run_id}.json'
+    headers = {'Content-Disposition': f'attachment; filename="{filename}"'}
+    return Response(
+        content=file_content,
+        media_type='application/json',
+        headers=headers,
+    )

pyfilament-0.0.0/src/filament/api/setup_logging.py

@@ -0,0 +1,25 @@
+import logging
+
+_is_setup = False
+
+
+class UnixTimeFormatter(logging.Formatter):
+    def formatTime(self, record, datefmt=None):
+        t = round(record.created, 1)
+        return f'{t:.1f}'
+
+
+def setup_logging():
+    global _is_setup
+    if _is_setup:
+        return
+    formatter = UnixTimeFormatter(fmt='%(asctime)s [%(levelname)s] %(name)s: %(message)s')
+    handler = logging.StreamHandler()
+    handler.setFormatter(formatter)
+    rootLogger = logging.getLogger()
+    rootLogger.setLevel(logging.DEBUG)
+    rootLogger.handlers = [handler]  # Replace any existing handlers
+    _is_setup = True
+
+
+setup_logging()

pyfilament-0.0.0/src/filament/api/types/task.py

@@ -0,0 +1,121 @@
+import json
+from datetime import datetime
+
+import strawberry
+from sqlalchemy import select
+from strawberry import Info
+
+from filament.db.models import TaskRun as TaskRunModel
+from filament.redis.client import r
+
+
+async def get_logs(
+    task_run: TaskRunModel, with_children: bool = True, max_depth: int = 3, max_num_children: int = 100
+) -> list[dict]:
+    logs = []
+    redis_key = f'filament_log:{(await task_run.awaitable_attrs.task_type).func_address}:{task_run.task_uuid}'
+    range_results = await r.lrange(redis_key, 0, -1)
+    for range_result in range_results:
+        logs.append(json.loads(range_result))
+    if with_children and max_depth > 0:
+        for child_task in (await task_run.awaitable_attrs.child_tasks)[:max_num_children]:
+            logs.extend(await get_logs(child_task, with_children, max_depth - 1, max_num_children))
+    return sorted(logs, key=lambda x: x['timestamp'])
+
+
+@strawberry.type
+class TaskRun:
+    id: int
+    task_uuid: str
+    name: str | None
+    created_at: datetime
+    state: str
+    state_since: datetime
+    heartbeat: datetime
+    run_count: int
+    parent_task_uuid: str | None
+    parameters_json: str | None
+    result_json: str | None
+
+    @strawberry.field
+    async def task_type(self) -> 'TaskType':
+        return await self.awaitable_attrs.task_type
+
+    @strawberry.field
+    async def state_transitions(self) -> list['TaskRunStateTransition']:
+        return await self.awaitable_attrs.state_transitions
+
+    @strawberry.field
+    async def child_tasks(self) -> list['TaskRun']:
+        return await self.awaitable_attrs.child_tasks
+
+    @strawberry.field
+    async def logs(
+        self, with_children: bool = True, max_depth: int = 3, max_num_children: int = 100
+    ) -> list['TaskRunLog']:
+        logs = await get_logs(self, with_children, max_depth, max_num_children)
+        return [TaskRunLog(**log) for log in logs]
+
+    @strawberry.field
+    async def task_runs_stack(self) -> list['TaskRun']:
+        current = self
+        task_runs_stack = []
+        task_runs_stack.append(current)
+        while await current.awaitable_attrs.parent_task:
+            task_runs_stack.append(await current.awaitable_attrs.parent_task)
+            current = await current.awaitable_attrs.parent_task
+        return list(reversed(task_runs_stack))
+
+
+@strawberry.type
+class TaskRunStateTransition:
+    id: int
+    task_uuid: str
+    from_state: str
+    to_state: str
+    state_since: datetime
+
+    @strawberry.field
+    async def task_run(self) -> TaskRun:
+        return await self.awaitable_attrs.task_run
+
+
+@strawberry.type
+class TaskRunLog:
+    timestamp: float
+    level: str
+    name: str
+    message: str
+
+
+@strawberry.type
+class TaskType:
+    id: int
+    func_address: str
+    name: str | None
+    parameters_spec: str | None
+    result_spec: str | None
+
+    @strawberry.field
+    async def task_runs(self, info: Info) -> list[TaskRun]:
+        session = info.context['session']
+        statement = (
+            select(TaskRunModel)
+            .where(TaskRunModel.task_type_id == self.id)
+            .order_by(TaskRunModel.created_at.desc())
+            .limit(99)
+        )
+        task_runs = (await session.execute(statement)).scalars().all()
+        return task_runs
+
+    @strawberry.field
+    async def latest_task_run(self, info: Info) -> TaskRun | None:
+        session = info.context['session']
+        statement = (
+            select(TaskRunModel)
+            .where(TaskRunModel.task_type_id == self.id)
+            .order_by(TaskRunModel.state_since.desc())
+            .limit(1)
+        )
+        task_run = (await session.execute(statement)).scalars().one_or_none()
+        return task_run

pyfilament-0.0.0/src/filament/cli/cli_cleanup.py

@@ -0,0 +1,86 @@
+import logging
+import sys
+from datetime import datetime, timedelta, timezone
+
+import anyio
+from plasma import Plasma
+from sqlalchemy import select
+
+from filament.db.models import TaskRun, TaskState
+from filament.db.session import async_session_scope
+from filament.state.task_run_state import cancel_task_run, delete_task_run
+
+
+def setup_logging():
+    logger = logging.getLogger()
+    handler = logging.StreamHandler(sys.stdout)
+    formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+    handler.setFormatter(formatter)
+    logger.addHandler(handler)
+    logger.setLevel(logging.INFO)
+
+
+setup_logging()
+
+logger = logging.getLogger(__name__)
+
+
+async def main(
+    stonith_max_heartbeat_seconds: int = 60 * 60, delete_old_task_runs_days: int = 30, skip_root_tasks: bool = True
+):
+    await stonith(max_heartbeat_seconds=stonith_max_heartbeat_seconds)
+    await delete_old_task_runs(days=delete_old_task_runs_days, skip_root_tasks=skip_root_tasks)
+
+
+async def stonith(max_heartbeat_seconds: int, batch_size: int = 100):
+    while True:
+        async with async_session_scope() as session:
+            now = datetime.now(timezone.utc)
+            heartbeat_threshold = now - timedelta(seconds=max_heartbeat_seconds)
+            task_runs_statement = (
+                select(TaskRun)
+                .where(~TaskRun.state.in_(TaskState.TERMINAL))
+                .where(TaskRun.heartbeat < heartbeat_threshold)
+                .order_by(TaskRun.heartbeat.desc())
+                .limit(batch_size)
+            )
+            task_runs = (await session.execute(task_runs_statement)).scalars().all()
+            logger.info(f'Found {len(task_runs)} task runs to STONITH')
+            any_age = None
+            for task_run in task_runs:
+                heartbeat_age = now - task_run.heartbeat
+                logger.debug(f'Cancelling task run {task_run.id} heartbeat_age={heartbeat_age}')
+                if any_age is None:
+                    any_age = heartbeat_age
+                await cancel_task_run(session, task_run)
+                await session.flush()
+            await session.commit()
+            logger.info(f'STONITHed {len(task_runs)} task runs, any heartbeat_age={any_age}')
+            await anyio.sleep(1)
+            if len(task_runs) < batch_size:
+                break
+
+
+async def delete_old_task_runs(days: int = 30, batch_size: int = 100, skip_root_tasks: bool = True):
+    while True:
+        async with async_session_scope() as session:
+            statement = select(TaskRun).where(TaskRun.created_at < datetime.now(timezone.utc) - timedelta(days=days))
+            if skip_root_tasks:
+                statement = statement.where(TaskRun.parent_task_uuid.is_not(None))
+            statement = statement.limit(batch_size)
+            task_runs = (await session.execute(statement)).scalars().all()
+            logger.info(f'Found {len(task_runs)} task runs to delete')
+            any_age = None
+            for task_run in task_runs:
+                await delete_task_run(session, task_run)
+                if any_age is None:
+                    any_age = task_run.created_at
+            await session.commit()
+            logger.info(f'Deleted {len(task_runs)} task runs, any_age={any_age}')
+            await anyio.sleep(1)
+            if len(task_runs) < batch_size:
+                break
+
+
+if __name__ == '__main__':
+    Plasma(main)