dbos 0.22.0a1__py3-none-any.whl → 0.22.0a4__py3-none-any.whl

dbos/_core.py

@@ -63,6 +63,7 @@ from ._registrations import (
     get_or_create_func_info,
     get_temp_workflow_type,
     set_dbos_func_name,
+    set_func_info,
     set_temp_workflow_type,
 )
 from ._roles import check_required_roles
@@ -286,6 +287,7 @@ def execute_workflow_by_id(
         ctx.request = (
             _serialization.deserialize(request) if request is not None else None
         )
+        # If this function belongs to a configured class, add that class instance as its first argument
         if status["config_name"] is not None:
             config_name = status["config_name"]
             class_name = status["class_name"]
@@ -295,28 +297,9 @@ def execute_workflow_by_id(
                     workflow_id,
                     f"Cannot execute workflow because instance '{iname}' is not registered",
                 )
-
-            if startNew:
-                return start_workflow(
-                    dbos,
-                    wf_func,
-                    status["queue_name"],
-                    True,
-                    dbos._registry.instance_info_map[iname],
-                    *inputs["args"],
-                    **inputs["kwargs"],
-                )
-            else:
-                with SetWorkflowID(workflow_id):
-                    return start_workflow(
-                        dbos,
-                        wf_func,
-                        status["queue_name"],
-                        True,
-                        dbos._registry.instance_info_map[iname],
-                        *inputs["args"],
-                        **inputs["kwargs"],
-                    )
+            class_instance = dbos._registry.instance_info_map[iname]
+            inputs["args"] = (class_instance,) + inputs["args"]
+        # If this function is a class method, add that class object as its first argument
         elif status["class_name"] is not None:
             class_name = status["class_name"]
             if class_name not in dbos._registry.class_info_map:
@@ -324,30 +307,20 @@ def execute_workflow_by_id(
                     workflow_id,
                     f"Cannot execute workflow because class '{class_name}' is not registered",
                 )
+            class_object = dbos._registry.class_info_map[class_name]
+            inputs["args"] = (class_object,) + inputs["args"]
 
-            if startNew:
-                return start_workflow(
-                    dbos,
-                    wf_func,
-                    status["queue_name"],
-                    True,
-                    dbos._registry.class_info_map[class_name],
-                    *inputs["args"],
-                    **inputs["kwargs"],
-                )
-            else:
-                with SetWorkflowID(workflow_id):
-                    return start_workflow(
-                        dbos,
-                        wf_func,
-                        status["queue_name"],
-                        True,
-                        dbos._registry.class_info_map[class_name],
-                        *inputs["args"],
-                        **inputs["kwargs"],
-                    )
+        if startNew:
+            return start_workflow(
+                dbos,
+                wf_func,
+                status["queue_name"],
+                True,
+                *inputs["args"],
+                **inputs["kwargs"],
+            )
         else:
-            if startNew:
+            with SetWorkflowID(workflow_id):
                 return start_workflow(
                     dbos,
                     wf_func,
@@ -356,16 +329,6 @@ def execute_workflow_by_id(
                     *inputs["args"],
                     **inputs["kwargs"],
                 )
-            else:
-                with SetWorkflowID(workflow_id):
-                    return start_workflow(
-                        dbos,
-                        wf_func,
-                        status["queue_name"],
-                        True,
-                        *inputs["args"],
-                        **inputs["kwargs"],
-                    )
 
 
 @overload
@@ -398,9 +361,12 @@ def start_workflow(
     *args: P.args,
     **kwargs: P.kwargs,
 ) -> "WorkflowHandle[R]":
+    # If the function has a class, add the class object as its first argument
     fself: Optional[object] = None
     if hasattr(func, "__self__"):
         fself = func.__self__
+    if fself is not None:
+        args = (fself,) + args  # type: ignore
 
     fi = get_func_info(func)
     if fi is None:
@@ -436,17 +402,13 @@ def start_workflow(
     new_wf_ctx.id_assigned_for_next_workflow = new_wf_ctx.assign_workflow_id()
     new_wf_id = new_wf_ctx.id_assigned_for_next_workflow
 
-    gin_args: Tuple[Any, ...] = args
-    if fself is not None:
-        gin_args = (fself,)
-
     status = _init_workflow(
         dbos,
         new_wf_ctx,
         inputs=inputs,
         wf_name=get_dbos_func_name(func),
-        class_name=get_dbos_class_name(fi, func, gin_args),
-        config_name=get_config_name(fi, func, gin_args),
+        class_name=get_dbos_class_name(fi, func, args),
+        config_name=get_config_name(fi, func, args),
         temp_wf_type=get_temp_workflow_type(func),
         queue=queue_name,
         max_recovery_attempts=fi.max_recovery_attempts,
@@ -464,27 +426,15 @@ def start_workflow(
         )
         return WorkflowHandlePolling(new_wf_id, dbos)
 
-    if fself is not None:
-        future = dbos._executor.submit(
-            cast(Callable[..., R], _execute_workflow_wthread),
-            dbos,
-            status,
-            func,
-            new_wf_ctx,
-            fself,
-            *args,
-            **kwargs,
-        )
-    else:
-        future = dbos._executor.submit(
-            cast(Callable[..., R], _execute_workflow_wthread),
-            dbos,
-            status,
-            func,
-            new_wf_ctx,
-            *args,
-            **kwargs,
-        )
+    future = dbos._executor.submit(
+        cast(Callable[..., R], _execute_workflow_wthread),
+        dbos,
+        status,
+        func,
+        new_wf_ctx,
+        *args,
+        **kwargs,
+    )
     return WorkflowHandleFuture(new_wf_id, future, dbos)
 
 
@@ -516,6 +466,8 @@ def workflow_wrapper(
 
     @wraps(func)
     def wrapper(*args: Any, **kwargs: Any) -> R:
+        fi = get_func_info(func)
+        assert fi is not None
         if dbosreg.dbos is None:
             raise DBOSException(
                 f"Function {func.__name__} invoked before DBOS initialized"
@@ -726,6 +678,8 @@ def decorate_transaction(
         set_temp_workflow_type(temp_wf, "transaction")
         dbosreg.register_wf_function(get_dbos_func_name(temp_wf), wrapped_wf)
         wrapper.__orig_func = temp_wf  # type: ignore
+        set_func_info(wrapped_wf, get_or_create_func_info(func))
+        set_func_info(temp_wf, get_or_create_func_info(func))
 
         return cast(F, wrapper)
 
@@ -875,6 +829,8 @@ def decorate_step(
         set_temp_workflow_type(temp_wf, "step")
         dbosreg.register_wf_function(get_dbos_func_name(temp_wf), wrapped_wf)
         wrapper.__orig_func = temp_wf  # type: ignore
+        set_func_info(wrapped_wf, get_or_create_func_info(func))
+        set_func_info(temp_wf, get_or_create_func_info(func))
 
         return cast(Callable[P, R], wrapper)
 

dbos/_dbos.py

@@ -997,6 +997,10 @@ def _dbos_exit_hook() -> None:
         )
         return
     if not _dbos_global_instance._launched:
+        if _dbos_global_instance.fastapi is not None:
+            # FastAPI lifespan middleware will call launch/destroy, so we can ignore this.
+            # This is likely to happen during fastapi dev runs, where the reloader loads the module multiple times.
+            return
         print("DBOS exiting; DBOS exists but launch() was not called")
         dbos_logger.warning("DBOS exiting; DBOS exists but launch() was not called")
         return

dbos/_queue.py

@@ -76,7 +76,9 @@ def queue_thread(stop_event: threading.Event, dbos: "DBOS") -> None:
                     execute_workflow_by_id(dbos, id)
             except OperationalError as e:
                 # Ignore serialization error
-                if not isinstance(e.orig, errors.SerializationFailure):
+                if not isinstance(
+                    e.orig, (errors.SerializationFailure, errors.LockNotAvailable)
+                ):
                     dbos.logger.warning(
                         f"Exception encountered in queue thread: {traceback.format_exc()}"
                     )

dbos/_registrations.py

@@ -1,4 +1,5 @@
 import inspect
+from dataclasses import dataclass
 from enum import Enum
 from types import FunctionType
 from typing import Any, Callable, List, Literal, Optional, Tuple, Type, cast
@@ -31,9 +32,9 @@ def set_temp_workflow_type(f: Any, name: TempWorkflowType) -> None:
     setattr(f, "dbos_temp_workflow_type", name)
 
 
+@dataclass
 class DBOSClassInfo:
-    def __init__(self) -> None:
-        self.def_required_roles: Optional[List[str]] = None
+    def_required_roles: Optional[List[str]] = None
 
 
 class DBOSFuncType(Enum):
@@ -44,12 +45,12 @@ class DBOSFuncType(Enum):
     Instance = 4
 
 
+@dataclass
 class DBOSFuncInfo:
-    def __init__(self) -> None:
-        self.class_info: Optional[DBOSClassInfo] = None
-        self.func_type: DBOSFuncType = DBOSFuncType.Unknown
-        self.required_roles: Optional[List[str]] = None
-        self.max_recovery_attempts = DEFAULT_MAX_RECOVERY_ATTEMPTS
+    class_info: Optional[DBOSClassInfo] = None
+    func_type: DBOSFuncType = DBOSFuncType.Unknown
+    required_roles: Optional[List[str]] = None
+    max_recovery_attempts: int = DEFAULT_MAX_RECOVERY_ATTEMPTS
 
 
 def get_or_create_class_info(cls: Type[Any]) -> DBOSClassInfo:
@@ -110,6 +111,10 @@ def get_or_create_func_info(func: Callable[..., Any]) -> DBOSFuncInfo:
     return fi
 
 
+def set_func_info(func: Callable[..., Any], fi: DBOSFuncInfo) -> None:
+    setattr(func, "dbos_func_decorator_info", fi)
+
+
 def get_class_info(cls: Type[Any]) -> Optional[DBOSClassInfo]:
     if hasattr(cls, "dbos_class_decorator_info"):
         ci: DBOSClassInfo = getattr(cls, "dbos_class_decorator_info")

dbos/_sys_db.py

@@ -189,6 +189,10 @@ class SystemDatabase:
             host=config["database"]["hostname"],
             port=config["database"]["port"],
             database="postgres",
+            # fills the "application_name" column in pg_stat_activity
+            query={
+                "application_name": f"dbos_transact_{os.environ.get('DBOS__VMID', 'local')}_{os.environ.get('DBOS__APPVERSION', '')}"
+            },
         )
         engine = sa.create_engine(postgres_db_url)
         with engine.connect() as conn:
@@ -207,6 +211,10 @@ class SystemDatabase:
             host=config["database"]["hostname"],
             port=config["database"]["port"],
             database=sysdb_name,
+            # fills the "application_name" column in pg_stat_activity
+            query={
+                "application_name": f"dbos_transact_{os.environ.get('DBOS__VMID', 'local')}_{os.environ.get('DBOS__APPVERSION', '')}"
+            },
         )
 
         # Create a connection pool for the system database
@@ -1307,6 +1315,55 @@ class SystemDatabase:
             # Dequeue functions eligible for this worker and ordered by the time at which they were enqueued.
             # If there is a global or local concurrency limit N, select only the N oldest enqueued
             # functions, else select all of them.
+
+            # First lets figure out how many tasks the worker can dequeue
+            running_tasks_query = (
+                sa.select(
+                    SystemSchema.workflow_queue.c.executor_id,
+                    sa.func.count().label("task_count"),
+                )
+                .where(SystemSchema.workflow_queue.c.queue_name == queue.name)
+                .where(
+                    SystemSchema.workflow_queue.c.executor_id.isnot(
+                        None
+                    )  # Task is dequeued
+                )
+                .where(
+                    SystemSchema.workflow_queue.c.completed_at_epoch_ms.is_(
+                        None
+                    )  # Task is not completed
+                )
+                .group_by(SystemSchema.workflow_queue.c.executor_id)
+            )
+            running_tasks_result = c.execute(running_tasks_query).fetchall()
+            running_tasks_result_dict = {row[0]: row[1] for row in running_tasks_result}
+            running_tasks_for_this_worker = running_tasks_result_dict.get(
+                executor_id, 0
+            )  # Get count for current executor
+
+            max_tasks = float("inf")
+            if queue.worker_concurrency is not None:
+                # Worker local concurrency limit should always be >= running_tasks_for_this_worker
+                # This should never happen but a check + warning doesn't hurt
+                if running_tasks_for_this_worker > queue.worker_concurrency:
+                    dbos_logger.warning(
+                        f"Number of tasks on this worker ({running_tasks_for_this_worker}) exceeds the worker concurrency limit ({queue.worker_concurrency})"
+                    )
+                max_tasks = max(
+                    0, queue.worker_concurrency - running_tasks_for_this_worker
+                )
+            if queue.concurrency is not None:
+                total_running_tasks = sum(running_tasks_result_dict.values())
+                # Queue global concurrency limit should always be >= running_tasks_count
+                # This should never happen but a check + warning doesn't hurt
+                if total_running_tasks > queue.concurrency:
+                    dbos_logger.warning(
+                        f"Total running tasks ({total_running_tasks}) exceeds the global concurrency limit ({queue.concurrency})"
+                    )
+                available_tasks = max(0, queue.concurrency - total_running_tasks)
+                max_tasks = min(max_tasks, available_tasks)
+
+            # Lookup tasks
             query = (
                 sa.select(
                     SystemSchema.workflow_queue.c.workflow_uuid,
@@ -1315,29 +1372,25 @@ class SystemDatabase:
                 )
                 .where(SystemSchema.workflow_queue.c.queue_name == queue.name)
                 .where(SystemSchema.workflow_queue.c.completed_at_epoch_ms == None)
-                .where(
-                    # Only select functions that have not been started yet or have been started by this worker
-                    or_(
-                        SystemSchema.workflow_queue.c.executor_id == None,
-                        SystemSchema.workflow_queue.c.executor_id == executor_id,
-                    )
-                )
+                .where(SystemSchema.workflow_queue.c.executor_id == None)
                 .order_by(SystemSchema.workflow_queue.c.created_at_epoch_ms.asc())
+                .with_for_update(nowait=True)  # Error out early
             )
-            # Set a dequeue limit if necessary
-            if queue.worker_concurrency is not None:
-                query = query.limit(queue.worker_concurrency)
-            elif queue.concurrency is not None:
-                query = query.limit(queue.concurrency)
+            # Apply limit only if max_tasks is finite
+            if max_tasks != float("inf"):
+                query = query.limit(int(max_tasks))
 
             rows = c.execute(query).fetchall()
 
-            # Now, get the workflow IDs of functions that have not yet been started
-            dequeued_ids: List[str] = [row[0] for row in rows if row[1] is None]
+            # Get the workflow IDs
+            dequeued_ids: List[str] = [row[0] for row in rows]
+            if len(dequeued_ids) > 0:
+                dbos_logger.debug(
+                    f"[{queue.name}] dequeueing {len(dequeued_ids)} task(s)"
+                )
             ret_ids: list[str] = []
-            dbos_logger.debug(f"[{queue.name}] dequeueing {len(dequeued_ids)} task(s)")
-            for id in dequeued_ids:
 
+            for id in dequeued_ids:
                 # If we have a limiter, stop starting functions when the number
                 # of functions started this period exceeds the limit.
                 if queue.limiter is not None:

{dbos-0.22.0a1.dist-info → dbos-0.22.0a4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: dbos
-Version: 0.22.0a1
+Version: 0.22.0a4
 Summary: Ultra-lightweight durable execution in Python
 Author-Email: "DBOS, Inc." <contact@dbos.dev>
 License: MIT

{dbos-0.22.0a1.dist-info → dbos-0.22.0a4.dist-info}/RECORD

@@ -1,7 +1,7 @@
-dbos-0.22.0a1.dist-info/METADATA,sha256=DBJiSN2oOdnkb8P0DBogQjH4Wjh79UQ9gI6bDVdbbP8,5309
-dbos-0.22.0a1.dist-info/WHEEL,sha256=thaaA2w1JzcGC48WYufAs8nrYZjJm8LqNfnXFOFyCC4,90
-dbos-0.22.0a1.dist-info/entry_points.txt,sha256=_QOQ3tVfEjtjBlr1jS4sHqHya9lI2aIEIWkz8dqYp14,58
-dbos-0.22.0a1.dist-info/licenses/LICENSE,sha256=VGZit_a5-kdw9WT6fY5jxAWVwGQzgLFyPWrcVVUhVNU,1067
+dbos-0.22.0a4.dist-info/METADATA,sha256=3hz0YdIdTpfM3YyJqBeoKy5ADFI91UA3j-MW4lNhBfs,5309
+dbos-0.22.0a4.dist-info/WHEEL,sha256=thaaA2w1JzcGC48WYufAs8nrYZjJm8LqNfnXFOFyCC4,90
+dbos-0.22.0a4.dist-info/entry_points.txt,sha256=_QOQ3tVfEjtjBlr1jS4sHqHya9lI2aIEIWkz8dqYp14,58
+dbos-0.22.0a4.dist-info/licenses/LICENSE,sha256=VGZit_a5-kdw9WT6fY5jxAWVwGQzgLFyPWrcVVUhVNU,1067
 dbos/__init__.py,sha256=CxRHBHEthPL4PZoLbZhp3rdm44-KkRTT2-7DkK9d4QQ,724
 dbos/_admin_server.py,sha256=PJgneZG9-64TapZrPeJtt73puAswRImCE5uce2k2PKU,4750
 dbos/_app_db.py,sha256=_tv2vmPjjiaikwgxH3mqxgJ4nUUcG2-0uMXKWCqVu1c,5509
@@ -10,10 +10,10 @@ dbos/_cloudutils/authentication.py,sha256=V0fCWQN9stCkhbuuxgPTGpvuQcDqfU3KAxPAh0
 dbos/_cloudutils/cloudutils.py,sha256=5e3CW1deSW-dI5G3QN0XbiVsBhyqT8wu7fuV2f8wtGU,7688
 dbos/_cloudutils/databases.py,sha256=x4187Djsyoa-QaG3Kog8JT2_GERsnqa93LIVanmVUmg,8393
 dbos/_context.py,sha256=FHB_fpE4fQt4fIJvAmMMsbY4xHwH77gsW01cFsRZjsE,17779
-dbos/_core.py,sha256=nGiXyYgV8H5TRRZG0e8HCd5IZimufYQLmKNr7nBbwbo,36564
+dbos/_core.py,sha256=PvII-lcUgii0WqaDb9-56gfQ_69VoWthKAkKexRebuk,35461
 dbos/_croniter.py,sha256=hbhgfsHBqclUS8VeLnJ9PSE9Z54z6mi4nnrr1aUXn0k,47561
 dbos/_db_wizard.py,sha256=xgKLna0_6Xi50F3o8msRosXba8NScHlpJR5ICVCkHDQ,7534
-dbos/_dbos.py,sha256=wAjdlUgDSIC_Q8D_GZYDoiKaxjtr6KNHeq6DDuUh9do,36340
+dbos/_dbos.py,sha256=klpuF1Wj05_qsWiKnLSWfei-16-K-g4BQ5R0MMfTuNE,36620
 dbos/_dbos_config.py,sha256=DfiqVVxNqnafkocSzLqBp1Ig5vCviDTDK_GO3zTtQqI,8298
 dbos/_error.py,sha256=vtaSsG0QW6cRlwfZ4zzZWy_IHCZlomwSlrDyGWuyn8c,4337
 dbos/_fastapi.py,sha256=ke03vqsSYDnO6XeOtOVFXj0-f-v1MGsOxa9McaROvNc,3616
@@ -31,9 +31,9 @@ dbos/_migrations/versions/d76646551a6b_job_queue_limiter.py,sha256=8PyFi8rd6CN-m
 dbos/_migrations/versions/d76646551a6c_workflow_queue.py,sha256=G942nophZ2uC2vc4hGBC02Ptng1715roTjY3xiyzZU4,729
 dbos/_migrations/versions/eab0cc1d9a14_job_queue.py,sha256=uvhFOtqbBreCePhAxZfIT0qCAI7BiZTou9wt6QnbY7c,1412
 dbos/_outcome.py,sha256=FDMgWVjZ06vm9xO-38H17mTqBImUYQxgKs_bDCSIAhE,6648
-dbos/_queue.py,sha256=o_aczwualJTMoXb0XXL-Y5QH77OEukWzuerogbWi2ho,2779
+dbos/_queue.py,sha256=eZiapBcyn70-viW0y9fo7u09V6_VF5ACNGJxD-U_dNM,2844
 dbos/_recovery.py,sha256=rek9rm2CaENbbl_vu3To-BdXop7tMEyGvtoNiJLVxjQ,2772
-dbos/_registrations.py,sha256=mei6q6_3R5uei8i_Wo_TqGZs85s10shOekDX41sFYD0,6642
+dbos/_registrations.py,sha256=_zy6k944Ll8QwqU12Kr3OP23ukVtm8axPNN1TS_kJRc,6717
 dbos/_request.py,sha256=cX1B3Atlh160phgS35gF1VEEV4pD126c9F3BDgBmxZU,929
 dbos/_roles.py,sha256=iOsgmIAf1XVzxs3gYWdGRe1B880YfOw5fpU7Jwx8_A8,2271
 dbos/_scheduler.py,sha256=0I3e8Y-OIBG3wiUCIskShd-Sk_eUFCFyRB5u4L7IHXI,1940
@@ -41,7 +41,7 @@ dbos/_schemas/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 dbos/_schemas/application_database.py,sha256=KeyoPrF7hy_ODXV7QNike_VFSD74QBRfQ76D7QyE9HI,966
 dbos/_schemas/system_database.py,sha256=rwp4EvCSaXcUoMaRczZCvETCxGp72k3-hvLyGUDkih0,5163
 dbos/_serialization.py,sha256=YCYv0qKAwAZ1djZisBC7khvKqG-5OcIv9t9EC5PFIog,1743
-dbos/_sys_db.py,sha256=B8I-MUOHcXalrtOewomKuxbyRvYdItxQhHsfJ1IjunA,59719
+dbos/_sys_db.py,sha256=pH4vfPAAFVqu0KAfUFSUDKp5T-4OrwaLWJylMrzDFL4,62421
 dbos/_templates/dbos-db-starter/README.md,sha256=GhxhBj42wjTt1fWEtwNriHbJuKb66Vzu89G4pxNHw2g,930
 dbos/_templates/dbos-db-starter/__package/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 dbos/_templates/dbos-db-starter/__package/main.py,sha256=eI0SS9Nwj-fldtiuSzIlIG6dC91GXXwdRsoHxv6S_WI,2719
@@ -60,4 +60,4 @@ dbos/cli/cli.py,sha256=_tXw2IQrWW7fV_h51f_R99vEBSi6aMLz-vCOxKaENiQ,14155
 dbos/dbos-config.schema.json,sha256=X5TpXNcARGceX0zQs0fVgtZW_Xj9uBbY5afPt9Rz9yk,5741
 dbos/py.typed,sha256=QfzXT1Ktfk3Rj84akygc7_42z0lRpCq0Ilh8OXI6Zas,44
 version/__init__.py,sha256=L4sNxecRuqdtSFdpUGX3TtBi9KL3k7YsZVIvv-fv9-A,1678
-dbos-0.22.0a1.dist-info/RECORD,,
+dbos-0.22.0a4.dist-info/RECORD,,