Repository: AgnostiqHQ/covalent
Pull request #1982 (2025-03-26 20:28:27)

Move in-memory dispatcher state to the main DB
The dispatcher currently maintains a small amount of in-memory state consisting of several quantities that would be too costly (`O(num_electrons)`) to compute de novo from the database each time they are needed (`O(1)` times per electron). These are:
- Number of unresolved tasks in the dispatch (`num_submitted - num_teminal)`
- Number of dependencies of a task group that have yet to complete.
- The list of topologically sorted electrons within a task group.

Remembering these quantities over the course of a dispatch amortizes the initial computational cost over all of the electrons in the workflow.

This PR allows one to store the state in a real SQL DB by setting the connection string in the environment variable `COVALENT_CACHE_DB_URL`. By default, the cache db will live in an ephemeral sqlite DB that will be cleaned up when the server shuts down. If `/tmp` is a tmpfs mount like on many recent Linux distributions, the ephemeral DB will live purely in memory.

Benchmarks:

1000 trivial electrons, workflow processing time dominated by dispatcher I/O
```
import covalent as ct

@ct.electron
def task(x):
    return x + 1

@ct.lattice
def workflow(x):
    res = task(x)
    for _ in range(999):
        res = task(res)
    return res

dispatch_id = ct.dispatch(workflow)(0)
print(dispatch_id)
res = ct.get_result(dispatch_id, wait=True)
print(res)
```

With default settings (cache DB on tmpfs):
```
(venv) casey@datarobot covalent (bye-redis)$ time python benchmark.py 

real    1m46.539s
user    0m4.526s
sys     0m1.121s
```

With `COVALENT_CACHE_DB_URL` set to the connection string for the main dispatch DB (on SSD)
```
(venv) casey@datarobot covalent (bye-redis)$ time python benchmark.py 

real    2m55.202s
user    0m4.548s
sys     0m1.165s
```

Current develop (dispatcher state in python dictionary):
```
(venv) casey@datarobot covalent (develop)$ time python benchmark.py 

real    1m40.423s
user    0m4.521s
sys     0m1.117s
```


<!--
⚠️ Your pull request title should be short detailed and understandable for all.
⚠️ Add a note to /CHANGELOG.md summarizing the changes.
⚠️ If your pull request fixes an open issue, please link to the issue.
⚠️ Ensure that your branch is up-to-date with the base branch.
-->

- [ ] I have added the tests to cover my changes.
- [ ] I have updated the documentation and CHANGELOG accordingly.
- [ ] I have read the CONTRIBUTING document.


cjao: The external dep `pytest-asyncio` is somehow failing the license check.
cjao: This one is less time sensitive than #1981  and #1883  @wjcunningham7 

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 554d0544..54d69b39 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [UNRELEASED]
 
+### Changed
+
+- Moved in-memory dispatcher state to main DB
+
 ## [0.238.0-rc.0] - 2025-03-05
 
 ### Authors
diff --git a/covalent_dispatcher/_core/dispatcher.py b/covalent_dispatcher/_core/dispatcher.py
index b0ecd27b..14ae116f 100644
--- a/covalent_dispatcher/_core/dispatcher.py
+++ b/covalent_dispatcher/_core/dispatcher.py
@@ -34,7 +34,7 @@ from . import data_manager as datasvc
 from . import runner_ng
 from .data_modules import graph as tg_utils
 from .data_modules import job_manager as jbmgr
-from .dispatcher_modules.caches import _pending_parents, _sorted_task_groups, _unresolved_tasks
+from .dispatcher_modules.caches import _task_group_cache, _workflow_run_cache
 from .runner_modules.cancel import cancel_tasks
 
 app_log = logger.app_log
@@ -95,7 +95,7 @@ async def _handle_completed_node(dispatch_id: str, node_id: int):
         gid = child["task_group_id"]
         app_log.debug(f"dispatch {dispatch_id}: parent gid {parent_gid}, child gid {gid}")
         if parent_gid != gid:
-            now_pending = await _pending_parents.decrement(dispatch_id, gid)
+            now_pending = await _task_group_cache.decrement(dispatch_id, gid)
             if now_pending < 1:
                 app_log.debug(f"Queuing task group {gid} for execution")
                 next_task_groups.append(gid)
@@ -414,13 +414,12 @@ async def _finalize_dispatch(dispatch_id: str):
 
 
 async def _initialize_caches(dispatch_id, pending_parents, sorted_task_groups):
-    for gid, indegree in pending_parents.items():
-        await _pending_parents.set_pending(dispatch_id, gid, indegree)
-
-    for gid, sorted_nodes in sorted_task_groups.items():
-        await _sorted_task_groups.set_task_group(dispatch_id, gid, sorted_nodes)
+    for gid in pending_parents:
+        indegree = pending_parents[gid]
+        sorted_nodes = sorted_task_groups[gid]
+        await _task_group_cache.set(dispatch_id, gid, indegree, sorted_nodes)
 
-    await _unresolved_tasks.set_unresolved(dispatch_id, 0)
+    await _workflow_run_cache.set_unresolved(dispatch_id, 0)
 
 
 async def _submit_initial_tasks(dispatch_id: str):
@@ -442,7 +441,7 @@ async def _submit_initial_tasks(dispatch_id: str):
     for gid in initial_groups:
         sorted_nodes = sorted_task_groups[gid]
         app_log.debug(f"Sorted nodes group group {gid}: {sorted_nodes}")
-        await _unresolved_tasks.increment(dispatch_id, len(sorted_nodes))
+        await _workflow_run_cache.increment(dispatch_id, len(sorted_nodes))
 
     for gid in initial_groups:
         sorted_nodes = sorted_task_groups[gid]
@@ -469,8 +468,8 @@ async def _handle_node_status_update(dispatch_id, node_id, node_status, detail):
     if node_status == RESULT_STATUS.COMPLETED:
         next_task_groups = await _handle_completed_node(dispatch_id, node_id)
         for gid in next_task_groups:
-            sorted_nodes = await _sorted_task_groups.get_task_group(dispatch_id, gid)
-            await _unresolved_tasks.increment(dispatch_id, len(sorted_nodes))
+            sorted_nodes = await _task_group_cache.get_task_group(dispatch_id, gid)
+            await _workflow_run_cache.increment(dispatch_id, len(sorted_nodes))
             await _submit_task_group(dispatch_id, sorted_nodes, gid)
 
     if node_status == RESULT_STATUS.FAILED:
@@ -481,7 +480,7 @@ async def _handle_node_status_update(dispatch_id, node_id, node_status, detail):
 
     # Decrement after any increments to avoid race with
     # finalize_dispatch()
-    await _unresolved_tasks.decrement(dispatch_id)
+    await _workflow_run_cache.decrement(dispatch_id)
 
 
 async def _handle_dispatch_exception(dispatch_id: str, ex: Exception) -> RESULT_STATUS:
@@ -530,7 +529,7 @@ async def _handle_event(msg: Dict):
             fut.set_result(dispatch_status)
         return dispatch_status
 
-    unresolved = await _unresolved_tasks.get_unresolved(dispatch_id)
+    unresolved = await _workflow_run_cache.get_unresolved(dispatch_id)
     if unresolved < 1:
         app_log.debug("Finalizing dispatch")
         try:
@@ -549,7 +548,7 @@ async def _handle_event(msg: Dict):
 
 async def _clear_caches(dispatch_id: str):
     """Clean up all keys in caches."""
-    await _unresolved_tasks.remove(dispatch_id)
+    await _workflow_run_cache.remove(dispatch_id)
 
     g_node_link = await tg_utils.get_nodes_links(dispatch_id)
     g = nx.readwrite.node_link_graph(g_node_link)
@@ -558,5 +557,4 @@ async def _clear_caches(dispatch_id: str):
 
     for gid in task_groups:
         # Clean up no longer referenced keys
-        await _pending_parents.remove(dispatch_id, gid)
-        await _sorted_task_groups.remove(dispatch_id, gid)
+        await _task_group_cache.remove(dispatch_id, gid)
diff --git a/covalent_dispatcher/_core/dispatcher_modules/caches.py b/covalent_dispatcher/_core/dispatcher_modules/caches.py
index ce7b53e5..bca8db9d 100644
--- a/covalent_dispatcher/_core/dispatcher_modules/caches.py
+++ b/covalent_dispatcher/_core/dispatcher_modules/caches.py
@@ -18,84 +18,199 @@
 Helper classes for the dispatcher
 """
 
-from .store import _DictStore, _KeyValueBase
+import json
+import os
+import tempfile
 
+from covalent_dispatcher._core.data_modules.utils import run_in_executor
+from covalent_dispatcher._dal.base import workflow_db
+from covalent_dispatcher._dal.dispatcher_state import TaskGroupState, WorkflowState
+from covalent_dispatcher._db.datastore import DataStore
 
-def _pending_parents_key(dispatch_id: str, node_id: int):
-    return f"pending-parents-{dispatch_id}:{node_id}"
 
+class _WorkflowRunState:
 
-def _unresolved_tasks_key(dispatch_id: str):
-    return f"unresolved-{dispatch_id}"
+    def __init__(self, db: DataStore):
+        self.db = db
 
-
-def _task_groups_key(dispatch_id: str, task_group_id: int):
-    return f"task-groups-{dispatch_id}:{task_group_id}"
-
-
-class _UnresolvedTasksCache:
-    def __init__(self, store: _KeyValueBase = _DictStore()):
-        self._store = store
+    def _get_unresolved(self, dispatch_id: str):
+        with self.db.session() as session:
+            records = WorkflowState.get(
+                session,
+                fields=["num_unresolved_tasks"],
+                equality_filters={"dispatch_id": dispatch_id},
+                membership_filters={},
+            )
+            return records[0].num_unresolved_tasks
 
     async def get_unresolved(self, dispatch_id: str):
-        key = _unresolved_tasks_key(dispatch_id)
-        return await self._store.get(key)
+        return await run_in_executor(self._get_unresolved, dispatch_id)
+
+    def _set_unresolved(self, dispatch_id: str, val: int):
+        with self.db.session() as session:
+            WorkflowState.create(
+                session,
+                insert_kwargs={
+                    "dispatch_id": dispatch_id,
+                    "num_unresolved_tasks": val,
+                },
+            )
+            session.commit()
 
     async def set_unresolved(self, dispatch_id: str, val: int):
-        key = _unresolved_tasks_key(dispatch_id)
-        await self._store.insert(key, val)
+        return await run_in_executor(self._set_unresolved, dispatch_id, val)
+
+    def _increment(self, dispatch_id: str, interval: int = 1):
+        with self.db.session() as session:
+            WorkflowState.incr_bulk(
+                session=session,
+                increments={"num_unresolved_tasks": interval},
+                equality_filters={"dispatch_id": dispatch_id},
+                membership_filters={},
+            )
+            records = WorkflowState.get(
+                session,
+                fields=["num_unresolved_tasks"],
+                equality_filters={"dispatch_id": dispatch_id},
+                membership_filters={},
+            )
+            session.commit()
+            return records[0].num_unresolved_tasks
 
     async def increment(self, dispatch_id: str, interval: int = 1):
-        key = _unresolved_tasks_key(dispatch_id)
-        return await self._store.increment(key, interval)
+        return await run_in_executor(self._increment, dispatch_id, interval)
+
+    def _decrement(self, dispatch_id: str):
+        with self.db.session() as session:
+            WorkflowState.incr_bulk(
+                session=session,
+                increments={"num_unresolved_tasks": -1},
+                equality_filters={"dispatch_id": dispatch_id},
+                membership_filters={},
+            )
+            records = WorkflowState.get(
+                session,
+                fields=["num_unresolved_tasks"],
+                equality_filters={"dispatch_id": dispatch_id},
+                membership_filters={},
+            )
+            session.commit()
+            return records[0].num_unresolved_tasks
 
     async def decrement(self, dispatch_id: str):
-        key = _unresolved_tasks_key(dispatch_id)
-        return await self._store.increment(key, -1)
+        return await run_in_executor(self._decrement, dispatch_id)
+
+    def _remove(self, dispatch_id: str):
+        with self.db.session() as session:
+            WorkflowState.delete_bulk(
+                session=session,
+                equality_filters={"dispatch_id": dispatch_id},
+                membership_filters={},
+            )
+            session.commit()
 
     async def remove(self, dispatch_id: str):
-        key = _unresolved_tasks_key(dispatch_id)
-        await self._store.remove(key)
+        await run_in_executor(self._remove, dispatch_id)
 
 
-class _PendingParentsCache:
-    def __init__(self, store: _KeyValueBase = _DictStore()):
-        self._store = store
+class TaskGroupRunState:
 
-    async def get_pending(self, dispatch_id: str, task_group_id: int):
-        key = _pending_parents_key(dispatch_id, task_group_id)
-        return await self._store.get(key)
+    def __init__(self, db):
+        self.db = db
+
+    def _get_pending(self, dispatch_id: str, task_group_id: int):
+        with self.db.session() as session:
+            records = TaskGroupState.get(
+                session=session,
+                fields=["num_pending_parents"],
+                equality_filters={"dispatch_id": dispatch_id, "task_group_id": task_group_id},
+                membership_filters={},
+            )
+            return records[0].num_pending_parents
 
-    async def set_pending(self, dispatch_id: str, task_group_id: int, val: int):
-        key = _pending_parents_key(dispatch_id, task_group_id)
-        await self._store.insert(key, val)
+    async def get_pending(self, dispatch_id: str, task_group_id: int):
+        return await run_in_executor(self._get_pending, dispatch_id, task_group_id)
+
+    def _set(self, dispatch_id: str, task_group_id: int, num_pending: int, sorted_nodes):
+        with self.db.session() as session:
+            TaskGroupState.create(
+                session=session,
+                insert_kwargs={
+                    "dispatch_id": dispatch_id,
+                    "task_group_id": task_group_id,
+                    "num_pending_parents": num_pending,
+                    "sorted_tasks": json.dumps(sorted_nodes),
+                },
+            )
+            session.commit()
+
+    async def set(self, dispatch_id: str, task_group_id: int, num_pending: int, sorted_nodes):
+        return await run_in_executor(
+            self._set, dispatch_id, task_group_id, num_pending, sorted_nodes
+        )
+
+    def _decrement(self, dispatch_id: str, task_group_id):
+        with self.db.session() as session:
+            TaskGroupState.incr_bulk(
+                session=session,
+                increments={"num_pending_parents": -1},
+                equality_filters={"dispatch_id": dispatch_id, "task_group_id": task_group_id},
+                membership_filters={},
+            )
+            records = TaskGroupState.get(
+                session,
+                fields=["num_pending_parents"],
+                equality_filters={"dispatch_id": dispatch_id, "task_group_id": task_group_id},
+                membership_filters={},
+            )
+            session.commit()
+            return records[0].num_pending_parents
 
     async def decrement(self, dispatch_id: str, task_group_id: int):
-        key = _pending_parents_key(dispatch_id, task_group_id)
-        return await self._store.increment(key, -1)
+        return await run_in_executor(self._decrement, dispatch_id, task_group_id)
 
     async def remove(self, dispatch_id: str, task_group_id: int):
-        key = _pending_parents_key(dispatch_id, task_group_id)
-        await self._store.remove(key)
-
-
-class _SortedTaskGroups:
-    def __init__(self, store: _KeyValueBase = _DictStore()):
-        self._store = store
+        pass
+
+    def _get_task_group(self, dispatch_id: str, task_group_id: int):
+        with self.db.session() as session:
+            records = TaskGroupState.get(
+                session=session,
+                fields=["sorted_tasks"],
+                equality_filters={"dispatch_id": dispatch_id, "task_group_id": task_group_id},
+                membership_filters={},
+            )
+            return json.loads(records[0].sorted_tasks)
 
     async def get_task_group(self, dispatch_id: str, task_group_id: int):
-        key = _task_groups_key(dispatch_id, task_group_id)
-        return await self._store.get(key)
+        return await run_in_executor(self._get_task_group, dispatch_id, task_group_id)
 
-    async def set_task_group(self, dispatch_id: str, task_group_id: int, sorted_nodes: list):
-        key = _task_groups_key(dispatch_id, task_group_id)
-        await self._store.insert(key, sorted_nodes)
+    def _remove(self, dispatch_id: str, task_group_id: int):
+        with self.db.session() as session:
+            TaskGroupState.delete_bulk(
+                session=session,
+                equality_filters={"dispatch_id": dispatch_id, "task_group_id": task_group_id},
+                membership_filters={},
+            )
+            session.commit()
 
     async def remove(self, dispatch_id: str, task_group_id: int):
-        key = _task_groups_key(dispatch_id, task_group_id)
-        await self._store.remove(key)
+        await run_in_executor(self._remove, dispatch_id, task_group_id)
+
+
+# Default to tmpfs backed file
+cache_db_file = tempfile.NamedTemporaryFile(
+    mode="w+b", prefix="covalent-dispatcher-cache-", suffix=".db"
+)
+cache_db_URL = os.environ.get("COVALENT_CACHE_DB_URL", f"sqlite+pysqlite:///{cache_db_file.name}")
+initialize_db = True
+
+# If we want to store dispatcher state in the main DB, let the alembic migrations
+# create the tables
+if cache_db_URL == workflow_db.db_URL:
+    initialize_db = False
 
+cache_db = DataStore(db_URL=cache_db_URL, initialize_db=initialize_db)
 
-_pending_parents = _PendingParentsCache()
-_unresolved_tasks = _UnresolvedTasksCache()
-_sorted_task_groups = _SortedTaskGroups()
+_task_group_cache = TaskGroupRunState(db=cache_db)
+_workflow_run_cache = _WorkflowRunState(db=cache_db)
diff --git a/covalent_dispatcher/_core/dispatcher_modules/store.py b/covalent_dispatcher/_core/dispatcher_modules/store.py
deleted file mode 100644
index a27fe587..00000000
--- a/covalent_dispatcher/_core/dispatcher_modules/store.py
+++ /dev/null
@@ -1,66 +0,0 @@
-# Copyright 2021 Agnostiq Inc.
-#
-# This file is part of Covalent.
-#
-# Licensed under the Apache License 2.0 (the "License"). A copy of the
-# License may be obtained with this software package or at
-#
-#     https://www.apache.org/licenses/LICENSE-2.0
-#
-# Use of this file is prohibited except in compliance with the License.
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""
-Simple Key-Value store base
-"""
-
-
-class _KeyValueBase:
-    async def get(self, key):
-        raise NotImplementedError
-
-    async def insert(self, key, val):
-        raise NotImplementedError
-
-    async def belongs(self, key):
-        raise NotImplementedError
-
-    async def remove(self, key):
-        raise NotImplementedError
-
-    async def increment(self, key: str, delta: int) -> int:
-        """Increments value for `key` by amount `delta`
-
-        Parameters:
-            key: the value to change
-            delta: the amount to change (can be negative)
-        Returns:
-            The new value
-        """
-
-        raise NotImplementedError
-
-
-class _DictStore(_KeyValueBase):
-    def __init__(self):
-        self._store = {}
-
-    async def get(self, key):
-        return self._store[key]
-
-    async def insert(self, key, val):
-        self._store[key] = val
-
-    async def belongs(self, key):
-        return key in self._store
-
-    async def remove(self, key):
-        del self._store[key]
-
-    async def increment(self, key, delta: int):
-        self._store[key] += delta
-        return self._store[key]
diff --git a/covalent_dispatcher/_dal/controller.py b/covalent_dispatcher/_dal/controller.py
index 53928a2f..b429e42a 100644
--- a/covalent_dispatcher/_dal/controller.py
+++ b/covalent_dispatcher/_dal/controller.py
@@ -19,7 +19,7 @@ from __future__ import annotations
 
 from typing import Generic, List, Optional, Sequence, Type, TypeVar, Union
 
-from sqlalchemy import select, update
+from sqlalchemy import delete, select, update
 from sqlalchemy.engine import Row
 from sqlalchemy.orm import Session
 from sqlalchemy.sql.expression import Select, desc
@@ -182,6 +182,29 @@ class Record(Generic[T]):
             stmt = stmt.where(getattr(cls.model, attr).in_(vals))
         session.execute(stmt)
 
+    @classmethod
+    def delete_bulk(
+        cls,
+        session: Session,
+        *,
+        equality_filters: dict,
+        membership_filters: dict,
+    ):
+        """Bulk delete.
+
+        Args:
+            session: SQLAlchemy session
+            equality_filters: Dict{field_name: value}
+            membership_filters: Dict{field_name: value_list}
+        """
+
+        stmt = delete(cls.model)
+        for attr, val in equality_filters.items():
+            stmt = stmt.where(getattr(cls.model, attr) == val)
+        for attr, vals in membership_filters.items():
+            stmt = stmt.where(getattr(cls.model, attr).in_(vals))
+        session.execute(stmt)
+
     def update(self, session: Session, *, values: dict):
         """Update the corresponding DB record."""
 
diff --git a/covalent_dispatcher/_dal/dispatcher_state.py b/covalent_dispatcher/_dal/dispatcher_state.py
new file mode 100644
index 00000000..65af292e
--- /dev/null
+++ b/covalent_dispatcher/_dal/dispatcher_state.py
@@ -0,0 +1,29 @@
+# Copyright 2024 Agnostiq Inc.
+#
+# This file is part of Covalent.
+#
+# Licensed under the Apache License 2.0 (the "License"). A copy of the
+# License may be obtained with this software package or at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Use of this file is prohibited except in compliance with the License.
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from __future__ import annotations
+
+from .._db import models
+from .controller import Record
+
+
+class WorkflowState(Record[models.WorkflowState]):
+    model = models.WorkflowState
+
+
+class TaskGroupState(Record[models.TaskGroupState]):
+    model = models.TaskGroupState
diff --git a/covalent_dispatcher/_db/datastore.py b/covalent_dispatcher/_db/datastore.py
index 2fe99015..32328b8c 100644
--- a/covalent_dispatcher/_db/datastore.py
+++ b/covalent_dispatcher/_db/datastore.py
@@ -17,7 +17,7 @@
 from contextlib import contextmanager
 from os import environ, path
 from pathlib import Path
-from typing import Generator, Optional
+from typing import Any, Generator, Optional
 
 from alembic import command
 from alembic.config import Config
@@ -40,6 +40,7 @@ class DataStore:
         self,
         db_URL: Optional[str] = None,
         initialize_db: bool = False,
+        declarative_base: Any = models.Base,
         **kwargs,
     ):
         if db_URL:
@@ -57,7 +58,7 @@ class DataStore:
 
         # flag should only be used in pytest - tables should be generated using migrations
         if initialize_db:
-            models.Base.metadata.create_all(self.engine)
+            declarative_base.metadata.create_all(self.engine)
 
     @staticmethod
     def factory():
diff --git a/covalent_dispatcher/_db/models.py b/covalent_dispatcher/_db/models.py
index a8d10073..59534243 100644
--- a/covalent_dispatcher/_db/models.py
+++ b/covalent_dispatcher/_db/models.py
@@ -39,7 +39,10 @@ Base = declarative_base()
 
 class Lattice(Base):
     __tablename__ = "lattices"
-    __table_args__ = (UniqueConstraint("dispatch_id", name="u_dispatch_id"),)
+    __table_args__ = (
+        UniqueConstraint("dispatch_id", name="u_dispatch_id"),
+        UniqueConstraint("electron_id", name="u_electron_id"),
+    )
 
     id = Column(Integer, primary_key=True)
     dispatch_id = Column(String(64), nullable=False)
@@ -292,3 +295,34 @@ class Asset(Base):
 
     # Size in bytes
     size = Column(Integer, nullable=True)
+
+
+# Hold state of running workflows to avoid expensive computation per electron
+
+
+class WorkflowState(Base):
+    __tablename__ = "workflowstate"
+
+    dispatch_id = Column(Text, primary_key=True)
+
+    # Counter for dispatcher to track the number of unresolved tasks (submitted - terminal)
+    num_unresolved_tasks = Column(Integer, nullable=False)
+
+
+class TaskGroupState(Base):
+    __tablename__ = "taskgroupstate"
+    __table_args__ = (Index("task_grp_idx", "dispatch_id", "task_group_id"),)
+
+    id = Column(Integer, primary_key=True)
+
+    dispatch_id = Column(Text, nullable=False)
+
+    task_group_id = Column(Integer, nullable=False)
+
+    # Number of parent tasks that have yet to complete
+    # Used by dispatcher
+    num_pending_parents = Column(Integer, nullable=False)
+
+    # JSON list of node ids in topological order
+    # Used by dispatcher
+    sorted_tasks = Column(Text, nullable=False)
diff --git a/covalent_dispatcher/_service/app.py b/covalent_dispatcher/_service/app.py
index 3d582a7b..bf02b66a 100644
--- a/covalent_dispatcher/_service/app.py
+++ b/covalent_dispatcher/_service/app.py
@@ -20,11 +20,10 @@
 import asyncio
 import json
 from contextlib import asynccontextmanager
-from typing import List, Union
+from typing import Annotated, List, Union
 
 from fastapi import APIRouter, FastAPI, HTTPException, Query
 from fastapi.responses import JSONResponse
-from typing_extensions import Annotated
 
 import covalent_dispatcher.entry_point as dispatcher
 from covalent._shared_files import logger
diff --git a/covalent_migrations/versions/7e9fb153ecfb_add_task_group_table.py b/covalent_migrations/versions/7e9fb153ecfb_add_task_group_table.py
new file mode 100644
index 00000000..80cda084
--- /dev/null
+++ b/covalent_migrations/versions/7e9fb153ecfb_add_task_group_table.py
@@ -0,0 +1,73 @@
+# Copyright 2021 Agnostiq Inc.
+#
+# This file is part of Covalent.
+#
+# Licensed under the Apache License 2.0 (the "License"). A copy of the
+# License may be obtained with this software package or at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Use of this file is prohibited except in compliance with the License.
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Add task_group table
+
+Revision ID: 7e9fb153ecfb
+Revises: 7b40a55c0752
+Create Date: 2025-03-25 16:36:49.408467
+
+"""
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+# pragma: allowlist nextline secret
+revision = "7e9fb153ecfb"
+# pragma: allowlist nextline secret
+down_revision = "7b40a55c0752"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table(
+        "taskgroupstate",
+        sa.Column("id", sa.Integer(), nullable=False),
+        sa.Column("dispatch_id", sa.Text(), nullable=False),
+        sa.Column("task_group_id", sa.Integer(), nullable=False),
+        sa.Column("num_pending_parents", sa.Integer(), nullable=False),
+        sa.Column("sorted_tasks", sa.Text(), nullable=False),
+        sa.PrimaryKeyConstraint("id"),
+    )
+    op.create_table(
+        "workflowstate",
+        sa.Column("dispatch_id", sa.Text(), nullable=False),
+        sa.Column("num_unresolved_tasks", sa.Integer(), nullable=False),
+        sa.PrimaryKeyConstraint("dispatch_id"),
+    )
+
+    with op.batch_alter_table("taskgroupstate", schema=None) as batch_op:
+        batch_op.create_index("task_grp_idx", ["dispatch_id", "task_group_id"], unique=False)
+
+    with op.batch_alter_table("lattices", schema=None) as batch_op:
+        batch_op.create_unique_constraint("u_electron_id", ["electron_id"])
+
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table("lattices", schema=None) as batch_op:
+        batch_op.drop_constraint("u_electron_id", type_="unique")
+
+    with op.batch_alter_table("taskgroupstate", schema=None) as batch_op:
+        batch_op.drop_index("task_grp_idx")
+
+    op.drop_table("workflowstate")
+    op.drop_table("taskgroupstate")
+    # ### end Alembic commands ###
diff --git a/covalent_ui/api/v1/models/dispatch_model.py b/covalent_ui/api/v1/models/dispatch_model.py
index 27b7c91f..fba817bd 100644
--- a/covalent_ui/api/v1/models/dispatch_model.py
+++ b/covalent_ui/api/v1/models/dispatch_model.py
@@ -18,11 +18,10 @@
 
 
 from datetime import datetime
-from typing import List, Optional, Union
+from typing import Annotated, List, Optional, Union
 from uuid import UUID
 
 from pydantic import BaseModel, ConfigDict, Field
-from typing_extensions import Annotated
 
 from covalent_ui.api.v1.utils.models_helper import SortBy, SortDirection
 from covalent_ui.api.v1.utils.status import Status
diff --git a/covalent_ui/api/v1/models/logs_model.py b/covalent_ui/api/v1/models/logs_model.py
index a8a35610..394be075 100644
--- a/covalent_ui/api/v1/models/logs_model.py
+++ b/covalent_ui/api/v1/models/logs_model.py
@@ -17,10 +17,9 @@
 """Logs response model"""
 
 
-from typing import List, Optional, Union
+from typing import Annotated, List, Optional, Union
 
 from pydantic import BaseModel, ConfigDict, Field
-from typing_extensions import Annotated
 
 from covalent_ui.api.v1.utils.models_helper import CaseInsensitiveEnum, SortDirection
 
diff --git a/covalent_ui/api/v1/routes/end_points/summary_routes.py b/covalent_ui/api/v1/routes/end_points/summary_routes.py
index 68ae33df..8f554db8 100644
--- a/covalent_ui/api/v1/routes/end_points/summary_routes.py
+++ b/covalent_ui/api/v1/routes/end_points/summary_routes.py
@@ -16,12 +16,11 @@
 
 """Summary Routes"""
 
-from typing import Optional
+from typing import Annotated, Optional
 
 from fastapi import APIRouter, Query
 from pydantic import Field
 from sqlalchemy.orm import Session
-from typing_extensions import Annotated
 
 import covalent_ui.api.v1.database.config.db as db
 from covalent_ui.api.v1.data_layer.summary_dal import Summary
diff --git a/requirements.txt b/requirements.txt
index 15ecc909..a77471aa 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -12,13 +12,11 @@ natsort>=8.4.0
 networkx>=2.8.6
 psutil>=5.9.0
 pydantic>=2.1.1
-python-multipart>=0.0.6
 python-socketio>=5.7.1
 requests>=2.24.0
 rich>=12.0.0,<=13.3.5
 sqlalchemy>=1.4.37,<2.0.0
 sqlalchemy_utils>=0.38.3
 toml>=0.10.2
-typing-extensions>=4.8.0
 uvicorn[standard]
 watchdog>=2.2.1


Test patch:
diff --git a/tests/covalent_dispatcher_tests/_core/dispatcher_db_integration_test.py b/tests/covalent_dispatcher_tests/_core/dispatcher_db_integration_test.py
index 1544b753..31c06205 100644
--- a/tests/covalent_dispatcher_tests/_core/dispatcher_db_integration_test.py
+++ b/tests/covalent_dispatcher_tests/_core/dispatcher_db_integration_test.py
@@ -253,57 +253,67 @@ async def test_get_abstract_task_inputs(mocker, test_db):
 
 
 @pytest.mark.asyncio
-async def test_handle_completed_node(mocker, test_db):
+async def test_handle_completed_node(mocker):
     """Unit test for completed node handler"""
 
-    from covalent_dispatcher._core.dispatcher import _initialize_caches, _pending_parents
+    import tempfile
 
-    mocker.patch("covalent_dispatcher._db.write_result_to_db.workflow_db", test_db)
-    mocker.patch("covalent_dispatcher._db.upsert.workflow_db", test_db)
-    mocker.patch("covalent_dispatcher._dal.base.workflow_db", test_db)
+    from covalent_dispatcher._core.dispatcher import _initialize_caches, _task_group_cache
 
-    pending_parents = {}
-    sorted_task_groups = {}
-    sdkres = get_mock_result()
-    result_object = get_mock_srvresult(sdkres, test_db)
+    # Use a file-backed test DB so that it can be accessed by multiple threads
+    with tempfile.NamedTemporaryFile(mode="w+b") as tmp_file:
+        path = tmp_file.name
+        test_db = DataStore(
+            db_URL=f"sqlite+pysqlite:///{path}",
+            initialize_db=True,
+        )
 
-    async def get_node_successors(dispatch_id: str, node_id: int):
-        return result_object.lattice.transport_graph.get_successors(node_id, ["task_group_id"])
+        mocker.patch("covalent_dispatcher._db.write_result_to_db.workflow_db", test_db)
+        mocker.patch("covalent_dispatcher._db.upsert.workflow_db", test_db)
+        mocker.patch("covalent_dispatcher._dal.base.workflow_db", test_db)
 
-    async def electron_get(dispatch_id, node_id, keys):
-        return {keys[0]: node_id}
+        pending_parents = {}
+        sorted_task_groups = {}
+        sdkres = get_mock_result()
+        result_object = get_mock_srvresult(sdkres, test_db)
 
-    mocker.patch(
-        "covalent_dispatcher._core.dispatcher.tg_utils.get_node_successors",
-        get_node_successors,
-    )
+        async def get_node_successors(dispatch_id: str, node_id: int):
+            return result_object.lattice.transport_graph.get_successors(node_id, ["task_group_id"])
 
-    mocker.patch(
-        "covalent_dispatcher._core.data_manager.electron.get",
-        electron_get,
-    )
+        async def electron_get(dispatch_id, node_id, keys):
+            return {keys[0]: node_id}
+
+        mocker.patch(
+            "covalent_dispatcher._core.dispatcher.tg_utils.get_node_successors",
+            get_node_successors,
+        )
+
+        mocker.patch(
+            "covalent_dispatcher._core.data_manager.electron.get",
+            electron_get,
+        )
 
-    # tg edges are (1, 0), (0, 2)
-    pending_parents[0] = 1
-    pending_parents[1] = 0
-    pending_parents[2] = 1
-    sorted_task_groups[0] = [0]
-    sorted_task_groups[1] = [1]
-    sorted_task_groups[2] = [2]
+        # tg edges are (1, 0), (0, 2)
+        pending_parents[0] = 1
+        pending_parents[1] = 0
+        pending_parents[2] = 1
+        sorted_task_groups[0] = [0]
+        sorted_task_groups[1] = [1]
+        sorted_task_groups[2] = [2]
 
-    await _initialize_caches(result_object.dispatch_id, pending_parents, sorted_task_groups)
+        await _initialize_caches(result_object.dispatch_id, pending_parents, sorted_task_groups)
 
-    node_result = {"node_id": 1, "status": Result.COMPLETED}
-    assert await _pending_parents.get_pending(result_object.dispatch_id, 0) == 1
-    assert await _pending_parents.get_pending(result_object.dispatch_id, 1) == 0
-    assert await _pending_parents.get_pending(result_object.dispatch_id, 2) == 1
+        node_result = {"node_id": 1, "status": Result.COMPLETED}
+        assert await _task_group_cache.get_pending(result_object.dispatch_id, 0) == 1
+        assert await _task_group_cache.get_pending(result_object.dispatch_id, 1) == 0
+        assert await _task_group_cache.get_pending(result_object.dispatch_id, 2) == 1
 
-    next_nodes = await _handle_completed_node(result_object.dispatch_id, 1)
-    assert next_nodes == [0]
+        next_nodes = await _handle_completed_node(result_object.dispatch_id, 1)
+        assert next_nodes == [0]
 
-    assert await _pending_parents.get_pending(result_object.dispatch_id, 0) == 0
-    assert await _pending_parents.get_pending(result_object.dispatch_id, 1) == 0
-    assert await _pending_parents.get_pending(result_object.dispatch_id, 2) == 1
+        assert await _task_group_cache.get_pending(result_object.dispatch_id, 0) == 0
+        assert await _task_group_cache.get_pending(result_object.dispatch_id, 1) == 0
+        assert await _task_group_cache.get_pending(result_object.dispatch_id, 2) == 1
 
 
 @pytest.mark.asyncio
diff --git a/tests/covalent_dispatcher_tests/_core/dispatcher_test.py b/tests/covalent_dispatcher_tests/_core/dispatcher_test.py
index 6f373b6d..e055d163 100644
--- a/tests/covalent_dispatcher_tests/_core/dispatcher_test.py
+++ b/tests/covalent_dispatcher_tests/_core/dispatcher_test.py
@@ -215,18 +215,18 @@ async def test_handle_completed_node_update(mocker):
         "covalent_dispatcher._core.dispatcher._handle_completed_node", return_value=next_groups
     )
     mock_decrement = mocker.patch(
-        "covalent_dispatcher._core.dispatcher._unresolved_tasks.decrement"
+        "covalent_dispatcher._core.dispatcher._workflow_run_cache.decrement"
     )
 
     mock_increment = mocker.patch(
-        "covalent_dispatcher._core.dispatcher._unresolved_tasks.increment"
+        "covalent_dispatcher._core.dispatcher._workflow_run_cache.increment"
     )
 
     async def get_task_group(dispatch_id, gid):
         return [gid]
 
     mock_get_sorted_task_groups = mocker.patch(
-        "covalent_dispatcher._core.dispatcher._sorted_task_groups.get_task_group",
+        "covalent_dispatcher._core.dispatcher._task_group_cache.get_task_group",
         get_task_group,
     )
     mock_submit_task_group = mocker.patch(
@@ -251,7 +251,7 @@ async def test_handle_cancelled_node_update(mocker):
         "covalent_dispatcher._core.dispatcher._handle_cancelled_node",
     )
     mock_decrement = mocker.patch(
-        "covalent_dispatcher._core.dispatcher._unresolved_tasks.decrement"
+        "covalent_dispatcher._core.dispatcher._workflow_run_cache.decrement"
     )
 
     await _handle_node_status_update(dispatch_id, node_id, status, detail)
@@ -271,7 +271,7 @@ async def test_run_handle_failed_node_update(mocker):
         "covalent_dispatcher._core.dispatcher._handle_failed_node",
     )
     mock_decrement = mocker.patch(
-        "covalent_dispatcher._core.dispatcher._unresolved_tasks.decrement"
+        "covalent_dispatcher._core.dispatcher._workflow_run_cache.decrement"
     )
 
     await _handle_node_status_update(dispatch_id, node_id, status, detail)
@@ -293,7 +293,7 @@ async def test_run_handle_sublattice_node_update(mocker):
         "covalent_dispatcher._core.dispatcher.run_dispatch",
     )
     mock_decrement = mocker.patch(
-        "covalent_dispatcher._core.dispatcher._unresolved_tasks.decrement"
+        "covalent_dispatcher._core.dispatcher._workflow_run_cache.decrement"
     )
     await _handle_node_status_update(dispatch_id, node_id, status, detail)
     mock_run_dispatch.assert_called_with("sub_dispatch")
@@ -320,7 +320,7 @@ async def test_handle_event(mocker, unresolved_count):
     )
 
     mock_get_unresolved = mocker.patch(
-        "covalent_dispatcher._core.dispatcher._unresolved_tasks.get_unresolved",
+        "covalent_dispatcher._core.dispatcher._workflow_run_cache.get_unresolved",
         return_value=unresolved_count,
     )
 
@@ -361,7 +361,7 @@ async def test_handle_event_exception(mocker):
     )
 
     mock_get_unresolved = mocker.patch(
-        "covalent_dispatcher._core.dispatcher._unresolved_tasks.get_unresolved",
+        "covalent_dispatcher._core.dispatcher._workflow_run_cache.get_unresolved",
         return_value=2,
     )
 
@@ -407,7 +407,7 @@ async def test_handle_event_finalize_exception(mocker):
     )
 
     mock_get_unresolved = mocker.patch(
-        "covalent_dispatcher._core.dispatcher._unresolved_tasks.get_unresolved",
+        "covalent_dispatcher._core.dispatcher._workflow_run_cache.get_unresolved",
         return_value=0,
     )
 
@@ -497,7 +497,7 @@ async def test_submit_initial_tasks(mocker):
         "covalent_dispatcher._core.dispatcher._initialize_caches",
     )
 
-    mock_inc = mocker.patch("covalent_dispatcher._core.dispatcher._unresolved_tasks.increment")
+    mock_inc = mocker.patch("covalent_dispatcher._core.dispatcher._workflow_run_cache.increment")
     mock_submit_task_group = mocker.patch(
         "covalent_dispatcher._core.dispatcher._submit_task_group",
     )
@@ -662,21 +662,16 @@ async def test_clear_caches(mocker):
     mocker.patch("covalent_dispatcher._core.dispatcher.tg_utils.get_nodes_links")
     mocker.patch("networkx.readwrite.node_link_graph", return_value=g)
     mock_unresolved_remove = mocker.patch(
-        "covalent_dispatcher._core.dispatcher._unresolved_tasks.remove"
+        "covalent_dispatcher._core.dispatcher._workflow_run_cache.remove"
     )
-    mock_pending_remove = mocker.patch(
-        "covalent_dispatcher._core.dispatcher._pending_parents.remove"
-    )
-
-    mock_groups_remove = mocker.patch(
-        "covalent_dispatcher._core.dispatcher._sorted_task_groups.remove"
+    mock_task_groups_remove = mocker.patch(
+        "covalent_dispatcher._core.dispatcher._task_group_cache.remove"
     )
 
     await _clear_caches("dispatch")
 
     assert mock_unresolved_remove.await_count == 1
-    assert mock_pending_remove.await_count == 2
-    assert mock_groups_remove.await_count == 2
+    assert mock_task_groups_remove.await_count == 2
 
 
 @pytest.mark.asyncio