Conductor Improvements (#617)

- Support for a "private mode" where workflow data (as opposed to
metadata) is never sent to Conductor
- Add the ability to "tag" executors with metadata for display in the
dashboard
- Add a new workflow aggregates API

Author: kraftp

dbos/_conductor/conductor.py

@@ -123,6 +123,7 @@ def run(self) -> None:
                                 hostname=socket.gethostname(),
                                 language="python",
                                 dbos_version=GlobalParams.dbos_version,
+                                executor_metadata=self.dbos._conductor_executor_metadata,
                             )
                             websocket.send(info_response.to_json())
                             self.dbos.logger.info("Connected to DBOS conductor")
@@ -285,8 +286,8 @@ def run(self) -> None:
                                     workflow_id_prefix=body.get(
                                         "workflow_id_prefix", None
                                     ),
-                                    load_input=body.get("load_input", False),
-                                    load_output=body.get("load_output", False),
+                                    load_input=body.get("load_input", True),
+                                    load_output=body.get("load_output", True),
                                     executor_id=body.get("executor_id", None),
                                     queues_only=body.get("queues_only", False),
                                 )
@@ -330,8 +331,8 @@ def run(self) -> None:
                                     workflow_id_prefix=q_body.get(
                                         "workflow_id_prefix", None
                                     ),
-                                    load_input=q_body.get("load_input", False),
-                                    load_output=q_body.get("load_output", False),
+                                    load_input=q_body.get("load_input", True),
+                                    load_output=q_body.get("load_output", True),
                                     executor_id=q_body.get("executor_id", None),
                                     queues_only=True,
                                 )
@@ -358,7 +359,10 @@ def run(self) -> None:
                             info = None
                             try:
                                 info = get_workflow(
-                                    self.dbos._sys_db, get_workflow_message.workflow_id
+                                    self.dbos._sys_db,
+                                    get_workflow_message.workflow_id,
+                                    load_input=get_workflow_message.load_input,
+                                    load_output=get_workflow_message.load_output,
                                 )
                             except Exception as e:
                                 error_message = f"Exception encountered when getting workflow {get_workflow_message.workflow_id}: {traceback.format_exc()}"
@@ -478,7 +482,8 @@ def run(self) -> None:
                             step_info = None
                             try:
                                 step_info = self.dbos._sys_db.list_workflow_steps(
-                                    list_steps_message.workflow_id
+                                    list_steps_message.workflow_id,
+                                    load_output=list_steps_message.load_output,
                                 )
                             except Exception as e:
                                 error_message = f"Exception encountered when getting workflow {list_steps_message.workflow_id}: {traceback.format_exc()}"
@@ -645,9 +650,12 @@ def run(self) -> None:
                             sched_body = list_sched_msg.body
                             schedules: list[p.ScheduleOutput] = []
                             try:
+                                load_context = sched_body.get("load_context", True)
                                 schedules = [
                                     p.ScheduleOutput.from_schedule(
-                                        s, self.dbos._sys_db.serializer
+                                        s,
+                                        self.dbos._sys_db.serializer,
+                                        load_context=load_context,
                                     )
                                     for s in self.dbos._sys_db.list_schedules(
                                         status=sched_body.get("status", None),
@@ -679,7 +687,9 @@ def run(self) -> None:
                                 )
                                 if sched is not None:
                                     output = p.ScheduleOutput.from_schedule(
-                                        sched, self.dbos._sys_db.serializer
+                                        sched,
+                                        self.dbos._sys_db.serializer,
+                                        load_context=get_sched_msg.load_context,
                                     )
                             except Exception:
                                 error_message = f"Exception encountered when getting schedule '{get_sched_msg.schedule_name}': {traceback.format_exc()}"
@@ -814,6 +824,55 @@ def run(self) -> None:
                                     error_message=error_message,
                                 ).to_json()
                             )
+                        elif msg_type == p.MessageType.GET_WORKFLOW_AGGREGATES:
+                            agg_message = p.GetWorkflowAggregatesRequest.from_json(
+                                message
+                            )
+                            agg_body = agg_message.body
+                            agg_output: list[p.WorkflowAggregateOutput] = []
+                            try:
+                                agg_rows = self.dbos._sys_db.get_workflow_aggregates(
+                                    group_by_status=agg_body.get(
+                                        "group_by_status", False
+                                    ),
+                                    group_by_name=agg_body.get("group_by_name", False),
+                                    group_by_queue_name=agg_body.get(
+                                        "group_by_queue_name", False
+                                    ),
+                                    group_by_executor_id=agg_body.get(
+                                        "group_by_executor_id", False
+                                    ),
+                                    group_by_application_version=agg_body.get(
+                                        "group_by_application_version", False
+                                    ),
+                                    status=agg_body.get("status", None),
+                                    start_time=agg_body.get("start_time", None),
+                                    end_time=agg_body.get("end_time", None),
+                                    name=agg_body.get("name", None),
+                                    app_version=agg_body.get("app_version", None),
+                                    executor_id=agg_body.get("executor_id", None),
+                                    queue_name=agg_body.get("queue_name", None),
+                                    workflow_id_prefix=agg_body.get(
+                                        "workflow_id_prefix", None
+                                    ),
+                                )
+                                agg_output = [
+                                    p.WorkflowAggregateOutput(
+                                        group=r["group"], count=r["count"]
+                                    )
+                                    for r in agg_rows
+                                ]
+                            except Exception:
+                                error_message = f"Exception encountered when getting workflow aggregates: {traceback.format_exc()}"
+                                self.dbos.logger.error(error_message)
+                            websocket.send(
+                                p.GetWorkflowAggregatesResponse(
+                                    type=p.MessageType.GET_WORKFLOW_AGGREGATES,
+                                    request_id=base_message.request_id,
+                                    output=agg_output,
+                                    error_message=error_message,
+                                ).to_json()
+                            )
                         else:
                             self.dbos.logger.warning(
                                 f"Unexpected message type: {msg_type}"

dbos/_conductor/protocol.py

@@ -42,6 +42,7 @@ class MessageType(str, Enum):
     GET_WORKFLOW_EVENTS = "get_workflow_events"
     GET_WORKFLOW_NOTIFICATIONS = "get_workflow_notifications"
     GET_WORKFLOW_STREAMS = "get_workflow_streams"
+    GET_WORKFLOW_AGGREGATES = "get_workflow_aggregates"
 
 
 T = TypeVar("T", bound="BaseMessage")
@@ -88,6 +89,7 @@ class ExecutorInfoResponse(BaseMessage):
     hostname: Optional[str]
     language: Optional[str]
     dbos_version: Optional[str]
+    executor_metadata: Optional[Dict[str, object]] = None
     error_message: Optional[str] = None
 
 
@@ -335,6 +337,8 @@ class ListQueuedWorkflowsResponse(BaseMessage):
 @dataclass
 class GetWorkflowRequest(BaseMessage):
     workflow_id: str
+    load_input: bool = True
+    load_output: bool = True
 
 
 @dataclass
@@ -358,6 +362,7 @@ class ExistPendingWorkflowsResponse(BaseMessage):
 @dataclass
 class ListStepsRequest(BaseMessage):
     workflow_id: str
+    load_output: bool = True
 
 
 @dataclass
@@ -470,16 +475,22 @@ class ScheduleOutput:
     workflow_class_name: Optional[str]
     schedule: str
     status: str
-    context: str
+    context: Optional[str]
     last_fired_at: Optional[str]
     automatic_backfill: bool
     cron_timezone: Optional[str]
 
     @classmethod
     def from_schedule(
-        cls, s: WorkflowSchedule, serializer: Serializer
+        cls,
+        s: WorkflowSchedule,
+        serializer: Serializer,
+        *,
+        load_context: bool = True,
     ) -> "ScheduleOutput":
-        context_str = str(serializer.deserialize(s["context"]))
+        context_str = (
+            str(serializer.deserialize(s["context"])) if load_context else None
+        )
         return cls(
             schedule_id=s["schedule_id"],
             schedule_name=s["schedule_name"],
@@ -498,6 +509,7 @@ class ListSchedulesBody(TypedDict, total=False):
     status: Optional[Union[str, List[str]]]
     workflow_name: Optional[Union[str, List[str]]]
     schedule_name_prefix: Optional[Union[str, List[str]]]
+    load_context: bool
 
 
 @dataclass
@@ -514,6 +526,7 @@ class ListSchedulesResponse(BaseMessage):
 @dataclass
 class GetScheduleRequest(BaseMessage):
     schedule_name: str
+    load_context: bool = True
 
 
 @dataclass
@@ -678,3 +691,36 @@ class GetWorkflowStreamsRequest(BaseMessage):
 class GetWorkflowStreamsResponse(BaseMessage):
     streams: Optional[List[StreamEntryOutput]]
     error_message: Optional[str] = None
+
+
+class GetWorkflowAggregatesBody(TypedDict, total=False):
+    group_by_status: bool
+    group_by_name: bool
+    group_by_queue_name: bool
+    group_by_executor_id: bool
+    group_by_application_version: bool
+    status: Optional[List[str]]
+    start_time: Optional[str]
+    end_time: Optional[str]
+    name: Optional[List[str]]
+    app_version: Optional[List[str]]
+    executor_id: Optional[List[str]]
+    queue_name: Optional[List[str]]
+    workflow_id_prefix: Optional[List[str]]
+
+
+@dataclass
+class GetWorkflowAggregatesRequest(BaseMessage):
+    body: GetWorkflowAggregatesBody
+
+
+@dataclass
+class WorkflowAggregateOutput:
+    group: Dict[str, Optional[str]]
+    count: int
+
+
+@dataclass
+class GetWorkflowAggregatesResponse(BaseMessage):
+    output: List[WorkflowAggregateOutput]
+    error_message: Optional[str] = None

dbos/_dbos.py

@@ -4,6 +4,7 @@
 import copy
 import hashlib
 import inspect
+import json
 import os
 import sys
 import threading
@@ -408,6 +409,16 @@ def __init__(
         self._alert_handler: Optional[Callable[[str, str, Dict[str, str]], None]] = None
         serializer = config.get("serializer")
         self._serializer: Serializer = serializer if serializer else DefaultSerializer()
+        self._conductor_executor_metadata: Optional[Dict[str, Any]] = config.get(
+            "conductor_executor_metadata"
+        )
+        if self._conductor_executor_metadata is not None:
+            try:
+                json.dumps(self._conductor_executor_metadata)
+            except Exception as e:
+                raise DBOSException(
+                    f"conductor_executor_metadata must be JSON-serializable: {e}"
+                )
 
         # Globally set the application version and executor ID.
         # In DBOS Cloud, instead use the values supplied through environment variables.

dbos/_dbos_config.py

@@ -42,6 +42,7 @@ class DBOSConfig(TypedDict, total=False):
         system_database_engine (sa.Engine): A custom system database engine. If provided, DBOS will not create an engine but use this instead.
         conductor_key (str): An API key for DBOS Conductor. Pass this in to connect your process to Conductor.
         conductor_url (str): The websockets URL for your DBOS Conductor service. Only set if you're self-hosting Conductor.
+        conductor_executor_metadata (Dict[str, Any]): Metadata associated with this executor that may be used to identify an executor on the Conductor dashboard. Must be JSON-serializable.
         serializer (Serializer): A custom serializer and deserializer DBOS uses when storing program data in the system database
         use_listen_notify (bool): Whether to use LISTEN/NOTIFY or polling to listen for notifications and events.  Defaults to True. As this affects migrations, may not be changed after the system database is first created.
         notification_listener_polling_interval_sec (float): Polling interval in seconds for the notification listener background process. Defaults to 1.0. Minimum value is 0.001. Lower values can speed up test execution.
@@ -69,6 +70,7 @@ class DBOSConfig(TypedDict, total=False):
     system_database_engine: Optional[sa.Engine]
     conductor_key: Optional[str]
     conductor_url: Optional[str]
+    conductor_executor_metadata: Optional[Dict[str, Any]]
     serializer: Optional[Serializer]
     enable_patching: Optional[bool]
     use_listen_notify: Optional[bool]