Fix job status update

CHANGELOG.md

@@ -42,6 +42,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Fix the bug in the update diff check that replaces uuids
 - Fix snowflake vector search issues.
 - Fix crontab drop method.
+- Fix job status update.
+- Fix a random bug caused by queue thread-safety issues when OSS used crontab.
 
 ## [0.4.0](https://github.com/superduper-io/superduper/compare/0.4.0...0.3.0])    (2024-Nov-02)
 

plugins/mongodb/superduper_mongodb/__init__.py

@@ -4,7 +4,7 @@
 from .query import MongoQuery
 from .vector_search import MongoAtlasVectorSearcher as VectorSearcher
 
-__version__ = "0.4.4"
+__version__ = "0.4.5"
 
 __all__ = [
     "ArtifactStore",

plugins/mongodb/superduper_mongodb/metadata.py

@@ -424,3 +424,8 @@ def set_component_status(self, uuid, status: Status):
         return self.component_collection.update_one(
             {'uuid': uuid}, {'$set': {'status': status}}
         )
+
+    def _get_component_status(self, uuid):
+        """Get status of component."""
+        data = self.component_collection.find_one({'uuid': uuid}, {'status': 1})
+        return data['status'] if data else None

plugins/sqlalchemy/superduper_sqlalchemy/__init__.py

@@ -1,5 +1,5 @@
 from .metadata import SQLAlchemyMetadata as MetaDataStore
 
-__version__ = "0.4.3"
+__version__ = "0.4.4"
 
 __all__ = ['MetaDataStore']

plugins/sqlalchemy/superduper_sqlalchemy/metadata.py

@@ -553,6 +553,19 @@ def set_component_status(self, uuid, status: Status):
             )
             session.execute(stmt)
 
+    def _get_component_status(self, uuid):
+        """Get status of component."""
+        with self.session_context() as session:
+            stmt = (
+                select(self.component_table)
+                .where(self.component_table.c.uuid == uuid)
+                .limit(1)
+            )
+            res = self.query_results(self.component_table, stmt, session)
+            if not res:
+                return None
+            return res[0]['status']
+
     def _show_cdcs(self, table):
         """Show all triggers in the database.
 
@@ -634,9 +647,7 @@ def get_job(self, job_id: str):
         """
         with self.session_context() as session:
             stmt = (
-                select(self.job_table)
-                .where(self.job_table.c.identifier == job_id)
-                .limit(1)
+                select(self.job_table).where(self.job_table.c.job_id == job_id).limit(1)
             )
             res = self.query_results(self.job_table, stmt, session)
             return res[0] if res else None

superduper/backends/base/metadata.py

@@ -6,7 +6,6 @@
 
 if t.TYPE_CHECKING:
     from superduper.backends.base.query import Select
-    from superduper.components.component import Status
 
 
 class NonExistentMetadataError(Exception):
@@ -158,7 +157,7 @@ def drop(self, force: bool = False):
         pass
 
     @abstractmethod
-    def set_component_status(self, uuid: str, status: 'Status'):
+    def set_component_status(self, uuid: str, status: str):
         """
         Set the status of a component.
 
@@ -167,6 +166,15 @@ def set_component_status(self, uuid: str, status: 'Status'):
         """
         pass
 
+    def get_component_status(self, uuid: str):
+        """
+        Get the status of a component.
+
+        :param uuid: ``Component.uuid``
+        """
+        status = self._get_component_status(uuid)
+        return status
+
     # ------------------ JOBS ------------------
 
     @abstractmethod

superduper/backends/base/queue.py

@@ -156,6 +156,7 @@ def _consume_event_type(event_type, ids, table, db: 'Datalayer'):
     context = str(uuid.uuid4())
     jobs: t.List[Job] = []
     job_lookup: t.DefaultDict = defaultdict(dict)
+    logging.info(f'Consuming {event_type} events on {table}')
 
     from superduper.components.cdc import build_streaming_graph
 
@@ -186,7 +187,7 @@ def _consume_event_type(event_type, ids, table, db: 'Datalayer'):
         for job in sub_jobs:
             job_lookup[component.uuid][job.method] = job.job_id
         jobs += sub_jobs
-        print(f'Streaming with {component.type_id}:{component.identifier}')
+        logging.info(f'Streaming with {component.type_id}:{component.identifier}')
 
     for job in jobs:
         job.execute(db)

superduper/backends/local/compute.py

@@ -63,8 +63,9 @@ def submit(self, job: Job) -> str:
         """
         args, kwargs = job.get_args_kwargs(self.futures[job.context])
 
+        assert job.job_id is not None
         component = self.db.load(uuid=job.uuid)
-        self.db.metadata.update_job(job.identifier, 'status', 'running')
+        self.db.metadata.update_job(job.job_id, 'status', 'running')
 
         try:
             logging.debug(
@@ -73,10 +74,10 @@ def submit(self, job: Job) -> str:
             method = getattr(component, job.method)
             output = method(*args, **kwargs)
         except Exception as e:
-            self.db.metadata.update_job(job.identifier, 'status', 'failed')
+            self.db.metadata.update_job(job.job_id, 'status', 'failed')
             raise e
 
-        self.db.metadata.update_job(job.identifier, 'status', 'success')
+        self.db.metadata.update_job(job.job_id, 'status', 'success')
         self.futures[job.context][job.job_id] = output
         assert job.job_id is not None
         return job.job_id

superduper/backends/local/queue.py

@@ -1,3 +1,4 @@
+import threading
 import typing as t
 
 from superduper import logging
@@ -27,6 +28,7 @@ def __init__(self, uri: t.Optional[str] = None):
         super().__init__(uri=uri)
         self.consumer = self.build_consumer()
         self._component_uuid_mapping: t.Dict = {}
+        self.lock = threading.Lock()
 
     def list(self):
         """List all components."""
@@ -46,7 +48,8 @@ def initialize(self):
             identifier = component_data['identifier']
             r = self.db.show(type_id=type_id, identifier=identifier, version=-1)
             if r.get('trigger'):
-                self.queue[type_id, identifier] = []
+                with self.lock:
+                    self.queue[type_id, identifier] = []
 
     def _put(self, component):
         msg = 'Table name "_apply" collides with Superduper namespace'
@@ -77,9 +80,10 @@ def publish(self, events: t.List[Event]):
 
         :param events: list of events
         """
-        for event in events:
-            self.queue[event.queue].append(event)
-        self.consumer.consume(db=self.db, queue=self.queue)
+        with self.lock:
+            for event in events:
+                self.queue[event.queue].append(event)
+            self.consumer.consume(db=self.db, queue=self.queue)
 
 
 class LocalQueueConsumer(BaseQueueConsumer):

superduper/components/component.py

@@ -199,12 +199,12 @@ def handle_update_or_same(self, other):
         other.uuid = self.uuid
         other.version = self.version
 
-    def set_db(self, value):
+    def set_db(self, value: "Datalayer"):
         """Set the db for the component.
 
         :param value: The value to set the db to.
         """
-        self.db = value
+        self.db: "Datalayer" = value
         for child in self.get_children():
             child.set_db(value)