lablink_allocator_service.main

lablink_allocator_service.main

Attributes

AGENT_TOKEN = secrets.token_urlsafe(32) module-attribute

ENVIRONMENT = os.getenv('ENVIRONMENT', 'prod').strip().lower().replace(' ', '-') module-attribute

REGISTER_TOKEN = secrets.token_urlsafe(32) module-attribute

TOFU_DIR = (Path(__file__).parent / 'terraform').resolve() module-attribute

admin_session_expiry_service = None module-attribute

allocator_ip = os.getenv('ALLOCATOR_PUBLIC_IP') module-attribute

app = Flask(__name__) module-attribute

cfg = get_config() module-attribute

cloud_init_output_log_group = os.getenv('CLOUD_INIT_LOG_GROUP') module-attribute

database = None module-attribute

key_name = os.getenv('ALLOCATOR_KEY_NAME') module-attribute

logger = logging.getLogger(__name__) module-attribute

metrics_db = None module-attribute

operations_db = None module-attribute

operations_worker = None module-attribute

reboot_service = None module-attribute

schedule_db = None module-attribute

scheduler_service = None module-attribute

users = {cfg.app.admin_user: generate_password_hash(cfg.app.admin_password)} module-attribute

Classes

Functions:

init_database()

Initialize the database connection.

Source code in packages/allocator/src/lablink_allocator_service/main.py

def init_database():
    """Initialize the database connection."""
    global database
    database = VmDatabase(
        dbname=DB_NAME,
        user=DB_USER,
        password=cfg.db.password,
        host=DB_HOST,
        port=DB_PORT,
        table_name=VM_TABLE_NAME,
    )
    global schedule_db
    schedule_db = ScheduleDatabase(pool=database.pool)
    global metrics_db
    metrics_db = MetricsDatabase(pool=database.pool, table_name=VM_TABLE_NAME)
    # Expose the underlying psycopg2 pool to blueprints (e.g. /desktop,
    # /internal/proxy_auth) that need a raw connection for the signed-cookie
    # helpers, without coupling them to the VmDatabase wrapper.
    app.config["DB_POOL"] = database._pool
    app.config["VM_TABLE_NAME"] = VM_TABLE_NAME
    # Persist the deployment register-token as an argon2 hash at rest
    # (SR-F14). Validation reads this back via settings (Option A).
    database.set_setting("register_token_hash", hash_secret(REGISTER_TOKEN))

main()

Main entry point for the allocator service.

Source code in packages/allocator/src/lablink_allocator_service/main.py

def main():
    """Main entry point for the allocator service."""
    global scheduler_service, reboot_service, admin_session_expiry_service
    global operations_worker, operations_db
    global _startup_time

    verify_secrets_resolved()

    try:
        _startup_time = time.monotonic()
        with app.app_context():
            init_database()

        # Initialize scheduler service
        logger.info("Initializing scheduler service...")
        db_url = (
            f"postgresql://{DB_USER}:{cfg.db.password}"
            f"@{DB_HOST}:{DB_PORT}/{DB_NAME}"
        )
        scheduler_service = ScheduledDestructionService(
            schedule_db=schedule_db,
            db_url=db_url,
        )
        scheduler_service.start()
        atexit.register(scheduler_service.stop)
        logger.info("Scheduler service started successfully")

        # Initialize auto-reboot service
        logger.info("Initializing auto-reboot service...")
        reboot_service = AutoRebootService(
            database=database,
            region=cfg.app.region,
            tofu_dir=str(TOFU_DIR),
            provider=app.config.get("LABLINK_PROVIDER"),
        )
        reboot_service.start()
        atexit.register(reboot_service.stop)
        logger.info("Auto-reboot service started successfully")

        # Initialize admin-session expiry sweep
        logger.info("Initializing admin-session expiry service...")
        admin_session_expiry_service = AdminSessionExpiryService(
            database=database,
            timeout_minutes=cfg.app.admin_session_timeout_minutes,
        )
        admin_session_expiry_service.start()
        atexit.register(admin_session_expiry_service.stop)
        logger.info("Admin-session expiry service started successfully")

        # Initialize operations worker (on-demand apply/destroy jobs).
        # No atexit registration: see the module-level comment on
        # operations_worker — there's no background loop to stop.
        logger.info("Initializing operations worker...")
        operations_db = OperationsDatabase(pool=database.pool)
        operations_worker = OperationsWorker(database=operations_db)
        operations_worker.start()
        logger.info("Operations worker started successfully")

        # OpenTofu initialization — gated on the provider's capability flag
        # (mirrors the policy at module top: branch on capability, not type).
        # Manual/BYO providers don't provision hosts, so `tofu init` is
        # irrelevant and the binary may not even be present in the image.
        provider = app.config["LABLINK_PROVIDER"]
        if not provider.can_provision_hosts:
            logger.info(
                "Skipping tofu init: provider %s does not provision hosts.",
                getattr(provider, "name", type(provider).__name__),
            )
        elif not (TOFU_DIR / "terraform.runtime.tfvars").exists():
            logger.info("Initializing OpenTofu...")
            if ENVIRONMENT not in ["prod", "test", "ci-test"]:
                (TOFU_DIR / "backend.tf").unlink(missing_ok=True)
                subprocess.run(
                    ["tofu", "init"],
                    cwd=TOFU_DIR,
                    check=True,
                )
            else:
                # Use bucket_name from config for client VM tofu state
                default_bucket = "tf-state-lablink-allocator-bucket"
                bucket_name = (
                    cfg.bucket_name if hasattr(cfg, "bucket_name") else default_bucket
                )
                # Derive deployment_name for state key scoping
                deployment_name = (
                    cfg.deployment_name
                    if hasattr(cfg, "deployment_name") and cfg.deployment_name
                    else "lablink"
                )
                state_key = f"{deployment_name}/{ENVIRONMENT}/client/terraform.tfstate"
                logger.info(
                    f"Initializing OpenTofu with S3 backend: {bucket_name} "
                    f"(key: {state_key})"
                )
                subprocess.run(
                    [
                        "tofu",
                        "init",
                        f"-backend-config=backend-client-{ENVIRONMENT}.hcl",
                        f"-backend-config=key={state_key}",
                        f"-backend-config=bucket={bucket_name}",
                        f"-backend-config=region={cfg.app.region}",
                    ],
                    cwd=TOFU_DIR,
                    check=True,
                )

        logger.info("Auto-generated API token for machine-to-machine auth")
        logger.info("Starting Flask application...")
        flask_host = os.environ.get("FLASK_HOST", "127.0.0.1")
        flask_port = int(os.environ.get("FLASK_PORT", "8000"))
        app.run(host=flask_host, port=flask_port, threaded=True)

    except Exception as e:
        logger.error(f"Failed to start allocator service: {e}", exc_info=True)

        # Clean up services if they were initialized
        if reboot_service is not None:
            try:
                logger.info("Stopping auto-reboot service due to startup failure...")
                reboot_service.stop()
            except Exception as cleanup_error:
                logger.error(
                    f"Error stopping reboot service during cleanup: {cleanup_error}"
                )

        if scheduler_service is not None:
            try:
                logger.info("Stopping scheduler service due to startup failure...")
                scheduler_service.stop()
            except Exception as cleanup_error:
                logger.error(
                    f"Error stopping scheduler during cleanup: {cleanup_error}"
                )

        if admin_session_expiry_service is not None:
            try:
                logger.info(
                    "Stopping admin-session expiry service due to startup failure..."
                )
                admin_session_expiry_service.stop()
            except Exception as cleanup_error:
                logger.error(
                    f"Error stopping admin-session expiry service during "
                    f"cleanup: {cleanup_error}"
                )

        # Re-raise the exception to exit with error code
        raise

verify_secrets_resolved()

Refuse to start when a credential is still an unfilled placeholder.

Called from main(), not at import: starting the service is the only thing this must block, while importing this module is also how CI smoke-tests the shipped image and how the tests build a Flask client — neither of which supplies real credentials. Config-time validation cannot do the job either, since lablink-template validates config.yaml before substituting its PLACEHOLDER_* sentinels; the placeholders are legitimate there and only suspicious once the allocator is about to serve them.

Raises:

Type	Description
SystemExit	If any required credential is still a sentinel.

Source code in packages/allocator/src/lablink_allocator_service/main.py

def verify_secrets_resolved() -> None:
    """Refuse to start when a credential is still an unfilled placeholder.

    Called from ``main()``, not at import: starting the service is the only
    thing this must block, while *importing* this module is also how CI
    smoke-tests the shipped image and how the tests build a Flask client —
    neither of which supplies real credentials. Config-time validation cannot
    do the job either, since lablink-template validates config.yaml *before*
    substituting its PLACEHOLDER_* sentinels; the placeholders are legitimate
    there and only suspicious once the allocator is about to serve them.

    Raises:
        SystemExit: If any required credential is still a sentinel.
    """
    missing = [
        name
        for name, value in (
            ("app.admin_user", cfg.app.admin_user),
            ("app.admin_password", cfg.app.admin_password),
            ("db.password", cfg.db.password),
        )
        if is_unresolved_secret(value)
    ]
    if missing:
        raise SystemExit(
            f"FATAL: Required secrets not configured: {', '.join(missing)}.\n"
            "These still hold a placeholder, so no real credential was ever "
            "supplied. Refusing to start rather than serve an admin UI whose "
            "password is a literal published in the deployment template.\n"
            "Deploying via GitHub Actions substitutes them from the "
            "ADMIN_PASSWORD/DB_PASSWORD repository secrets; a local "
            "`tofu apply` does not — set real values in config.yaml first."
        )