Workflow cancel , restart, resume as admin API (#187)

Author: manojdbos

dbos/_admin_server.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import json
+import re
 import threading
 from functools import partial
 from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
@@ -15,6 +16,9 @@
 _health_check_path = "/dbos-healthz"
 _workflow_recovery_path = "/dbos-workflow-recovery"
 _deactivate_path = "/deactivate"
+# /workflows/:workflow_id/cancel
+# /workflows/:workflow_id/resume
+# /workflows/:workflow_id/restart
 
 
 class AdminServer:
@@ -79,12 +83,51 @@ def do_POST(self) -> None:
             self._end_headers()
             self.wfile.write(json.dumps(workflow_ids).encode("utf-8"))
         else:
-            self.send_response(404)
-            self._end_headers()
+
+            restart_match = re.match(
+                r"^/workflows/(?P<workflow_id>[^/]+)/restart$", self.path
+            )
+            resume_match = re.match(
+                r"^/workflows/(?P<workflow_id>[^/]+)/resume$", self.path
+            )
+            cancel_match = re.match(
+                r"^/workflows/(?P<workflow_id>[^/]+)/cancel$", self.path
+            )
+
+            if restart_match:
+                workflow_id = restart_match.group("workflow_id")
+                self._handle_restart(workflow_id)
+            elif resume_match:
+                workflow_id = resume_match.group("workflow_id")
+                self._handle_resume(workflow_id)
+            elif cancel_match:
+                workflow_id = cancel_match.group("workflow_id")
+                self._handle_cancel(workflow_id)
+            else:
+                self.send_response(404)
+                self._end_headers()
 
     def log_message(self, format: str, *args: Any) -> None:
         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()
+
+    def _handle_resume(self, workflow_id: str) -> None:
+        print("Resuming workflow", workflow_id)
+        self.dbos.resume_workflow(workflow_id)
+        self.send_response(204)
+        self._end_headers()
+
+    def _handle_cancel(self, workflow_id: str) -> None:
+        print("Cancelling workflow", workflow_id)
+        self.dbos.cancel_workflow(workflow_id)
+        self.send_response(204)
+        self._end_headers()
+
 
 # Be consistent with DBOS-TS response.
 class PerfUtilization(TypedDict):

dbos/_core.py

@@ -266,7 +266,9 @@ def _execute_workflow_wthread(
                 raise
 
 
-def execute_workflow_by_id(dbos: "DBOS", workflow_id: str) -> "WorkflowHandle[Any]":
+def execute_workflow_by_id(
+    dbos: "DBOS", workflow_id: str, startNew: bool = False
+) -> "WorkflowHandle[Any]":
     status = dbos._sys_db.get_workflow_status(workflow_id)
     if not status:
         raise DBOSRecoveryError(workflow_id, "Workflow status not found")
@@ -293,7 +295,8 @@ def execute_workflow_by_id(dbos: "DBOS", workflow_id: str) -> "WorkflowHandle[An
                     workflow_id,
                     f"Cannot execute workflow because instance '{iname}' is not registered",
                 )
-            with SetWorkflowID(workflow_id):
+
+            if startNew:
                 return start_workflow(
                     dbos,
                     wf_func,
@@ -303,14 +306,26 @@ def execute_workflow_by_id(dbos: "DBOS", workflow_id: str) -> "WorkflowHandle[An
                     *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"],
+                    )
         elif status["class_name"] is not None:
             class_name = status["class_name"]
             if class_name not in dbos._registry.class_info_map:
                 raise DBOSWorkflowFunctionNotFoundError(
                     workflow_id,
                     f"Cannot execute workflow because class '{class_name}' is not registered",
                 )
-            with SetWorkflowID(workflow_id):
+
+            if startNew:
                 return start_workflow(
                     dbos,
                     wf_func,
@@ -320,8 +335,19 @@ def execute_workflow_by_id(dbos: "DBOS", workflow_id: str) -> "WorkflowHandle[An
                     *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"],
+                    )
         else:
-            with SetWorkflowID(workflow_id):
+            if startNew:
                 return start_workflow(
                     dbos,
                     wf_func,
@@ -330,6 +356,16 @@ def execute_workflow_by_id(dbos: "DBOS", workflow_id: str) -> "WorkflowHandle[An
                     *inputs["args"],
                     **inputs["kwargs"],
                 )
+            else:
+                with SetWorkflowID(workflow_id):
+                    return start_workflow(
+                        dbos,
+                        wf_func,
+                        status["queue_name"],
+                        True,
+                        *inputs["args"],
+                        **inputs["kwargs"],
+                    )
 
 
 @overload

dbos/_dbos.py

@@ -56,6 +56,7 @@
 )
 from ._roles import default_required_roles, required_roles
 from ._scheduler import ScheduledWorkflow, scheduled
+from ._sys_db import WorkflowStatusString
 from ._tracer import dbos_tracer
 
 if TYPE_CHECKING:
@@ -231,6 +232,7 @@ def __new__(
                         f"DBOS configured multiple times with conflicting information"
                     )
                 config = _dbos_global_registry.config
+
             _dbos_global_instance = super().__new__(cls)
             _dbos_global_instance.__init__(fastapi=fastapi, config=config, flask=flask)  # type: ignore
         else:
@@ -767,13 +769,30 @@ def execute_workflow_id(cls, workflow_id: str) -> WorkflowHandle[Any]:
         """Execute a workflow by ID (for recovery)."""
         return execute_workflow_by_id(_get_dbos_instance(), workflow_id)
 
+    @classmethod
+    def restart_workflow(cls, workflow_id: str) -> None:
+        """Execute a workflow by ID (for recovery)."""
+        execute_workflow_by_id(_get_dbos_instance(), workflow_id, True)
+
     @classmethod
     def recover_pending_workflows(
         cls, executor_ids: List[str] = ["local"]
     ) -> List[WorkflowHandle[Any]]:
         """Find all PENDING workflows and execute them."""
         return recover_pending_workflows(_get_dbos_instance(), executor_ids)
 
+    @classmethod
+    def cancel_workflow(cls, workflow_id: str) -> None:
+        """Cancel a workflow by ID."""
+        _get_dbos_instance()._sys_db.set_workflow_status(
+            workflow_id, WorkflowStatusString.CANCELLED, False
+        )
+
+    @classmethod
+    def resume_workflow(cls, workflow_id: str) -> None:
+        """Resume a workflow by ID."""
+        execute_workflow_by_id(_get_dbos_instance(), workflow_id, False)
+
     @classproperty
     def logger(cls) -> Logger:
         """Return the DBOS `Logger` for the current context."""

dbos/_workflow_commands.py

@@ -1,3 +1,6 @@
+import importlib
+import os
+import sys
 from typing import Any, List, Optional, cast
 
 import typer
@@ -6,6 +9,7 @@
 from dbos import DBOS
 
 from . import _serialization, load_config
+from ._core import execute_workflow_by_id
 from ._dbos_config import ConfigFile, _is_valid_app_name
 from ._sys_db import (
     GetWorkflowsInput,
@@ -123,11 +127,6 @@ def _cancel_workflow(config: ConfigFile, uuid: str) -> None:
             sys_db.destroy()
 
 
-def _reattempt_workflow(uuid: str, startNewWorkflow: bool) -> None:
-    print(f"Reattempt workflow info for {uuid} not implemented")
-    return
-
-
 def _get_workflow_info(
     sys_db: SystemDatabase, workflowUUID: str, getRequest: bool
 ) -> Optional[WorkflowInformation]:

dbos/cli.py

@@ -9,6 +9,7 @@
 from typing import Any
 
 import jsonpickle  # type: ignore
+import requests
 import sqlalchemy as sa
 import tomlkit
 import typer
@@ -22,12 +23,7 @@
 from ._app_db import ApplicationDatabase
 from ._dbos_config import _is_valid_app_name
 from ._sys_db import SystemDatabase
-from ._workflow_commands import (
-    _cancel_workflow,
-    _get_workflow,
-    _list_workflows,
-    _reattempt_workflow,
-)
+from ._workflow_commands import _cancel_workflow, _get_workflow, _list_workflows
 
 app = typer.Typer()
 workflow = typer.Typer()
@@ -432,5 +428,49 @@ def cancel(
     print(f"Workflow {uuid} has been cancelled")
 
 
+@workflow.command(help="Resume a workflow that has been cancelled")
+def resume(
+    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,
+) -> None:
+    response = requests.post(
+        f"http://{host}:{port}/workflows/{uuid}/resume", json=[], timeout=5
+    )
+
+    if response.status_code == 200:
+        print(f"Workflow {uuid} has been resumed")
+    else:
+        print(f"Failed to resume workflow {uuid}. Status code: {response.status_code}")
+
+
+@workflow.command(help="Restart a workflow from the beginning with a new id")
+def restart(
+    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,
+) -> None:
+    response = requests.post(
+        f"http://{host}:{port}/workflows/{uuid}/restart", json=[], timeout=5
+    )
+
+    if response.status_code == 200:
+        print(f"Workflow {uuid} has been restarted")
+    else:
+        print(f"Failed to resume workflow {uuid}. Status code: {response.status_code}")
+
+
 if __name__ == "__main__":
     app()