Migrate to SQLAlchemy and Alembic for database migrations. (#19)

* Migrate to SQLAlchemy and Alembic for database migrations.
* remove db.init() call from gradio app
* run migrations at app startup
* fix run migrations
* Convert SQLAlchemy rows to simple lists
* add missing config key

alembic.ini

@@ -0,0 +1,44 @@
+# A generic, single database configuration.
+
+[alembic]
+
+# database URL.  This is consumed by the user-maintained env.py script only.
+# other means of configuring database URLs may be customized within the env.py
+# file.
+sqlalchemy.url = driver://user:pass@localhost/dbname
+
+
+# Logging configuration
+[loggers]
+keys = root,sqlalchemy,alembic
+
+[handlers]
+keys = console
+
+[formatters]
+keys = generic
+
+[logger_root]
+level = WARNING
+handlers = console
+qualname =
+
+[logger_sqlalchemy]
+level = WARNING
+handlers =
+qualname = sqlalchemy.engine
+
+[logger_alembic]
+level = INFO
+handlers =
+qualname = alembic
+
+[handler_console]
+class = StreamHandler
+args = (sys.stderr,)
+level = NOTSET
+formatter = generic
+
+[formatter_generic]
+format = %(levelname)-5.5s [%(name)s] %(message)s
+datefmt = %H:%M:%S

alembic/README

@@ -0,0 +1 @@
+pyproject configuration, based on the generic configuration.

alembic/env.py

@@ -0,0 +1,78 @@
+from logging.config import fileConfig
+
+from sqlalchemy import engine_from_config
+from sqlalchemy import pool
+
+from alembic import context
+from blossomtune_gradio import config as cfg
+from blossomtune_gradio.database import Base
+
+
+# this is the Alembic Config object, which provides
+# access to the values within the .ini file in use.
+config = context.config
+config.set_main_option("sqlalchemy.url", cfg.SQLALCHEMY_URL)
+
+# Interpret the config file for Python logging.
+# This line sets up loggers basically.
+if config.config_file_name is not None:
+    fileConfig(config.config_file_name)
+
+# add your model's MetaData object here
+# for 'autogenerate' support
+target_metadata = Base.metadata
+
+# other values from the config, defined by the needs of env.py,
+# can be acquired:
+# my_important_option = config.get_main_option("my_important_option")
+# ... etc.
+
+
+def run_migrations_offline() -> None:
+    """Run migrations in 'offline' mode.
+
+    This configures the context with just a URL
+    and not an Engine, though an Engine is acceptable
+    here as well.  By skipping the Engine creation
+    we don't even need a DBAPI to be available.
+
+    Calls to context.execute() here emit the given string to the
+    script output.
+
+    """
+    url = config.get_main_option("sqlalchemy.url")
+    context.configure(
+        url=url,
+        target_metadata=target_metadata,
+        literal_binds=True,
+        dialect_opts={"paramstyle": "named"},
+    )
+
+    with context.begin_transaction():
+        context.run_migrations()
+
+
+def run_migrations_online() -> None:
+    """Run migrations in 'online' mode.
+
+    In this scenario we need to create an Engine
+    and associate a connection with the context.
+
+    """
+    connectable = engine_from_config(
+        config.get_section(config.config_ini_section, {}),
+        prefix="sqlalchemy.",
+        poolclass=pool.NullPool,
+    )
+
+    with connectable.connect() as connection:
+        context.configure(connection=connection, target_metadata=target_metadata)
+
+        with context.begin_transaction():
+            context.run_migrations()
+
+
+if context.is_offline_mode():
+    run_migrations_offline()
+else:
+    run_migrations_online()

alembic/script.py.mako

@@ -0,0 +1,28 @@
+"""${message}
+
+Revision ID: ${up_revision}
+Revises: ${down_revision | comma,n}
+Create Date: ${create_date}
+
+"""
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+${imports if imports else ""}
+
+# revision identifiers, used by Alembic.
+revision: str = ${repr(up_revision)}
+down_revision: Union[str, Sequence[str], None] = ${repr(down_revision)}
+branch_labels: Union[str, Sequence[str], None] = ${repr(branch_labels)}
+depends_on: Union[str, Sequence[str], None] = ${repr(depends_on)}
+
+
+def upgrade() -> None:
+    """Upgrade schema."""
+    ${upgrades if upgrades else "pass"}
+
+
+def downgrade() -> None:
+    """Downgrade schema."""
+    ${downgrades if downgrades else "pass"}

alembic/versions/0b01b44ab005_initial.py

@@ -0,0 +1,56 @@
+"""initial
+
+Revision ID: 0b01b44ab005
+Revises:
+Create Date: 2025-10-08 12:13:39.666969
+
+"""
+
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision: str = "0b01b44ab005"
+down_revision: Union[str, Sequence[str], None] = None
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    """Upgrade schema."""
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table(
+        "config",
+        sa.Column("key", sa.String(), nullable=False),
+        sa.Column("value", sa.String(), nullable=False),
+        sa.PrimaryKeyConstraint("key"),
+    )
+    op.create_table(
+        "requests",
+        sa.Column("participant_id", sa.String(), nullable=False),
+        sa.Column("status", sa.String(), nullable=False),
+        sa.Column(
+            "timestamp",
+            sa.DateTime(),
+            server_default=sa.text("(CURRENT_TIMESTAMP)"),
+            nullable=False,
+        ),
+        sa.Column("partition_id", sa.Integer(), nullable=True),
+        sa.Column("email", sa.String(), nullable=True),
+        sa.Column("hf_handle", sa.String(), nullable=True),
+        sa.Column("activation_code", sa.String(), nullable=True),
+        sa.Column("is_activated", sa.Integer(), nullable=False),
+        sa.PrimaryKeyConstraint("participant_id"),
+    )
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    """Downgrade schema."""
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_table("requests")
+    op.drop_table("config")
+    # ### end Alembic commands ###

app.py

@@ -1,5 +0,0 @@
-from blossomtune_gradio.gradio_app import demo
-
-
-if __name__ == "__main__":
-    demo.launch()

blossomtune_gradio/__main__.py

@@ -1,5 +1,9 @@
+from blossomtune_gradio import config as cfg
+from blossomtune_gradio import database as db
 from blossomtune_gradio.gradio_app import demo
 
 
 if __name__ == "__main__":
+    if cfg.RUN_MIGRATIONS_ON_STARTUP:
+        db.run_migrations()
     demo.launch()

blossomtune_gradio/config.py

@@ -13,6 +13,7 @@ SPACE_OWNER = os.getenv("SPACE_OWNER", SPACE_ID.split("/")[0] if SPACE_ID else N
 DB_PATH = (
     "/data/db/federation.db" if os.path.isdir("/data/db") else "./data/db/federation.db"
 )
+SQLALCHEMY_URL = f"sqlite:///{os.path.abspath(DB_PATH)}"
 MAX_NUM_NODES = int(os.getenv("MAX_NUM_NODES", "20"))
 SMTP_SENDER = os.getenv("SMTP_SENDER", "hello@ethicalabs.ai")
 SMTP_SERVER = os.getenv("SMTP_SERVER", "localhost")
@@ -25,6 +26,9 @@ SUPERLINK_HOST = os.getenv("SUPERLINK_HOST", "127.0.0.1:9092")
 SUPERLINK_PORT = int(os.getenv("SUPERLINK_PORT", 9092))
 SUPERLINK_CONTROL_API_PORT = int(os.getenv("SUPERLINK_CONTROL_API_PORT", 9093))
 SUPERLINK_MODE = os.getenv("SUPERLINK_MODE", "internal").lower()  # Or external
+RUN_MIGRATIONS_ON_STARTUP = util.strtobool(
+    os.getenv("RUN_MIGRATIONS_ON_STARTUP", "true")
+)  # Set to false in prod.
 
 # TLS root cert path. For production only.
 TLS_CERT_DIR = os.getenv("TLS_CERT_DIR", "./certs/")

blossomtune_gradio/database.py

@@ -1,34 +1,65 @@
-import sqlite3
+from alembic import config
+from sqlalchemy import create_engine, Column, String, Integer, DateTime, func
+from sqlalchemy.orm import sessionmaker, declarative_base
+
+
 from blossomtune_gradio import config as cfg
 
 
-def init():
-    """Initializes the DB, adds columns, and creates config table."""
-    with sqlite3.connect(cfg.DB_PATH) as conn:
-        cursor = conn.cursor()
-        cursor.execute("""
-            CREATE TABLE IF NOT EXISTS requests (
-                participant_id TEXT PRIMARY KEY,
-                status TEXT NOT NULL,
-                timestamp TEXT NOT NULL,
-                partition_id INTEGER,
-                email TEXT,
-                hf_handle TEXT,
-                activation_code TEXT,
-                is_activated INTEGER DEFAULT 0
-            )
-        """)
-
-        # Create a simple key-value config table
-        cursor.execute("""
-            CREATE TABLE IF NOT EXISTS config (
-                key TEXT PRIMARY KEY,
-                value TEXT NOT NULL
-            )
-        """)
-        # Initialize default number of partitions if not present
-        cursor.execute(
-            "INSERT OR IGNORE INTO config (key, value) VALUES (?, ?)",
-            ("num_partitions", "10"),
+Base = declarative_base()
+engine = create_engine(cfg.SQLALCHEMY_URL)
+
+# The sessionmaker factory generates new Session objects when called.
+SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
+
+
+class Request(Base):
+    """
+    SQLAlchemy model for the 'requests' table.
+    This table stores information about participants wanting to join the federation.
+    """
+
+    __tablename__ = "requests"
+
+    participant_id = Column(String, primary_key=True)
+    status = Column(String, nullable=False, default="pending")
+    timestamp = Column(DateTime, nullable=False, server_default=func.now())
+    partition_id = Column(Integer, nullable=True)
+    email = Column(String, nullable=True)
+    hf_handle = Column(String, nullable=True)
+    activation_code = Column(String, nullable=True)
+    is_activated = Column(Integer, nullable=False, default=0)
+
+    def __repr__(self):
+        return (
+            f"<Request(participant_id='{self.participant_id}', status='{self.status}')>"
         )
-        conn.commit()
+
+
+class Config(Base):
+    """
+    SQLAlchemy model for the 'config' table.
+    A simple key-value store for application settings.
+    """
+
+    __tablename__ = "config"
+
+    key = Column(String, primary_key=True)
+    value = Column(String, nullable=False)
+
+    def __repr__(self):
+        return f"<Config(key='{self.key}', value='{self.value}')>"
+
+
+def run_migrations():
+    """
+    Applies any pending Alembic migrations to the database.
+    This should be called on application startup.
+    """
+    print("Running database migrations...")
+    alembicArgs = [
+        "--raiseerr",
+        "upgrade",
+        "head",
+    ]
+    config.main(argv=alembicArgs)

blossomtune_gradio/federation.py

@@ -1,13 +1,11 @@
 import string
 import secrets
-import sqlite3
-
-from datetime import datetime
 
 from blossomtune_gradio import config as cfg
 from blossomtune_gradio import mail
 from blossomtune_gradio import util
 from blossomtune_gradio.settings import settings
+from blossomtune_gradio.database import SessionLocal, Request, Config
 
 
 def generate_participant_id(length=6):
@@ -23,105 +21,99 @@ def generate_activation_code(length=8):
 
 
 def check_participant_status(pid_to_check: str, email: str, activation_code: str):
-    """Handles a participant's request to join, activate, or check status."""
-    with sqlite3.connect(cfg.DB_PATH) as conn:
-        cursor = conn.cursor()
+    """
+    Handles a participant's request to join, activate, or check status using SQLAlchemy.
+    Returns a tuple: (is_approved: bool, message: str, data: any | None)
+    The 'is_approved' boolean is True ONLY when the participant's final status is 'approved'.
+    """
+    with SessionLocal() as db:
+        query = db.query(Request).filter(
+            Request.hf_handle == pid_to_check, Request.email == email
+        )
         if activation_code:
-            cursor.execute(
-                "SELECT participant_id, status, partition_id, is_activated, activation_code FROM requests WHERE hf_handle = ? AND email = ? AND activation_code = ?",
-                (pid_to_check, email, activation_code),
-            )
-        else:
-            cursor.execute(
-                "SELECT participant_id, status, partition_id, is_activated, activation_code FROM requests WHERE hf_handle = ? AND email = ?",
-                (pid_to_check, email),
-            )
-        result = cursor.fetchone()
+            query = query.filter(Request.activation_code == activation_code)
 
-        cursor.execute("SELECT value FROM config WHERE key = 'num_partitions'")
-        num_partitions_res = cursor.fetchone()
-        num_partitions = num_partitions_res[0] if num_partitions_res else "10"
+        request = query.first()
 
-    # Case 1: New user registration
-    if result is None:
-        if activation_code:
-            return (False, settings.get_text("activation_invalid_md"), None)
-        if not util.validate_email(email):
-            return (False, settings.get_text("invalid_email_md"), None)
-
-        with sqlite3.connect(cfg.DB_PATH) as conn:
-            cursor = conn.cursor()
-            cursor.execute("SELECT COUNT(*) FROM requests WHERE status = 'approved'")
-            approved_count = cursor.fetchone()[0]
+        num_partitions_config = (
+            db.query(Config).filter(Config.key == "num_partitions").first()
+        )
+        num_partitions = num_partitions_config.value if num_partitions_config else "10"
+
+        # Case 1: New user registration
+        if request is None:
+            if activation_code:
+                return (False, settings.get_text("activation_invalid_md"), None)
+            if not util.validate_email(email):
+                return (False, settings.get_text("invalid_email_md"), None)
+
+            approved_count = (
+                db.query(Request).filter(Request.status == "approved").count()
+            )
             if approved_count >= cfg.MAX_NUM_NODES:
                 return (False, settings.get_text("federation_full_md"), None)
 
-        participant_id = generate_participant_id()
-        new_activation_code = generate_activation_code()
-        mail_sent, message = mail.send_activation_email(email, new_activation_code)
-        if mail_sent:
-            with sqlite3.connect(cfg.DB_PATH) as conn:
-                conn.execute(
-                    "INSERT INTO requests (participant_id, status, timestamp, hf_handle, email, activation_code, is_activated) VALUES (?, ?, ?, ?, ?, ?, ?)",
-                    (
-                        participant_id,
-                        "pending",
-                        datetime.utcnow().isoformat(),
-                        pid_to_check,
-                        email,
-                        new_activation_code,
-                        0,
-                    ),
-                )
-            return (True, settings.get_text("registration_submitted_md"), None)
-        else:
-            return (False, message, None)
-
-    # Existing user
-    participant_id, status, partition_id, is_activated, stored_code = result
-
-    # Case 2: User is activating their account
-    if not is_activated:
-        if activation_code == stored_code:
-            with sqlite3.connect(cfg.DB_PATH) as conn:
-                conn.execute(
-                    "UPDATE requests SET is_activated = 1 WHERE hf_handle = ?",
-                    (pid_to_check,),
+            participant_id = generate_participant_id()
+            new_activation_code = generate_activation_code()
+            mail_sent, message = mail.send_activation_email(email, new_activation_code)
+
+            if mail_sent:
+                new_request = Request(
+                    participant_id=participant_id,
+                    hf_handle=pid_to_check,
+                    email=email,
+                    activation_code=new_activation_code,
                 )
-            return (True, settings.get_text("activation_successful_md"), None)
-        else:
-            return (False, settings.get_text("activation_invalid_md"), None)
-    else:
+                db.add(new_request)
+                db.commit()
+                # A successful registration step, but not yet approved for federation.
+                return (False, settings.get_text("registration_submitted_md"), None)
+            else:
+                return (False, message, None)
+
+        # Case 2: User is activating their account
+        if not request.is_activated:
+            if activation_code == request.activation_code:
+                request.is_activated = 1
+                db.commit()
+                # A successful activation step, but not yet approved.
+                return (False, settings.get_text("activation_successful_md"), None)
+            else:
+                return (False, settings.get_text("activation_invalid_md"), None)
+
+        # At this point, user is activated.
+        # They must provide the activation code to check their final status.
         if not activation_code:
-            return (False, settings.get_text("missing_activation_code_md"))
-
-    # Case 3: Activated user is checking their status
-    if status == "approved":
-        hostname = (
-            "localhost"
-            if not cfg.SPACE_ID
-            else f"{cfg.SPACE_ID.split('/')[1]}-{cfg.SPACE_ID.split('/')[0]}.hf.space"
-        )
-        superlink_hostname = cfg.SUPERLINK_HOST or hostname
-
-        connection_string = settings.get_text(
-            "status_approved_md",
-            participant_id=participant_id,
-            partition_id=partition_id,
-            superlink_hostname=superlink_hostname,
-            superlink_port=cfg.SUPERLINK_PORT,
-            num_partitions=num_partitions,
-        )
-        # TODO: build and provide .blossomfile for download
-        return (True, connection_string, cfg.BLOSSOMTUNE_TLS_CERT_PATH)
-    elif status == "pending":
-        return (False, settings.get_text("status_pending_md"))
-    else:  # Denied
-        return (
-            False,
-            settings.get_text("status_denied_md", participant_id=participant_id),
-            None,
-        )
+            return (False, settings.get_text("missing_activation_code_md"), None)
+
+        # Case 3: Activated user is checking their final status
+        if request.status == "approved":
+            hostname = (
+                "localhost"
+                if not cfg.SPACE_ID
+                else f"{cfg.SPACE_ID.split('/')[1]}-{cfg.SPACE_ID.split('/')[0]}.hf.space"
+            )
+            superlink_hostname = cfg.SUPERLINK_HOST or hostname
+
+            connection_string = settings.get_text(
+                "status_approved_md",
+                participant_id=request.participant_id,
+                partition_id=request.partition_id,
+                superlink_hostname=superlink_hostname,
+                num_partitions=num_partitions,
+            )
+            # The user is fully approved. Return success and the cert path.
+            return (True, connection_string, cfg.BLOSSOMTUNE_TLS_CERT_PATH)
+        elif request.status == "pending":
+            return (False, settings.get_text("status_pending_md"), None)
+        else:  # Denied
+            return (
+                False,
+                settings.get_text(
+                    "status_denied_md", participant_id=request.participant_id
+                ),
+                None,
+            )
 
 
 def manage_request(participant_id: str, partition_id: str, action: str):
@@ -129,29 +121,31 @@ def manage_request(participant_id: str, partition_id: str, action: str):
     if not participant_id:
         return False, "Please select a participant from the pending requests table."
 
-    if action == "approve":
-        if not partition_id or not partition_id.isdigit():
-            return False, "Please provide a valid integer for the Partition ID."
+    with SessionLocal() as db:
+        request = (
+            db.query(Request).filter(Request.participant_id == participant_id).first()
+        )
+        if not request:
+            return False, "Participant not found."
 
-        p_id_int = int(partition_id)
-        with sqlite3.connect(cfg.DB_PATH) as conn:
-            cursor = conn.cursor()
-            cursor.execute(
-                "SELECT is_activated FROM requests WHERE participant_id = ?",
-                (participant_id,),
-            )
-            is_activated_res = cursor.fetchone()
-            if not is_activated_res or not is_activated_res[0]:
+        if action == "approve":
+            if not partition_id or not partition_id.isdigit():
+                return False, "Please provide a valid integer for the Partition ID."
+
+            p_id_int = int(partition_id)
+            if not request.is_activated:
                 return (
                     False,
                     settings.get_text("participant_not_activated_warning_md"),
                 )
 
-            cursor.execute(
-                "SELECT 1 FROM requests WHERE partition_id = ? AND status = 'approved'",
-                (p_id_int,),
+            existing_participant = (
+                db.query(Request)
+                .filter(Request.partition_id == p_id_int, Request.status == "approved")
+                .first()
             )
-            if cursor.fetchone():
+
+            if existing_participant:
                 return (
                     False,
                     settings.get_text(
@@ -159,20 +153,17 @@ def manage_request(participant_id: str, partition_id: str, action: str):
                     ),
                 )
 
-            conn.execute(
-                "UPDATE requests SET status = ?, partition_id = ? WHERE participant_id = ?",
-                ("approved", p_id_int, participant_id),
-            )
+            request.status = "approved"
+            request.partition_id = p_id_int
+            db.commit()
             return (
                 True,
                 f"Participant {participant_id} is allowed to join the federation.",
             )
-    else:  # Deny
-        with sqlite3.connect(cfg.DB_PATH) as conn:
-            conn.execute(
-                "UPDATE requests SET status = ?, partition_id = NULL WHERE participant_id = ?",
-                ("denied", participant_id),
-            )
+        else:  # Deny
+            request.status = "denied"
+            request.partition_id = None
+            db.commit()
             return (
                 True,
                 f"Participant {participant_id} is not allowed to join the federation.",
@@ -180,12 +171,14 @@ def manage_request(participant_id: str, partition_id: str, action: str):
 
 
 def get_next_partion_id() -> int:
-    with sqlite3.connect(cfg.DB_PATH) as conn:
-        cursor = conn.cursor()
-        cursor.execute(
-            "SELECT partition_id FROM requests WHERE status = 'approved' AND partition_id IS NOT NULL"
+    """Finds the lowest available partition ID."""
+    with SessionLocal() as db:
+        used_ids_query = (
+            db.query(Request.partition_id)
+            .filter(Request.status == "approved", Request.partition_id.isnot(None))
+            .all()
         )
-        used_ids = {row[0] for row in cursor.fetchall()}
+        used_ids = {row[0] for row in used_ids_query}
 
     next_id = 0
     while next_id in used_ids:

blossomtune_gradio/gradio_app.py

@@ -1,12 +1,10 @@
 import gradio as gr
 
-from blossomtune_gradio import database as db
 from blossomtune_gradio.ui import components
 from blossomtune_gradio.ui import callbacks
 
 
 with gr.Blocks(theme=gr.themes.Soft(), title="Flower Superlink & Runner") as demo:
-    db.init()
     gr.Markdown("BlossomTune 🌸 Flower Superlink & Runner")
     with gr.Row():
         login_button = gr.LoginButton()

blossomtune_gradio/processing.py

@@ -1,11 +1,12 @@
 import os
 import shutil
-import sqlite3
 import threading
 import subprocess
 
 from blossomtune_gradio.logs import log
 from blossomtune_gradio import config as cfg
+from blossomtune_gradio import util
+from blossomtune_gradio.database import SessionLocal, Config
 
 
 # In-memory store for background processes and logs
@@ -44,7 +45,9 @@ def start_superlink():
 
     if process_store["superlink"] and process_store["superlink"].poll() is None:
         return False, "Superlink process is already running."
-    command = [shutil.which("flower-superlink"), "--insecure"]
+
+    # The command needs to be adapted for TLS if it's not insecure
+    command = [shutil.which("flower-superlink"), "--insecure"]  # Placeholder
     threading.Thread(
         target=run_process, args=(command, "superlink"), daemon=True
     ).start()
@@ -58,31 +61,36 @@ def start_runner(
 ):
     if process_store["runner"] and process_store["runner"].poll() is None:
         return False, "A Runner process is already running."
-    if (
-        not (process_store["superlink"] and process_store["superlink"].poll() is None)
-        and not cfg.SUPERLINK_MODE == "external"
-    ):
+
+    # Check if the Superlink is running, respecting the configured mode
+    if cfg.SUPERLINK_MODE == "external":
+        if not util.is_port_open(cfg.SUPERLINK_HOST, cfg.SUPERLINK_PORT):
+            return False, "External Superlink is not running or unreachable."
+    elif not (process_store["superlink"] and process_store["superlink"].poll() is None):
         return (
             False,
-            "Superlink is not running. Please start it before starting the runner.",
+            "Internal Superlink is not running. Please start it before starting the runner.",
         )
-    # TODO: check if external superlink is running
+
     if not all([runner_app, run_id, num_partitions]):
         return False, "Please provide a Runner App, Run ID, and Total Partitions."
     if not num_partitions.isdigit() or int(num_partitions) <= 0:
         return False, "Total Partitions must be a positive integer."
 
-    with sqlite3.connect(cfg.DB_PATH) as conn:
-        conn.execute(
-            "UPDATE config SET value = ? WHERE key = 'num_partitions'",
-            (num_partitions,),
-        )
-        conn.commit()
+    # Update the number of partitions in the database using SQLAlchemy
+    with SessionLocal() as db:
+        config_entry = db.query(Config).filter(Config.key == "num_partitions").first()
+        if config_entry:
+            config_entry.value = num_partitions
+        else:
+            db.add(Config(key="num_partitions", value=num_partitions))
+        db.commit()
 
     runner_app_path = runner_app.replace(".", os.path.sep)
     if not os.path.exists(runner_app_path):
         return False, f"Unable to find app path '{runner_app_path}'."
 
+    # Construct the command for a TLS-enabled runner
     command = [
         shutil.which("flwr"),
         "run",

blossomtune_gradio/ui/callbacks.py

@@ -1,5 +1,4 @@
 import time
-import sqlite3
 import gradio as gr
 import pandas as pd
 
@@ -9,6 +8,7 @@ from blossomtune_gradio import federation as fed
 from blossomtune_gradio import processing
 from blossomtune_gradio.settings import settings
 from blossomtune_gradio import util
+from blossomtune_gradio.database import SessionLocal, Request
 
 from . import components
 from . import auth
@@ -46,16 +46,31 @@ def get_full_status_update(
     else:
         auth_status = settings.get_text("auth_status_local_mode_md")
 
-    with sqlite3.connect(cfg.DB_PATH) as conn:
-        pending_rows = conn.execute(
-            "SELECT participant_id, hf_handle, email FROM requests WHERE status = 'pending' AND is_activated = 1 ORDER BY timestamp ASC"
-        ).fetchall()
-        approved_rows = conn.execute(
-            "SELECT participant_id, hf_handle, email, partition_id FROM requests WHERE status = 'approved' ORDER BY timestamp DESC"
-        ).fetchall()
+    with SessionLocal() as db:
+        pending_results = (
+            db.query(Request.participant_id, Request.hf_handle, Request.email)
+            .filter(Request.status == "pending", Request.is_activated == 1)
+            .order_by(Request.timestamp.asc())
+            .all()
+        )
+        approved_results = (
+            db.query(
+                Request.participant_id,
+                Request.hf_handle,
+                Request.email,
+                Request.partition_id,
+            )
+            .filter(Request.status == "approved")
+            .order_by(Request.timestamp.desc())
+            .all()
+        )
+
+    # Convert SQLAlchemy rows to simple lists
+    pending_rows = [list(row) for row in pending_results]
+    approved_rows = [list(row) for row in approved_results]
 
     # Superlink Status Logic
-    superlink_btn_update = gr.update()  # Default empty update
+    superlink_btn_update = gr.update()
 
     if cfg.SUPERLINK_MODE == "internal":
         superlink_is_running = (
@@ -78,7 +93,6 @@ def get_full_status_update(
         else:
             is_open = util.is_port_open(cfg.SUPERLINK_HOST, cfg.SUPERLINK_PORT)
             superlink_status = "🟢 Running" if is_open else "🔴 Not Running"
-        # Disable the button in external mode
         superlink_btn_update = gr.update(value="Managed Externally", interactive=False)
     else:
         superlink_status = "⚠️ Invalid Mode"
@@ -126,7 +140,6 @@ def toggle_superlink(
 ):
     """Toggles the Superlink process on or off."""
     if not auth.is_space_owner(profile, oauth_token):
-        # Hardcode warning text as it's not in the schema
         gr.Warning("You are not authorized to perform this operation.")
         return
     if (
@@ -147,7 +160,6 @@ def toggle_runner(
 ):
     """Toggles the Runner process on or off."""
     if not auth.is_space_owner(profile, oauth_token):
-        # Hardcode warning text as it's not in the schema
         gr.Warning("You are not authorized to perform this operation.")
         return
     if (
@@ -204,8 +216,8 @@ def on_check_participant_status(
     pid_to_check = user_hf_handle.strip()
     email_to_add = email.strip()
     activation_code_to_check = activation_code.strip()
-    # The federation module is responsible for getting the correct text from settings
-    _, message, download = fed.check_participant_status(
+
+    approved, message, download = fed.check_participant_status(
         pid_to_check, email_to_add, activation_code_to_check
     )
     return {
@@ -217,7 +229,6 @@ def on_check_participant_status(
 
 
 def on_manage_fed_request(participant_id: str, partition_id: str, action: str):
-    # The federation module is responsible for getting the correct text from settings
     result, message = fed.manage_request(participant_id, partition_id, action)
     if result:
         gr.Info(message)

flower_apps/quickstart_huggingface/pyproject.toml

@@ -60,3 +60,4 @@ options.backend.client-resources.num-gpus = 0.0 # at most 4 ClientApp will run i
 [tool.flwr.federations.local-deployment]
 address = "0.0.0.0:9093"
 insecure = true
+root-certificate = ""

pyproject.toml

@@ -16,6 +16,8 @@ dependencies = [
     "markupsafe==2.1.3",
     "jinja2>=3.1.6",
     "mlx[cpu]>=0.29.2",
+    "alembic>=1.16.5",
+    "sqlalchemy>=2.0.43",
 ]
 
 [tool.uv.sources]
@@ -90,3 +92,87 @@ build-backend = "setuptools.build_meta"
 where = ["."]
 include = ["blossomtune_gradio", "flower_apps"]
 exclude = ["results", "data"]
+
+
+[tool.alembic]
+
+# path to migration scripts.
+# this is typically a path given in POSIX (e.g. forward slashes)
+# format, relative to the token %(here)s which refers to the location of this
+# ini file
+script_location = "%(here)s/alembic"
+
+# template used to generate migration file names; The default value is %%(rev)s_%%(slug)s
+# Uncomment the line below if you want the files to be prepended with date and time
+# see https://alembic.sqlalchemy.org/en/latest/tutorial.html#editing-the-ini-file
+# for all available tokens
+# file_template = "%%(year)d_%%(month).2d_%%(day).2d_%%(hour).2d%%(minute).2d-%%(rev)s_%%(slug)s"
+
+# additional paths to be prepended to sys.path. defaults to the current working directory.
+prepend_sys_path = [
+    "."
+]
+
+# timezone to use when rendering the date within the migration file
+# as well as the filename.
+# If specified, requires the python>=3.9 or backports.zoneinfo library and tzdata library.
+# Any required deps can installed by adding `alembic[tz]` to the pip requirements
+# string value is passed to ZoneInfo()
+# leave blank for localtime
+# timezone =
+
+# max length of characters to apply to the "slug" field
+# truncate_slug_length = 40
+
+# set to 'true' to run the environment during
+# the 'revision' command, regardless of autogenerate
+# revision_environment = false
+
+# set to 'true' to allow .pyc and .pyo files without
+# a source .py file to be detected as revisions in the
+# versions/ directory
+# sourceless = false
+
+# version location specification; This defaults
+# to <script_location>/versions.  When using multiple version
+# directories, initial revisions must be specified with --version-path.
+# version_locations = [
+#    "%(here)s/alembic/versions",
+#    "%(here)s/foo/bar"
+# ]
+
+
+# set to 'true' to search source files recursively
+# in each "version_locations" directory
+# new in Alembic version 1.10
+# recursive_version_locations = false
+
+# the output encoding used when revision files
+# are written from script.py.mako
+# output_encoding = "utf-8"
+
+# This section defines scripts or Python functions that are run
+# on newly generated revision scripts.  See the documentation for further
+# detail and examples
+# [[tool.alembic.post_write_hooks]]
+# format using "black" - use the console_scripts runner,
+# against the "black" entrypoint
+# name = "black"
+# type = "console_scripts"
+# entrypoint = "black"
+# options = "-l 79 REVISION_SCRIPT_FILENAME"
+#
+# [[tool.alembic.post_write_hooks]]
+# lint with attempts to fix using "ruff" - use the module runner, against the "ruff" module
+# name = "ruff"
+# type = "module"
+# module = "ruff"
+# options = "check --fix REVISION_SCRIPT_FILENAME"
+#
+# [[tool.alembic.post_write_hooks]]
+# Alternatively, use the exec runner to execute a binary found on your PATH
+# name = "ruff"
+# type = "exec"
+# executable = "ruff"
+# options = "check --fix REVISION_SCRIPT_FILENAME"
+

tests/conftest.py

@@ -1,18 +1,60 @@
 import pytest
-import sqlite3
+from unittest.mock import MagicMock
 
-from blossomtune_gradio import database
+from alembic.config import Config
+from alembic import command
+from sqlalchemy import create_engine
+from sqlalchemy.orm import sessionmaker
+
+from blossomtune_gradio import config
+
+
+@pytest.fixture(scope="session")
+def alembic_config():
+    """Fixture to create a valid Alembic Config object."""
+    return Config("alembic.ini")
 
 
 @pytest.fixture
-def in_memory_db(mocker):
+def db_session(mocker, tmp_path):
     """
-    Fixture to set up and tear down an in-memory SQLite database for tests.
-    It ensures that the same connection object is used for both schema
-    initialization and the test execution.
+    Fixture to set up a clean file-based SQLite database for each test function.
+    It creates the database schema using Alembic programmatically and ensures all
+    modules use the same test database session.
     """
-    con = sqlite3.connect(":memory:")
-    mocker.patch("sqlite3.connect", return_value=con)
-    database.init()
-    yield con
-    con.close()
+    db_file = tmp_path / "test_federation.db"
+    db_url = f"sqlite:///{db_file}"
+    mocker.patch.object(config, "SQLALCHEMY_URL", db_url)
+
+    # Create an Alembic Config object and point it to the temp database.
+    alembic_cfg = Config()
+    alembic_cfg.set_main_option("script_location", "alembic")
+    alembic_cfg.set_main_option("sqlalchemy.url", db_url)
+
+    # Apply the migrations to create the schema in the temporary database.
+    command.upgrade(alembic_cfg, "head")
+
+    # Set up the SQLAlchemy engine and session factory for the tests to use.
+    engine = create_engine(db_url)
+    TestingSessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
+    session = TestingSessionLocal()
+
+    # Mock the SessionLocal factory in each module where it is imported and used.
+    mocker.patch("blossomtune_gradio.federation.SessionLocal", return_value=session)
+    mocker.patch("blossomtune_gradio.processing.SessionLocal", return_value=session)
+    mocker.patch("blossomtune_gradio.ui.callbacks.SessionLocal", return_value=session)
+
+    yield session
+
+    session.close()
+
+
+@pytest.fixture
+def mock_settings(mocker):
+    """Fixture to mock the settings module, available to all tests."""
+    mock_get = MagicMock(
+        side_effect=lambda key, **kwargs: f"mock_{key}".format(**kwargs)
+    )
+    mocker.patch("blossomtune_gradio.federation.settings.get_text", mock_get)
+    mocker.patch("blossomtune_gradio.ui.callbacks.settings.get_text", mock_get)
+    return mock_get

tests/test_federation.py

@@ -1,18 +1,8 @@
 import pytest
-from unittest.mock import MagicMock
+from datetime import datetime
 
 from blossomtune_gradio import federation as fed
-
-
-@pytest.fixture
-def mock_settings(mocker):
-    """Fixture to mock the settings module."""
-    # The lambda returns a formatted string to simulate Jinja2's behavior
-    mock_get = MagicMock(
-        side_effect=lambda key, **kwargs: f"mock_{key}".format(**kwargs)
-    )
-    mocker.patch("blossomtune_gradio.federation.settings.get_text", mock_get)
-    return mock_get
+from blossomtune_gradio.database import Request
 
 
 @pytest.fixture
@@ -38,208 +28,202 @@ def test_generate_activation_code():
 
 
 class TestCheckParticipantStatus:
-    """Test suite for the check_participant_status function."""
+    """Test suite for the check_participant_status function using SQLAlchemy."""
 
-    def test_new_user_registration_success(
-        self, in_memory_db, mock_settings, mock_mail
-    ):
+    def test_new_user_registration_success(self, db_session, mock_settings, mock_mail):
         """Verify successful registration for a new user."""
         mock_mail.return_value = (True, "")
-        success, message, download = fed.check_participant_status(
+        approved, message, download = fed.check_participant_status(
             "new_user", "new@example.com", ""
         )
-        assert success is True
+        assert approved is False
         assert download is None
         assert message == "mock_registration_submitted_md"
 
         # Verify the user was added to the database
-        cursor = in_memory_db.cursor()
-        cursor.execute("SELECT hf_handle FROM requests WHERE hf_handle = 'new_user'")
-        assert cursor.fetchone() is not None
+        request = db_session.query(Request).filter_by(hf_handle="new_user").first()
+        assert request is not None
+        assert request.email == "new@example.com"
 
-    def test_new_user_invalid_email(self, in_memory_db, mock_settings):
+    def test_new_user_invalid_email(self, db_session, mock_settings):
         """Verify registration fails with an invalid email."""
-        success, message, download = fed.check_participant_status(
+        approved, message, download = fed.check_participant_status(
             "user", "invalid-email", ""
         )
-        assert success is False
+        assert approved is False
         assert download is None
         assert message == "mock_invalid_email_md"
 
-    def test_new_user_federation_full(self, in_memory_db, mock_settings, mocker):
+    def test_new_user_federation_full(self, db_session, mock_settings, mocker):
         """Verify registration fails when the federation is full."""
-        mocker.patch("blossomtune_gradio.federation.cfg.MAX_NUM_NODES", 0)
-        success, message, download = fed.check_participant_status(
+        mocker.patch("blossomtune_gradio.config.MAX_NUM_NODES", 0)
+        approved, message, download = fed.check_participant_status(
             "another_user", "another@example.com", ""
         )
-        assert success is False
+        assert approved is False
         assert download is None
         assert message == "mock_federation_full_md"
 
-    def test_user_activation_success(self, in_memory_db, mock_settings):
+    def test_user_activation_success(self, db_session, mock_settings):
         """Verify a user can successfully activate their account."""
         # Setup: Add a pending, non-activated user
-        cursor = in_memory_db.cursor()
-        cursor.execute(
-            "INSERT INTO requests (participant_id, status, timestamp, hf_handle, email, activation_code, is_activated) VALUES (?, ?, ?, ?, ?, ?, ?)",
-            ("PID123", "pending", "now", "test_user", "test@example.com", "ABCDEF", 0),
+        pending_user = Request(
+            participant_id="PID123",
+            hf_handle="test_user",
+            email="test@example.com",
+            activation_code="ABCDEF",
+            is_activated=0,
         )
-        in_memory_db.commit()
+        db_session.add(pending_user)
+        db_session.commit()
 
-        success, message, download = fed.check_participant_status(
+        approved, message, download = fed.check_participant_status(
             "test_user", "test@example.com", "ABCDEF"
         )
-        assert success is True
+        assert approved is False
         assert download is None
         assert message == "mock_activation_successful_md"
         # Verify the user is now activated
-        cursor.execute(
-            "SELECT is_activated FROM requests WHERE hf_handle = 'test_user'"
+        activated_user = (
+            db_session.query(Request).filter_by(hf_handle="test_user").first()
         )
-        assert cursor.fetchone()[0] == 1
+        assert activated_user.is_activated == 1
 
-    def test_user_activation_invalid_code(self, in_memory_db, mock_settings):
+    def test_user_activation_invalid_code(self, db_session, mock_settings):
         """Verify activation fails with an invalid code."""
-        # Setup: Add a pending, non-activated user
-        cursor = in_memory_db.cursor()
-        cursor.execute(
-            "INSERT INTO requests (participant_id, status, timestamp, hf_handle, email, activation_code, is_activated) VALUES (?, ?, ?, ?, ?, ?, ?)",
-            ("PID123", "pending", "now", "test_user", "test@example.com", "ABCDEF", 0),
+        pending_user = Request(
+            participant_id="PID123",
+            hf_handle="test_user",
+            email="test@example.com",
+            activation_code="ABCDEF",
+            is_activated=0,
         )
-        in_memory_db.commit()
+        db_session.add(pending_user)
+        db_session.commit()
 
-        success, message, download = fed.check_participant_status(
+        approved, message, download = fed.check_participant_status(
             "test_user", "test@example.com", "WRONGCODE"
         )
-        assert success is False
+        assert approved is False
         assert download is None
         assert message == "mock_activation_invalid_md"
 
-    def test_status_check_approved(self, in_memory_db, mock_settings):
+    def test_status_check_approved(self, db_session, mock_settings):
         """Verify the status check for an approved user."""
-        cursor = in_memory_db.cursor()
-        cursor.execute(
-            "INSERT INTO requests (participant_id, status, timestamp, hf_handle, email, activation_code, is_activated, partition_id) VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
-            (
-                "PID456",
-                "approved",
-                "now",
-                "approved_user",
-                "approved@example.com",
-                "GHIJKL",
-                1,
-                5,
-            ),
-        )
-        in_memory_db.commit()
-        success, message, download = fed.check_participant_status(
+        approved_user = Request(
+            participant_id="PID456",
+            status="approved",
+            hf_handle="approved_user",
+            email="approved@example.com",
+            activation_code="GHIJKL",
+            is_activated=1,
+            partition_id=5,
+        )
+        db_session.add(approved_user)
+        db_session.commit()
+
+        approved, message, download = fed.check_participant_status(
             "approved_user", "approved@example.com", "GHIJKL"
         )
-        assert success is True
+        assert approved is True
         assert download is not None
         assert "mock_status_approved_md" in message
 
 
 class TestManageRequest:
-    """Test suite for the manage_request function."""
+    """Test suite for the manage_request function using SQLAlchemy."""
 
-    def test_approve_success(self, in_memory_db):
+    def test_approve_success(self, db_session):
         """Verify successful approval of a participant."""
-        # Setup: Add an activated, pending user
-        cursor = in_memory_db.cursor()
-        cursor.execute(
-            "INSERT INTO requests (participant_id, status, timestamp, hf_handle, email, activation_code, is_activated) VALUES (?, ?, ?, ?, ?, ?, ?)",
-            (
-                "PENDING1",
-                "pending",
-                "now",
-                "pending_user",
-                "pending@example.com",
-                "CODE",
-                1,
-            ),
-        )
-        in_memory_db.commit()
+        pending_user = Request(
+            participant_id="PENDING1",
+            status="pending",
+            hf_handle="pending_user",
+            email="pending@example.com",
+            is_activated=1,
+        )
+        db_session.add(pending_user)
+        db_session.commit()
 
         success, message = fed.manage_request("PENDING1", "10", "approve")
         assert success is True
         assert "is allowed to join" in message
 
         # Verify status in DB
-        cursor.execute(
-            "SELECT status, partition_id FROM requests WHERE participant_id = 'PENDING1'"
+        updated_user = (
+            db_session.query(Request).filter_by(participant_id="PENDING1").first()
         )
-        status, partition_id = cursor.fetchone()
-        assert status == "approved"
-        assert partition_id == 10
+        assert updated_user.status == "approved"
+        assert updated_user.partition_id == 10
 
-    def test_approve_not_activated(self, in_memory_db, mock_settings):
+    def test_approve_not_activated(self, db_session, mock_settings):
         """Verify approval fails if the user is not activated."""
-        cursor = in_memory_db.cursor()
-        cursor.execute(
-            "INSERT INTO requests (participant_id, status, timestamp, hf_handle, email, activation_code, is_activated) VALUES (?, ?, ?, ?, ?, ?, ?)",
-            (
-                "PENDING2",
-                "pending",
-                "now",
-                "pending_user2",
-                "pending2@example.com",
-                "CODE",
-                0,
-            ),
-        )
-        in_memory_db.commit()
+        pending_user = Request(
+            participant_id="PENDING2",
+            status="pending",
+            hf_handle="pending_user2",
+            is_activated=0,
+        )
+        db_session.add(pending_user)
+        db_session.commit()
         success, message = fed.manage_request("PENDING2", "11", "approve")
         assert success is False
         assert message == "mock_participant_not_activated_warning_md"
 
-    def test_deny_success(self, in_memory_db):
+    def test_deny_success(self, db_session):
         """Verify successful denial of a participant."""
-        cursor = in_memory_db.cursor()
-        cursor.execute(
-            "INSERT INTO requests (participant_id, status, timestamp, hf_handle, email, activation_code, is_activated) VALUES (?, ?, ?, ?, ?, ?, ?)",
-            (
-                "PENDING3",
-                "pending",
-                "now",
-                "pending_user3",
-                "pending3@example.com",
-                "CODE",
-                1,
-            ),
-        )
-        in_memory_db.commit()
+        pending_user = Request(
+            participant_id="PENDING3",
+            status="pending",
+            hf_handle="pending_user3",
+            is_activated=1,
+        )
+        db_session.add(pending_user)
+        db_session.commit()
         success, message = fed.manage_request("PENDING3", "", "deny")
         assert success is True
         assert "is not allowed to join" in message
 
         # Verify status in DB
-        cursor.execute("SELECT status FROM requests WHERE participant_id = 'PENDING3'")
-        assert cursor.fetchone()[0] == "denied"
+        updated_user = (
+            db_session.query(Request).filter_by(participant_id="PENDING3").first()
+        )
+        assert updated_user.status == "denied"
 
 
-def test_get_next_partition_id(in_memory_db):
+def test_get_next_partition_id(db_session):
     """Verify the logic for finding the next available partition ID."""
-    cursor = in_memory_db.cursor()
     # No approved users yet
     assert fed.get_next_partion_id() == 0
 
-    # Add some approved users with assigned partitions, including the required timestamp
-    cursor.execute(
-        "INSERT INTO requests (participant_id, status, timestamp, partition_id) VALUES (?, ?, ?, ?)",
-        ("P1", "approved", "now", 0),
+    # Add some approved users with assigned partitions
+    db_session.add(
+        Request(
+            participant_id="P1",
+            status="approved",
+            partition_id=0,
+            timestamp=datetime.utcnow(),
+        )
     )
-    cursor.execute(
-        "INSERT INTO requests (participant_id, status, timestamp, partition_id) VALUES (?, ?, ?, ?)",
-        ("P2", "approved", "now", 1),
+    db_session.add(
+        Request(
+            participant_id="P2",
+            status="approved",
+            partition_id=1,
+            timestamp=datetime.utcnow(),
+        )
     )
-    in_memory_db.commit()
+    db_session.commit()
     assert fed.get_next_partion_id() == 2
 
     # Add another user, skipping a partition ID
-    cursor.execute(
-        "INSERT INTO requests (participant_id, status, timestamp, partition_id) VALUES (?, ?, ?, ?)",
-        ("P3", "approved", "now", 3),
+    db_session.add(
+        Request(
+            participant_id="P3",
+            status="approved",
+            partition_id=3,
+            timestamp=datetime.utcnow(),
+        )
     )
-    in_memory_db.commit()
+    db_session.commit()
     assert fed.get_next_partion_id() == 2

tests/test_processing.py

@@ -2,6 +2,7 @@ import pytest
 from unittest.mock import MagicMock, patch
 
 from blossomtune_gradio import processing
+from blossomtune_gradio.database import Config
 
 
 @pytest.fixture(autouse=True)
@@ -12,8 +13,6 @@ def reset_process_store():
     """
     processing.process_store = {"superlink": None, "runner": None}
     yield
-    # Teardown is not strictly necessary as it's reset at the start,
-    # but it's good practice.
     processing.process_store = {"superlink": None, "runner": None}
 
 
@@ -29,18 +28,15 @@ def test_start_superlink_success(mock_which, mock_thread):
     assert success is True
     assert message == "Superlink process started."
     mock_thread.assert_called_once()
-    # Check that the thread is targeting the run_process function
     call_args = mock_thread.call_args
     assert call_args.kwargs["target"] == processing.run_process
-    # Check that the command is correct
     assert call_args.kwargs["args"][0] == ["/fake/path/flower-superlink", "--insecure"]
 
 
 def test_start_superlink_already_running(mocker):
     """Verify that start_superlink returns False if a process is already running."""
-    # Mock a running process
     mock_process = MagicMock()
-    mock_process.poll.return_value = None  # None means it's running
+    mock_process.poll.return_value = None
     processing.process_store["superlink"] = mock_process
 
     success, message = processing.start_superlink()
@@ -49,33 +45,60 @@ def test_start_superlink_already_running(mocker):
 
 
 @patch("blossomtune_gradio.processing.os.path.exists", return_value=True)
-@patch("blossomtune_gradio.processing.sqlite3.connect")
 @patch("blossomtune_gradio.processing.threading.Thread")
 @patch("blossomtune_gradio.processing.shutil.which", return_value="/fake/path/flwr")
-def test_start_runner_success(mock_which, mock_thread, mock_sqlite, mock_exists):
-    """Verify that start_runner successfully starts a thread when conditions are met."""
-    # Mock a running superlink process
+def test_start_runner_success_internal_superlink(
+    mock_which, mock_thread, mock_exists, db_session
+):
+    """Verify start_runner succeeds with an internal superlink and updates the DB."""
+    # Arrange: Mock a running internal superlink process
     mock_superlink = MagicMock()
     mock_superlink.poll.return_value = None
     processing.process_store["superlink"] = mock_superlink
 
+    # Act
+    success, message = processing.start_runner("app.main", "run1", "15")
+
+    # Assert
+    assert success is True
+    assert message == "Federation Run is starting...."
+    mock_thread.assert_called_once()
+
+    # Verify DB was updated using SQLAlchemy
+    config_entry = db_session.query(Config).filter_by(key="num_partitions").first()
+    assert config_entry is not None
+    assert config_entry.value == "15"
+
+
+@patch("blossomtune_gradio.processing.os.path.exists", return_value=True)
+@patch("blossomtune_gradio.processing.threading.Thread")
+@patch("blossomtune_gradio.processing.shutil.which", return_value="/fake/path/flwr")
+def test_start_runner_success_external_superlink(
+    mock_which, mock_thread, mock_exists, db_session, mocker
+):
+    """Verify start_runner succeeds with an external superlink."""
+    # Arrange
+    mocker.patch("blossomtune_gradio.config.SUPERLINK_MODE", "external")
+    mocker.patch("blossomtune_gradio.util.is_port_open", return_value=True)
+
+    # Act
     success, message = processing.start_runner("app.main", "run1", "10")
 
+    # Assert
     assert success is True
     assert message == "Federation Run is starting...."
     mock_thread.assert_called_once()
-    mock_sqlite.assert_called_once()  # Check if DB was updated
 
 
-def test_start_runner_superlink_not_running():
-    """Verify start_runner fails if superlink is not running."""
+def test_start_runner_internal_superlink_not_running(db_session):
+    """Verify start_runner fails if internal superlink is not running."""
     processing.process_store["superlink"] = None
     success, message = processing.start_runner("app.main", "run1", "10")
     assert success is False
-    assert "Superlink is not running" in message
+    assert "Internal Superlink is not running" in message
 
 
-def test_start_runner_missing_args():
+def test_start_runner_missing_args(db_session):
     """Verify start_runner fails if arguments are missing."""
     mock_superlink = MagicMock()
     mock_superlink.poll.return_value = None
@@ -87,7 +110,7 @@ def test_start_runner_missing_args():
 
 
 @patch("blossomtune_gradio.processing.os.path.exists", return_value=False)
-def test_start_runner_app_path_not_found(mock_exists, in_memory_db):
+def test_start_runner_app_path_not_found(mock_exists, db_session):
     """Verify start_runner fails if the app path doesn't exist."""
     mock_superlink = MagicMock()
     mock_superlink.poll.return_value = None
@@ -101,7 +124,7 @@ def test_start_runner_app_path_not_found(mock_exists, in_memory_db):
 def test_stop_process_running():
     """Verify stop_process terminates a running process."""
     mock_process = MagicMock()
-    mock_process.poll.return_value = None  # Running
+    mock_process.poll.return_value = None
     processing.process_store["superlink"] = mock_process
 
     processing.stop_process("superlink")
@@ -118,7 +141,6 @@ def test_stop_process_not_running(mocker):
 
     processing.stop_process("superlink")
 
-    # Check that the specific "no process was running" log was called
     log_mock.assert_any_call(
         "[Superlink] Stop command received, but no process was running."
     )