dbos 0.26.0a13__tar.gz → 0.26.0a15__tar.gz

{dbos-0.26.0a13 → dbos-0.26.0a15}/PKG-INFO

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

{dbos-0.26.0a13 → dbos-0.26.0a15}/dbos/_admin_server.py

@@ -7,6 +7,7 @@ from functools import partial
 from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
 from typing import TYPE_CHECKING, Any, List, TypedDict
 
+from ._error import DBOSException
 from ._logger import dbos_logger
 from ._recovery import recover_pending_workflows
 from ._utils import GlobalParams
@@ -22,6 +23,7 @@ _workflow_queues_metadata_path = "/dbos-workflow-queues-metadata"
 # /workflows/:workflow_id/resume
 # /workflows/:workflow_id/restart
 # /workflows/:workflow_id/steps
+# /workflows/:workflow_id/fork
 
 
 class AdminServer:
@@ -123,6 +125,9 @@ class AdminRequestHandler(BaseHTTPRequestHandler):
             restart_match = re.match(
                 r"^/workflows/(?P<workflow_id>[^/]+)/restart$", self.path
             )
+            fork_match = re.match(
+                r"^/workflows/(?P<workflow_id>[^/]+)/fork$", self.path
+            )
             resume_match = re.match(
                 r"^/workflows/(?P<workflow_id>[^/]+)/resume$", self.path
             )
@@ -130,7 +135,23 @@ class AdminRequestHandler(BaseHTTPRequestHandler):
                 r"^/workflows/(?P<workflow_id>[^/]+)/cancel$", self.path
             )
 
-            if restart_match:
+            if fork_match:
+                workflow_id = fork_match.group("workflow_id")
+                try:
+                    data = json.loads(post_data.decode("utf-8"))
+                    start_step: int = data.get("start_step", 1)
+                    self._handle_fork(workflow_id, start_step)
+                except (json.JSONDecodeError, AttributeError) as e:
+                    self.send_response(500)
+                    self.send_header("Content-Type", "application/json")
+                    self.end_headers()
+                    self.wfile.write(
+                        json.dumps({"error": f"Invalid JSON input: {str(e)}"}).encode(
+                            "utf-8"
+                        )
+                    )
+                    return
+            elif restart_match:
                 workflow_id = restart_match.group("workflow_id")
                 self._handle_restart(workflow_id)
             elif resume_match:
@@ -147,10 +168,33 @@ class AdminRequestHandler(BaseHTTPRequestHandler):
         return  # Disable admin server request logging
 
     def _handle_restart(self, workflow_id: str) -> None:
-        self.dbos.restart_workflow(workflow_id)
-        print("Restarting workflow", workflow_id)
-        self.send_response(204)
-        self._end_headers()
+        try:
+            print(f"Restarting workflow {workflow_id}")
+            self.dbos.restart_workflow(workflow_id)
+            self.send_response(204)
+            self._end_headers()
+        except DBOSException as e:
+            print(f"Error restarting workflow: {e}")
+            self.send_response(500)
+            response_body = json.dumps({"error": str(e)}).encode("utf-8")
+            self.send_header("Content-Type", "application/json")
+            self.send_header("Content-Length", str(len(response_body)))
+            self.end_headers()
+            self.wfile.write(response_body)
+
+    def _handle_fork(self, workflow_id: str, start_step: int) -> None:
+        try:
+            self.dbos.fork_workflow(workflow_id, start_step)
+            self.send_response(204)
+            self._end_headers()
+        except DBOSException as e:
+            print(f"Error forking workflow: {e}")
+            self.send_response(500)
+            response_body = json.dumps({"error": str(e)}).encode("utf-8")
+            self.send_header("Content-Type", "application/json")
+            self.send_header("Content-Length", str(len(response_body)))
+            self.end_headers()
+            self.wfile.write(response_body)
 
     def _handle_resume(self, workflow_id: str) -> None:
         print("Resuming workflow", workflow_id)

{dbos-0.26.0a13 → dbos-0.26.0a15}/dbos/_app_db.py

@@ -228,3 +228,59 @@ class ApplicationDatabase:
             )
             for row in rows
         ]
+
+    def get_max_function_id(self, workflow_uuid: str) -> Optional[int]:
+        with self.engine.begin() as conn:
+            max_function_id_row = conn.execute(
+                sa.select(
+                    sa.func.max(ApplicationSchema.transaction_outputs.c.function_id)
+                ).where(
+                    ApplicationSchema.transaction_outputs.c.workflow_uuid
+                    == workflow_uuid
+                )
+            ).fetchone()
+
+            max_function_id = max_function_id_row[0] if max_function_id_row else None
+
+            return max_function_id
+
+    def clone_workflow_transactions(
+        self, src_workflow_id: str, forked_workflow_id: str, start_step: int
+    ) -> None:
+        """
+        Copies all steps from dbos.transctions_outputs where function_id < input function_id
+        into a new workflow_uuid. Returns the new workflow_uuid.
+        """
+
+        with self.engine.begin() as conn:
+
+            insert_stmt = sa.insert(ApplicationSchema.transaction_outputs).from_select(
+                [
+                    "workflow_uuid",
+                    "function_id",
+                    "output",
+                    "error",
+                    "txn_id",
+                    "txn_snapshot",
+                    "executor_id",
+                    "function_name",
+                ],
+                sa.select(
+                    sa.literal(forked_workflow_id).label("workflow_uuid"),
+                    ApplicationSchema.transaction_outputs.c.function_id,
+                    ApplicationSchema.transaction_outputs.c.output,
+                    ApplicationSchema.transaction_outputs.c.error,
+                    ApplicationSchema.transaction_outputs.c.txn_id,
+                    ApplicationSchema.transaction_outputs.c.txn_snapshot,
+                    ApplicationSchema.transaction_outputs.c.executor_id,
+                    ApplicationSchema.transaction_outputs.c.function_name,
+                ).where(
+                    (
+                        ApplicationSchema.transaction_outputs.c.workflow_uuid
+                        == src_workflow_id
+                    )
+                    & (ApplicationSchema.transaction_outputs.c.function_id < start_step)
+                ),
+            )
+
+            conn.execute(insert_stmt)

{dbos-0.26.0a13 → dbos-0.26.0a15}/dbos/_core.py

@@ -365,7 +365,9 @@ def _execute_workflow_wthread(
                 if isinstance(result, Immediate):
                     return cast(Immediate[R], result)()
                 else:
-                    return asyncio.run(cast(Pending[R], result)())
+                    return dbos._background_event_loop.submit_coroutine(
+                        cast(Pending[R], result)()
+                    )
             except Exception:
                 dbos.logger.error(
                     f"Exception encountered in asynchronous workflow: {traceback.format_exc()}"

{dbos-0.26.0a13 → dbos-0.26.0a15}/dbos/_dbos.py

@@ -4,7 +4,6 @@ import asyncio
 import atexit
 import hashlib
 import inspect
-import json
 import os
 import sys
 import threading
@@ -31,7 +30,6 @@ from typing import (
 
 from opentelemetry.trace import Span
 
-from dbos import _serialization
 from dbos._conductor.conductor import ConductorWebsocket
 from dbos._utils import INTERNAL_QUEUE_NAME, GlobalParams
 from dbos._workflow_commands import (
@@ -112,6 +110,7 @@ from ._error import (
     DBOSException,
     DBOSNonExistentWorkflowError,
 )
+from ._event_loop import BackgroundEventLoop
 from ._logger import add_otlp_to_all_loggers, config_logger, dbos_logger, init_logger
 from ._sys_db import SystemDatabase
 from ._workflow_commands import WorkflowStatus, get_workflow
@@ -341,6 +340,7 @@ class DBOS:
         self.conductor_url: Optional[str] = conductor_url
         self.conductor_key: Optional[str] = conductor_key
         self.conductor_websocket: Optional[ConductorWebsocket] = None
+        self._background_event_loop: BackgroundEventLoop = BackgroundEventLoop()
 
         init_logger()
 
@@ -451,6 +451,7 @@ class DBOS:
             dbos_logger.info(f"Executor ID: {GlobalParams.executor_id}")
             dbos_logger.info(f"Application version: {GlobalParams.app_version}")
             self._executor_field = ThreadPoolExecutor(max_workers=64)
+            self._background_event_loop.start()
             self._sys_db_field = SystemDatabase(
                 self.config["database"], debug_mode=debug_mode
             )
@@ -568,6 +569,7 @@ class DBOS:
         self._initialized = False
         for event in self.stop_events:
             event.set()
+        self._background_event_loop.stop()
         if self._sys_db_field is not None:
             self._sys_db_field.destroy()
             self._sys_db_field = None
@@ -958,14 +960,45 @@ class DBOS:
     def restart_workflow(cls, workflow_id: str) -> WorkflowHandle[Any]:
         """Restart a workflow with a new workflow ID"""
 
+        return cls.fork_workflow(workflow_id, 1)
+
+    @classmethod
+    def fork_workflow(
+        cls, workflow_id: str, start_step: int = 1
+    ) -> WorkflowHandle[Any]:
+        """Restart a workflow with a new workflow ID"""
+
+        def get_max_function_id(workflow_uuid: str) -> int:
+            max_transactions = (
+                _get_dbos_instance()._app_db.get_max_function_id(workflow_uuid) or 0
+            )
+            max_operations = (
+                _get_dbos_instance()._sys_db.get_max_function_id(workflow_uuid) or 0
+            )
+            return max(max_transactions, max_operations)
+
+        max_function_id = get_max_function_id(workflow_id)
+        if max_function_id > 0 and start_step > max_function_id:
+            raise DBOSException(
+                f"Cannot fork workflow {workflow_id} at step {start_step}. The workflow has  {max_function_id} steps."
+            )
+
         def fn() -> str:
-            dbos_logger.info(f"Restarting workflow: {workflow_id}")
-            return _get_dbos_instance()._sys_db.fork_workflow(workflow_id)
+            forked_workflow_id = str(uuid.uuid4())
+            dbos_logger.info(f"Forking workflow: {workflow_id} from step {start_step}")
+
+            _get_dbos_instance()._app_db.clone_workflow_transactions(
+                workflow_id, forked_workflow_id, start_step
+            )
+
+            return _get_dbos_instance()._sys_db.fork_workflow(
+                workflow_id, forked_workflow_id, start_step
+            )
 
-        forked_workflow_id = _get_dbos_instance()._sys_db.call_function_as_step(
-            fn, "DBOS.restartWorkflow"
+        new_id = _get_dbos_instance()._sys_db.call_function_as_step(
+            fn, "DBOS.forkWorkflow"
         )
-        return cls.retrieve_workflow(forked_workflow_id)
+        return cls.retrieve_workflow(new_id)
 
     @classmethod
     def list_workflows(

dbos-0.26.0a15/dbos/_event_loop.py

@@ -0,0 +1,67 @@
+import asyncio
+import threading
+from typing import Any, Coroutine, Optional, TypeVar
+
+
+class BackgroundEventLoop:
+    """
+    This is the event loop to which DBOS submits any coroutines that are not started from within an event loop.
+    In particular, coroutines submitted to queues (such as from scheduled workflows) run on this event loop.
+    """
+
+    def __init__(self) -> None:
+        self._loop: Optional[asyncio.AbstractEventLoop] = None
+        self._thread: Optional[threading.Thread] = None
+        self._running = False
+        self._ready = threading.Event()
+
+    def start(self) -> None:
+        if self._running:
+            return
+
+        self._thread = threading.Thread(target=self._run_event_loop, daemon=True)
+        self._thread.start()
+        self._ready.wait()  # Wait until the loop is running
+
+    def stop(self) -> None:
+        if not self._running or self._loop is None or self._thread is None:
+            return
+
+        asyncio.run_coroutine_threadsafe(self._shutdown(), self._loop)
+        self._thread.join()
+        self._running = False
+
+    def _run_event_loop(self) -> None:
+        self._loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(self._loop)
+
+        self._running = True
+        self._ready.set()  # Signal that the loop is ready
+
+        try:
+            self._loop.run_forever()
+        finally:
+            self._loop.close()
+
+    async def _shutdown(self) -> None:
+        if self._loop is None:
+            raise RuntimeError("Event loop not started")
+        tasks = [
+            task
+            for task in asyncio.all_tasks(self._loop)
+            if task is not asyncio.current_task(self._loop)
+        ]
+
+        for task in tasks:
+            task.cancel()
+
+        await asyncio.gather(*tasks, return_exceptions=True)
+        self._loop.stop()
+
+    T = TypeVar("T")
+
+    def submit_coroutine(self, coro: Coroutine[Any, Any, T]) -> T:
+        """Submit a coroutine to the background event loop"""
+        if self._loop is None:
+            raise RuntimeError("Event loop not started")
+        return asyncio.run_coroutine_threadsafe(coro, self._loop).result()

{dbos-0.26.0a13 → dbos-0.26.0a15}/dbos/_sys_db.py

@@ -489,15 +489,29 @@ class SystemDatabase:
                 .values(status=WorkflowStatusString.ENQUEUED.value, recovery_attempts=0)
             )
 
-    def fork_workflow(self, original_workflow_id: str) -> str:
+    def get_max_function_id(self, workflow_uuid: str) -> Optional[int]:
+        with self.engine.begin() as conn:
+            max_function_id_row = conn.execute(
+                sa.select(
+                    sa.func.max(SystemSchema.operation_outputs.c.function_id)
+                ).where(SystemSchema.operation_outputs.c.workflow_uuid == workflow_uuid)
+            ).fetchone()
+
+            max_function_id = max_function_id_row[0] if max_function_id_row else None
+
+            return max_function_id
+
+    def fork_workflow(
+        self, original_workflow_id: str, forked_workflow_id: str, start_step: int = 1
+    ) -> str:
+
         status = self.get_workflow_status(original_workflow_id)
         if status is None:
             raise Exception(f"Workflow {original_workflow_id} not found")
         inputs = self.get_workflow_inputs(original_workflow_id)
         if inputs is None:
             raise Exception(f"Workflow {original_workflow_id} not found")
-        # Generate a random ID for the forked workflow
-        forked_workflow_id = str(uuid.uuid4())
+
         with self.engine.begin() as c:
             # Create an entry for the forked workflow with the same
             # initial values as the original.
@@ -524,6 +538,37 @@ class SystemDatabase:
                     inputs=_serialization.serialize_args(inputs),
                 )
             )
+
+            if start_step > 1:
+
+                # Copy the original workflow's outputs into the forked workflow
+                insert_stmt = sa.insert(SystemSchema.operation_outputs).from_select(
+                    [
+                        "workflow_uuid",
+                        "function_id",
+                        "output",
+                        "error",
+                        "function_name",
+                        "child_workflow_id",
+                    ],
+                    sa.select(
+                        sa.literal(forked_workflow_id).label("workflow_uuid"),
+                        SystemSchema.operation_outputs.c.function_id,
+                        SystemSchema.operation_outputs.c.output,
+                        SystemSchema.operation_outputs.c.error,
+                        SystemSchema.operation_outputs.c.function_name,
+                        SystemSchema.operation_outputs.c.child_workflow_id,
+                    ).where(
+                        (
+                            SystemSchema.operation_outputs.c.workflow_uuid
+                            == original_workflow_id
+                        )
+                        & (SystemSchema.operation_outputs.c.function_id < start_step)
+                    ),
+                )
+
+                c.execute(insert_stmt)
+
             # Enqueue the forked workflow on the internal queue
             c.execute(
                 pg.insert(SystemSchema.workflow_queue).values(

{dbos-0.26.0a13 → dbos-0.26.0a15}/dbos/_workflow_commands.py

@@ -1,4 +1,5 @@
 import json
+import uuid
 from typing import Any, List, Optional
 
 from . import _serialization

{dbos-0.26.0a13 → dbos-0.26.0a15}/dbos/cli/cli.py

@@ -433,13 +433,60 @@ def restart(
     ] = 3001,
 ) -> None:
     response = requests.post(
-        f"http://{host}:{port}/workflows/{uuid}/restart", json=[], timeout=5
+        f"http://{host}:{port}/workflows/{uuid}/restart",
+        json=[],
+        timeout=5,
     )
 
     if response.status_code == 204:
         print(f"Workflow {uuid} has been restarted")
     else:
-        print(f"Failed to resume workflow {uuid}. Status code: {response.status_code}")
+        error_message = response.json().get("error", "Unknown error")
+        print(
+            f"Failed to restart workflow {uuid}. "
+            f"Status code: {response.status_code}. "
+            f"Error: {error_message}"
+        )
+
+
+@workflow.command(
+    help="fork a workflow from the beginning with a new id and from a step"
+)
+def fork(
+    uuid: Annotated[str, typer.Argument()],
+    host: Annotated[
+        typing.Optional[str],
+        typer.Option("--host", "-H", help="Specify the admin host"),
+    ] = "localhost",
+    port: Annotated[
+        typing.Optional[int],
+        typer.Option("--port", "-p", help="Specify the admin port"),
+    ] = 3001,
+    step: Annotated[
+        typing.Optional[int],
+        typer.Option(
+            "--step",
+            "-s",
+            help="Restart from this step (default: first step)",
+        ),
+    ] = 1,
+) -> None:
+    print(f"Forking workflow {uuid} from step {step}")
+    response = requests.post(
+        f"http://{host}:{port}/workflows/{uuid}/fork",
+        json={"start_step": step},
+        timeout=5,
+    )
+
+    if response.status_code == 204:
+        print(f"Workflow {uuid} has been forked")
+    else:
+        error_message = response.json().get("error", "Unknown error")
+        print(
+            f"Failed to fork workflow {uuid}. "
+            f"Status code: {response.status_code}. "
+            f"Error: {error_message}"
+        )
 
 
 @queue.command(name="list", help="List enqueued functions for your application")

{dbos-0.26.0a13 → dbos-0.26.0a15}/pyproject.toml

@@ -28,7 +28,7 @@ dependencies = [
 ]
 requires-python = ">=3.9"
 readme = "README.md"
-version = "0.26.0a13"
+version = "0.26.0a15"
 
 [project.license]
 text = "MIT"

{dbos-0.26.0a13 → dbos-0.26.0a15}/tests/test_async.py

@@ -7,8 +7,8 @@ import pytest
 import sqlalchemy as sa
 
 # Public API
-from dbos import DBOS, SetWorkflowID
-from dbos._dbos import WorkflowHandleAsync
+from dbos import DBOS, Queue, SetWorkflowID
+from dbos._dbos import WorkflowHandle, WorkflowHandleAsync
 from dbos._dbos_config import ConfigFile
 from dbos._error import DBOSException
 
@@ -56,6 +56,15 @@ async def test_async_workflow(dbos: DBOS) -> None:
     assert step_counter == 1
     assert txn_counter == 1
 
+    # Test DBOS.start_workflow_async
+    handle = await DBOS.start_workflow_async(test_workflow, "alice", "bob")
+    assert (await handle.get_result()) == "alicetxn21bobstep2"
+
+    # Test DBOS.start_workflow. Not recommended for async workflows,
+    # but needed for backwards compatibility.
+    sync_handle = DBOS.start_workflow(test_workflow, "alice", "bob")
+    assert sync_handle.get_result() == "alicetxn31bobstep3"  # type: ignore
+
 
 @pytest.mark.asyncio
 async def test_async_step(dbos: DBOS) -> None:
@@ -160,10 +169,11 @@ async def test_send_recv_async(dbos: DBOS) -> None:
     none_uuid = str(uuid.uuid4())
     none_handle = None
     with SetWorkflowID(none_uuid):
-        none_handle = dbos.start_workflow(test_recv_timeout, 10.0)
+        none_handle = await dbos.start_workflow_async(test_recv_timeout, 10.0)
     await test_send_none(none_uuid)
     begin_time = time.time()
-    assert none_handle.get_result() is None
+    result = await none_handle.get_result()  # type: ignore
+    assert result is None
     duration = time.time() - begin_time
     assert duration < 1.0  # None is from the received message, not from the timeout.
 
@@ -400,3 +410,59 @@ async def test_retrieve_workflow_async(dbos: DBOS) -> None:
     wfstatus = await handle.get_status()
     assert wfstatus.status == "SUCCESS"
     assert wfstatus.workflow_id == wfuuid
+
+
+def test_unawaited_workflow(dbos: DBOS) -> None:
+    input = 5
+    child_id = str(uuid.uuid4())
+    queue = Queue("test_queue")
+
+    @DBOS.workflow()
+    async def child_workflow(x: int) -> int:
+        await asyncio.sleep(0.1)
+        return x
+
+    @DBOS.workflow()
+    async def parent_workflow(x: int) -> None:
+        with SetWorkflowID(child_id):
+            await DBOS.start_workflow_async(child_workflow, x)
+
+    assert queue.enqueue(parent_workflow, input).get_result() is None
+    handle: WorkflowHandle[int] = DBOS.retrieve_workflow(
+        child_id, existing_workflow=False
+    )
+    assert handle.get_result() == 5
+
+
+def test_unawaited_workflow_exception(dbos: DBOS) -> None:
+    child_id = str(uuid.uuid4())
+    queue = Queue("test_queue")
+
+    @DBOS.workflow()
+    async def child_workflow(s: str) -> int:
+        await asyncio.sleep(0.1)
+        raise Exception(s)
+
+    @DBOS.workflow()
+    async def parent_workflow(s: str) -> None:
+        with SetWorkflowID(child_id):
+            await DBOS.start_workflow_async(child_workflow, s)
+
+    # Verify the unawaited child properly throws an exception
+    input = "alice"
+    assert queue.enqueue(parent_workflow, input).get_result() is None
+    handle: WorkflowHandle[int] = DBOS.retrieve_workflow(
+        child_id, existing_workflow=False
+    )
+    with pytest.raises(Exception) as exc_info:
+        handle.get_result()
+    assert input in str(exc_info.value)
+
+    # Verify it works if run again
+    input = "bob"
+    child_id = str(uuid.uuid4())
+    assert queue.enqueue(parent_workflow, input).get_result() is None
+    handle = DBOS.retrieve_workflow(child_id, existing_workflow=False)
+    with pytest.raises(Exception) as exc_info:
+        handle.get_result()
+    assert input in str(exc_info.value)