feat: export write_trace in pdl.pdl (#966)

Signed-off-by: Louis Mandel <[email protected]>

Author: mandel

src/pdl/pdl.py

@@ -22,7 +22,10 @@
 from .pdl_lazy import PdlDict
 from .pdl_parser import parse_dict, parse_file, parse_str
 from .pdl_runner import exec_docker
-from .pdl_utils import validate_scope
+from .pdl_utils import (  # pylint: disable=unused-import # noqa: F401
+    validate_scope,
+    write_trace,
+)
 
 
 class InterpreterConfig(TypedDict, total=False):

src/pdl/pdl_interpreter.py

@@ -103,7 +103,6 @@
     deserialize,
     ensure_context,
 )
-from .pdl_dumper import as_json, block_to_dict  # noqa: E402
 from .pdl_lazy import PdlConst, PdlDict, PdlLazy, PdlList, lazy_apply  # noqa: E402
 from .pdl_llms import LitellmModel  # noqa: E402
 from .pdl_location_utils import append, get_loc_string  # noqa: E402
@@ -119,6 +118,7 @@
     replace_contribute_value,
     stringify,
     value_of_expr,
+    write_trace,
 )
 
 empty_scope: ScopeType = PdlDict({"pdl_context": DependentContext([])})
@@ -202,25 +202,6 @@ def generate(
     return 0
 
 
-def write_trace(
-    trace_file: str | Path,
-    trace: BlockType,
-):
-    """Write the execution trace into a file.
-
-    Args:
-        trace_file:  File to save the execution trace.
-        trace: Execution trace.
-    """
-    try:
-        d: Any = block_to_dict(trace, json_compatible=True)
-        d = as_json(d)
-        with open(trace_file, "w", encoding="utf-8") as fp:
-            json.dump(d, fp)
-    except Exception as e:
-        print(f"Failure generating the trace: {str(e)}", file=sys.stderr)
-
-
 def process_prog(
     state: InterpreterState,
     scope: ScopeType,

src/pdl/pdl_utils.py

@@ -1,15 +1,19 @@
 import fnmatch
 import json
+import sys
+from pathlib import Path
 from typing import Any, Generator, Generic, Sequence, TypeVar
 
 from .pdl_ast import (
+    BlockType,
     ContributeTarget,
     ContributeValue,
     ExpressionType,
     FunctionBlock,
     LocalizedExpression,
     get_sampling_defaults,
 )
+from .pdl_dumper import as_json, block_to_dict
 
 GeneratorWrapperYieldT = TypeVar("GeneratorWrapperYieldT")
 GeneratorWrapperSendT = TypeVar("GeneratorWrapperSendT")
@@ -196,3 +200,22 @@ def validate_pdl_model_defaults(model_defaults: list[dict[str, dict[str, Any]]])
                     f"invalid defaults {glob_defaults} for model matcher {model_glob}"
                 )
             assert isinstance(glob_defaults, dict)
+
+
+def write_trace(
+    trace_file: str | Path,
+    trace: BlockType,
+):
+    """Write the execution trace into a file.
+
+    Args:
+        trace_file:  File to save the execution trace.
+        trace: Execution trace.
+    """
+    try:
+        d: Any = block_to_dict(trace, json_compatible=True)
+        d = as_json(d)
+        with open(trace_file, "w", encoding="utf-8") as fp:
+            json.dump(d, fp)
+    except Exception as e:
+        print(f"Failure generating the trace: {str(e)}", file=sys.stderr)