Adds partition key to logger + UI

This adds the partition key in a new metadata.json file in the
application directory. This defaults to None (valid) if empty, so we're
backwards compatible.

burr/cli/demo_data.py

@@ -87,12 +87,13 @@ def _run_conversation(app_id, prompts):
 
 def generate_counter_data(data_dir: str = "~/.burr"):
     counts = [1, 10, 100, 50, 42]
-    for count in counts:
+    for i, count in enumerate(counts):
         app = counter_application.application(
             count_up_to=count,
             app_id=f"count-to-{count}",
             storage_dir=data_dir,
             hooks=[ProgressHook()],
+            partition_key=f"user_{i}",
         )
         app.run(halt_after=["result"])
 

burr/core/application.py

@@ -259,6 +259,7 @@ def __init__(
             state=self._state,
             application_graph=self._graph,
             app_id=self._uid,
+            partition_key=self._partition_key,
         )
         # TODO -- consider adding global inputs + global input factories to the builder
         self.dependency_factory = {

burr/lifecycle/base.py

@@ -131,14 +131,16 @@ def post_application_create(
         self,
         *,
         app_id: str,
-        partition_key: str,
+        partition_key: Optional[str],
         state: "State",
         application_graph: "ApplicationGraph",
         **future_kwargs: Any,
     ):
         """Runs after an "application" object is instantiated. This is run by the Application, in its constructor,
         as the last step.
 
+        :param app_id: Application ID
+        :param partition_key: Partition key of application
         :param state: Current state of the application
         :param application_graph: Application graph of the application, representing the state machine
         :param future_kwargs: Future keyword arguments for backwards compatibility

burr/tracking/client.py

@@ -16,6 +16,7 @@
     PreStartSpanHook,
 )
 from burr.tracking.common.models import (
+    ApplicationMetadataModel,
     ApplicationModel,
     BeginEntryModel,
     BeginSpanModel,
@@ -57,6 +58,7 @@ class LocalTrackingClient(
     """
 
     GRAPH_FILENAME = "graph.json"
+    METADATA_FILENAME = "metadata.json"
     LOG_FILENAME = "log.jsonl"
     DEFAULT_STORAGE_DIR = "~/.burr"
 
@@ -154,20 +156,31 @@ def post_application_create(
         self,
         *,
         app_id: str,
+        partition_key: Optional[str],
         state: "State",
         application_graph: "ApplicationGraph",
         **future_kwargs: Any,
     ):
         self._ensure_dir_structure(app_id)
         self.f = open(os.path.join(self.storage_dir, app_id, self.LOG_FILENAME), "a")
-        path = os.path.join(self.storage_dir, app_id, self.GRAPH_FILENAME)
-        if os.path.exists(path):
-            logger.info(f"Graph already exists at {path}. Not overwriting.")
+        graph_path = os.path.join(self.storage_dir, app_id, self.GRAPH_FILENAME)
+        if os.path.exists(graph_path):
+            logger.info(f"Graph already exists at {graph_path}. Not overwriting.")
             return
         graph = ApplicationModel.from_application_graph(application_graph).model_dump()
-        with open(path, "w") as f:
+        with open(graph_path, "w") as f:
             json.dump(graph, f)
 
+        metadata_path = os.path.join(self.storage_dir, app_id, self.METADATA_FILENAME)
+        if os.path.exists(metadata_path):
+            logger.info(f"Metadata already exists at {metadata_path}. Not overwriting.")
+            return
+        metadata = ApplicationMetadataModel(
+            partition_key=partition_key,
+        ).model_dump()
+        with open(metadata_path, "w") as f:
+            json.dump(metadata, f)
+
     def _append_write_line(self, model: pydantic.BaseModel):
         self.f.write(model.model_dump_json() + "\n")
         self.f.flush()

burr/tracking/common/models.py

@@ -86,6 +86,14 @@ def from_application_graph(application_graph: ApplicationGraph) -> "ApplicationM
         )
 
 
+class ApplicationMetadataModel(IdentifyingModel):
+    """Pydantic model that represents metadata for an application.
+    We will want to add tags here when we have them."""
+
+    partition_key: Optional[str]
+    type: str = "application_metadata"
+
+
 INPUT_FILTERLIST = {"__tracer"}
 
 

burr/tracking/server/backend.py

@@ -7,6 +7,7 @@
 import aiofiles.os as aiofilesos
 import fastapi
 
+from burr.tracking.common import models
 from burr.tracking.common.models import BeginEntryModel, BeginSpanModel, EndEntryModel, EndSpanModel
 from burr.tracking.server import schema
 from burr.tracking.server.schema import ApplicationLogs, ApplicationSummary
@@ -119,11 +120,20 @@ async def list_apps(
                 # skip hidden files/directories
                 continue
             full_path = os.path.join(project_filepath, entry)
+            metadata_path = os.path.join(full_path, "metadata.json")
             log_path = os.path.join(full_path, "log.jsonl")
             if os.path.isdir(full_path):
+                if os.path.exists(metadata_path):
+                    async with aiofiles.open(metadata_path) as f:
+                        metadata = models.ApplicationMetadataModel.parse_obj(
+                            json.loads(await f.read())
+                        )
+                else:
+                    metadata = models.ApplicationMetadataModel()  # backwards compatibility
                 out.append(
                     schema.ApplicationSummary(
                         app_id=entry,
+                        partition_key=metadata.partition_key,
                         first_written=await aiofilesos.path.getctime(full_path),
                         last_written=await aiofilesos.path.getmtime(full_path),
                         num_steps=await self.get_max_sequence_id(log_path),

burr/tracking/server/schema.py

@@ -23,6 +23,7 @@ class Project(pydantic.BaseModel):
 
 class ApplicationSummary(pydantic.BaseModel):
     app_id: str
+    partition_key: Optional[str]
     first_written: datetime.datetime
     last_written: datetime.datetime
     num_steps: int

examples/conversational_rag/application.py

@@ -95,7 +95,7 @@ def application(
         )
         .with_entrypoint("human_converse")
         .with_tracker(project="demo:conversational-rag", params={"storage_dir": storage_dir})
-        .with_identifiers(app_id=app_id)
+        .with_identifiers(app_id=app_id, partition_key="sample_user")
         .with_hooks(*hooks if hooks else [])
         .build()
     )

telemetry/ui/src/api/models/ApplicationSummary.ts

@@ -3,10 +3,10 @@
 /* tslint:disable */
 /* eslint-disable */
 export type ApplicationSummary = {
-    app_id: string;
-    first_written: string;
-    last_written: string;
-    num_steps: number;
-    tags: Record<string, string>;
+  app_id: string;
+  partition_key: string | null;
+  first_written: string;
+  last_written: string;
+  num_steps: number;
+  tags: Record<string, string>;
 };
-

telemetry/ui/src/components/routes/AppList.tsx

@@ -44,6 +44,7 @@ export const AppListTable = (props: { apps: ApplicationSummary[]; projectId: str
     <Table>
       <TableHead>
         <TableRow>
+          <TableHeader>Partition Key</TableHeader>
           <TableHeader>ID</TableHeader>
           <TableHeader>First Seen</TableHeader>
           <TableHeader>Last Run</TableHeader>
@@ -64,6 +65,7 @@ export const AppListTable = (props: { apps: ApplicationSummary[]; projectId: str
             className="hover:bg-gray-50 cursor-pointer"
             href={`/project/${props.projectId}/${app.app_id}`}
           >
+            <TableCell className="text-gray-600 font-sans">{app.partition_key}</TableCell>
             <TableCell className="font-semibold text-gray-700">{app.app_id}</TableCell>
             <TableCell>
               <DateTimeDisplay date={app.first_written} mode="long" />