lablink_allocator_service.scheduler

lablink_allocator_service.scheduler

Attributes

logger = logging.getLogger(__name__) module-attribute

Classes

ScheduledDestructionService(schedule_db, db_url, tofu_dir=None)

Initialize the scheduler service.

Parameters:

Name	Type	Description	Default
schedule_db	ScheduleDatabase	ScheduleDatabase instance	required
db_url	str	PostgreSQL connection URL for APScheduler job store	required
tofu_dir	Optional[str]	Path to OpenTofu directory (optional, auto-detected if None)	None

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

def __init__(
    self,
    schedule_db: ScheduleDatabase,
    db_url: str,
    tofu_dir: Optional[str] = None,
):
    """Initialize the scheduler service.

    Args:
        schedule_db: ScheduleDatabase instance
        db_url: PostgreSQL connection URL for APScheduler job store
        tofu_dir: Path to OpenTofu directory (optional, auto-detected if None)
    """
    self.schedule_db: ScheduleDatabase = schedule_db
    self.db_url = db_url

    self.tofu_dir = tofu_dir or os.path.join(
        os.path.dirname(__file__), "terraform"
    )

    # Configure APScheduler with SQLAlchemy job store
    jobstores = {"default": SQLAlchemyJobStore(url=db_url)}

    executors = {"default": ThreadPoolExecutor(max_workers=2)}

    job_defaults = {
        "coalesce": False,
        "max_instances": 1,
        "misfire_grace_time": 300,
    }

    self.scheduler = BackgroundScheduler(
        jobstores=jobstores,
        executors=executors,
        job_defaults=job_defaults,
        timezone="UTC",
    )

Attributes

db_url = db_url instance-attribute

schedule_db = schedule_db instance-attribute

scheduler = BackgroundScheduler(jobstores=jobstores, executors=executors, job_defaults=job_defaults, timezone='UTC') instance-attribute

tofu_dir = tofu_dir or os.path.join(os.path.dirname(__file__), 'terraform') instance-attribute

Methods:

cancel_scheduled_destruction(schedule_id)

Cancel a scheduled destruction.

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

def cancel_scheduled_destruction(self, schedule_id: int) -> None:
    """Cancel a scheduled destruction."""

    # Remove from APScheduler
    job_id = f"destruction_{schedule_id}"
    if self.scheduler.get_job(job_id):
        self.scheduler.remove_job(job_id)

    # Update database
    self.schedule_db.cancel_scheduled_destruction(schedule_id)

    logger.info(f"Cancelled scheduled destruction ID: {schedule_id}")

schedule_destruction(schedule_name, destruction_time, recurrence_rule=None, created_by=None, notification_enabled=True, notification_hours_before=1)

Schedule a new destruction job.

Parameters:

Name	Type	Description	Default
schedule_name	str	Unique name for the schedule	required
destruction_time	datetime	Initial destruction time (UTC)	required
recurrence_rule	Optional[str]	Optional recurrence rule in RRULE format	None
created_by	Optional[str]	User who created the schedule	None
notification_enabled	bool	Whether to enable notifications	True
notification_hours_before	int	Hours before destruction to notify	1

Returns:

Type	Description
int	ID of the created schedule in the database

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

def schedule_destruction(
    self,
    schedule_name: str,
    destruction_time: datetime,
    recurrence_rule: Optional[str] = None,
    created_by: Optional[str] = None,
    notification_enabled: bool = True,
    notification_hours_before: int = 1,
) -> int:
    """Schedule a new destruction job.

    Args:
        schedule_name: Unique name for the schedule
        destruction_time: Initial destruction time (UTC)
        recurrence_rule: Optional recurrence rule in RRULE format
        created_by: User who created the schedule
        notification_enabled: Whether to enable notifications
        notification_hours_before: Hours before destruction to notify

    Returns:
        ID of the created schedule in the database
    """
    # Create database record
    # This will raise ValueError for duplicate names or RuntimeError for
    # other DB errors
    schedule_id = self.schedule_db.create_scheduled_destruction(
        schedule_name=schedule_name,
        destruction_time=destruction_time,
        recurrence_rule=recurrence_rule,
        created_by=created_by,
        notification_enabled=notification_enabled,
        notification_hours_before=notification_hours_before,
    )

    # Add job to APScheduler
    self._add_scheduler_job(
        schedule_id=schedule_id,
        destruction_time=destruction_time,
        recurrence_rule=recurrence_rule,
    )

    logger.info(
        f"Scheduled destruction '{schedule_name}' (ID: {schedule_id}) "
        f"for {destruction_time}"
    )
    return schedule_id

start()

Start the scheduler and load existing schedules

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

def start(self):
    """Start the scheduler and load existing schedules"""
    logger.info("Starting Scheduled Destruction Service...")
    self.scheduler.start()

    # Load existing schedules from the database
    self._load_scheduled_destructions()

stop()

Stop the scheduler.

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

def stop(self):
    """Stop the scheduler."""
    logger.info("Stopping scheduled destruction service")
    self.scheduler.shutdown(wait=True)

Functions:

execute_scheduled_destruction_job(schedule_id, tofu_dir)

Execute a scheduled destruction job.

This is a standalone function (not a method) to avoid pickling issues with APScheduler's SQLAlchemy job store.

Database credentials and provider configuration are read from the config system at runtime to avoid storing sensitive data in the APScheduler job store. Destruction is dispatched through provider.destroy_hosts (the same path as the /destroy HTTP route) so that all teardown logic lives in one place.

Parameters:

Name	Type	Description	Default
schedule_id	int	ID of the scheduled destruction	required
tofu_dir	str	Path to OpenTofu directory	required

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

def execute_scheduled_destruction_job(
    schedule_id: int,
    tofu_dir: str,
):
    """
    Execute a scheduled destruction job.

    This is a standalone function (not a method) to avoid pickling issues
    with APScheduler's SQLAlchemy job store.

    Database credentials and provider configuration are read from the config
    system at runtime to avoid storing sensitive data in the APScheduler job
    store.  Destruction is dispatched through ``provider.destroy_hosts``
    (the same path as the ``/destroy`` HTTP route) so that all teardown logic
    lives in one place.

    Args:
        schedule_id: ID of the scheduled destruction
        tofu_dir: Path to OpenTofu directory
    """
    from lablink_allocator_service.conf.structured_config import (
        DB_HOST,
        DB_NAME,
        DB_PORT,
        DB_USER,
        VM_TABLE_NAME,
    )
    from lablink_allocator_service.db.pool import make_pool
    from lablink_allocator_service.db.vms import VmDatabase
    from lablink_allocator_service.get_config import get_config
    from lablink_allocator_service.providers.registry import get_provider

    # Load config at runtime to get credentials (avoids storing passwords in job store)
    cfg = get_config()

    # Build a single pool for the whole job and derive every handle from
    # it — including the VmDatabase kept around for
    # clear_database() at the end of the job — instead of opening a
    # second pool internally. VmDatabase does not take ownership
    # of an injected pool (see its __init__ docstring), so `pool` remains
    # the sole owner and `finally` below is the only place that closes it.
    pool = make_pool(cfg.db.password)

    try:
        schedule_db = ScheduleDatabase(pool=pool)
        metrics_db = MetricsDatabase(pool=pool, table_name=VM_TABLE_NAME)
        database = VmDatabase(
            dbname=DB_NAME,
            user=DB_USER,
            password=cfg.db.password,
            host=DB_HOST,
            port=DB_PORT,
            table_name=VM_TABLE_NAME,
            pool=pool,
        )

        # Instantiate the provider from config (same as main.py at startup)
        provider = get_provider(
            cfg.provider,
            region=cfg.app.region,
            tofu_dir=tofu_dir,
        )

        logger.info(f"Executing scheduled destruction ID: {schedule_id}")

        # Mark as executing
        schedule_db.update_scheduled_destruction_status(
            schedule_id=schedule_id,
            status="executing",
        )

        from lablink_allocator_service.db.operations import OperationsDatabase

        operations_db = OperationsDatabase(pool=pool)
        if operations_db.get_in_progress_operation() is not None:
            logger.info(
                "Scheduled destruction %s skipped — an on-demand operation "
                "is in progress",
                schedule_id,
            )
            schedule_db.update_scheduled_destruction_status(
                schedule_id=schedule_id,
                status="failed",
                execution_result=(
                    "Skipped: an on-demand operation was in progress"
                ),
            )
            return

        if not provider.can_destroy_hosts:
            logger.info(
                "Scheduled destruction skipped — provider %s does not support destroy.",
                getattr(provider, "name", type(provider).__name__),
            )
        else:
            # Dispatch through provider.destroy_hosts (mirrors /destroy route)
            # by way of run_scheduled_destroy, which also seals any open
            # session-metrics rows so they don't survive the tear-down.
            handles = provider.list_hosts()
            run_scheduled_destroy(handles, metrics_db, provider)

        # Clear database
        logger.info("Clearing all VMs from database")
        database.clear_database()

        # Mark as completed
        schedule_db.update_scheduled_destruction_status(
            schedule_id=schedule_id,
            status="completed",
            execution_result="All VMs destroyed successfully",
        )

        logger.info(f"Scheduled destruction {schedule_id} completed successfully")

    except subprocess.CalledProcessError as e:
        error_msg = f"OpenTofu destroy failed: {e.stderr}"
        logger.error(error_msg)

        schedule_db.update_scheduled_destruction_status(
            schedule_id=schedule_id,
            status="failed",
            execution_result=error_msg,
        )

    except Exception as e:
        error_msg = f"Destruction failed: {str(e)}"
        logger.error(error_msg)

        schedule_db.update_scheduled_destruction_status(
            schedule_id=schedule_id,
            status="failed",
            execution_result=error_msg,
        )

    finally:
        pool.closeall()
        logger.debug("Scheduled-job connection pool closed.")

run_scheduled_destroy(handles, metrics_db, provider)

Seal session-metrics rows, then tear down the VMs.

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

def run_scheduled_destroy(handles: list, metrics_db, provider) -> None:
    """Seal session-metrics rows, then tear down the VMs."""
    logger.info("Bulk-sealing session metrics before destroy")
    sealed = metrics_db.bulk_seal_session_metrics()
    logger.info("Sealed %d session-metrics rows", sealed)
    logger.info("Running provider.destroy_hosts via scheduled job")
    result = provider.destroy_hosts(handles)
    logger.info("Scheduled destroy succeeded:\n%s", result.stdout)