langfun 0.1.2.dev202501080804__py3-none-any.whl → 0.1.2.dev202501240804__py3-none-any.whl

langfun/core/concurrent.py

@@ -15,6 +15,7 @@
 
 import abc
 import collections
+from collections.abc import Mapping
 import concurrent.futures
 import dataclasses
 import io
@@ -22,10 +23,8 @@ import random
 import sys
 import threading
 import time
-from typing import Any, Callable, Iterable, Iterator, Literal, Sequence, Tuple, Type, Union
+from typing import Annotated, Any, Callable, Iterable, Iterator, Literal, Sequence, Tuple, Type, Union
 
-from langfun.core import component
-from langfun.core import text_formatting
 import pyglove as pg
 
 
@@ -39,18 +38,6 @@ except ImportError:
   tqdm = None
 
 
-def with_context_access(func: Callable[..., Any]) -> Callable[..., Any]:
-  """Derives a user function with the access to the current context."""
-  with component.context() as current_context:
-    pass
-
-  def _func(*args, **kwargs) -> Any:
-    with component.context(**current_context):
-      return func(*args, **kwargs)
-
-  return _func
-
-
 class RetryError(RuntimeError):
   """Retry error."""
 
@@ -144,43 +131,33 @@ def with_retry(
     A function with the same signature of the input function, with the retry
     capability.
   """
-  rand = random if seed is None else random.Random(seed)
-
-  def _func(*args, **kwargs) -> Any:
-    def base_interval() -> int:
-      if isinstance(retry_interval, tuple):
-        return rand.randint(retry_interval[0], retry_interval[1])
-      else:
-        assert isinstance(retry_interval, int)
-        return retry_interval
 
-    def next_wait_interval(attempt: int) -> float:
-      if not exponential_backoff:
-        attempt = 1
-      return min(max_retry_interval, base_interval() * (2 ** (attempt - 1)))
-
-    wait_intervals = []
-    errors = []
-    while True:
-      with pg.catch_errors(retry_on_errors) as error_context:
-        return func(*args, **kwargs)
+  def _func(*args, **kwargs):
+    job = Job(
+        func,
+        args,
+        kwargs,
+        retry_on_errors=retry_on_errors,
+        max_attempts=max_attempts,
+        retry_interval=retry_interval,
+        exponential_backoff=exponential_backoff,
+        max_retry_interval=max_retry_interval,
+        seed=seed,
+    )
+    job()
+    if job.error:
+      raise job.error
+    return job.result
 
-      # Branch when errors are met for retry.
-      errors.append(error_context.error)
-      if len(errors) < max_attempts:
-        wait_interval = next_wait_interval(len(errors))
-        wait_intervals.append(wait_interval)
+  return _func
 
-        pg.logging.warning(
-            f'Calling {func!r} encountered {error_context.error!r} '
-            f'(attempts={len(errors)}), retrying in {wait_interval} seconds...'
-        )
 
-        time.sleep(wait_interval)
-      else:
-        raise RetryError(func, errors, wait_intervals)
+class RetryEntry(pg.Object):
+  """Retry entry."""
 
-  return _func
+  call_interval: float
+  error: BaseException | None = None
+  wait_interval: float = 0.
 
 
 def concurrent_execute(
@@ -198,6 +175,7 @@ def concurrent_execute(
     retry_interval: int | tuple[int, int] = (5, 60),
     exponential_backoff: bool = True,
     max_retry_interval: int = 300,
+    return_jobs: bool = False,
 ) -> list[Any]:
   """Executes a function concurrently under current component context.
 
@@ -221,32 +199,53 @@ def concurrent_execute(
     max_retry_interval: The max retry interval in seconds. This is useful when
       the retry interval is exponential, to avoid the wait time to grow
       exponentially.
+    return_jobs: If True, return a list of `Job` objects. Otherwise, return a
+      list of outputs.
 
   Returns:
     A list of ouputs. Each is the return value of `func` based on the input
       value. Order is preserved.
   """
-  if retry_on_errors is not None:
-    func = with_retry(
-        func,
-        retry_on_errors,
-        max_attempts=max_attempts,
-        retry_interval=retry_interval,
-        exponential_backoff=exponential_backoff,
-        max_retry_interval=max_retry_interval,
+  jobs = []
+  for inputs in parallel_inputs:
+    jobs.append(
+        Job(
+            func,
+            (inputs,),
+            retry_on_errors=retry_on_errors,
+            max_attempts=max_attempts,
+            retry_interval=retry_interval,
+            exponential_backoff=exponential_backoff,
+            max_retry_interval=max_retry_interval,
+        )
     )
 
   # NOTE(daiyip): when executor is not specified and max_worker is 1,
   # we don't need to create a executor pool. Instead, the inputs will be
   # processed by the user function in sequence within the current thread.
   if executor is None and max_workers == 1:
-    return [func(i) for i in parallel_inputs]
+    for job in jobs:
+      job()
+      if job.error:
+        raise job.error
+    return jobs if return_jobs else [job.result for job in jobs]
 
   shutdown_after_finish = executor is None
   executor = _executor_pool.executor_from(executor, max_workers=max_workers)
 
   try:
-    return list(executor.map(with_context_access(func), parallel_inputs))
+    executed_jobs = list(
+        executor.map(
+            lambda job: job(),
+            [pg.with_contextual_override(job) for job in jobs]
+        )
+    )
+    for job in executed_jobs:
+      if job.error:
+        raise job.error
+    return (
+        executed_jobs if return_jobs else [job.result for job in executed_jobs]
+    )
   finally:
     if shutdown_after_finish:
       # Do not wait threads to finish if they are timed out.
@@ -258,9 +257,61 @@ class Job:
   """Thread pool job."""
 
   func: Callable[[Any], Any]
-  arg: Any
+  args: Sequence[Any] = ()
+  kwargs: Mapping[str, Any] = dataclasses.field(default_factory=dict)
+  _: dataclasses.KW_ONLY
+
   result: Any = pg.MISSING_VALUE
-  error: BaseException | None = None
+  error: Annotated[
+      BaseException | None,
+      'The non-retryable error encountered during the job execution.',
+  ] = None
+  retry_entries: Annotated[
+      Sequence[RetryEntry], 'Records of retry attempts.'
+  ] = dataclasses.field(default_factory=list)
+
+  retry_on_errors: Annotated[
+      Sequence[Type[BaseException] | str],
+      (
+          'A sequence of exception types or tuples of exception type and error '
+          'messages (described in regular expression) as the desired exception '
+          'types to retry.'
+      ),
+  ] = ()
+  max_attempts: Annotated[
+      int, 'Max number of attempts if an error to retry is encountered.'
+  ] = 5
+  retry_interval: Annotated[
+      int | tuple[int, int],
+      (
+          'The (base) retry interval in seconds. If a tuple, the retry '
+          'interval will be randomly chosen between the first and the second '
+          'element of the tuple.'
+      ),
+  ] = (5, 60)
+  exponential_backoff: Annotated[
+      bool,
+      (
+          'If True, exponential wait time will be applied on top of the base '
+          'retry interval.'
+      ),
+  ] = True
+  max_retry_interval: Annotated[
+      int,
+      (
+          'The max retry interval in seconds. This is useful when the retry '
+          'interval is exponential, to avoid the wait time to grow '
+          'exponentially.'
+      ),
+  ] = 300
+  seed: Annotated[
+      int | None,
+      (
+          'Random seed to generate retry interval. If None, the seed will be'
+          ' determined based on current time.'
+      ),
+  ] = None
+
   timeit: pg.object_utils.TimeIt = dataclasses.field(
       default_factory=lambda: pg.object_utils.TimeIt('job')
   )
@@ -270,14 +321,70 @@ class Job:
     """Returns the running time in seconds since the job get started."""
     return self.timeit.elapse
 
-  def __call__(self) -> Any:
+  def _retry_call(self) -> 'Job':
+    """Retries func call on args."""
+    rand = random if self.seed is None else random.Random(self.seed)
+
+    def base_interval() -> int:
+      if isinstance(self.retry_interval, tuple):
+        return rand.randint(*self.retry_interval)
+      else:
+        assert isinstance(self.retry_interval, int)
+        return self.retry_interval
+
+    def next_wait_interval(attempt: int) -> float:
+      if not self.exponential_backoff:
+        attempt = 1
+      return min(
+          self.max_retry_interval, base_interval() * (2 ** (attempt - 1))
+      )
+
+    retry_entries = []
+    wait_interval = 0
+    while True:
+      with pg.catch_errors(self.retry_on_errors) as error_context:
+        begin_time = time.time()
+        self.result = self.func(*self.args, **self.kwargs)
+
+      end_time = time.time()
+      retry_entries.append(RetryEntry(
+          call_interval=end_time - begin_time,
+          wait_interval=wait_interval,
+          error=error_context.error,
+      ))
+      if error_context.error is None:
+        self.retry_entries = retry_entries
+        return self
+
+      # Branch when errors are met for retry.
+      if len(retry_entries) < self.max_attempts:
+        wait_interval = next_wait_interval(len(retry_entries))
+
+        pg.logging.warning(
+            f'Calling {self.func!r} encountered {error_context.error!r} '
+            f'(attempts={len(retry_entries)}), retrying in '
+            f'{wait_interval} seconds...'
+        )
+
+        time.sleep(wait_interval)
+      else:
+        errors = [e.error for e in retry_entries]
+        # First wait interval is 0.
+        wait_intervals = [e.wait_interval for e in retry_entries[1:]]
+        raise RetryError(self.func, errors, wait_intervals)
+
+  def __call__(self) -> 'Job':
+    if getattr(self, '_has_call', False):
+      raise ValueError('Job can only be called once.')
+    self._has_call = True
     try:
       with self.timeit:
-        self.result = self.func(self.arg)
-        return self.result
+        if self.retry_on_errors:
+          return self._retry_call()
+        self.result = self.func(*self.args, **self.kwargs)
     except BaseException as e:  # pylint: disable=broad-exception-caught
       self.error = e
-      return e
+    return self
 
   def mark_canceled(self, error: BaseException) -> None:
     """Marks the job as canceled."""
@@ -538,7 +645,8 @@ def concurrent_map(
     max_attempts: int = 5,
     retry_interval: int | tuple[int, int] = (5, 60),
     exponential_backoff: bool = True,
-) -> Iterator[tuple[Any, Any, BaseException | None]]:
+    return_jobs: bool = False,
+) -> Iterator[Any]:
   """Maps inputs to outptus via func concurrently under current context.
 
   Args:
@@ -581,9 +689,10 @@ def concurrent_map(
       of the tuple.
     exponential_backoff: If True, exponential wait time will be applied on top
       of the base retry interval.
+    return_jobs: If True, the returned iterator will emit `Job` objects.
 
   Yields:
-    An iterator of (input, output, error).
+    An iterator of (input, output, error) or Job object.
 
   Raises:
     Exception: Errors that are not in `silence_on_errors` or `retry_on_errors`,
@@ -593,15 +702,6 @@ def concurrent_map(
   """
     # Internal usage logging.
 
-  if retry_on_errors:
-    func = with_retry(
-        func,
-        retry_on_errors,
-        max_attempts=max_attempts,
-        retry_interval=retry_interval,
-        exponential_backoff=exponential_backoff,
-    )
-
   status_fn = status_fn or (lambda p: {   # pylint: disable=g-long-lambda
       'Succeeded': '%.2f%% (%d/%d)' % (
           p.success_rate * 100, p.succeeded, p.completed),
@@ -616,10 +716,15 @@ def concurrent_map(
   pending_futures = []
   total = 0
   for inputs in parallel_inputs:
-    job = Job(func, inputs)
-    future = executor.submit(
-        with_context_access(job),
+    job = Job(
+        func,
+        (inputs,),
+        retry_on_errors=retry_on_errors,
+        max_attempts=max_attempts,
+        retry_interval=retry_interval,
+        exponential_backoff=exponential_backoff,
     )
+    future = executor.submit(pg.with_contextual_override(job))
     pending_futures.append(future)
     future_to_job[future] = job
     total += 1
@@ -669,7 +774,7 @@ def concurrent_map(
                 silence_on_errors and isinstance(job.error, silence_on_errors)):
               raise job.error   # pylint: disable=g-doc-exception
 
-            yield job.arg, job.result, job.error
+            yield job if return_jobs else job.args[0], job.result, job.error
             progress.update(job)
             update_progress_bar(progress)
             ProgressBar.refresh()
@@ -690,7 +795,7 @@ def concurrent_map(
             if job.error is not None and not (
                 silence_on_errors and isinstance(job.error, silence_on_errors)):
               raise job.error   # pylint: disable=g-doc-exception
-            yield job.arg, job.result, job.error
+            yield job if return_jobs else job.args[0], job.result, job.error
             progress.update(job)
             update_progress_bar(progress)
             completed_batch.add(future)
@@ -713,7 +818,7 @@ def concurrent_map(
                       and isinstance(job.error, silence_on_errors)):
                 raise job.error  # pylint: disable=g-doc-exception
 
-            yield job.arg, job.result, job.error
+            yield job.args[0], job.result, job.error
             progress.update(job)
             update_progress_bar(progress)
           else:
@@ -844,10 +949,10 @@ class _ConsoleProgressControl(_ProgressControl):
   def refresh(self):
     s = io.StringIO()
     if self.label is not None:
-      s.write(text_formatting.colored(self.label, 'red', styles=['bold']))
+      s.write(pg.colored(self.label, 'red', styles=['bold']))
       s.write(': ')
     s.write(
-        text_formatting.colored(
+        pg.colored(
             '%d%% (%d/%d)' %
             (
                 self._progress * 100 // self.total,

langfun/core/concurrent_test.py

@@ -29,22 +29,6 @@ class A(component.Component):
   y: int = component.contextual()
 
 
-class WithContextAccessTest(unittest.TestCase):
-
-  def test_context_access(self):
-    inputs = [A(1), A(2)]
-    with futures.ThreadPoolExecutor() as executor:
-      with component.context(y=3):
-        self.assertEqual(
-            list(
-                executor.map(
-                    concurrent.with_context_access(lambda x: x.y), inputs
-                )
-            ),
-            [3, 3],
-        )
-
-
 class RetryErrorTest(unittest.TestCase):
 
   def test_basics(self):
@@ -94,7 +78,7 @@ class RetryErrorTest(unittest.TestCase):
     )
 
 
-class WithRetryTest(unittest.TestCase):
+class RetryTest(unittest.TestCase):
 
   def assert_retry(self, func, expected_attempts, expected_wait_intervals):
     with pg.catch_errors(concurrent.RetryError) as error_context:
@@ -162,6 +146,31 @@ class WithRetryTest(unittest.TestCase):
     with self.assertRaises(ValueError):
       foo_with_retry()
 
+  def test_retry_with_job(self):
+    count = 0
+
+    def foo():
+      nonlocal count
+      count += 1
+      if count < 3:
+        raise ValueError('Foo temporary error.')
+      return 'Success'
+
+    job = concurrent.Job(
+        foo,
+        retry_on_errors=ValueError,
+        retry_interval=1,
+    )
+    job()
+    self.assertEqual(job.result, 'Success')
+    self.assertEqual(
+        [retry_entry.wait_interval for retry_entry in job.retry_entries],
+        [0, 1, 2],
+    )
+    self.assertIsInstance(job.retry_entries[0].error, ValueError)
+    self.assertIsInstance(job.retry_entries[1].error, ValueError)
+    self.assertIsNone(job.retry_entries[2].error)
+
 
 class ConcurrentExecuteTest(unittest.TestCase):
 
@@ -217,8 +226,8 @@ class ProgressTest(unittest.TestCase):
     def fun2(unused_x):
       raise ValueError('Intentional error.')
 
-    job1 = concurrent.Job(fun, 1)
-    job2 = concurrent.Job(fun2, 2)
+    job1 = concurrent.Job(fun, (1,))
+    job2 = concurrent.Job(fun2, (2,))
     job1()
     job2()
 

langfun/core/console.py

@@ -15,7 +15,7 @@
 
 import sys
 from typing import Any
-from langfun.core.text_formatting import colored
+import pyglove as pg
 
 
 def write(
@@ -42,10 +42,14 @@ def write(
   """
   # Print title if present.
   if title is not None:
-    print(colored(title, styles=['bold']))
+    print(pg.colored(title, styles=['bold']))
 
   # Print body.
-  print(colored(str(value), color=color, background=background, styles=styles))
+  print(
+      pg.colored(
+          str(value), color=color, background=background, styles=styles
+      )
+  )
 
 
 try:

langfun/core/eval/base.py

@@ -1298,7 +1298,7 @@ class Evaluation(Evaluable):
             id=self.id,
             dir=self.dir,
             model=self.lm.model_id,
-            prompt_template=lf.text_formatting.decolored(str(self.prompt)),
+            prompt_template=pg.decolor(str(self.prompt)),
             method=self.method,
             schema_fn=str(self.schema_fn),
         ),
@@ -2110,8 +2110,7 @@ class Summary(pg.Object):
 
 def _format_error(error: Exception):
   """Formats an error into a string."""
-  return (f'({error.__class__.__name__}) '
-          + lf.text_formatting.decolored(str(error)))
+  return (f'({error.__class__.__name__}) ' + pg.decolor(str(error)))
 
 
 def _error_key(error: Exception) -> str:

langfun/core/eval/v2/evaluation.py

@@ -516,7 +516,9 @@ class Evaluation(experiment_lib.Experiment):
                       target='example-view',
                       css_classes=['example-link'],
                   )
-                  for dp in metric_value.data_points
+                  for dp in sorted(
+                      metric_value.data_points, key=lambda dp: dp.example_id
+                  )
               ]
           )
       )

langfun/core/eval/v2/reporting.py

@@ -51,6 +51,7 @@ class HtmlReporter(experiment_lib.Plugin):
     self._update_thread = None
     self._stop_update = False
     self._stop_update_experiment_ids = set()
+    self._summary_lock = None
     self._experiment_index_lock = None
 
   def on_run_start(
@@ -62,6 +63,7 @@ class HtmlReporter(experiment_lib.Plugin):
     self._last_experiment_report_time = {leaf.id: 0 for leaf in root.leaf_nodes}
     self._stop_update = False
     self._stop_update_experiment_ids = set()
+    self._summary_lock = threading.Lock()
     self._experiment_index_lock = {
         leaf.id: threading.Lock() for leaf in root.leaf_nodes
     }
@@ -137,21 +139,23 @@ class HtmlReporter(experiment_lib.Plugin):
     """Maybe update the summary of current run."""
     run = runner.current_run
     def _summary():
-      run.experiment.to_html(
+      html = run.experiment.to_html(
           collapse_level=None,
           extra_flags=dict(
               current_run=run, interactive=False, card_view=True,
           )
-      ).save(
-          run.output_path_for(run.experiment, _SUMMARY_FILE)
       )
+      with self._summary_lock:
+        html.save(
+            run.output_path_for(run.experiment, _SUMMARY_FILE)
+        )
 
     if force or (time.time() - self._last_summary_time > self.summary_interval):
+      self._last_summary_time = time.time()
       if background:
         runner.background_run(_summary)
       else:
         _summary()
-      self._last_summary_time = time.time()
 
   def _maybe_update_experiment_html(
       self,