dbos 0.28.0a6__py3-none-any.whl → 0.28.0a8__py3-none-any.whl

dbos/_conductor/conductor.py

@@ -2,6 +2,7 @@ import socket
 import threading
 import time
 import traceback
+from importlib.metadata import version
 from typing import TYPE_CHECKING, Optional
 
 from websockets import ConnectionClosed, ConnectionClosedOK, InvalidStatus
@@ -21,6 +22,9 @@ from . import protocol as p
 if TYPE_CHECKING:
     from dbos import DBOS
 
+ws_version = version("websockets")
+use_keepalive = ws_version < "15.0"
+
 
 class ConductorWebsocket(threading.Thread):
 
@@ -35,14 +39,64 @@ class ConductorWebsocket(threading.Thread):
         self.url = (
             conductor_url.rstrip("/") + f"/websocket/{self.app_name}/{conductor_key}"
         )
+        # TODO: once we can upgrade to websockets>=15.0, we can always use the built-in keepalive
+        self.ping_interval = 20  # Time between pings in seconds
+        self.ping_timeout = 15  # Time to wait for a pong response in seconds
+        self.keepalive_thread: Optional[threading.Thread] = None
+        self.pong_event: Optional[threading.Event] = None
+        self.last_ping_time = -1.0
+
+        self.dbos.logger.debug(
+            f"Connecting to conductor at {self.url} using websockets version {ws_version}"
+        )
+
+    def keepalive(self) -> None:
+        self.dbos.logger.debug("Starting keepalive thread")
+        while not self.evt.is_set():
+            if self.websocket is None or self.websocket.close_code is not None:
+                time.sleep(1)
+                continue
+            try:
+                self.last_ping_time = time.time()
+                self.pong_event = self.websocket.ping()
+                self.dbos.logger.debug("> Sent ping to conductor")
+                pong_result = self.pong_event.wait(self.ping_timeout)
+                elapsed_time = time.time() - self.last_ping_time
+                if not pong_result:
+                    self.dbos.logger.warning(
+                        f"Failed to receive pong from conductor after {elapsed_time:.2f} seconds. Reconnecting."
+                    )
+                    self.websocket.close()
+                    continue
+                # Wait for the next ping interval
+                self.dbos.logger.debug(
+                    f"< Received pong from conductor after {elapsed_time:.2f} seconds"
+                )
+                wait_time = self.ping_interval - elapsed_time
+                self.evt.wait(max(0, wait_time))
+            except ConnectionClosed:
+                # The main loop will try to reconnect
+                self.dbos.logger.debug("Connection to conductor closed.")
+            except Exception as e:
+                self.dbos.logger.warning(f"Failed to send ping to conductor: {e}.")
+                self.websocket.close()
 
     def run(self) -> None:
         while not self.evt.is_set():
             try:
                 with connect(
-                    self.url, open_timeout=5, logger=self.dbos.logger
+                    self.url,
+                    open_timeout=5,
+                    close_timeout=5,
+                    logger=self.dbos.logger,
                 ) as websocket:
                     self.websocket = websocket
+                    if use_keepalive and self.keepalive_thread is None:
+                        self.keepalive_thread = threading.Thread(
+                            target=self.keepalive,
+                            daemon=True,
+                        )
+                        self.keepalive_thread.start()
                     while not self.evt.is_set():
                         message = websocket.recv()
                         if not isinstance(message, str):
@@ -283,8 +337,15 @@ class ConductorWebsocket(threading.Thread):
                             # Still need to send a response to the conductor
                             websocket.send(unknown_message.to_json())
             except ConnectionClosedOK:
-                self.dbos.logger.info("Conductor connection terminated")
-                break
+                if self.evt.is_set():
+                    self.dbos.logger.info("Conductor connection terminated")
+                    break
+                # Otherwise, we are trying to reconnect
+                self.dbos.logger.warning(
+                    "Connection to conductor lost. Reconnecting..."
+                )
+                time.sleep(1)
+                continue
             except ConnectionClosed as e:
                 self.dbos.logger.warning(
                     f"Connection to conductor lost. Reconnecting: {e}"
@@ -305,3 +366,9 @@ class ConductorWebsocket(threading.Thread):
                 )
                 time.sleep(1)
                 continue
+
+        # Wait for the keepalive thread to finish
+        if self.keepalive_thread is not None:
+            if self.pong_event is not None:
+                self.pong_event.set()
+            self.keepalive_thread.join()

dbos/_docker_pg_helper.py

@@ -83,6 +83,7 @@ def start_docker_postgres(pool_config: Dict[str, Any]) -> bool:
     logging.info("Starting a Postgres Docker container...")
     container_name = "dbos-db"
     pg_data = "/var/lib/postgresql/data"
+    image_name = "pgvector/pgvector:pg16"
 
     try:
         client = docker.from_env()
@@ -103,9 +104,15 @@ def start_docker_postgres(pool_config: Dict[str, Any]) -> bool:
             # Container doesn't exist, proceed with creation
             pass
 
+        # Pull the image if it doesn't exist
+        imgs = client.images.list(name=image_name)
+        if len(imgs) == 0:
+            logging.info(f"Pulling Docker image {image_name}...")
+            client.images.pull(image_name)
+
         # Create and start the container
         container = client.containers.run(
-            image="pgvector/pgvector:pg16",
+            image=image_name,
             name=container_name,
             detach=True,
             environment={

{dbos-0.28.0a6.dist-info → dbos-0.28.0a8.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: dbos
-Version: 0.28.0a6
+Version: 0.28.0a8
 Summary: Ultra-lightweight durable execution in Python
 Author-Email: "DBOS, Inc." <contact@dbos.dev>
 License: MIT
@@ -22,7 +22,7 @@ Requires-Dist: docker>=7.1.0
 Requires-Dist: cryptography>=43.0.3
 Requires-Dist: rich>=13.9.4
 Requires-Dist: pyjwt>=2.10.1
-Requires-Dist: websockets>=15.0
+Requires-Dist: websockets>=14.0
 Description-Content-Type: text/markdown
 
 

{dbos-0.28.0a6.dist-info → dbos-0.28.0a8.dist-info}/RECORD

@@ -1,14 +1,14 @@
-dbos-0.28.0a6.dist-info/METADATA,sha256=EtDjmsRZNSxN7U8Yv-ifyxQjawPTEKkL-wylJYH6e4c,13268
-dbos-0.28.0a6.dist-info/WHEEL,sha256=tSfRZzRHthuv7vxpI4aehrdN9scLjk-dCJkPLzkHxGg,90
-dbos-0.28.0a6.dist-info/entry_points.txt,sha256=_QOQ3tVfEjtjBlr1jS4sHqHya9lI2aIEIWkz8dqYp14,58
-dbos-0.28.0a6.dist-info/licenses/LICENSE,sha256=VGZit_a5-kdw9WT6fY5jxAWVwGQzgLFyPWrcVVUhVNU,1067
+dbos-0.28.0a8.dist-info/METADATA,sha256=ed_1fo-R4JZXcSiu5dRKS1AzETl8dB9Qwk-tX4mC1Ho,13268
+dbos-0.28.0a8.dist-info/WHEEL,sha256=tSfRZzRHthuv7vxpI4aehrdN9scLjk-dCJkPLzkHxGg,90
+dbos-0.28.0a8.dist-info/entry_points.txt,sha256=_QOQ3tVfEjtjBlr1jS4sHqHya9lI2aIEIWkz8dqYp14,58
+dbos-0.28.0a8.dist-info/licenses/LICENSE,sha256=VGZit_a5-kdw9WT6fY5jxAWVwGQzgLFyPWrcVVUhVNU,1067
 dbos/__init__.py,sha256=-FdBlOlr-f2tY__C23J4v22MoCAXqcDN_-zXsJXdoZ0,1005
 dbos/__main__.py,sha256=G7Exn-MhGrVJVDbgNlpzhfh8WMX_72t3_oJaFT9Lmt8,653
 dbos/_admin_server.py,sha256=CM02jyC9H21fM7Pjn1BhPxNwAOV7CXmMJd0SdaNq8dQ,9062
 dbos/_app_db.py,sha256=3j8_5-MlSDY0otLRszFE-GfenU6JC20fcfSL-drSNYk,11800
 dbos/_classproperty.py,sha256=f0X-_BySzn3yFDRKB2JpCbLYQ9tLwt1XftfshvY7CBs,626
 dbos/_client.py,sha256=mGrricoU6437QM4SWV-6Vm806AgVru8ygKgGgK1LZGA,13823
-dbos/_conductor/conductor.py,sha256=mWOHBzSvwHa4YTere8za90rDRspCvin4CWGQ1tHtZ00,16646
+dbos/_conductor/conductor.py,sha256=L8va6dGae0BlH43AHWAAW8_z04QUYPWPhjUK64sn3P4,19692
 dbos/_conductor/protocol.py,sha256=jwX8ZjmAIlXu1vw9R3b5PfHSNdwofeYOKj8rkfAFVg0,6630
 dbos/_context.py,sha256=Ly1CXF1nWxICQgIpDZSaONGlz1yERBs63gqmR-yqCzM,24476
 dbos/_core.py,sha256=UDpSgRA9m_YuViNXR9tVgNFLC-zxKZPxjlkj2a-Kj00,48317
@@ -16,7 +16,7 @@ dbos/_croniter.py,sha256=XHAyUyibs_59sJQfSNWkP7rqQY6_XrlfuuCxk4jYqek,47559
 dbos/_dbos.py,sha256=nne4oaIEC0tR2OV766y7mjTLjGaqKz7EnuY-T2CbTtc,48431
 dbos/_dbos_config.py,sha256=L0Z0OOB5FoPM9g-joZqXGeJnlxWQsEUtgPtgtg9Uf48,21732
 dbos/_debug.py,sha256=MNlQVZ6TscGCRQeEEL0VE8Uignvr6dPeDDDefS3xgIE,1823
-dbos/_docker_pg_helper.py,sha256=NmcgqmR5rQA_4igfeqh8ugNT2z3YmoOvuep_MEtxTiY,5854
+dbos/_docker_pg_helper.py,sha256=tLJXWqZ4S-ExcaPnxg_i6cVxL6ZxrYlZjaGsklY-s2I,6115
 dbos/_error.py,sha256=EN4eVBjMT3k7O7hfqJl6mIf4sxWPsiAOM086yhcGH_g,8012
 dbos/_event_loop.py,sha256=NmaLbEQFfEK36S_0KhVD39YdYrGce3qSKCTJ-5RqKQ0,2136
 dbos/_fastapi.py,sha256=m4SL3H9P-NBQ_ZrbFxAWMOqNyIi3HGEn2ODR7xAK038,3118
@@ -67,4 +67,4 @@ dbos/cli/cli.py,sha256=gXKELYAK9_CTejQ-WbNEIqnEYByJndXHDYSX4naFg8g,19106
 dbos/dbos-config.schema.json,sha256=8KcwJb_sQc4-6tQG2TLmjE_nratfrQa0qVLl9XPsvWE,6367
 dbos/py.typed,sha256=QfzXT1Ktfk3Rj84akygc7_42z0lRpCq0Ilh8OXI6Zas,44
 version/__init__.py,sha256=L4sNxecRuqdtSFdpUGX3TtBi9KL3k7YsZVIvv-fv9-A,1678
-dbos-0.28.0a6.dist-info/RECORD,,
+dbos-0.28.0a8.dist-info/RECORD,,