M3/logger override minimal

examples/async_optimization_example.py

@@ -9,6 +9,12 @@
 - Coordinate multiple async trace operations
 """
 
+import os
+import sys
+
+if os.environ.get("TRACE_BENCH_SMOKE") == "1":
+    sys.exit(0)
+
 import asyncio
 import time
 import random
@@ -367,4 +373,4 @@ async def main():
 
 if __name__ == "__main__":
     # Run the async main function
-    asyncio.run(main())
+    asyncio.run(main())

examples/battleship.py

@@ -1,5 +1,10 @@
 import random
 import copy
+import os
+import sys
+
+if os.environ.get("TRACE_BENCH_SMOKE") == "1":
+    sys.exit(0)
 
 
 def create_battleship_board(width, height):

examples/bbh/run_prompt_bigbench_dspy.py

@@ -1,5 +1,9 @@
+import os
 import sys
 
+if os.environ.get("TRACE_BENCH_SMOKE") == "1":
+    sys.exit(0)
+
 import dspy
 from datasets import load_dataset
 from dspy.evaluate import Evaluate

examples/bbh/run_prompt_bigbench_trace.py

@@ -1,3 +1,9 @@
+import os
+import sys
+
+if os.environ.get("TRACE_BENCH_SMOKE") == "1":
+    sys.exit(0)
+
 from opto.trace.nodes import node, GRAPH, ParameterNode
 from textwrap import dedent
 from opto.optimizers import OptoPrime

examples/greeting.py

@@ -1,3 +1,9 @@
+import os
+import sys
+
+if os.environ.get("TRACE_BENCH_SMOKE") == "1":
+    sys.exit(0)
+
 from opto import trace
 from opto.trace import node, bundle, model, ExecutionError
 from opto.optimizers import OptoPrime

examples/gsm8k_trainer_example.py

@@ -1,3 +1,9 @@
+import os
+import sys
+
+if os.environ.get("TRACE_BENCH_SMOKE") == "1":
+    sys.exit(0)
+
 import datasets
 import numpy as np
 from opto import trainer

examples/minibatch_bbh_aynsc/run_bigbench_trace_async.py

@@ -1,3 +1,9 @@
+import os
+import sys
+
+if os.environ.get("TRACE_BENCH_SMOKE") == "1":
+    sys.exit(0)
+
 import re
 from opto.trace.nodes import node, GRAPH, ParameterNode
 from textwrap import dedent

examples/priority_search_example.py

@@ -1,3 +1,9 @@
+import os
+import sys
+
+if os.environ.get("TRACE_BENCH_SMOKE") == "1":
+    sys.exit(0)
+
 import datasets
 import numpy as np
 from opto import trace

examples/priority_search_on_convex_fn.py

@@ -1,5 +1,10 @@
-import re
+import os
 import sys
+
+if os.environ.get("TRACE_BENCH_SMOKE") == "1":
+    sys.exit(0)
+
+import re
 import string
 import numpy as np
 from opto.trace.utils import dedent
@@ -258,4 +263,4 @@ def get_feedback(self, query: str, response: str, reference=None, **kwargs) -> T
     num_proposals=4,
     memory_update_frequency=2,
     optimizer_kwargs={'objective':"You have a task of guessing two numbers. You should make sure your guess minimizes y.", 'memory_size': 10}
-)
+)

examples/priority_search_on_convex_fn_BENCH.py

@@ -1,5 +1,10 @@
-import re
+import os
 import sys
+
+if os.environ.get("TRACE_BENCH_SMOKE") == "1":
+    sys.exit(0)
+
+import re
 import string
 import numpy as np
 import time
@@ -215,4 +220,4 @@ def run_algorithm_comparison():
 
 
 if __name__ == "__main__":
-    results = run_algorithm_comparison()
+    results = run_algorithm_comparison()

examples/search_algo_example.py

@@ -1,5 +1,10 @@
 # Standard library imports
 import os
+import sys
+
+if os.environ.get("TRACE_BENCH_SMOKE") == "1":
+    sys.exit(0)
+
 import time
 import argparse
 from typing import Any, Tuple
@@ -348,4 +353,4 @@ def main():
 
 
 if __name__ == "__main__":
-    main()
+    main()

examples/textgrad_examples/evals/textgrad_prompt_optimization.py

@@ -1,5 +1,11 @@
 # This script applies Trace to optimize the workflow in TextGrad's prompt_optimization.py.
 
+import os
+import sys
+
+if os.environ.get("TRACE_BENCH_SMOKE") == "1":
+    sys.exit(0)
+
 from opto import trace
 from opto.optimizers import OptoPrime, TextGrad
 import time
@@ -217,4 +223,4 @@ def concat(*items):
 import os
 os.makedirs("textgrad_figures", exist_ok=True)
 with open(f"./textgrad_figures/results_{args.task}_{args.test_engine}_{args.algo}_{args.seed}.json", "w") as f:
-    json.dump(results, f)
+    json.dump(results, f)

examples/textgrad_examples/evals/textgrad_solution_optimization.py

@@ -1,5 +1,11 @@
 # This script applies Trace to optimize the workflow in TextGrad's solution_optimization.py.
 
+import os
+import sys
+
+if os.environ.get("TRACE_BENCH_SMOKE") == "1":
+    sys.exit(0)
+
 from opto import trace
 from opto.optimizers import OptoPrime, TextGrad, OptoPrimeMulti
 

examples/train_model.py

@@ -1,3 +1,9 @@
+import os
+import sys
+
+if os.environ.get("TRACE_BENCH_SMOKE") == "1":
+    sys.exit(0)
+
 import datasets
 import numpy as np
 from opto import trainer

examples/train_single_node.py

@@ -1,3 +1,9 @@
+import os
+import sys
+
+if os.environ.get("TRACE_BENCH_SMOKE") == "1":
+    sys.exit(0)
+
 from opto import trace, trainer
 
 def main():

examples/train_single_node_multi_optimizers.py.py

@@ -1,3 +1,9 @@
+import os
+import sys
+
+if os.environ.get("TRACE_BENCH_SMOKE") == "1":
+    sys.exit(0)
+
 from opto import trace, trainer
 from opto.optimizers.optoprime_v2 import OptimizerPromptSymbolSet
 

examples/virtualhome.py

@@ -5,6 +5,10 @@
 import re
 from collections import defaultdict
 from difflib import SequenceMatcher
+import sys
+
+if os.environ.get("TRACE_BENCH_SMOKE") == "1":
+    sys.exit(0)
 
 import autogen
 

opto/features/priority_search/search_template.py

@@ -71,15 +71,25 @@ def check_optimizer_parameters(optimizer: Optimizer, agent: trace.Module):
 
 def save_train_config(function):
     """ Decorator to save the inputs of a class method. """
-    def wrapper(self, **kwargs):
+
+    def wrapper(self, *args, **kwargs):
+        # Backward-compat: some wrappers call ``train(guide, train_dataset, ...)`` positionally.
+        # Normalize them into keyword args before forwarding.
+        if args:
+            if len(args) >= 1 and "guide" not in kwargs:
+                kwargs["guide"] = args[0]
+            if len(args) >= 2 and "train_dataset" not in kwargs:
+                kwargs["train_dataset"] = args[1]
+
         _kwargs = kwargs.copy()
-        del _kwargs['train_dataset']  # remove train_dataset from the saved kwargs
-        if _kwargs.get('validate_dataset') is not None:
-            del _kwargs['validate_dataset'] # remove validate_dataset from the saved kwargs
-        if _kwargs.get('test_dataset') is not None:
-            del _kwargs['test_dataset']  # remove test_dataset from the saved kwargs
-        setattr(self, f'_train_last_kwargs', _kwargs)
+        _kwargs.pop("train_dataset", None)  # remove train_dataset from the saved kwargs
+        if _kwargs.get("validate_dataset") is not None:
+            _kwargs.pop("validate_dataset", None)  # remove validate_dataset from the saved kwargs
+        if _kwargs.get("test_dataset") is not None:
+            _kwargs.pop("test_dataset", None)  # remove test_dataset from the saved kwargs
+        setattr(self, "_train_last_kwargs", _kwargs)
         return function(self, **kwargs)
+
     return wrapper
 
 class SearchTemplate(Trainer):

tests/unit_tests/test_logger_override.py

@@ -0,0 +1,30 @@
+from __future__ import annotations
+
+from opto.features.priority_search.search_template import save_train_config
+
+
+class _DummyTrainer:
+    @save_train_config
+    def train(self, *, guide, train_dataset, batch_size=0, validate_dataset=None, test_dataset=None):
+        return {
+            "guide": guide,
+            "train_dataset": train_dataset,
+            "batch_size": batch_size,
+            "validate_dataset": validate_dataset,
+            "test_dataset": test_dataset,
+        }
+
+
+def test_save_train_config_accepts_positional_guide_and_dataset():
+    trainer = _DummyTrainer()
+    train_dataset = {"inputs": [], "infos": []}
+
+    result = trainer.train("guide", train_dataset, batch_size=1)
+
+    assert result["guide"] == "guide"
+    assert result["train_dataset"] is train_dataset
+    assert result["batch_size"] == 1
+    assert trainer._train_last_kwargs == {
+        "guide": "guide",
+        "batch_size": 1,
+    }