Initial inference documentation pass (#3330)

In the recent meeting it was raised that the codebase is not documented
and hard to understand as a result. This adds some initial
documentation. For now it is quite basic but hopefully will allow a bit
more understanding of some parts of the code "at a glance", reducing the
friction somewhat.

Author: olliestanley

inference/safety/main.py

@@ -1,4 +1,8 @@
-# A FastAPI server to run the safety pipeline
+"""
+A simple FastAPI server which serves a `blade2blade2` safety model.
+
+See https://github.com/LAION-AI/blade2blade for context.
+"""
 
 import asyncio
 from concurrent.futures import ThreadPoolExecutor
@@ -40,6 +44,7 @@ async def load_pipeline():
 
 
 async def async_predict(pipeline: Blade2Blade, inputs: str):
+    """Run predictions in a separate thread for a small server parallelism benefit."""
     return await asyncio.get_event_loop().run_in_executor(executor, pipeline.predict, inputs)
 
 

inference/safety/settings.py

@@ -2,6 +2,7 @@
 
 
 class Settings(pydantic.BaseSettings):
+    # HuggingFace model ID for the model to load in blade2blade
     safety_model_name: str = "shahules786/blade2blade-t5-base"
 
 

inference/server/export.py

@@ -1,3 +1,5 @@
+"""Script to facilitate exporting chat data from the server database."""
+
 import argparse
 import asyncio
 import contextlib

inference/server/main.py

@@ -56,6 +56,7 @@ def terminate_server(signum, frame):
 
 @app.on_event("startup")
 async def alembic_upgrade():
+    """Upgrades database schema based on Alembic migration scripts."""
     signal.signal(signal.SIGINT, terminate_server)
     if not settings.update_alembic:
         logger.warning("Skipping alembic upgrade on startup (update_alembic is False)")
@@ -113,7 +114,7 @@ async def maybe_add_debug_api_keys():
 app.include_router(workers.router)
 app.include_router(configs.router)
 
-# mount plugins
+# mount builtin plugins to be hosted on this server
 for app_prefix, sub_app in plugins.plugin_apps.items():
     app.mount(path=settings.plugins_path_prefix + app_prefix, app=sub_app)
 

inference/server/oasst_inference_server/admin.py

@@ -1,3 +1,5 @@
+"""Logic related to admin actions."""
+
 import fastapi
 from loguru import logger
 from oasst_inference_server import database, models

inference/server/oasst_inference_server/auth.py

@@ -1,3 +1,5 @@
+"""Logic related to authorization actions."""
+
 import hashlib
 import json
 from datetime import datetime, timedelta

inference/server/oasst_inference_server/chat_repository.py

@@ -12,6 +12,8 @@
 
 
 class ChatRepository(pydantic.BaseModel):
+    """Wrapper around a database session providing functionality relating to chats."""
+
     session: database.AsyncSession
 
     class Config:
@@ -38,6 +40,10 @@ async def get_prompter_message_by_id(self, message_id: str) -> models.DbMessage:
     async def start_work(
         self, *, message_id: str, worker_id: str, worker_config: inference.WorkerConfig
     ) -> models.DbMessage:
+        """
+        Update an assistant message in the database to be allocated to a specific worker.
+        The message must be in `pending` state. An exception is raised if the message has timed out or was cancelled.
+        """
         logger.debug(f"Starting work on message {message_id}")
         message = await self.get_assistant_message_by_id(message_id)
 
@@ -65,6 +71,10 @@ async def start_work(
         return message
 
     async def reset_work(self, message_id: str) -> models.DbMessage:
+        """
+        Update an assistant message in the database which has already been allocated to a worker to remove the
+        allocation and reset the message state to `pending`.
+        """
         logger.warning(f"Resetting work on message {message_id}")
         message = await self.get_assistant_message_by_id(message_id)
         message.state = inference.MessageState.pending
@@ -78,6 +88,7 @@ async def reset_work(self, message_id: str) -> models.DbMessage:
         return message
 
     async def abort_work(self, message_id: str, reason: str) -> models.DbMessage:
+        """Update an assistant message in the database to mark it as having been aborted by the allocated worker."""
         logger.warning(f"Aborting work on message {message_id}")
         message = await self.get_assistant_message_by_id(message_id)
         message.state = inference.MessageState.aborted_by_worker
@@ -88,7 +99,13 @@ async def abort_work(self, message_id: str, reason: str) -> models.DbMessage:
         await self.session.refresh(message)
         return message
 
-    async def complete_work(self, message_id: str, content: str, used_plugin: inference.PluginUsed) -> models.DbMessage:
+    async def complete_work(
+        self, message_id: str, content: str, used_plugin: inference.PluginUsed | None
+    ) -> models.DbMessage:
+        """
+        Update an assistant message in the database to mark it as having been completed with the given content, also
+        updating the used plugin if one is specified.
+        """
         logger.debug(f"Completing work on message {message_id}")
         message = await self.get_assistant_message_by_id(message_id)
         message.state = inference.MessageState.complete

inference/server/oasst_inference_server/chat_utils.py

@@ -3,6 +3,7 @@
 
 
 def get_model_config(model_config_name: str) -> model_configs.ModelConfig:
+    """Get a `ModelConfig` by its name. See `oasst_shared.model_configs`."""
     if settings.allowed_model_config_names != "*":
         if model_config_name not in settings.allowed_model_config_names_list:
             raise ValueError(f"Model {model_config_name} not in allowed models: {settings.allowed_model_config_names}")

inference/server/oasst_inference_server/compliance.py

@@ -1,3 +1,5 @@
+"""Logic related to worker compliance checks, which seek to ensure workers do not produce malicious responses."""
+
 import datetime
 from typing import cast
 
@@ -14,6 +16,10 @@
 async def find_compliance_work_request_message(
     session: database.AsyncSession, worker_config: inference.WorkerConfig, worker_id: str
 ) -> models.DbMessage | None:
+    """
+    Find a suitable assistant message to carry out a worker compliance check for the given worker. Such a message must
+    have been generated by a different worker, but one with the same compatibility hash as the given worker.
+    """
     compat_hash = worker_config.compat_hash
     query = (
         sqlmodel.select(models.DbMessage)
@@ -30,6 +36,10 @@ async def find_compliance_work_request_message(
 
 
 async def should_do_compliance_check(session: database.AsyncSession, worker_id: str) -> bool:
+    """
+    Check whether we should carry out a compliance check for the given worker, based on time since last check.
+    Trusted workers are excluded.
+    """
     worker = await worker_utils.get_worker(worker_id, session)
     if worker.trusted:
         return False
@@ -43,6 +53,13 @@ async def should_do_compliance_check(session: database.AsyncSession, worker_id:
 
 
 async def run_compliance_check(websocket: fastapi.WebSocket, worker_id: str, worker_config: inference.WorkerConfig):
+    """
+    Run a compliance check for the given worker:
+    - Find a suitable compliance check assistant message
+    - Task the worker with generating a response with the same context
+    - Compare the respons against the existing completed message
+    - Update the database with the outcome
+    """
     async with deps.manual_create_session() as session:
         try:
             worker = await worker_utils.get_worker(worker_id, session)

inference/server/oasst_inference_server/database.py

@@ -75,6 +75,7 @@ async def get_async_session(autoflush=True):
 
 
 def alembic_upgrade(connection):
+    """Upgrades database schema based on Alembic migration scripts."""
     alembic_ini_path = Path(__file__).parent.parent / "alembic.ini"
     alembic_cfg = alembic.config.Config(str(alembic_ini_path))
     alembic_cfg.set_main_option("sqlalchemy.url", settings.database_uri)