Multiple schedulers (#156)

* implement multiple schedulers

* add dependencies

* add a comment

* remove debug log

* set scheduler in request_metadata

* update scheduler info on request_metadata

* typing

* qa

* changes cli

* refactor: streamline worker count retrieval and add Plackett-Luce shuffle for load balancing

* refactor: update Broker class to use 'schedulers' dictionary instead of 'clients' and remove unused 'scheduler_url' parameter

Author: francesconazzaro

alembic/versions/a4e8be715296_add_deleted_as_new_status.py

@@ -46,9 +46,13 @@ def upgrade() -> None:
 def downgrade() -> None:
     # Remove the new status from the enum
     # this doesn't work
-    #op.execute("ALTER TYPE status DELETE VALUE 'deleted'")
-    op.execute("CREATE TYPE status_old AS ENUM ('accepted','running','failed','successful','dismissed')")
+    # op.execute("ALTER TYPE status DELETE VALUE 'deleted'")
+    op.execute(
+        "CREATE TYPE status_old AS ENUM ('accepted','running','failed','successful','dismissed')"
+    )
     op.execute("DELETE FROM system_requests where status='deleted'")
-    op.execute("ALTER TABLE system_requests ALTER COLUMN status TYPE status_old USING (status::text::status_old)")
+    op.execute(
+        "ALTER TABLE system_requests ALTER COLUMN status TYPE status_old USING (status::text::status_old)"
+    )
     op.execute("DROP TYPE status")
-    op.execute("ALTER TYPE status_old RENAME TO status")
+    op.execute("ALTER TYPE status_old RENAME TO status")

cads_broker/database.py

@@ -719,6 +719,7 @@ def set_request_status(
     error_reason: str | None = None,
     resubmit: bool | None = None,
     priority: float | None = None,
+    scheduler: str | None = None,
 ) -> SystemRequest:
     """Set the status of a request."""
     statement = sa.select(SystemRequest).where(SystemRequest.request_uid == request_uid)
@@ -734,6 +735,9 @@ def set_request_status(
     if priority is not None:
         metadata.update({"priority": priority})
         request.request_metadata = metadata
+    if scheduler is not None:
+        metadata.update({"scheduler": scheduler})
+        request.request_metadata = metadata
     if status == "successful":
         request.finished_at = sa.func.now()
     elif status == "failed":