UncountablePythonSDK 0.0.115__py3-none-any.whl → 0.0.142.dev0__py3-none-any.whl

uncountable/integration/queue_runner/job_scheduler.py

@@ -1,19 +1,27 @@
 import asyncio
+import os
 import sys
+import threading
 import typing
 from concurrent.futures import ProcessPoolExecutor
 from dataclasses import dataclass
 
 from opentelemetry.trace import get_current_span
 
-from uncountable.integration.db.connect import IntegrationDBService, create_db_engine
-from uncountable.integration.db.session import get_session_maker
 from uncountable.integration.queue_runner.command_server import (
     CommandEnqueueJob,
     CommandEnqueueJobResponse,
     CommandQueue,
+    CommandRetryJob,
+    CommandRetryJobResponse,
     CommandTask,
 )
+from uncountable.integration.queue_runner.command_server.types import (
+    CommandCancelJob,
+    CommandCancelJobResponse,
+    CommandCancelJobStatus,
+    CommandVaccuumQueuedJobs,
+)
 from uncountable.integration.queue_runner.datastore import DatastoreSqlite
 from uncountable.integration.queue_runner.datastore.interface import Datastore
 from uncountable.integration.queue_runner.worker import Worker
@@ -21,7 +29,7 @@ from uncountable.integration.scan_profiles import load_profiles
 from uncountable.integration.telemetry import Logger
 from uncountable.types import job_definition_t, queued_job_t
 
-from .types import ResultQueue, ResultTask
+from .types import RESTART_EXIT_CODE, ResultQueue, ResultTask
 
 _MAX_JOB_WORKERS = 5
 
@@ -83,14 +91,11 @@ def _start_workers(
     return job_worker_lookup
 
 
-async def start_scheduler(command_queue: CommandQueue) -> None:
+async def start_scheduler(
+    command_queue: CommandQueue, datastore: DatastoreSqlite
+) -> None:
     logger = Logger(get_current_span())
     result_queue: ResultQueue = asyncio.Queue()
-    engine = create_db_engine(IntegrationDBService.RUNNER)
-    session_maker = get_session_maker(engine)
-
-    datastore = DatastoreSqlite(session_maker)
-    datastore.setup(engine)
 
     with ProcessPoolExecutor(max_workers=_MAX_JOB_WORKERS) as process_pool:
         job_worker_lookup = _start_workers(
@@ -104,7 +109,9 @@ async def start_scheduler(command_queue: CommandQueue) -> None:
                 worker = job_worker_lookup[queued_job.job_ref_name]
             except KeyError as e:
                 logger.log_exception(e)
-                datastore.remove_job_from_queue(queued_job.queued_job_uuid)
+                datastore.update_job_status(
+                    queued_job.queued_job_uuid, queued_job_t.JobStatus.FAILED
+                )
                 return
             await worker.listen_queue.put(queued_job)
 
@@ -119,17 +126,11 @@ async def start_scheduler(command_queue: CommandQueue) -> None:
                     queued_job_t.InvocationContextManual,
                 ),
             ):
-                existing_queued_jobs = datastore.load_job_queue()
-                duplicate_job = next(
-                    (
-                        job
-                        for job in existing_queued_jobs
-                        if job.job_ref_name == job_ref_name
-                    ),
-                    None,
+                existing_queued_job = datastore.get_next_queued_job_for_ref_name(
+                    job_ref_name=job_ref_name
                 )
-                if duplicate_job is not None:
-                    return duplicate_job.queued_job_uuid
+                if existing_queued_job is not None:
+                    return existing_queued_job.queued_job_uuid
             queued_job = datastore.add_job_to_queue(
                 job_payload=payload,
                 job_ref_name=job_ref_name,
@@ -146,6 +147,53 @@ async def start_scheduler(command_queue: CommandQueue) -> None:
                 CommandEnqueueJobResponse(queued_job_uuid=queued_job_uuid)
             )
 
+        async def _handle_cancel_job_command(command: CommandCancelJob) -> None:
+            queued_job = datastore.get_queued_job(uuid=command.queued_job_uuid)
+            if queued_job is None:
+                await command.response_queue.put(
+                    CommandCancelJobResponse(status=CommandCancelJobStatus.NO_JOB_FOUND)
+                )
+                return
+
+            if queued_job.status == queued_job_t.JobStatus.QUEUED:
+                datastore.remove_job_from_queue(command.queued_job_uuid)
+                await command.response_queue.put(
+                    CommandCancelJobResponse(
+                        status=CommandCancelJobStatus.CANCELLED_WITH_RESTART
+                    )
+                )
+
+                def delayed_exit() -> None:
+                    os._exit(RESTART_EXIT_CODE)
+
+                threading.Timer(interval=5, function=delayed_exit).start()
+
+            else:
+                await command.response_queue.put(
+                    CommandCancelJobResponse(
+                        status=CommandCancelJobStatus.JOB_ALREADY_COMPLETED
+                    )
+                )
+
+        async def _handle_retry_job_command(command: CommandRetryJob) -> None:
+            queued_job = datastore.retry_job(command.queued_job_uuid)
+            if queued_job is None:
+                await command.response_queue.put(
+                    CommandRetryJobResponse(queued_job_uuid=None)
+                )
+                return
+
+            await enqueue_queued_job(queued_job)
+            await command.response_queue.put(
+                CommandRetryJobResponse(queued_job_uuid=queued_job.queued_job_uuid)
+            )
+
+        def _handle_vaccuum_queued_jobs_command(
+            command: CommandVaccuumQueuedJobs,
+        ) -> None:
+            logger.log_info("Vaccuuming queued jobs...")
+            datastore.vaccuum_queued_jobs()
+
         for queued_job in queued_jobs:
             await enqueue_queued_job(queued_job)
 
@@ -162,10 +210,26 @@ async def start_scheduler(command_queue: CommandQueue) -> None:
                     match command:
                         case CommandEnqueueJob():
                             await _handle_enqueue_job_command(command=command)
+                        case CommandRetryJob():
+                            await _handle_retry_job_command(command=command)
+                        case CommandVaccuumQueuedJobs():
+                            _handle_vaccuum_queued_jobs_command(command=command)
+                        case CommandCancelJob():
+                            await _handle_cancel_job_command(command=command)
                         case _:
                             typing.assert_never(command)
                     command_task = asyncio.create_task(command_queue.get())
                 elif task == result_task:
                     queued_job_result = result_task.result()
-                    datastore.remove_job_from_queue(queued_job_result.queued_job_uuid)
+                    match queued_job_result.job_result.success:
+                        case True:
+                            datastore.update_job_status(
+                                queued_job_result.queued_job_uuid,
+                                queued_job_t.JobStatus.SUCCESS,
+                            )
+                        case False:
+                            datastore.update_job_status(
+                                queued_job_result.queued_job_uuid,
+                                queued_job_t.JobStatus.FAILED,
+                            )
                     result_task = asyncio.create_task(result_queue.get())

uncountable/integration/queue_runner/queue_runner.py

@@ -1,16 +1,24 @@
 import asyncio
 
+from uncountable.integration.db.connect import IntegrationDBService, create_db_engine
+from uncountable.integration.db.session import get_session_maker
 from uncountable.integration.queue_runner.command_server import serve
 from uncountable.integration.queue_runner.command_server.types import CommandQueue
+from uncountable.integration.queue_runner.datastore import DatastoreSqlite
 from uncountable.integration.queue_runner.job_scheduler import start_scheduler
 
 
 async def queue_runner_loop() -> None:
     command_queue: CommandQueue = asyncio.Queue()
+    engine = create_db_engine(IntegrationDBService.RUNNER)
+    session_maker = get_session_maker(engine)
 
-    command_server = asyncio.create_task(serve(command_queue))
+    datastore = DatastoreSqlite(session_maker)
+    datastore.setup(engine)
 
-    scheduler = asyncio.create_task(start_scheduler(command_queue))
+    command_server = asyncio.create_task(serve(command_queue, datastore))
+
+    scheduler = asyncio.create_task(start_scheduler(command_queue, datastore))
 
     await scheduler
     await command_server

uncountable/integration/queue_runner/types.py

@@ -5,3 +5,5 @@ from uncountable.types import queued_job_t
 ListenQueue = Queue[queued_job_t.QueuedJob]
 ResultQueue = Queue[queued_job_t.QueuedJobResult]
 ResultTask = Task[queued_job_t.QueuedJobResult]
+
+RESTART_EXIT_CODE = 147

uncountable/integration/queue_runner/worker.py

@@ -87,33 +87,32 @@ def run_queued_job(
             base_span=span,
             profile_metadata=job_details.profile_metadata,
             job_definition=job_details.job_definition,
+            queued_job_uuid=queued_job.queued_job_uuid,
         )
-        try:
-            client = construct_uncountable_client(
-                profile_meta=job_details.profile_metadata, logger=job_logger
-            )
-            batch_processor = AsyncBatchProcessor(client=client)
-
-            payload = _resolve_queued_job_payload(queued_job)
-
-            args = JobArguments(
-                job_definition=job_details.job_definition,
-                client=client,
-                batch_processor=batch_processor,
-                profile_metadata=job_details.profile_metadata,
-                logger=job_logger,
-                payload=payload,
-            )
-
-            return execute_job(
-                args=args,
-                profile_metadata=job_details.profile_metadata,
-                job_definition=job_details.job_definition,
-                job_uuid=queued_job.queued_job_uuid,
-            )
-        except Exception as e:
-            job_logger.log_exception(e)
-            return job_definition_t.JobResult(success=False)
-        except BaseException as e:
-            job_logger.log_exception(e)
-            raise e
+        with job_logger.resource_tracking():
+            try:
+                client = construct_uncountable_client(
+                    profile_meta=job_details.profile_metadata, logger=job_logger
+                )
+                batch_processor = AsyncBatchProcessor(client=client)
+
+                payload = _resolve_queued_job_payload(queued_job)
+
+                args = JobArguments(
+                    job_definition=job_details.job_definition,
+                    client=client,
+                    batch_processor=batch_processor,
+                    profile_metadata=job_details.profile_metadata,
+                    logger=job_logger,
+                    payload=payload,
+                    job_uuid=queued_job.queued_job_uuid,
+                )
+
+                return execute_job(
+                    args=args,
+                    profile_metadata=job_details.profile_metadata,
+                    job_definition=job_details.job_definition,
+                )
+            except BaseException as e:
+                job_logger.log_exception(e)
+                return job_definition_t.JobResult(success=False)

uncountable/integration/scheduler.py

@@ -5,6 +5,8 @@ import sys
 import time
 from dataclasses import dataclass
 from datetime import UTC
+from enum import StrEnum
+from typing import assert_never
 
 from opentelemetry.trace import get_current_span
 
@@ -15,15 +17,24 @@ from uncountable.integration.queue_runner.command_server import (
     check_health,
 )
 from uncountable.integration.queue_runner.queue_runner import start_queue_runner
+from uncountable.integration.queue_runner.types import RESTART_EXIT_CODE
 from uncountable.integration.telemetry import Logger
 
 SHUTDOWN_TIMEOUT_SECS = 30
 
+AnyProcess = multiprocessing.Process | subprocess.Popen[bytes]
+
+
+class ProcessName(StrEnum):
+    QUEUE_RUNNER = "queue_runner"
+    CRON_SERVER = "cron_server"
+    UWSGI = "uwsgi"
+
 
 @dataclass(kw_only=True)
 class ProcessInfo:
-    name: str
-    process: multiprocessing.Process | subprocess.Popen[bytes]
+    name: ProcessName
+    process: AnyProcess
 
     @property
     def is_alive(self) -> bool:
@@ -46,14 +57,27 @@ class ProcessInfo:
                 return self.process.poll()
 
 
-def handle_shutdown(logger: Logger, processes: list[ProcessInfo]) -> None:
+@dataclass(kw_only=True)
+class ProcessAlarmRestart:
+    process: ProcessInfo
+
+
+@dataclass(kw_only=True)
+class ProcessAlarmShutdownAll:
+    pass
+
+
+ProcessAlarm = ProcessAlarmRestart | ProcessAlarmShutdownAll
+
+
+def handle_shutdown(logger: Logger, processes: dict[ProcessName, ProcessInfo]) -> None:
     logger.log_info("received shutdown command, shutting down sub-processes")
-    for proc_info in processes:
+    for proc_info in processes.values():
         if proc_info.is_alive:
             proc_info.process.terminate()
 
     shutdown_start = time.time()
-    still_living_processes = processes
+    still_living_processes = list(processes.values())
     while (
         time.time() - shutdown_start < SHUTDOWN_TIMEOUT_SECS
         and len(still_living_processes) > 0
@@ -82,14 +106,59 @@ def handle_shutdown(logger: Logger, processes: list[ProcessInfo]) -> None:
         proc_info.process.kill()
 
 
-def check_process_alive(logger: Logger, processes: list[ProcessInfo]) -> None:
-    for proc_info in processes:
+def restart_process(
+    logger: Logger, proc_info: ProcessInfo, processes: dict[ProcessName, ProcessInfo]
+) -> None:
+    logger.log_error(
+        f"process {proc_info.name} shut down unexpectedly - exit code {proc_info.exitcode}. Restarting..."
+    )
+
+    match proc_info.name:
+        case ProcessName.QUEUE_RUNNER:
+            queue_proc = multiprocessing.Process(target=start_queue_runner)
+            queue_proc.start()
+            new_info = ProcessInfo(name=ProcessName.QUEUE_RUNNER, process=queue_proc)
+            processes[ProcessName.QUEUE_RUNNER] = new_info
+            try:
+                _wait_queue_runner_online()
+                logger.log_info("queue runner restarted successfully")
+            except Exception as e:
+                logger.log_exception(e)
+                logger.log_error(
+                    "queue runner failed to restart, shutting down scheduler"
+                )
+                handle_shutdown(logger, processes)
+                sys.exit(1)
+
+        case ProcessName.CRON_SERVER:
+            cron_proc = multiprocessing.Process(target=cron_target)
+            cron_proc.start()
+            new_info = ProcessInfo(name=ProcessName.CRON_SERVER, process=cron_proc)
+            processes[ProcessName.CRON_SERVER] = new_info
+            logger.log_info("cron server restarted successfully")
+
+        case ProcessName.UWSGI:
+            uwsgi_proc: AnyProcess = subprocess.Popen(["uwsgi", "--die-on-term"])
+            new_info = ProcessInfo(name=ProcessName.UWSGI, process=uwsgi_proc)
+            processes[ProcessName.UWSGI] = new_info
+            logger.log_info("uwsgi restarted successfully")
+
+
+def check_process_alarms(
+    logger: Logger, processes: dict[ProcessName, ProcessInfo]
+) -> ProcessAlarm | None:
+    for proc_info in processes.values():
         if not proc_info.is_alive:
+            if proc_info.exitcode == RESTART_EXIT_CODE:
+                logger.log_warning(
+                    f"process {proc_info.name} requested restart! restarting"
+                )
+                return ProcessAlarmRestart(process=proc_info)
             logger.log_error(
                 f"process {proc_info.name} shut down unexpectedly! shutting down scheduler; exit code is {proc_info.exitcode}"
             )
-            handle_shutdown(logger, processes)
-            sys.exit(1)
+            return ProcessAlarmShutdownAll()
+    return None
 
 
 def _wait_queue_runner_online() -> None:
@@ -113,36 +182,65 @@ def _wait_queue_runner_online() -> None:
 
 def main() -> None:
     logger = Logger(get_current_span())
-    processes: list[ProcessInfo] = []
+    processes: dict[ProcessName, ProcessInfo] = {}
+
+    multiprocessing.set_start_method("forkserver")
 
     def add_process(process: ProcessInfo) -> None:
-        processes.append(process)
+        processes[process.name] = process
         logger.log_info(f"started process {process.name}")
 
-    runner_process = multiprocessing.Process(target=start_queue_runner)
-    runner_process.start()
-    add_process(ProcessInfo(name="queue runner", process=runner_process))
+    def _start_queue_runner() -> None:
+        runner_process = multiprocessing.Process(target=start_queue_runner)
+        runner_process.start()
+        add_process(
+            ProcessInfo(
+                name=ProcessName.QUEUE_RUNNER,
+                process=runner_process,
+            )
+        )
 
-    try:
-        _wait_queue_runner_online()
-    except Exception as e:
-        logger.log_exception(e)
-        handle_shutdown(logger, processes=processes)
-        return
+        try:
+            _wait_queue_runner_online()
+        except Exception as e:
+            logger.log_exception(e)
+            handle_shutdown(logger, processes=processes)
+            return
+
+    _start_queue_runner()
 
     cron_process = multiprocessing.Process(target=cron_target)
     cron_process.start()
-    add_process(ProcessInfo(name="cron server", process=cron_process))
+    add_process(ProcessInfo(name=ProcessName.CRON_SERVER, process=cron_process))
 
     uwsgi_process = subprocess.Popen([
         "uwsgi",
         "--die-on-term",
     ])
-    add_process(ProcessInfo(name="uwsgi", process=uwsgi_process))
+    add_process(ProcessInfo(name=ProcessName.UWSGI, process=uwsgi_process))
 
     try:
         while True:
-            check_process_alive(logger, processes=processes)
+            process_alarm = check_process_alarms(logger, processes=processes)
+            match process_alarm:
+                case ProcessAlarmRestart():
+                    match process_alarm.process.name:
+                        case ProcessName.QUEUE_RUNNER:
+                            del processes[ProcessName.QUEUE_RUNNER]
+                            _start_queue_runner()
+                        case ProcessName.CRON_SERVER | ProcessName.UWSGI:
+                            raise NotImplementedError(
+                                f"restarting {process_alarm.process.name} not yet implemented"
+                            )
+                        case _:
+                            assert_never(process_alarm.process.name)
+                case ProcessAlarmShutdownAll():
+                    handle_shutdown(logger, processes)
+                    sys.exit(1)
+                case None:
+                    pass
+                case _:
+                    assert_never(process_alarm)
             time.sleep(1)
     except KeyboardInterrupt:
         handle_shutdown(logger, processes=processes)

uncountable/integration/server.py

@@ -11,7 +11,11 @@ from apscheduler.triggers.cron import CronTrigger
 from opentelemetry.trace import get_current_span
 from sqlalchemy.engine.base import Engine
 
+from uncountable.core.environment import get_local_admin_server_port
 from uncountable.integration.cron import CronJobArgs, cron_job_executor
+from uncountable.integration.queue_runner.command_server.command_client import (
+    send_vaccuum_queued_jobs_message,
+)
 from uncountable.integration.telemetry import Logger
 from uncountable.types import base_t, job_definition_t
 from uncountable.types.job_definition_t import (
@@ -21,6 +25,14 @@ from uncountable.types.job_definition_t import (
 
 _MAX_APSCHEDULER_CONCURRENT_JOBS = 1
 
+VACCUUM_QUEUED_JOBS_JOB_ID = "vacuum_queued_jobs"
+
+STATIC_JOB_IDS = {VACCUUM_QUEUED_JOBS_JOB_ID}
+
+
+def vaccuum_queued_jobs() -> None:
+    send_vaccuum_queued_jobs_message(port=get_local_admin_server_port())
+
 
 class IntegrationServer:
     _scheduler: BaseScheduler
@@ -36,11 +48,27 @@ class IntegrationServer:
         )
         self._server_logger = Logger(get_current_span())
 
+    def _register_static_jobs(self) -> None:
+        all_job_ids = {job.id for job in self._scheduler.get_jobs()}
+        if VACCUUM_QUEUED_JOBS_JOB_ID in all_job_ids:
+            self._scheduler.remove_job(VACCUUM_QUEUED_JOBS_JOB_ID)
+
+        self._scheduler.add_job(
+            vaccuum_queued_jobs,
+            max_instances=1,
+            coalesce=True,
+            trigger=CronTrigger.from_crontab("5 4 * * 4"),
+            name="Vaccuum queued jobs",
+            id=VACCUUM_QUEUED_JOBS_JOB_ID,
+            kwargs={},
+            misfire_grace_time=None,
+        )
+
     def register_jobs(self, profiles: list[job_definition_t.ProfileMetadata]) -> None:
-        valid_job_ids = []
+        valid_job_ids: set[str] = set()
         for profile_metadata in profiles:
             for job_defn in profile_metadata.jobs:
-                valid_job_ids.append(job_defn.id)
+                valid_job_ids.add(job_defn.id)
                 match job_defn:
                     case CronJobDefinition():
                         # Add to ap scheduler
@@ -90,10 +118,11 @@ class IntegrationServer:
                         pass
                     case _:
                         assert_never(job_defn)
-        all_jobs = self._scheduler.get_jobs()
-        for job in all_jobs:
-            if job.id not in valid_job_ids:
-                self._scheduler.remove_job(job.id)
+        all_job_ids = {job.id for job in self._scheduler.get_jobs()}
+        invalid_job_ids = all_job_ids.difference(valid_job_ids.union(STATIC_JOB_IDS))
+
+        for job_id in invalid_job_ids:
+            self._scheduler.remove_job(job_id)
 
     def serve_forever(self) -> None:
         signal.pause()
@@ -106,6 +135,7 @@ class IntegrationServer:
 
     def __enter__(self) -> "IntegrationServer":
         self._start_apscheduler()
+        self._register_static_jobs()
         return self
 
     def __exit__(