Merge pull request #870 from MolSSI/task_queue_idx

Small task queue fix & migration code location

qcfractal/qcfractal/alembic/versions/2022-10-12-981f69781d65_remove_keywords_index_and_add_hashing.py

@@ -6,24 +6,17 @@
 
 """
 
-import json
-from hashlib import sha256
-
 import sqlalchemy as sa
 from alembic import op
 
-from qcportal.serialization import _JSONEncoder
-
 # revision identifiers, used by Alembic.
 revision = "981f69781d65"
 down_revision = "148fef89c2ec"
 branch_labels = None
 depends_on = None
 
 
-def hash_dict(d):
-    j = json.dumps(d, ensure_ascii=True, sort_keys=True, cls=_JSONEncoder).encode("utf-8")
-    return sha256(j).hexdigest()
+from migration_helpers.hashing import hash_dict_1
 
 
 def create_hashes(table):
@@ -32,7 +25,7 @@ def create_hashes(table):
     all_kw = res.fetchall()
 
     for spec_id, kw in all_kw:
-        h = hash_dict(kw)
+        h = hash_dict_1(kw)
         op.execute(sa.text(f"""UPDATE {table} SET keywords_hash = '{h}' WHERE id = {spec_id};"""))
 
 

qcfractal/qcfractal/alembic/versions/2024-12-18-8263992eb6c8_fix_task_queue_sort.py

@@ -0,0 +1,30 @@
+"""Fix task queue sort
+
+Revision ID: 8263992eb6c8
+Revises: f8a7c273f18a
+Create Date: 2024-12-18 11:34:57.186942
+
+"""
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "8263992eb6c8"
+down_revision = "f8a7c273f18a"
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_index("ix_task_queue_sort", table_name="task_queue")
+    op.execute(
+        "CREATE INDEX ix_task_queue_sort ON task_queue (priority DESC, sort_date, id, tag) WHERE available = True;"
+    )
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    pass
+    # ### end Alembic commands ###

qcfractal/qcfractal/alembic/versions/migration_helpers/hashing.py

@@ -0,0 +1,48 @@
+import base64
+import json
+from hashlib import sha256
+from typing import Any
+
+import numpy as np
+
+try:
+    from pydantic.v1 import BaseModel
+    from pydantic.v1.json import pydantic_encoder
+except ImportError:
+    from pydantic import BaseModel
+    from pydantic.json import pydantic_encoder
+
+
+class _JSONEncoder_1(json.JSONEncoder):
+    def default(self, obj: Any) -> Any:
+        # JSON does not handle byte arrays
+        # So convert to base64
+        if isinstance(obj, bytes):
+            return {"_bytes_base64_": base64.b64encode(obj).decode("ascii")}
+
+        # Now do anything with pydantic, excluding unset fields
+        # Also always use aliases when serializing
+        if isinstance(obj, BaseModel):
+            return obj.dict(exclude_unset=True, by_alias=True)
+
+        # Let pydantic handle other things
+        try:
+            return pydantic_encoder(obj)
+        except TypeError:
+            pass
+
+        # Flatten numpy arrays
+        # This is mostly for Molecule class
+        # TODO - remove once all data in the database in converted
+        if isinstance(obj, np.ndarray):
+            if obj.shape:
+                return obj.ravel().tolist()
+            else:
+                return obj.tolist()
+
+        return json.JSONEncoder.default(self, obj)
+
+
+def hash_dict_1(d):
+    j = json.dumps(d, ensure_ascii=True, sort_keys=True, cls=_JSONEncoder_1).encode("utf-8")
+    return sha256(j).hexdigest()

qcfractal/qcfractal/components/tasks/db_models.py

@@ -48,7 +48,9 @@ class TaskQueueORM(BaseORM):
     __table_args__ = (
         Index("ix_task_queue_tag", "tag"),
         Index("ix_task_queue_required_programs", "required_programs", postgresql_using="gin"),
-        Index("ix_task_queue_sort", priority.desc(), sort_date.asc(), id.asc(), tag),
+        Index(
+            "ix_task_queue_sort", priority.desc(), sort_date.asc(), id.asc(), tag, postgresql_where=(available == True)
+        ),
         UniqueConstraint("record_id", name="ux_task_queue_record_id"),
         # WARNING - these are not autodetected by alembic
         CheckConstraint(