UncountablePythonSDK 0.0.83__py3-none-any.whl → 0.0.132__py3-none-any.whl

uncountable/integration/queue_runner/datastore/datastore_sqlite.py

@@ -1,7 +1,8 @@
+import datetime
 import uuid
-from datetime import datetime, timezone
+from datetime import UTC
 
-from sqlalchemy import delete, insert, select, update
+from sqlalchemy import delete, insert, or_, select, text, update
 from sqlalchemy.engine import Engine
 
 from pkgs.argument_parser import CachedParser
@@ -13,6 +14,8 @@ from uncountable.types import queued_job_t
 
 queued_job_payload_parser = CachedParser(queued_job_t.QueuedJobPayload)
 
+MAX_QUEUE_WINDOW_DAYS = 30
+
 
 class DatastoreSqlite(Datastore):
     def __init__(self, session_maker: DBSessionMaker) -> None:
@@ -22,6 +25,17 @@ class DatastoreSqlite(Datastore):
     @classmethod
     def setup(cls, engine: Engine) -> None:
         Base.metadata.create_all(engine)
+        with engine.connect() as connection:
+            if not bool(
+                connection.execute(
+                    text(
+                        "select exists (select 1 from pragma_table_info('queued_jobs') where name='status');"
+                    )
+                ).scalar()
+            ):
+                connection.execute(
+                    text("alter table queued_jobs add column status VARCHAR")
+                )
 
     def add_job_to_queue(
         self, job_payload: queued_job_t.QueuedJobPayload, job_ref_name: str
@@ -30,11 +44,12 @@ class DatastoreSqlite(Datastore):
             serialized_payload = serialize_for_storage(job_payload)
             queued_job_uuid = str(uuid.uuid4())
             num_attempts = 0
-            submitted_at = datetime.now(timezone.utc)
+            submitted_at = datetime.datetime.now(UTC)
             insert_stmt = insert(QueuedJob).values({
                 QueuedJob.id.key: queued_job_uuid,
                 QueuedJob.job_ref_name.key: job_ref_name,
                 QueuedJob.payload.key: serialized_payload,
+                QueuedJob.status.key: queued_job_t.JobStatus.QUEUED,
                 QueuedJob.num_attempts: num_attempts,
                 QueuedJob.submitted_at: submitted_at,
             })
@@ -43,10 +58,48 @@ class DatastoreSqlite(Datastore):
                 queued_job_uuid=queued_job_uuid,
                 job_ref_name=job_ref_name,
                 payload=job_payload,
+                status=queued_job_t.JobStatus.QUEUED,
                 submitted_at=submitted_at,
                 num_attempts=num_attempts,
             )
 
+    def retry_job(
+        self,
+        queued_job_uuid: str,
+    ) -> queued_job_t.QueuedJob | None:
+        with self.session_maker() as session:
+            select_stmt = select(
+                QueuedJob.id,
+                QueuedJob.payload,
+                QueuedJob.num_attempts,
+                QueuedJob.job_ref_name,
+                QueuedJob.status,
+                QueuedJob.submitted_at,
+            ).filter(QueuedJob.id == queued_job_uuid)
+            existing_job = session.execute(select_stmt).one_or_none()
+
+            if (
+                existing_job is None
+                or existing_job.status != queued_job_t.JobStatus.FAILED
+            ):
+                return None
+
+            update_stmt = (
+                update(QueuedJob)
+                .values({QueuedJob.status.key: queued_job_t.JobStatus.QUEUED})
+                .filter(QueuedJob.id == queued_job_uuid)
+            )
+            session.execute(update_stmt)
+
+            return queued_job_t.QueuedJob(
+                queued_job_uuid=existing_job.id,
+                job_ref_name=existing_job.job_ref_name,
+                num_attempts=existing_job.num_attempts,
+                status=queued_job_t.JobStatus.QUEUED,
+                submitted_at=existing_job.submitted_at,
+                payload=queued_job_payload_parser.parse_storage(existing_job.payload),
+            )
+
     def increment_num_attempts(self, queued_job_uuid: str) -> int:
         with self.session_maker() as session:
             update_stmt = (
@@ -56,7 +109,7 @@ class DatastoreSqlite(Datastore):
             )
             session.execute(update_stmt)
             session.flush()
-            # IMPROVE: python3.12's sqlite does not support the RETURNING clause
+            # IMPROVE: python3's sqlite does not support the RETURNING clause
             select_stmt = select(QueuedJob.num_attempts).filter(
                 QueuedJob.id == queued_job_uuid
             )
@@ -67,15 +120,103 @@ class DatastoreSqlite(Datastore):
             delete_stmt = delete(QueuedJob).filter(QueuedJob.id == queued_job_uuid)
             session.execute(delete_stmt)
 
+    def update_job_status(
+        self, queued_job_uuid: str, status: queued_job_t.JobStatus
+    ) -> None:
+        with self.session_maker() as session:
+            update_stmt = (
+                update(QueuedJob)
+                .values({QueuedJob.status.key: status})
+                .filter(QueuedJob.id == queued_job_uuid)
+            )
+            session.execute(update_stmt)
+
+    def list_queued_job_metadata(
+        self, offset: int = 0, limit: int | None = 100
+    ) -> list[queued_job_t.QueuedJobMetadata]:
+        with self.session_maker() as session:
+            select_statement = (
+                select(
+                    QueuedJob.id,
+                    QueuedJob.job_ref_name,
+                    QueuedJob.num_attempts,
+                    QueuedJob.status,
+                    QueuedJob.submitted_at,
+                )
+                .order_by(QueuedJob.submitted_at)
+                .offset(offset)
+                .limit(limit)
+            )
+
+            queued_job_metadata: list[queued_job_t.QueuedJobMetadata] = [
+                queued_job_t.QueuedJobMetadata(
+                    queued_job_uuid=row.id,
+                    job_ref_name=row.job_ref_name,
+                    num_attempts=row.num_attempts,
+                    status=row.status or queued_job_t.JobStatus.QUEUED,
+                    submitted_at=row.submitted_at,
+                )
+                for row in session.execute(select_statement)
+            ]
+
+            return queued_job_metadata
+
+    def get_next_queued_job_for_ref_name(
+        self, job_ref_name: str
+    ) -> queued_job_t.QueuedJob | None:
+        with self.session_maker() as session:
+            select_stmt = (
+                select(
+                    QueuedJob.id,
+                    QueuedJob.payload,
+                    QueuedJob.num_attempts,
+                    QueuedJob.job_ref_name,
+                    QueuedJob.status,
+                    QueuedJob.submitted_at,
+                )
+                .filter(QueuedJob.job_ref_name == job_ref_name)
+                .filter(
+                    or_(
+                        QueuedJob.status == queued_job_t.JobStatus.QUEUED,
+                        QueuedJob.status.is_(None),
+                    )
+                )
+                .limit(1)
+                .order_by(QueuedJob.submitted_at)
+            )
+
+            for row in session.execute(select_stmt):
+                parsed_payload = queued_job_payload_parser.parse_storage(row.payload)
+                return queued_job_t.QueuedJob(
+                    queued_job_uuid=row.id,
+                    job_ref_name=row.job_ref_name,
+                    num_attempts=row.num_attempts,
+                    status=row.status or queued_job_t.JobStatus.QUEUED,
+                    submitted_at=row.submitted_at,
+                    payload=parsed_payload,
+                )
+
+            return None
+
     def load_job_queue(self) -> list[queued_job_t.QueuedJob]:
         with self.session_maker() as session:
-            select_stmt = select(
-                QueuedJob.id,
-                QueuedJob.payload,
-                QueuedJob.num_attempts,
-                QueuedJob.job_ref_name,
-                QueuedJob.submitted_at,
-            ).order_by(QueuedJob.submitted_at)
+            select_stmt = (
+                select(
+                    QueuedJob.id,
+                    QueuedJob.payload,
+                    QueuedJob.num_attempts,
+                    QueuedJob.job_ref_name,
+                    QueuedJob.status,
+                    QueuedJob.submitted_at,
+                )
+                .filter(
+                    or_(
+                        QueuedJob.status == queued_job_t.JobStatus.QUEUED,
+                        QueuedJob.status.is_(None),
+                    )
+                )
+                .order_by(QueuedJob.submitted_at)
+            )
 
             queued_jobs: list[queued_job_t.QueuedJob] = []
             for row in session.execute(select_stmt):
@@ -85,9 +226,25 @@ class DatastoreSqlite(Datastore):
                         queued_job_uuid=row.id,
                         job_ref_name=row.job_ref_name,
                         num_attempts=row.num_attempts,
+                        status=row.status or queued_job_t.JobStatus.QUEUED,
                         submitted_at=row.submitted_at,
                         payload=parsed_payload,
                     )
                 )
 
             return queued_jobs
+
+    def vaccuum_queued_jobs(self) -> None:
+        with self.session_maker() as session:
+            delete_stmt = (
+                delete(QueuedJob)
+                .filter(QueuedJob.status == queued_job_t.JobStatus.QUEUED)
+                .filter(
+                    QueuedJob.submitted_at
+                    <= (
+                        datetime.datetime.now(UTC)
+                        - datetime.timedelta(days=MAX_QUEUE_WINDOW_DAYS)
+                    )
+                )
+            )
+            session.execute(delete_stmt)

uncountable/integration/queue_runner/datastore/interface.py

@@ -17,3 +17,13 @@ class Datastore(ABC):
 
     @abstractmethod
     def load_job_queue(self) -> list[queued_job_t.QueuedJob]: ...
+
+    @abstractmethod
+    def get_next_queued_job_for_ref_name(
+        self, job_ref_name: str
+    ) -> queued_job_t.QueuedJob | None: ...
+
+    @abstractmethod
+    def list_queued_job_metadata(
+        self, offset: int, limit: int | None
+    ) -> list[queued_job_t.QueuedJobMetadata]: ...

uncountable/integration/queue_runner/datastore/model.py

@@ -1,7 +1,9 @@
-from sqlalchemy import JSON, BigInteger, Column, DateTime, Text
+from sqlalchemy import JSON, BigInteger, Column, DateTime, Enum, Text
 from sqlalchemy.orm import declarative_base
 from sqlalchemy.sql import func
 
+from uncountable.types import queued_job_t
+
 Base = declarative_base()
 
 
@@ -15,3 +17,8 @@ class QueuedJob(Base):
     )
     payload = Column(JSON, nullable=False)
     num_attempts = Column(BigInteger, nullable=False, default=0, server_default="0")
+    status = Column(
+        Enum(queued_job_t.JobStatus, length=None),
+        default=queued_job_t.JobStatus.QUEUED,
+        nullable=True,
+    )

uncountable/integration/queue_runner/job_scheduler.py

@@ -1,18 +1,22 @@
 import asyncio
+import sys
 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 (
+    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
@@ -34,6 +38,10 @@ class JobListenerKey:
 def _get_job_worker_key(
     job_definition: job_definition_t.JobDefinition, profile_name: str
 ) -> JobListenerKey:
+    if job_definition.subqueue_name is not None:
+        return JobListenerKey(
+            profile_name=profile_name, subqueue_name=job_definition.subqueue_name
+        )
     return JobListenerKey(profile_name=profile_name)
 
 
@@ -41,9 +49,12 @@ def on_worker_crash(
     worker_key: JobListenerKey,
 ) -> typing.Callable[[asyncio.Task], None]:
     def hook(task: asyncio.Task) -> None:
-        raise Exception(
-            f"worker {worker_key.profile_name}_{worker_key.subqueue_name} crashed unexpectedly"
+        Logger(get_current_span()).log_exception(
+            Exception(
+                f"worker {worker_key.profile_name}_{worker_key.subqueue_name} crashed unexpectedly"
+            )
         )
+        sys.exit(1)
 
     return hook
 
@@ -75,14 +86,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(
@@ -96,7 +104,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)
 
@@ -106,19 +116,16 @@ async def start_scheduler(command_queue: CommandQueue) -> None:
         ) -> str:
             if isinstance(
                 payload.invocation_context,
-                queued_job_t.InvocationContextCron,
+                (
+                    queued_job_t.InvocationContextCron,
+                    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,
@@ -135,6 +142,25 @@ async def start_scheduler(command_queue: CommandQueue) -> None:
                 CommandEnqueueJobResponse(queued_job_uuid=queued_job_uuid)
             )
 
+        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)
 
@@ -151,10 +177,24 @@ 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 _:
                             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/worker.py

@@ -90,7 +90,7 @@ def run_queued_job(
         )
         try:
             client = construct_uncountable_client(
-                profile_meta=job_details.profile_metadata, job_logger=job_logger
+                profile_meta=job_details.profile_metadata, logger=job_logger
             )
             batch_processor = AsyncBatchProcessor(client=client)
 
@@ -103,6 +103,7 @@ def run_queued_job(
                 profile_metadata=job_details.profile_metadata,
                 logger=job_logger,
                 payload=payload,
+                job_uuid=queued_job.queued_job_uuid,
             )
 
             return execute_job(
@@ -110,9 +111,6 @@ def run_queued_job(
                 profile_metadata=job_details.profile_metadata,
                 job_definition=job_details.job_definition,
             )
-        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
+            return job_definition_t.JobResult(success=False)

uncountable/integration/scan_profiles.py

@@ -26,7 +26,7 @@ def load_profiles() -> list[job_definition_t.ProfileMetadata]:
         profile_name = profile_file.name
         try:
             definition = profile_parser.parse_yaml_resource(
-                package=".".join([profiles_module, profile_name]),
+                package=f"{profiles_module}.{profile_name}",
                 resource="profile.yaml",
             )
             for job in definition.jobs:

uncountable/integration/scheduler.py

@@ -1,9 +1,11 @@
+import datetime
 import multiprocessing
 import subprocess
 import sys
 import time
 from dataclasses import dataclass
-from datetime import datetime, timezone
+from datetime import UTC
+from enum import StrEnum
 
 from opentelemetry.trace import get_current_span
 
@@ -18,11 +20,19 @@ 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:
@@ -45,14 +55,14 @@ class ProcessInfo:
                 return self.process.poll()
 
 
-def handle_shutdown(logger: Logger, processes: list[ProcessInfo]) -> None:
+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
@@ -81,46 +91,84 @@ 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_alive(
+    logger: Logger, processes: dict[ProcessName, ProcessInfo]
+) -> None:
+    for proc_info in processes.values():
         if not proc_info.is_alive:
-            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)
+            restart_process(logger, proc_info, processes)
 
 
 def _wait_queue_runner_online() -> None:
-    _MAX_QUEUE_RUNNER_HEALTH_CHECKS = 10
-    _QUEUE_RUNNER_HEALTH_CHECK_DELAY_SECS = 1
+    MAX_QUEUE_RUNNER_HEALTH_CHECKS = 10
+    QUEUE_RUNNER_HEALTH_CHECK_DELAY_SECS = 1
 
     num_attempts = 0
-    before = datetime.now(timezone.utc)
-    while num_attempts < _MAX_QUEUE_RUNNER_HEALTH_CHECKS:
+    before = datetime.datetime.now(UTC)
+    while num_attempts < MAX_QUEUE_RUNNER_HEALTH_CHECKS:
         try:
             if check_health(port=get_local_admin_server_port()):
                 return
         except CommandServerTimeout:
             pass
         num_attempts += 1
-        time.sleep(_QUEUE_RUNNER_HEALTH_CHECK_DELAY_SECS)
-    after = datetime.now(timezone.utc)
+        time.sleep(QUEUE_RUNNER_HEALTH_CHECK_DELAY_SECS)
+    after = datetime.datetime.now(UTC)
     duration_secs = (after - before).seconds
     raise Exception(f"queue runner failed to come online after {duration_secs} seconds")
 
 
 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))
+    add_process(ProcessInfo(name=ProcessName.QUEUE_RUNNER, process=runner_process))
 
     try:
         _wait_queue_runner_online()
@@ -131,13 +179,13 @@ def main() -> None:
 
     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:
@@ -147,4 +195,5 @@ def main() -> None:
         handle_shutdown(logger, processes=processes)
 
 
-main()
+if __name__ == "__main__":
+    main()

uncountable/integration/secret_retrieval/retrieve_secret.py

@@ -54,7 +54,7 @@ def _load_secret_overrides(profile_name: str) -> dict[SecretRetrieval, str]:
     profiles_module = os.environ["UNC_PROFILES_MODULE"]
     try:
         overrides = overrides_parser.parse_yaml_resource(
-            package=".".join([profiles_module, profile_name]),
+            package=f"{profiles_module}.{profile_name}",
             resource="local_overrides.yaml",
         )
         return {