From e4474dde90e254613a19573281485165e842e4f3 Mon Sep 17 00:00:00 2001 From: Nick Mitchell Date: Tue, 4 Mar 2025 12:39:14 -0500 Subject: [PATCH] fix: correct end_nanos accounting for async model calls MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This at least fixes the block timestamps for async model calls. i don’t think it’s possible to do this generally (at least with my understanding of python futures), because the pdl_lazy logic obscures the original asyncio.Future… but this future is available at the point of fire for the async model calls… so i’ve fixed that part. It’s imperfect, because e.g. an outer Text block… doesn’t know that there are Futures under the hood. but still this seems like an improvement over what we had. and that might be something we can patch over in the UI. This includes said patch to the UI (see src/view/timeline/model.ts, the comment that references #683) Signed-off-by: Nick Mitchell --- pdl-live-react/src/view/timeline/model.ts | 25 +++++++++++++++++------ src/pdl/pdl_dumper.py | 8 ++++++-- src/pdl/pdl_llms.py | 10 +++++++++ 3 files changed, 35 insertions(+), 8 deletions(-) diff --git a/pdl-live-react/src/view/timeline/model.ts b/pdl-live-react/src/view/timeline/model.ts index aae3e3b3c..9a74e6d13 100644 --- a/pdl-live-react/src/view/timeline/model.ts +++ b/pdl-live-react/src/view/timeline/model.ts @@ -75,12 +75,25 @@ function computeModelIter( block, } - return [ - ...(ignoreRoot ? [] : [root]), - ...childrenOf(block) - .filter(nonNullable) - .flatMap((child) => computeModelIter(child, root)), - ].filter(nonNullable) + const childrenModel = childrenOf(block) + .filter(nonNullable) + .flatMap((child) => computeModelIter(child, root)) + + // Correct for anomalies in the trace where a child may have an + // earlier end timestamp than its children. See + // https://github.com/IBM/prompt-declaration-language/pull/683 + if (root) { + const maxEnd = childrenModel.reduce( + (maxEnd, child) => Math.max(maxEnd, child.block.pdl__timing.end_nanos), + 0, + ) + root.block.pdl__timing.end_nanos = Math.max( + root.block.pdl__timing.end_nanos, + maxEnd, + ) + } + + return [...(ignoreRoot ? [] : [root]), ...childrenModel].filter(nonNullable) } export function childrenOf(block: NonScalarPdlBlock) { diff --git a/src/pdl/pdl_dumper.py b/src/pdl/pdl_dumper.py index 8f619469f..2a8767c80 100644 --- a/src/pdl/pdl_dumper.py +++ b/src/pdl/pdl_dumper.py @@ -97,8 +97,6 @@ def block_to_dict( # noqa: C901 context = block.context if len(context) > 0: d["context"] = context - if block.pdl__timing is not None: - d["pdl__timing"] = timing_to_dict(block.pdl__timing) if block.description is not None: d["description"] = block.description if block.role is not None: @@ -239,6 +237,12 @@ def block_to_dict( # noqa: C901 # d["location"] = location_to_dict(block.location) if block.fallback is not None: d["fallback"] = block_to_dict(block.fallback, json_compatible) + # Warning: remember to update timing here at the end! this ensures + # that any logic that updates timestamps when futures + # finish... has a chance to do its work before we record the + # timestamps to the trace + if block.pdl__timing is not None: + d["pdl__timing"] = timing_to_dict(block.pdl__timing) return d diff --git a/src/pdl/pdl_llms.py b/src/pdl/pdl_llms.py index 940b9297f..6c386559b 100644 --- a/src/pdl/pdl_llms.py +++ b/src/pdl/pdl_llms.py @@ -100,6 +100,16 @@ def generate_text( pdl_future: PdlLazy[tuple[dict[str, Any], Any]] = PdlConst(future) message = lazy_apply((lambda x: x[0]), pdl_future) response = lazy_apply((lambda x: x[1]), pdl_future) + + # update the end timestamp when the future is done + def update_end_nanos(future): + import time + + if block.pdl__timing is not None: + block.pdl__timing.end_nanos = time.time_ns() + + future.add_done_callback(update_end_nanos) + return message, response @staticmethod