From 38912bdd6043cc19d7871b4c4854ccfca10eec04 Mon Sep 17 00:00:00 2001
From: Ori Kabeli <ori@futurehouse.org>
Date: Tue, 18 Feb 2025 12:57:00 -0600
Subject: [PATCH 01/21] Rollout manager with exception counter & chunking fix;
 .gitignore update

---
 .gitignore                                    |  2 +
 ldp/alg/rollout.py                            | 53 ++++++++++++++++---
 ldp/graph/async_torch.py                      |  2 +-
 ldp/nn/__init__.py                            |  3 ++
 ....jinja => llama3_chat_template_test.jinja} |  0
 ldp/nn/handlers/chunking.py                   |  5 +-
 6 files changed, 55 insertions(+), 10 deletions(-)
 rename ldp/nn/chat_templates/{llama3_chat_template_ori.jinja => llama3_chat_template_test.jinja} (100%)

diff --git a/.gitignore b/.gitignore
index 01a52fbd..57d27823 100644
--- a/.gitignore
+++ b/.gitignore
@@ -297,3 +297,5 @@ cython_debug/
 
 # Version files made by setuptools_scm
 **/version.py
+
+.vscode/
\ No newline at end of file
diff --git a/ldp/alg/rollout.py b/ldp/alg/rollout.py
index ff06f266..227982f8 100644
--- a/ldp/alg/rollout.py
+++ b/ldp/alg/rollout.py
@@ -2,6 +2,7 @@
 import itertools
 import logging
 import uuid
+from collections import Counter
 from collections.abc import Callable, Iterator, Sequence
 from contextlib import contextmanager, nullcontext
 from typing import Any, TypeVar, overload
@@ -24,6 +25,7 @@ class CaughtError(Exception):
     """Base class for reraised exceptions when catching is enabled."""
 
     def __init__(self, original_exc: Exception):
+        super().__init__(str(original_exc))
         self.original_exc = original_exc
 
     exc_type = "undefined"
@@ -39,12 +41,13 @@ class EnvError(CaughtError):
 
 @contextmanager
 def reraise_exc_as(reraise: type[CaughtError], enabled: bool) -> Iterator[None]:
+    """Context manager that reraises exceptions as a custom CaughtError type if enabled."""
     try:
         yield
     except Exception as e:
         if enabled:
-            error_details = format_error_details(e)
-            logger.exception(f"Caught {reraise.exc_type} exception:\n{error_details}")
+            # Minimal logging instead of spamming. Detailed error stored in the trajectory's metadata.
+            logger.debug(f"Reraising {reraise.exc_type} exception.")
             raise reraise(e) from None
         raise
 
@@ -193,14 +196,50 @@ async def _sample_trajectories_from_envs(
         max_steps: int | None = None,
     ) -> list[Trajectory]:
         self.traj_buffer.clear()
+        exception_counter = Counter()
 
-        traj_ids = [uuid.uuid4().hex for _ in range(len(environments))]
-        await asyncio.gather(
-            *(
-                self._rollout(*args, max_steps=max_steps)
-                for args in zip(traj_ids, environments, strict=True)
+        traj_ids = [uuid.uuid4().hex for _ in environments]
+
+        # Create all tasks first
+        tasks = [
+            asyncio.create_task(
+                self._rollout(traj_id, env, max_steps=max_steps)
             )
+            for traj_id, env in zip(traj_ids, environments, strict=True)
+        ]
+
+        # Use a single line bar_format to avoid multiline spam.
+        from tqdm import tqdm
+        bar_format = (
+            "{l_bar}{bar} {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}]"
+            " {postfix}"
         )
+
+        with tqdm(
+            total=len(tasks),
+            desc="Rollouts",
+            unit="rollout",
+            bar_format=bar_format,
+        ) as pbar:
+            for task in asyncio.as_completed(tasks):
+                trajectory = await task
+                pbar.update(1)
+                # Check if this trajectory ended with an exception
+                if trajectory.steps:
+                    last_step = trajectory.steps[-1]
+                    if last_step.metadata.get("exception"):
+                        # We'll keep it short but still have something to categorize
+                        exc_str = last_step.metadata["exception"][:500].replace('"', "'")
+                        exception_counter[exc_str] += 1
+                        num_exceptions = sum(exception_counter.values())
+                        pbar.set_postfix({"num_exceptions": num_exceptions})
+
+        # Final summary of exceptions (if any)
+        if exception_counter:
+            logger.info("Caught exceptions:")
+            logger.info("{:<6} {:<50}".format("Count", "Exception"))
+            for exc, count in exception_counter.items():
+                logger.info("{:<6} {:<50}".format(count, exc))
         return [self.traj_buffer[traj_id] for traj_id in traj_ids]
 
     async def _rollout(
diff --git a/ldp/graph/async_torch.py b/ldp/graph/async_torch.py
index 55ef0e9b..c614ab0b 100644
--- a/ldp/graph/async_torch.py
+++ b/ldp/graph/async_torch.py
@@ -127,7 +127,7 @@ async def _maybe_process_batch(self) -> None:
 
         if (
             len(self._work_buffer) >= self.batch_size
-            or now - self._work_buffer[0][0] > self.timeout
+            or (now - self._work_buffer[0][0] > self.timeout) and len(self._work_buffer) > 0
         ):
             # if we're over batch size or have at least one input waiting for
             # more than timeout, pull out a batch to run
diff --git a/ldp/nn/__init__.py b/ldp/nn/__init__.py
index 46e6dcd0..73b300bf 100644
--- a/ldp/nn/__init__.py
+++ b/ldp/nn/__init__.py
@@ -1,6 +1,7 @@
 from .agent.simple_local_agent import AgentLMConfig, SimpleLocalLLMAgent
 from .graph.llm_call_op import LocalLLMCallOp
 from .handlers.chunking import TensorChunker
+from .handlers.module_handler import AsyncModuleHandler, ModuleExecutionInterface
 from .handlers.transformer_handler import (
     AsyncTransformer,
     AsyncTransformerInterface,
@@ -20,12 +21,14 @@
 
 __all__ = [
     "AgentLMConfig",
+    "AsyncModuleHandler",
     "AsyncTransformer",
     "AsyncTransformerInterface",
     "ExecutionMode",
     "LMConfig",
     "LMType",
     "LocalLLMCallOp",
+    "ModuleExecutionInterface",
     "ParallelAsyncTransformer",
     "ParallelModeConfig",
     "ParallelTransformerHandler",
diff --git a/ldp/nn/chat_templates/llama3_chat_template_ori.jinja b/ldp/nn/chat_templates/llama3_chat_template_test.jinja
similarity index 100%
rename from ldp/nn/chat_templates/llama3_chat_template_ori.jinja
rename to ldp/nn/chat_templates/llama3_chat_template_test.jinja
diff --git a/ldp/nn/handlers/chunking.py b/ldp/nn/handlers/chunking.py
index 38fdbe96..65771b5a 100644
--- a/ldp/nn/handlers/chunking.py
+++ b/ldp/nn/handlers/chunking.py
@@ -159,8 +159,9 @@ def _split_value(self, value):
             for i in range(self.num_chunks):
                 if i >= len(chunks):
                     # Chunk 0 will always exist, and we need only a batch of one ([:1])
-                    # to activate the model
-                    chunks.append(torch.full_like(chunks[0][:1], self.dummy_value))
+                    # to activate the model. 
+                    # We use real data to avoid errors in the model expecting certain token structure.
+                    chunks.append(chunks[0][:1])
                     dummy_chunk_flags.append(True)
                 else:
                     dummy_chunk_flags.append(False)

From 13c962224ad7147310b718d912503aebd9ea4290 Mon Sep 17 00:00:00 2001
From: Ori Kabeli <ori@futurehouse.org>
Date: Tue, 18 Feb 2025 14:13:45 -0600
Subject: [PATCH 02/21] nits

---
 .gitignore                             |  2 +-
 ldp/alg/rollout.py                     | 18 ++++++++----------
 ldp/graph/async_torch.py               | 12 +++++++++---
 ldp/nn/handlers/chunking.py            |  2 +-
 ldp/nn/handlers/transformer_handler.py | 15 +++++++++++++++
 5 files changed, 34 insertions(+), 15 deletions(-)

diff --git a/.gitignore b/.gitignore
index 57d27823..5293aadb 100644
--- a/.gitignore
+++ b/.gitignore
@@ -298,4 +298,4 @@ cython_debug/
 # Version files made by setuptools_scm
 **/version.py
 
-.vscode/
\ No newline at end of file
+.vscode/
diff --git a/ldp/alg/rollout.py b/ldp/alg/rollout.py
index 227982f8..a0c2a883 100644
--- a/ldp/alg/rollout.py
+++ b/ldp/alg/rollout.py
@@ -8,10 +8,10 @@
 from typing import Any, TypeVar, overload
 
 from aviary.core import Environment, Message
+from tqdm import tqdm
 
 from ldp.agent import Agent
 from ldp.data_structures import Trajectory, Transition
-from ldp.utils import format_error_details
 
 from .callbacks import Callback
 
@@ -196,20 +196,16 @@ async def _sample_trajectories_from_envs(
         max_steps: int | None = None,
     ) -> list[Trajectory]:
         self.traj_buffer.clear()
-        exception_counter = Counter()
+        exception_counter: Counter = Counter()
 
         traj_ids = [uuid.uuid4().hex for _ in environments]
 
         # Create all tasks first
         tasks = [
-            asyncio.create_task(
-                self._rollout(traj_id, env, max_steps=max_steps)
-            )
+            asyncio.create_task(self._rollout(traj_id, env, max_steps=max_steps))
             for traj_id, env in zip(traj_ids, environments, strict=True)
         ]
 
-        # Use a single line bar_format to avoid multiline spam.
-        from tqdm import tqdm
         bar_format = (
             "{l_bar}{bar} {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}]"
             " {postfix}"
@@ -229,7 +225,8 @@ async def _sample_trajectories_from_envs(
                     last_step = trajectory.steps[-1]
                     if last_step.metadata.get("exception"):
                         # We'll keep it short but still have something to categorize
-                        exc_str = last_step.metadata["exception"][:500].replace('"', "'")
+                        exc_str: str = last_step.metadata["exception"][:500]
+                        exc_str = exc_str.replace('"', "'")
                         exception_counter[exc_str] += 1
                         num_exceptions = sum(exception_counter.values())
                         pbar.set_postfix({"num_exceptions": num_exceptions})
@@ -237,9 +234,10 @@ async def _sample_trajectories_from_envs(
         # Final summary of exceptions (if any)
         if exception_counter:
             logger.info("Caught exceptions:")
-            logger.info("{:<6} {:<50}".format("Count", "Exception"))
+            logger.info("%-6s %-50s", "Count", "Exception")
             for exc, count in exception_counter.items():
-                logger.info("{:<6} {:<50}".format(count, exc))
+                logger.info("%-6d %-50s", count, exc)
+
         return [self.traj_buffer[traj_id] for traj_id in traj_ids]
 
     async def _rollout(
diff --git a/ldp/graph/async_torch.py b/ldp/graph/async_torch.py
index c614ab0b..d9adea65 100644
--- a/ldp/graph/async_torch.py
+++ b/ldp/graph/async_torch.py
@@ -120,14 +120,20 @@ async def _maybe_process_batch(self) -> None:
 
         If neither condition is met, do nothing.
         """
+        # Technically should not happen, but if a coroutine crashes, it could release
+        # self._lock before placing results in _results_buffer and additional process
+        # coming inside will crash.
+        if not self._work_buffer:
+            return
+
         now = time.time()
 
         # sort by oldest requests first
         self._work_buffer.sort(key=operator.itemgetter(0))
 
-        if (
-            len(self._work_buffer) >= self.batch_size
-            or (now - self._work_buffer[0][0] > self.timeout) and len(self._work_buffer) > 0
+        if len(self._work_buffer) >= self.batch_size or (
+            (now - self._work_buffer[0][0] > self.timeout)
+            and len(self._work_buffer) > 0
         ):
             # if we're over batch size or have at least one input waiting for
             # more than timeout, pull out a batch to run
diff --git a/ldp/nn/handlers/chunking.py b/ldp/nn/handlers/chunking.py
index 65771b5a..fe34da68 100644
--- a/ldp/nn/handlers/chunking.py
+++ b/ldp/nn/handlers/chunking.py
@@ -159,7 +159,7 @@ def _split_value(self, value):
             for i in range(self.num_chunks):
                 if i >= len(chunks):
                     # Chunk 0 will always exist, and we need only a batch of one ([:1])
-                    # to activate the model. 
+                    # to activate the model.
                     # We use real data to avoid errors in the model expecting certain token structure.
                     chunks.append(chunks[0][:1])
                     dummy_chunk_flags.append(True)
diff --git a/ldp/nn/handlers/transformer_handler.py b/ldp/nn/handlers/transformer_handler.py
index 8c53ede5..0bb84378 100644
--- a/ldp/nn/handlers/transformer_handler.py
+++ b/ldp/nn/handlers/transformer_handler.py
@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import atexit
 import logging
 import os
 import socket
@@ -193,6 +194,14 @@ async def __call__(  # type: ignore[override]
     @staticmethod
     def model_generate(model: PreTrainedModel, *args, **kwargs):
         """A method that can be used as module_call_fn to sample from an LLM."""
+        if dist.get_world_size() > 1:
+            synced_gpus = kwargs.pop("synced_gpus", None)
+            if synced_gpus is None:
+                logger.debug("synced_gpus not defined, defaulting to True.")
+            elif not synced_gpus:
+                raise ValueError("synced_gpus must be True when using FSDP.")
+            kwargs["synced_gpus"] = True
+
         # Summoning params per https://github.com/pytorch/pytorch/issues/100069
         # If model is not FSDP, this context manager is a no-op.
         with FullyShardedDataParallel.summon_full_params(model, recurse=False):
@@ -463,6 +472,8 @@ def __init__(self, config: TransformerHandlerConfig):
 
         self._initialized = True
 
+        atexit.register(self.teardown)
+
         # don't call AsyncTorchModule.__init__ because we don't need to set up module[_call_fn]
         AsyncBufferedWorker.__init__(
             self,
@@ -484,6 +495,10 @@ def _init_local_cluster(
         # lazy import since dask-cuda only works on Linux machines
         from dask_cuda import LocalCUDACluster
 
+        # This uses NVIDIA's NVML layer instead of native CUDA, which is more robust in GPU detection
+        # post initialization. This prevents issues with forked processes wrongly detecting the
+        # default GPU as cuda:0
+        os.environ["PYTORCH_NVML_BASED_CUDA_CHECK"] = "1"
         self.cluster = LocalCUDACluster(
             n_workers=parallel_mode_config.num_workers,
             threads_per_worker=parallel_mode_config.num_cpus_per_worker,

From a7c1bf2736790536f38a42a1cbdd31d1dde24b2e Mon Sep 17 00:00:00 2001
From: Ori Kabeli <ori@futurehouse.org>
Date: Tue, 18 Feb 2025 14:15:27 -0600
Subject: [PATCH 03/21] nits

---
 ldp/alg/rollout.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/ldp/alg/rollout.py b/ldp/alg/rollout.py
index a0c2a883..975de0db 100644
--- a/ldp/alg/rollout.py
+++ b/ldp/alg/rollout.py
@@ -225,7 +225,7 @@ async def _sample_trajectories_from_envs(
                     last_step = trajectory.steps[-1]
                     if last_step.metadata.get("exception"):
                         # We'll keep it short but still have something to categorize
-                        exc_str: str = last_step.metadata["exception"][:500]
+                        exc_str: str = str(last_step.metadata["exception"])[:500]
                         exc_str = exc_str.replace('"', "'")
                         exception_counter[exc_str] += 1
                         num_exceptions = sum(exception_counter.values())

From ae71669f05f5cfb0f6c9e3b5920cfd05e8f1fc03 Mon Sep 17 00:00:00 2001
From: Ori Kabeli <ori@futurehouse.org>
Date: Sun, 23 Feb 2025 08:47:29 -0600
Subject: [PATCH 04/21] nits

---
 ldp/alg/rollout.py                     |  1 +
 ldp/graph/async_torch.py               |  1 -
 ldp/nn/agent/simple_local_agent.py     | 20 ++++++++++++++++++++
 ldp/nn/handlers/transformer_handler.py |  9 ++++++++-
 4 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/ldp/alg/rollout.py b/ldp/alg/rollout.py
index 975de0db..31d9f9e0 100644
--- a/ldp/alg/rollout.py
+++ b/ldp/alg/rollout.py
@@ -245,6 +245,7 @@ async def _rollout(
         traj_id: str,
         env: Environment,
         max_steps: int | None,
+        max_tokens: int | None = None,  # <-- new argument
     ) -> Trajectory:
         trajectory = Trajectory(traj_id=traj_id)
 
diff --git a/ldp/graph/async_torch.py b/ldp/graph/async_torch.py
index d9adea65..f0b2f2df 100644
--- a/ldp/graph/async_torch.py
+++ b/ldp/graph/async_torch.py
@@ -133,7 +133,6 @@ async def _maybe_process_batch(self) -> None:
 
         if len(self._work_buffer) >= self.batch_size or (
             (now - self._work_buffer[0][0] > self.timeout)
-            and len(self._work_buffer) > 0
         ):
             # if we're over batch size or have at least one input waiting for
             # more than timeout, pull out a batch to run
diff --git a/ldp/nn/agent/simple_local_agent.py b/ldp/nn/agent/simple_local_agent.py
index 9954b03e..f59fbf3a 100644
--- a/ldp/nn/agent/simple_local_agent.py
+++ b/ldp/nn/agent/simple_local_agent.py
@@ -2,6 +2,7 @@
 
 import torch
 import torch.distributed as dist
+from litellm import token_counter
 from aviary.core import Message, Tool, ToolRequestMessage
 from pydantic import Field, field_validator
 
@@ -41,6 +42,11 @@ class AgentLMConfig(_LMConfig):
         "are better defaults than HF's.",
         validate_default=True,
     )
+    
+    max_traj_token_count: int | None = Field(
+        default=None,
+        description="If set, raise an error if the total tokens in the trajectory exceed this value."
+    )
 
     @field_validator("llm_call_kwargs")
     @classmethod
@@ -110,6 +116,20 @@ async def get_asv(
 
         # Update state messages with result and return the new state
         next_state.messages = [*next_state.messages, result.value]
+        
+          
+        import ipdb; ipdb.set_trace()              
+        if self.llm_model.max_traj_token_count is not None:
+            total_tokens = token_counter(
+                model=self.llm_model.llm_for_sft,  # or any field referencing the model name
+                messages=next_state.messages,
+                tools=next_state.tools,
+            )
+            if total_tokens > self.llm_model.max_traj_token_count:
+                raise ValueError(
+                    f"Token limit exceeded for trajectory: {total_tokens} > {self.llm_model.max_traj_token_count}"
+                )
+        
         return cast(OpResult[ToolRequestMessage], result), next_state, 0.0
 
     # TODO: maybe remove these recomputation methods. I added them to debug some things. But idk,
diff --git a/ldp/nn/handlers/transformer_handler.py b/ldp/nn/handlers/transformer_handler.py
index 0bb84378..47f5253f 100644
--- a/ldp/nn/handlers/transformer_handler.py
+++ b/ldp/nn/handlers/transformer_handler.py
@@ -47,6 +47,7 @@
 else:
     from typing_extensions import overload  # noqa: UP035
 
+logger = logging.getLogger(__name__)
 
 config.set({
     # We have no use for rebooting workers in aviary for now, and rebooting workers
@@ -60,7 +61,10 @@
     "distributed.comm.timeouts.tcp": "300s",
 })
 
-logger = logging.getLogger(__name__)
+compression = os.getenv("USE_DASK_COMPRESSION")
+if compression is not None:
+    config.set({"distributed.comm.compression": compression})
+    logger.info(f"Setting Dask compression to {compression}")
 
 TReturn = TypeVar("TReturn")
 TParams = ParamSpec("TParams")
@@ -201,6 +205,9 @@ def model_generate(model: PreTrainedModel, *args, **kwargs):
             elif not synced_gpus:
                 raise ValueError("synced_gpus must be True when using FSDP.")
             kwargs["synced_gpus"] = True
+            if os.getenv("USE_DASK_BARRIER"):
+                logger.info("Waiting for all workers to reach this point.")
+                dist.barrier()
 
         # Summoning params per https://github.com/pytorch/pytorch/issues/100069
         # If model is not FSDP, this context manager is a no-op.

From 39b0fc05b83c5bcb1bace29a37e76d0dd5d59c12 Mon Sep 17 00:00:00 2001
From: Ori Kabeli <ori@futurehouse.org>
Date: Sun, 23 Feb 2025 09:17:04 -0600
Subject: [PATCH 05/21] nits

---
 ldp/alg/rollout.py          | 10 ++--------
 ldp/nn/handlers/chunking.py |  3 ++-
 2 files changed, 4 insertions(+), 9 deletions(-)

diff --git a/ldp/alg/rollout.py b/ldp/alg/rollout.py
index 31d9f9e0..dc9acb5f 100644
--- a/ldp/alg/rollout.py
+++ b/ldp/alg/rollout.py
@@ -206,16 +206,11 @@ async def _sample_trajectories_from_envs(
             for traj_id, env in zip(traj_ids, environments, strict=True)
         ]
 
-        bar_format = (
-            "{l_bar}{bar} {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}]"
-            " {postfix}"
-        )
-
         with tqdm(
             total=len(tasks),
             desc="Rollouts",
             unit="rollout",
-            bar_format=bar_format,
+            ncols=0,
         ) as pbar:
             for task in asyncio.as_completed(tasks):
                 trajectory = await task
@@ -225,8 +220,7 @@ async def _sample_trajectories_from_envs(
                     last_step = trajectory.steps[-1]
                     if last_step.metadata.get("exception"):
                         # We'll keep it short but still have something to categorize
-                        exc_str: str = str(last_step.metadata["exception"])[:500]
-                        exc_str = exc_str.replace('"', "'")
+                        exc_str: str = str(last_step.metadata["exception"])[:500].replace('"', "'")
                         exception_counter[exc_str] += 1
                         num_exceptions = sum(exception_counter.values())
                         pbar.set_postfix({"num_exceptions": num_exceptions})
diff --git a/ldp/nn/handlers/chunking.py b/ldp/nn/handlers/chunking.py
index fe34da68..147b6dd8 100644
--- a/ldp/nn/handlers/chunking.py
+++ b/ldp/nn/handlers/chunking.py
@@ -160,7 +160,8 @@ def _split_value(self, value):
                 if i >= len(chunks):
                     # Chunk 0 will always exist, and we need only a batch of one ([:1])
                     # to activate the model.
-                    # We use real data to avoid errors in the model expecting certain token structure.
+                    # We use the first element of the existing chunks as real data to avoid
+                    # errors in the model that may expect a specific token structure.
                     chunks.append(chunks[0][:1])
                     dummy_chunk_flags.append(True)
                 else:

From da92fbf9b87f89b09b2a270dfec7d0a463602e3e Mon Sep 17 00:00:00 2001
From: Ori Kabeli <ori@futurehouse.org>
Date: Sun, 23 Feb 2025 09:28:46 -0600
Subject: [PATCH 06/21] nits

---
 src/ldp/graph/async_torch.py                  |   7 +-
 .../llama3_chat_template_test.jinja           | 113 ------------------
 2 files changed, 4 insertions(+), 116 deletions(-)
 delete mode 100644 src/ldp/nn/chat_templates/llama3_chat_template_test.jinja

diff --git a/src/ldp/graph/async_torch.py b/src/ldp/graph/async_torch.py
index f0b2f2df..f6597bb9 100644
--- a/src/ldp/graph/async_torch.py
+++ b/src/ldp/graph/async_torch.py
@@ -122,7 +122,7 @@ async def _maybe_process_batch(self) -> None:
         """
         # Technically should not happen, but if a coroutine crashes, it could release
         # self._lock before placing results in _results_buffer and additional process
-        # coming inside will crash.
+        # coming inside this func will crash as self._work_buffer will be empty.
         if not self._work_buffer:
             return
 
@@ -131,8 +131,9 @@ async def _maybe_process_batch(self) -> None:
         # sort by oldest requests first
         self._work_buffer.sort(key=operator.itemgetter(0))
 
-        if len(self._work_buffer) >= self.batch_size or (
-            (now - self._work_buffer[0][0] > self.timeout)
+        if (
+            len(self._work_buffer) >= self.batch_size
+            or now - self._work_buffer[0][0] > self.timeout
         ):
             # if we're over batch size or have at least one input waiting for
             # more than timeout, pull out a batch to run
diff --git a/src/ldp/nn/chat_templates/llama3_chat_template_test.jinja b/src/ldp/nn/chat_templates/llama3_chat_template_test.jinja
deleted file mode 100644
index e2807a55..00000000
--- a/src/ldp/nn/chat_templates/llama3_chat_template_test.jinja
+++ /dev/null
@@ -1,113 +0,0 @@
-{{- bos_token }}
-{%- if custom_tools is defined %}
-    {%- set tools = custom_tools %}
-{%- endif %}
-{%- if not tools_in_user_message is defined %}
-    {%- set tools_in_user_message = true %}
-{%- endif %}
-{%- if not date_string is defined %}
-    {%- set date_string = "26 Jul 2024" %}
-{%- endif %}
-{%- if not tools is defined %}
-    {%- set tools = none %}
-{%- endif %}
-
-{#- This block extracts the system message, so we can slot it into the right place. #}
-{%- if messages[0]['role'] == 'system' %}
-    {%- set system_message = messages[0]['content']|trim %}
-    {%- set messages = messages[1:] %}
-{%- else %}
-    {%- set system_message = "" %}
-{%- endif %}
-
-{#- System message + builtin tools #}
-{{- "<|start_header_id|>system<|end_header_id|>\n\n" }}
-{%- if builtin_tools is defined or tools is not none %}
-    {{- "Environment: ipython\n" }}
-{%- endif %}
-{%- if builtin_tools is defined %}
-    {{- "Tools: " + builtin_tools | reject('equalto', 'code_interpreter') | join(", ") + "\n\n"}}
-{%- endif %}
-{{- "Cutting Knowledge Date: December 2023\n" }}
-{{- "Today Date: " + date_string + "\n\n" }}
-{%- if tools is not none and not tools_in_user_message %}
-    {{- "You have access to the following functions. To call a function, please respond with JSON for a function call." }}
-    {{- 'Respond in the format {"name": function name, "parameters": dictionary of argument name and its value,  "thought": optional succeint reasoning processes leading to calling this tool, }.' }}
-    {{- "Do not use variables.\n\n" }}
-    {%- for t in tools %}
-        {{- t | tojson(indent=4) }}
-        {{- "\n\n" }}
-    {%- endfor %}
-{%- endif %}
-{{- system_message }}
-{{- "<|eot_id|>" }}
-
-{#- Custom tools are passed in a user message with some extra guidance #}
-{%- if tools_in_user_message and not tools is none %}
-    {#- Extract the first user message so we can plug it in here #}
-    {%- if messages | length != 0 %}
-        {%- set first_user_message = messages[0]['content']|trim %}
-        {%- set messages = messages[1:] %}
-    {%- else %}
-        {{- raise_exception("Cannot put tools in the first user message when there's no first user message!") }}
-{%- endif %}
-    {{- '<|start_header_id|>user<|end_header_id|>\n\n' -}}
-    {{- "Given the following functions, please respond with a JSON for a function call " }}
-    {{- "with its proper arguments that best answers the given prompt.\n\n" }}
-    {{- 'Respond in the format {"name": function name, "parameters": dictionary of argument name and its value, "thought": succeint reasoning processes leading to calling this tool, }.' }}
-    {{- "Do not use variables.\n\n" }}
-    {%- for t in tools %}
-        {{- t | tojson(indent=4) }}
-        {{- "\n\n" }}
-    {%- endfor %}
-    {{- first_user_message + "<|eot_id|>"}}
-{%- endif %}
-
-{%- for message in messages %}
-    {%- if not (message.role == 'ipython' or message.role == 'tool' or 'tool_calls' in message) %}
-        {%- if message['role'] == 'assistant' %}
-            {% generation %}{{- '<|start_header_id|>' + message['role'] + '<|end_header_id|>\n\n' + message['content'] | trim + '<|eot_id|>' }}{% endgeneration %}
-        {%- else %}
-            {{- '<|start_header_id|>' + message['role'] + '<|end_header_id|>\n\n' + message['content'] | trim + '<|eot_id|>' }}
-        {%- endif %}
-    {%- elif 'tool_calls' in message %}
-        {%- if not message.tool_calls|length == 1 %}
-            {{- raise_exception("This model only supports single tool-calls at once!") }}
-        {%- endif %}
-        {%- set tool_call = message.tool_calls[0].function %}
-        {% generation %}{%- if builtin_tools is defined and tool_call.name in builtin_tools %}
-            {{- '<|start_header_id|>assistant<|end_header_id|>\n\n' -}}
-            {{- "<|python_tag|>" + tool_call.name + ".call(" }}
-            {%- for arg_name, arg_val in tool_call.arguments | items %}
-                {{- arg_name + '="' + arg_val + '"' }}
-                {%- if not loop.last %}
-                    {{- ", " }}
-                {%- endif %}
-            {%- endfor %}
-            {{- ")" }}
-        {%- else  %}
-            {{- '<|start_header_id|>assistant<|end_header_id|>\n\n' -}}
-            {{- '{"name": "' + tool_call.name + '", ' }}
-            {{- '"parameters": ' }}
-            {{- tool_call.arguments | tojson }}
-            {{- "}" }}
-        {%- endif %}
-        {%- if builtin_tools is defined %}
-            {#- This means we're in ipython mode #}
-            {{- "<|eom_id|>" }}
-        {%- else %}
-            {{- "<|eot_id|>" }}
-        {%- endif %}{% endgeneration %}
-    {%- elif message.role == "tool" or message.role == "ipython" %}
-        {{- "<|start_header_id|>ipython<|end_header_id|>\n\n" }}
-        {%- if message.content is mapping or message.content is iterable %}
-            {{- message.content | tojson }}
-        {%- else %}
-            {{- message.content }}
-        {%- endif %}
-        {{- "<|eot_id|>" }}
-    {%- endif %}
-{%- endfor %}
-{%- if add_generation_prompt %}
-    {% generation %}{{- '<|start_header_id|>assistant<|end_header_id|>\n\n' }}{% endgeneration %}
-{%- endif %}

From 2ba58981a2171a438e6511d690b0f146a53c1574 Mon Sep 17 00:00:00 2001
From: Ori Kabeli <ori@futurehouse.org>
Date: Mon, 24 Feb 2025 06:43:51 -0600
Subject: [PATCH 07/21] nits

---
 src/ldp/alg/rollout.py                     |  4 ++-
 src/ldp/nn/agent/simple_local_agent.py     | 31 +++++++++++++++-------
 src/ldp/nn/handlers/transformer_handler.py |  6 +----
 3 files changed, 25 insertions(+), 16 deletions(-)

diff --git a/src/ldp/alg/rollout.py b/src/ldp/alg/rollout.py
index dc9acb5f..2b1a1b1d 100644
--- a/src/ldp/alg/rollout.py
+++ b/src/ldp/alg/rollout.py
@@ -220,7 +220,9 @@ async def _sample_trajectories_from_envs(
                     last_step = trajectory.steps[-1]
                     if last_step.metadata.get("exception"):
                         # We'll keep it short but still have something to categorize
-                        exc_str: str = str(last_step.metadata["exception"])[:500].replace('"', "'")
+                        exc_str: str = str(last_step.metadata["exception"])[
+                            :500
+                        ].replace('"', "'")
                         exception_counter[exc_str] += 1
                         num_exceptions = sum(exception_counter.values())
                         pbar.set_postfix({"num_exceptions": num_exceptions})
diff --git a/src/ldp/nn/agent/simple_local_agent.py b/src/ldp/nn/agent/simple_local_agent.py
index 011a4d48..f30da850 100644
--- a/src/ldp/nn/agent/simple_local_agent.py
+++ b/src/ldp/nn/agent/simple_local_agent.py
@@ -2,8 +2,8 @@
 
 import torch
 import torch.distributed as dist
-from litellm import token_counter
 from aviary.core import Message, Tool, ToolRequestMessage
+from litellm import token_counter
 from pydantic import Field, field_validator
 
 from ldp.agent import Agent, SimpleAgentState
@@ -41,10 +41,9 @@ class AgentLMConfig(_LMConfig):
         "are better defaults than HF's.",
         validate_default=True,
     )
-    
     max_traj_token_count: int | None = Field(
         default=None,
-        description="If set, raise an error if the total tokens in the trajectory exceed this value."
+        description="If set, raise an error if the total tokens in the trajectory exceed this value.",
     )
 
     @field_validator("llm_call_kwargs")
@@ -115,20 +114,32 @@ async def get_asv(
 
         # Update state messages with result and return the new state
         next_state.messages = [*next_state.messages, result.value]
-        
-          
-        import ipdb; ipdb.set_trace()              
+
         if self.llm_model.max_traj_token_count is not None:
+            messages_for_tokenizer = self._llm_call_op.prep_messages_for_tokenizer(
+                next_state.messages
+            )
+            tools_for_tokenizer = self._llm_call_op.prep_tools_for_tokenizer(
+                next_state.tools
+            )
             total_tokens = token_counter(
-                model=self.llm_model.llm_for_sft,  # or any field referencing the model name
-                messages=next_state.messages,
-                tools=next_state.tools,
+                model=self.llm_model.model,  # or any field referencing the model name
+                messages=messages_for_tokenizer,
+                tools=tools_for_tokenizer,
+            )
+            # TODO remove
+            print(
+                "The traj size is %d tokens, with a limit of %d tokens"
+                % (total_tokens, self.llm_model.max_traj_token_count)
             )
             if total_tokens > self.llm_model.max_traj_token_count:
+                import ipdb
+
+                ipdb.set_trace()  # TODO remove
                 raise ValueError(
                     f"Token limit exceeded for trajectory: {total_tokens} > {self.llm_model.max_traj_token_count}"
                 )
-        
+
         return cast(OpResult[ToolRequestMessage], result), next_state, 0.0
 
     # TODO: maybe remove these recomputation methods. I added them to debug some things. But idk,
diff --git a/src/ldp/nn/handlers/transformer_handler.py b/src/ldp/nn/handlers/transformer_handler.py
index fb3a5d46..6fcc07e3 100644
--- a/src/ldp/nn/handlers/transformer_handler.py
+++ b/src/ldp/nn/handlers/transformer_handler.py
@@ -60,11 +60,6 @@
     "distributed.comm.timeouts.tcp": "300s",
 })
 
-compression = os.getenv("USE_DASK_COMPRESSION")
-if compression is not None:
-    config.set({"distributed.comm.compression": compression})
-    logger.info(f"Setting Dask compression to {compression}")
-
 TReturn = TypeVar("TReturn")
 TParams = ParamSpec("TParams")
 
@@ -204,6 +199,7 @@ def model_generate(model: PreTrainedModel, *args, **kwargs):
             elif not synced_gpus:
                 raise ValueError("synced_gpus must be True when using FSDP.")
             kwargs["synced_gpus"] = True
+            # TODO remove
             if os.getenv("USE_DASK_BARRIER"):
                 logger.info("Waiting for all workers to reach this point.")
                 dist.barrier()

From 62977864e4755dac5587d19d8ee290d18c863d0f Mon Sep 17 00:00:00 2001
From: Ori Kabeli <ori@futurehouse.org>
Date: Mon, 3 Mar 2025 14:01:04 -0600
Subject: [PATCH 08/21] nits

---
 src/ldp/nn/agent/simple_local_agent.py     | 47 +++++++++++-----------
 src/ldp/nn/handlers/transformer_handler.py | 34 +++++++++++-----
 2 files changed, 48 insertions(+), 33 deletions(-)

diff --git a/src/ldp/nn/agent/simple_local_agent.py b/src/ldp/nn/agent/simple_local_agent.py
index f30da850..50ee80a5 100644
--- a/src/ldp/nn/agent/simple_local_agent.py
+++ b/src/ldp/nn/agent/simple_local_agent.py
@@ -1,3 +1,4 @@
+import logging
 from typing import cast
 
 import torch
@@ -18,6 +19,8 @@
 )
 from ldp.nn.lm_config import LMConfig as _LMConfig
 
+logger = logging.getLogger(__name__)
+
 
 class AgentLMConfig(_LMConfig):
     """Adds some additional configuration options for running an LM in an Op."""
@@ -94,6 +97,8 @@ async def get_asv(
             else next_state.messages
         )
 
+        self._validate_token_count(messages, next_state.tools)
+
         # Execute the LLM operation call
         result = cast(
             OpResult[Message | ToolRequestMessage],
@@ -114,33 +119,27 @@ async def get_asv(
 
         # Update state messages with result and return the new state
         next_state.messages = [*next_state.messages, result.value]
+        self._validate_token_count(next_state.messages, next_state.tools)
 
-        if self.llm_model.max_traj_token_count is not None:
-            messages_for_tokenizer = self._llm_call_op.prep_messages_for_tokenizer(
-                next_state.messages
-            )
-            tools_for_tokenizer = self._llm_call_op.prep_tools_for_tokenizer(
-                next_state.tools
-            )
-            total_tokens = token_counter(
-                model=self.llm_model.model,  # or any field referencing the model name
-                messages=messages_for_tokenizer,
-                tools=tools_for_tokenizer,
+        return cast(OpResult[ToolRequestMessage], result), next_state, 0.0
+
+    def _validate_token_count(self, messages: list[Message], tools: list[Tool]):
+        if self.llm_model.max_traj_token_count is None:
+            return
+        messages_for_tokenizer = self._llm_call_op.prep_messages_for_tokenizer(messages)
+        tools_for_tokenizer = self._llm_call_op.prep_tools_for_tokenizer(tools)
+        total_tokens = token_counter(
+            model=self.llm_model.model,
+            messages=messages_for_tokenizer,
+            tools=tools_for_tokenizer,
+        )
+        if total_tokens > self.llm_model.max_traj_token_count:
+            logger.error(
+                f"Token limit exceeded for trajectory: {total_tokens} > {self.llm_model.max_traj_token_count}"
             )
-            # TODO remove
-            print(
-                "The traj size is %d tokens, with a limit of %d tokens"
-                % (total_tokens, self.llm_model.max_traj_token_count)
+            raise ValueError(
+                f"Token limit exceeded for trajectory: {total_tokens} > {self.llm_model.max_traj_token_count}"
             )
-            if total_tokens > self.llm_model.max_traj_token_count:
-                import ipdb
-
-                ipdb.set_trace()  # TODO remove
-                raise ValueError(
-                    f"Token limit exceeded for trajectory: {total_tokens} > {self.llm_model.max_traj_token_count}"
-                )
-
-        return cast(OpResult[ToolRequestMessage], result), next_state, 0.0
 
     # TODO: maybe remove these recomputation methods. I added them to debug some things. But idk,
     # maybe they'll come in handy later.
diff --git a/src/ldp/nn/handlers/transformer_handler.py b/src/ldp/nn/handlers/transformer_handler.py
index 6fcc07e3..1edcb54c 100644
--- a/src/ldp/nn/handlers/transformer_handler.py
+++ b/src/ldp/nn/handlers/transformer_handler.py
@@ -17,7 +17,7 @@
 import torch.distributed as dist
 import tree
 from dask import config
-from dask.distributed import Client
+from dask.distributed import Client, as_completed, wait
 from pydantic import BaseModel, ConfigDict, Field, field_validator
 from torch import nn
 from torch.cuda import nccl
@@ -196,13 +196,10 @@ def model_generate(model: PreTrainedModel, *args, **kwargs):
             synced_gpus = kwargs.pop("synced_gpus", None)
             if synced_gpus is None:
                 logger.debug("synced_gpus not defined, defaulting to True.")
+                kwargs["synced_gpus"] = True
             elif not synced_gpus:
                 raise ValueError("synced_gpus must be True when using FSDP.")
-            kwargs["synced_gpus"] = True
-            # TODO remove
-            if os.getenv("USE_DASK_BARRIER"):
-                logger.info("Waiting for all workers to reach this point.")
-                dist.barrier()
+           
 
         # Summoning params per https://github.com/pytorch/pytorch/issues/100069
         # If model is not FSDP, this context manager is a no-op.
@@ -585,7 +582,7 @@ def get_cuda_visible_devices() -> int | None:
             futures.append(future_op)
             worker_ids.append(worker_id)
 
-        self.handlers = self.client.gather(futures)
+        self.handlers = self.client_gather(futures)
         self.worker_ids = worker_ids
 
     async def __call__(
@@ -656,7 +653,7 @@ def _submit_and_gather(
                 self.handlers, self.worker_ids, split_args, split_kwargs, strict=True
             )
         ]
-        results = self.client.gather(futures)
+        results = self.client_gather(futures)
         results = cast(list[TReturn], [res.result().result() for res in results])
 
         if split_data:
@@ -767,13 +764,32 @@ def save_checkpoint(self, ckpt: os.PathLike | str, **kwargs) -> None:
 
     def teardown(self) -> None:
         if self._initialized:
-            self.client.close()
+            self.client.shutdown()
             self.cluster.close()
             self._initialized = False
 
     def __del__(self) -> None:
         self.teardown()
 
+    def client_gather(self, futures):
+        """Gather results from futures, propagating exceptions as they arrive.
+
+        Unlike client.gather() which waits for all futures to complete before raising
+        any exceptions, this method processes futures as they complete and raises
+        exceptions immediately. This is crucial when using FSDP where workers may
+        be stuck waiting for each other where one worker crashes, causing long hangs.
+        """
+        # Initialize a list to hold results
+        results = [None] * len(futures)
+        for completed_future, result in as_completed(
+            futures, with_results=True, raise_errors=True
+        ):
+            # Find the index of the completed future
+            index = futures.index(completed_future)
+            # Store the result directly from as_completed
+            results[index] = result
+        return results
+
 
 # Helpers
 

From e4ccb45fb2f0dda7c3ebad7e94de3c762e4dae17 Mon Sep 17 00:00:00 2001
From: Ori Kabeli <ori@futurehouse.org>
Date: Tue, 4 Mar 2025 06:14:58 -0600
Subject: [PATCH 09/21] nits

---
 src/ldp/alg/rollout.py          | 71 ++++++++++++++++++++++++---------
 src/ldp/graph/async_torch.py    |  6 ---
 src/ldp/nn/handlers/chunking.py |  3 +-
 tests/test_nn_models.py         | 17 ++++----
 4 files changed, 60 insertions(+), 37 deletions(-)

diff --git a/src/ldp/alg/rollout.py b/src/ldp/alg/rollout.py
index 2b1a1b1d..210437de 100644
--- a/src/ldp/alg/rollout.py
+++ b/src/ldp/alg/rollout.py
@@ -79,6 +79,8 @@ async def sample_trajectories(  # noqa: D418
         environment_factory: Callable[[], TEnv],
         batch_size: int = 1,
         max_steps: int | None = None,
+        *,
+        log_exceptions_immediately: bool = False,
     ) -> list[tuple[Trajectory, TEnv]]:
         """Run rollouts in parallel, using a factory to construct environments.
 
@@ -92,6 +94,8 @@ async def sample_trajectories(  # noqa: D418
                 an environment instance
             batch_size (int, optional): Defaults to 1.
             max_steps (int | None, optional): Max steps per rollout. Defaults to None (see above).
+            log_exceptions_immediately (bool, optional): Whether to log exceptions as they occur
+                or only at the end of the rollouts. Defaults to False.
 
         Returns:
             list[tuple[Trajectory, Environment]]: A list of (trajectory, environment) tuples: one per rollout.
@@ -102,6 +106,8 @@ async def sample_trajectories(  # noqa: D418
         self,
         environments: Sequence[Environment],
         max_steps: int | None = None,
+        *,
+        log_exceptions_immediately: bool = False,
     ) -> list[Trajectory]:
         """Run rollouts in parallel on a list of provided environments.
 
@@ -109,26 +115,35 @@ async def sample_trajectories(  # noqa: D418
             environments: A list of environments to run rollouts on.
             max_steps: Max steps per rollout. Defaults to None, in which case the rollouts are run
                 until environment returns done.
+            log_exceptions_immediately (bool, optional): Whether to log exceptions as they occur
+                or only at the end of the rollouts. Defaults to False.
         """
 
-    async def sample_trajectories(self, **kwargs):
-        if "environment_factory" in kwargs:
-            assert "environments" not in kwargs, (
-                "Cannot use environment_factory with environments"
-            )
-
+    async def sample_trajectories(
+        self,
+        environment_factory: Callable[[], Environment] | None = None,
+        environments: Sequence[Environment] | None = None,
+        batch_size: int = 1,
+        max_steps: int | None = None,
+        *,
+        log_exceptions_immediately: bool = False,
+    ) -> list[tuple[Trajectory, Environment]] | list[Trajectory]:
+        """Sample trajectories from environments, either via factory or pre-created."""
+        if environment_factory is not None:
+            assert environments is None, "Cannot use environment_factory with environments"
             return await self._sample_trajectories_from_env_factory(
-                kwargs["environment_factory"],
-                kwargs.get("batch_size", 1),
-                kwargs.get("max_steps"),
+                environment_factory,
+                batch_size,
+                max_steps,
+                log_exceptions_immediately=log_exceptions_immediately,
             )
 
-        if "environments" in kwargs:
-            assert "environment_factory" not in kwargs, (
-                "Cannot use environments with environment_factory"
-            )
+        if environments is not None:
+            assert environment_factory is None, "Cannot use environments with environment_factory"
             return await self._sample_trajectories_from_envs(
-                kwargs["environments"], kwargs.get("max_steps")
+                environments, 
+                max_steps,
+                log_exceptions_immediately=log_exceptions_immediately,
             )
 
         raise TypeError(
@@ -141,6 +156,8 @@ async def _sample_trajectories_from_env_factory(
         environment_factory: Callable[[], Environment],
         batch_size: int = 1,
         max_steps: int | None = None,
+        *,
+        log_exceptions_immediately: bool = False,
     ) -> list[tuple[Trajectory, Environment]]:
         self.traj_buffer.clear()
 
@@ -156,6 +173,7 @@ async def rollout_with_args(idx: int, **rollout_kwargs):
                     traj_id=uuid.uuid4().hex,
                     env=environment_factory(),
                     max_steps=max_steps,
+                    log_exceptions_immediately=log_exceptions_immediately,
                 )
             )
             for idx in range(batch_size)
@@ -182,6 +200,7 @@ async def rollout_with_args(idx: int, **rollout_kwargs):
                             traj_id=uuid.uuid4().hex,
                             env=environment_factory(),
                             max_steps=remaining_steps,
+                            log_exceptions_immediately=log_exceptions_immediately,
                         )
                     )
                     new_tasks.append(new_task)
@@ -194,6 +213,8 @@ async def _sample_trajectories_from_envs(
         self,
         environments: Sequence[Environment],
         max_steps: int | None = None,
+        *,
+        log_exceptions_immediately: bool = False,
     ) -> list[Trajectory]:
         self.traj_buffer.clear()
         exception_counter: Counter = Counter()
@@ -202,7 +223,14 @@ async def _sample_trajectories_from_envs(
 
         # Create all tasks first
         tasks = [
-            asyncio.create_task(self._rollout(traj_id, env, max_steps=max_steps))
+            asyncio.create_task(
+                self._rollout(
+                    traj_id, 
+                    env, 
+                    max_steps=max_steps,
+                    log_exceptions_immediately=log_exceptions_immediately
+                )
+            )
             for traj_id, env in zip(traj_ids, environments, strict=True)
         ]
 
@@ -229,10 +257,12 @@ async def _sample_trajectories_from_envs(
 
         # Final summary of exceptions (if any)
         if exception_counter:
-            logger.info("Caught exceptions:")
-            logger.info("%-6s %-50s", "Count", "Exception")
-            for exc, count in exception_counter.items():
-                logger.info("%-6d %-50s", count, exc)
+            summary = ["Caught exceptions:", "Count  Exception"]
+            summary.extend(
+                f"{count:<6d} {exc:<50s}"
+                for exc, count in exception_counter.items()
+            )
+            logger.info("\n".join(summary))
 
         return [self.traj_buffer[traj_id] for traj_id in traj_ids]
 
@@ -294,6 +324,9 @@ async def store_step(step: Transition):
         except CaughtError as e:
             # NOTE: This trajectory should not be used for regular training.
             # We save the last transition here for debugging, etc.
+            if log_exceptions_immediately:
+                logger.exception(f"Exception in rollout {traj_id}: {e.original_exc}")
+                
             await store_step(
                 Transition(
                     timestep=len(trajectory.steps),
diff --git a/src/ldp/graph/async_torch.py b/src/ldp/graph/async_torch.py
index f6597bb9..55ef0e9b 100644
--- a/src/ldp/graph/async_torch.py
+++ b/src/ldp/graph/async_torch.py
@@ -120,12 +120,6 @@ async def _maybe_process_batch(self) -> None:
 
         If neither condition is met, do nothing.
         """
-        # Technically should not happen, but if a coroutine crashes, it could release
-        # self._lock before placing results in _results_buffer and additional process
-        # coming inside this func will crash as self._work_buffer will be empty.
-        if not self._work_buffer:
-            return
-
         now = time.time()
 
         # sort by oldest requests first
diff --git a/src/ldp/nn/handlers/chunking.py b/src/ldp/nn/handlers/chunking.py
index 147b6dd8..9369d6d3 100644
--- a/src/ldp/nn/handlers/chunking.py
+++ b/src/ldp/nn/handlers/chunking.py
@@ -9,9 +9,8 @@
 class TensorChunker:
     """Splits tensors into chunks and adds dummy chunks as needed for parallel processing frameworks like FSDP."""
 
-    def __init__(self, num_chunks: int, dummy_value: int = 0):
+    def __init__(self, num_chunks: int):
         self.num_chunks = num_chunks
-        self.dummy_value = dummy_value
 
     def chunkify(self, *args, **kwargs) -> tuple[list[tuple], list[dict], list[bool]]:
         """Splits the args into self.num_chunks chunks, adding dummy chunks as needed.
diff --git a/tests/test_nn_models.py b/tests/test_nn_models.py
index aa717b48..e872d9f1 100644
--- a/tests/test_nn_models.py
+++ b/tests/test_nn_models.py
@@ -27,11 +27,10 @@ class TestTensorChunker:
     def test_chunkify_add_dummy_chunks(self):
         batch_size = 3
         num_chunks = 5
-        dummy_value = 0
 
         sample_tensor = torch.arange(1, batch_size * 10 + 1).reshape(batch_size, 10)
 
-        chunker = ldp.nn.TensorChunker(num_chunks=num_chunks, dummy_value=dummy_value)
+        chunker = ldp.nn.TensorChunker(num_chunks=num_chunks)
         split_args, split_kwargs, dummy_chunk_flags = chunker.chunkify(sample_tensor)
 
         assert len(split_args) == num_chunks
@@ -41,20 +40,19 @@ def test_chunkify_add_dummy_chunks(self):
         assert torch.equal(split_args[1][0], sample_tensor[1:2])
         assert torch.equal(split_args[2][0], sample_tensor[2:3])
         assert torch.equal(
-            split_args[3][0], torch.full_like(sample_tensor[:1], dummy_value)
+            split_args[3][0], sample_tensor[:1]
         )
         assert torch.equal(
-            split_args[4][0], torch.full_like(sample_tensor[:1], dummy_value)
+            split_args[4][0], sample_tensor[:1]
         )
 
     def test_chunkify_no_dummy_chunks(self):
         batch_size = 9
         num_chunks = 5
-        dummy_value = 0
 
         sample_tensor = torch.arange(1, batch_size * 10 + 1).reshape(batch_size, 10)
 
-        chunker = ldp.nn.TensorChunker(num_chunks=num_chunks, dummy_value=dummy_value)
+        chunker = ldp.nn.TensorChunker(num_chunks=num_chunks)
         split_args, split_kwargs, dummy_chunk_flags = chunker.chunkify(sample_tensor)
 
         assert len(split_args) == num_chunks
@@ -69,7 +67,6 @@ def test_chunkify_no_dummy_chunks(self):
     def test_chunkify_with_args_and_kwargs(self):
         batch_size = 2
         num_chunks = 3
-        dummy_value = 0
 
         sample_tensor = torch.arange(1, batch_size * 10 + 1).reshape(batch_size, 10)
         sample_tensor_kwarg = torch.arange(1, batch_size * 5 + 1).reshape(batch_size, 5)
@@ -78,7 +75,7 @@ def test_chunkify_with_args_and_kwargs(self):
             "key2": "Not split",
         }
 
-        chunker = ldp.nn.TensorChunker(num_chunks=num_chunks, dummy_value=dummy_value)
+        chunker = ldp.nn.TensorChunker(num_chunks=num_chunks)
         split_args, split_kwargs, dummy_chunk_flags = chunker.chunkify(
             sample_tensor, **sample_kwargs
         )
@@ -89,13 +86,13 @@ def test_chunkify_with_args_and_kwargs(self):
         assert torch.equal(split_args[0][0], sample_tensor[:1])
         assert torch.equal(split_args[1][0], sample_tensor[1:2])
         assert torch.equal(
-            split_args[2][0], torch.full_like(sample_tensor[:1], dummy_value)
+            split_args[2][0], sample_tensor[:1]
         )
         assert torch.equal(split_kwargs[0]["key1"], sample_tensor_kwarg[:1])
         assert torch.equal(split_kwargs[1]["key1"], sample_tensor_kwarg[1:2])
         assert torch.equal(
             split_kwargs[2]["key1"],
-            torch.full_like(sample_tensor_kwarg[:1], dummy_value),
+            sample_tensor_kwarg[:1]
         )
         assert all(split_kwargs[i]["key2"] == "Not split" for i in range(num_chunks))
 

From 6ed30bd6a23cfecd60ffc07ad1cf13e11736d588 Mon Sep 17 00:00:00 2001
From: Ori Kabeli <ori@futurehouse.org>
Date: Tue, 4 Mar 2025 13:39:32 -0600
Subject: [PATCH 10/21] nits

---
 src/ldp/alg/rollout.py                     | 94 ++++++++++------------
 src/ldp/nn/agent/simple_local_agent.py     |  5 +-
 src/ldp/nn/handlers/transformer_handler.py |  5 +-
 tests/test_nn_models.py                    | 17 +---
 4 files changed, 51 insertions(+), 70 deletions(-)

diff --git a/src/ldp/alg/rollout.py b/src/ldp/alg/rollout.py
index 210437de..17c96a10 100644
--- a/src/ldp/alg/rollout.py
+++ b/src/ldp/alg/rollout.py
@@ -5,7 +5,7 @@
 from collections import Counter
 from collections.abc import Callable, Iterator, Sequence
 from contextlib import contextmanager, nullcontext
-from typing import Any, TypeVar, overload
+from typing import Any, TypeVar
 
 from aviary.core import Environment, Message
 from tqdm import tqdm
@@ -73,64 +73,52 @@ def __init__(
         self.traj_buffer: dict[str, Trajectory] = {}
         self.callbacks = callbacks or []
 
-    @overload
-    async def sample_trajectories(  # noqa: D418
+    async def sample_trajectories(
         self,
-        environment_factory: Callable[[], TEnv],
+        environment_factory: Callable[[], TEnv] | None = None,
+        environments: Sequence[TEnv] | None = None,
         batch_size: int = 1,
         max_steps: int | None = None,
         *,
         log_exceptions_immediately: bool = False,
-    ) -> list[tuple[Trajectory, TEnv]]:
-        """Run rollouts in parallel, using a factory to construct environments.
+    ) -> list[tuple[Trajectory, Environment]] | list[Trajectory]:
+        """Sample trajectories from environments, either via factory or pre-created.
 
-        We will construct `batch_size` environments and run rollouts on each of them.
-        If `max_steps` is set, rollouts will be truncated at this value. If a rollout
-        has fewer than `max_steps`, then a new environment will be constructed and another
-        rollout will be started until `max_steps` is reached.
+        There are two main ways to use this method:
 
-        Args:
-            environment_factory: A no-argument callable that returns
-                an environment instance
-            batch_size (int, optional): Defaults to 1.
-            max_steps (int | None, optional): Max steps per rollout. Defaults to None (see above).
-            log_exceptions_immediately (bool, optional): Whether to log exceptions as they occur
-                or only at the end of the rollouts. Defaults to False.
+        1. Using an environment factory:
+            Run rollouts in parallel, using a factory to construct environments.
+            We will construct `batch_size` environments and run rollouts on each of them.
+            If `max_steps` is set, rollouts will be truncated at this value. If a rollout
+            has fewer than `max_steps`, then a new environment will be constructed and another
+            rollout will be started until `max_steps` is reached.
 
-        Returns:
-            list[tuple[Trajectory, Environment]]: A list of (trajectory, environment) tuples: one per rollout.
-        """
+            In this case, returns a list of (trajectory, environment) tuples.
 
-    @overload
-    async def sample_trajectories(  # noqa: D418
-        self,
-        environments: Sequence[Environment],
-        max_steps: int | None = None,
-        *,
-        log_exceptions_immediately: bool = False,
-    ) -> list[Trajectory]:
-        """Run rollouts in parallel on a list of provided environments.
+        2. Using a sequence of environments:
+            Run rollouts in parallel on a list of provided environments.
+            In this case, returns a list of trajectories.
 
         Args:
-            environments: A list of environments to run rollouts on.
+            environment_factory: A no-argument callable that returns an environment instance
+            environments: A list of environments to run rollouts on
+            batch_size: Number of parallel environments to run when using environment_factory. Defaults to 1.
             max_steps: Max steps per rollout. Defaults to None, in which case the rollouts are run
                 until environment returns done.
-            log_exceptions_immediately (bool, optional): Whether to log exceptions as they occur
+            log_exceptions_immediately: Whether to log exceptions as they occur
                 or only at the end of the rollouts. Defaults to False.
-        """
 
-    async def sample_trajectories(
-        self,
-        environment_factory: Callable[[], Environment] | None = None,
-        environments: Sequence[Environment] | None = None,
-        batch_size: int = 1,
-        max_steps: int | None = None,
-        *,
-        log_exceptions_immediately: bool = False,
-    ) -> list[tuple[Trajectory, Environment]] | list[Trajectory]:
-        """Sample trajectories from environments, either via factory or pre-created."""
+        Returns:
+            Either list[tuple[Trajectory, Environment]] or list[Trajectory] depending on whether
+            environment_factory or environments is provided.
+
+        Raises:
+            TypeError: If neither environment_factory nor environments is provided.
+        """
         if environment_factory is not None:
-            assert environments is None, "Cannot use environment_factory with environments"
+            assert environments is None, (
+                "Cannot use environment_factory with environments"
+            )
             return await self._sample_trajectories_from_env_factory(
                 environment_factory,
                 batch_size,
@@ -139,9 +127,11 @@ async def sample_trajectories(
             )
 
         if environments is not None:
-            assert environment_factory is None, "Cannot use environments with environment_factory"
+            assert environment_factory is None, (
+                "Cannot use environments with environment_factory"
+            )
             return await self._sample_trajectories_from_envs(
-                environments, 
+                environments,
                 max_steps,
                 log_exceptions_immediately=log_exceptions_immediately,
             )
@@ -225,10 +215,10 @@ async def _sample_trajectories_from_envs(
         tasks = [
             asyncio.create_task(
                 self._rollout(
-                    traj_id, 
-                    env, 
+                    traj_id,
+                    env,
                     max_steps=max_steps,
-                    log_exceptions_immediately=log_exceptions_immediately
+                    log_exceptions_immediately=log_exceptions_immediately,
                 )
             )
             for traj_id, env in zip(traj_ids, environments, strict=True)
@@ -259,8 +249,7 @@ async def _sample_trajectories_from_envs(
         if exception_counter:
             summary = ["Caught exceptions:", "Count  Exception"]
             summary.extend(
-                f"{count:<6d} {exc:<50s}"
-                for exc, count in exception_counter.items()
+                f"{count:<6d} {exc:<50s}" for exc, count in exception_counter.items()
             )
             logger.info("\n".join(summary))
 
@@ -271,7 +260,8 @@ async def _rollout(
         traj_id: str,
         env: Environment,
         max_steps: int | None,
-        max_tokens: int | None = None,  # <-- new argument
+        *,
+        log_exceptions_immediately: bool = False,
     ) -> Trajectory:
         trajectory = Trajectory(traj_id=traj_id)
 
@@ -326,7 +316,7 @@ async def store_step(step: Transition):
             # We save the last transition here for debugging, etc.
             if log_exceptions_immediately:
                 logger.exception(f"Exception in rollout {traj_id}: {e.original_exc}")
-                
+
             await store_step(
                 Transition(
                     timestep=len(trajectory.steps),
diff --git a/src/ldp/nn/agent/simple_local_agent.py b/src/ldp/nn/agent/simple_local_agent.py
index 222ee7ce..00957d0d 100644
--- a/src/ldp/nn/agent/simple_local_agent.py
+++ b/src/ldp/nn/agent/simple_local_agent.py
@@ -4,7 +4,7 @@
 import torch
 import torch.distributed as dist
 from aviary.core import Message, Tool, ToolRequestMessage
-from litellm import token_counter
+from litellm.utils import token_counter
 from pydantic import Field, field_validator
 
 from ldp.agent import Agent, SimpleAgentState
@@ -131,10 +131,11 @@ def _validate_token_count(self, messages: list[Message], tools: list[Tool]):
             return
         messages_for_tokenizer = self._llm_call_op.prep_messages_for_tokenizer(messages)
         tools_for_tokenizer = self._llm_call_op.prep_tools_for_tokenizer(tools)
+
         total_tokens = token_counter(
             model=self.llm_model.model,
             messages=messages_for_tokenizer,
-            tools=tools_for_tokenizer,
+            tools=tools_for_tokenizer,  # type: ignore[arg-type]
         )
         if total_tokens > self.llm_model.max_traj_token_count:
             logger.error(
diff --git a/src/ldp/nn/handlers/transformer_handler.py b/src/ldp/nn/handlers/transformer_handler.py
index 4535c9fd..537a9129 100644
--- a/src/ldp/nn/handlers/transformer_handler.py
+++ b/src/ldp/nn/handlers/transformer_handler.py
@@ -17,7 +17,7 @@
 import torch.distributed as dist
 import tree
 from dask import config
-from dask.distributed import Client, as_completed, wait
+from dask.distributed import Client, as_completed
 from pydantic import BaseModel, ConfigDict, Field, field_validator
 from torch import nn
 from torch.cuda import nccl
@@ -192,14 +192,13 @@ async def __call__(  # type: ignore[override]
     @staticmethod
     def model_generate(model: PreTrainedModel, *args, **kwargs):
         """A method that can be used as module_call_fn to sample from an LLM."""
-        if dist.get_world_size() > 1:
+        if int(os.environ.get("WORLD_SIZE", "1")) > 1:
             synced_gpus = kwargs.pop("synced_gpus", None)
             if synced_gpus is None:
                 logger.debug("synced_gpus not defined, defaulting to True.")
                 kwargs["synced_gpus"] = True
             elif not synced_gpus:
                 raise ValueError("synced_gpus must be True when using FSDP.")
-           
 
         # Summoning params per https://github.com/pytorch/pytorch/issues/100069
         # If model is not FSDP, this context manager is a no-op.
diff --git a/tests/test_nn_models.py b/tests/test_nn_models.py
index e872d9f1..55917ce3 100644
--- a/tests/test_nn_models.py
+++ b/tests/test_nn_models.py
@@ -39,12 +39,8 @@ def test_chunkify_add_dummy_chunks(self):
         assert torch.equal(split_args[0][0], sample_tensor[:1])
         assert torch.equal(split_args[1][0], sample_tensor[1:2])
         assert torch.equal(split_args[2][0], sample_tensor[2:3])
-        assert torch.equal(
-            split_args[3][0], sample_tensor[:1]
-        )
-        assert torch.equal(
-            split_args[4][0], sample_tensor[:1]
-        )
+        assert torch.equal(split_args[3][0], sample_tensor[:1])
+        assert torch.equal(split_args[4][0], sample_tensor[:1])
 
     def test_chunkify_no_dummy_chunks(self):
         batch_size = 9
@@ -85,15 +81,10 @@ def test_chunkify_with_args_and_kwargs(self):
         assert dummy_chunk_flags == [False, False, True]
         assert torch.equal(split_args[0][0], sample_tensor[:1])
         assert torch.equal(split_args[1][0], sample_tensor[1:2])
-        assert torch.equal(
-            split_args[2][0], sample_tensor[:1]
-        )
+        assert torch.equal(split_args[2][0], sample_tensor[:1])
         assert torch.equal(split_kwargs[0]["key1"], sample_tensor_kwarg[:1])
         assert torch.equal(split_kwargs[1]["key1"], sample_tensor_kwarg[1:2])
-        assert torch.equal(
-            split_kwargs[2]["key1"],
-            sample_tensor_kwarg[:1]
-        )
+        assert torch.equal(split_kwargs[2]["key1"], sample_tensor_kwarg[:1])
         assert all(split_kwargs[i]["key2"] == "Not split" for i in range(num_chunks))
 
     def test_dechunkify(self):

From d245e3de0789202f818e802fa7f24f20c68971f2 Mon Sep 17 00:00:00 2001
From: Ori Kabeli <ori@futurehouse.org>
Date: Tue, 4 Mar 2025 13:53:35 -0600
Subject: [PATCH 11/21] nits

---
 src/ldp/alg/rollout.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/ldp/alg/rollout.py b/src/ldp/alg/rollout.py
index 17c96a10..9f34cf57 100644
--- a/src/ldp/alg/rollout.py
+++ b/src/ldp/alg/rollout.py
@@ -246,7 +246,7 @@ async def _sample_trajectories_from_envs(
                         pbar.set_postfix({"num_exceptions": num_exceptions})
 
         # Final summary of exceptions (if any)
-        if exception_counter:
+        if exception_counter and not log_exceptions_immediately:
             summary = ["Caught exceptions:", "Count  Exception"]
             summary.extend(
                 f"{count:<6d} {exc:<50s}" for exc, count in exception_counter.items()

From 51210dfcefa242350cfccc81609329f2425a742c Mon Sep 17 00:00:00 2001
From: Ori Kabeli <ori@futurehouse.org>
Date: Tue, 4 Mar 2025 14:16:33 -0600
Subject: [PATCH 12/21] nits

---
 src/ldp/alg/rollout.py | 88 ++++++++++++++++++++++++++++--------------
 1 file changed, 59 insertions(+), 29 deletions(-)

diff --git a/src/ldp/alg/rollout.py b/src/ldp/alg/rollout.py
index 9f34cf57..0ab2b7b6 100644
--- a/src/ldp/alg/rollout.py
+++ b/src/ldp/alg/rollout.py
@@ -80,7 +80,7 @@ async def sample_trajectories(
         batch_size: int = 1,
         max_steps: int | None = None,
         *,
-        log_exceptions_immediately: bool = False,
+        log_exceptions_immediately: bool = True,
     ) -> list[tuple[Trajectory, Environment]] | list[Trajectory]:
         """Sample trajectories from environments, either via factory or pre-created.
 
@@ -106,7 +106,7 @@ async def sample_trajectories(
             max_steps: Max steps per rollout. Defaults to None, in which case the rollouts are run
                 until environment returns done.
             log_exceptions_immediately: Whether to log exceptions as they occur
-                or only at the end of the rollouts. Defaults to False.
+                or only at the end of the rollouts.
 
         Returns:
             Either list[tuple[Trajectory, Environment]] or list[Trajectory] depending on whether
@@ -147,14 +147,17 @@ async def _sample_trajectories_from_env_factory(
         batch_size: int = 1,
         max_steps: int | None = None,
         *,
-        log_exceptions_immediately: bool = False,
+        log_exceptions_immediately: bool = True,
     ) -> list[tuple[Trajectory, Environment]]:
         self.traj_buffer.clear()
+        exception_counter: Counter = Counter()
 
         async def rollout_with_args(idx: int, **rollout_kwargs):
             return idx, await self._rollout(**rollout_kwargs), rollout_kwargs
 
         accumulated_steps = [0] * batch_size
+        total_trajectories = 0  # Counter for completed trajectories
+
         # submit initial batch of tasks
         tasks = [
             asyncio.create_task(
@@ -170,32 +173,59 @@ async def rollout_with_args(idx: int, **rollout_kwargs):
         ]
 
         results = []
-        while tasks:
-            done, pending = await asyncio.wait(
-                tasks, return_when=asyncio.FIRST_COMPLETED
-            )
-            new_tasks = []
-            for task in done:
-                idx, traj, kwargs = await task
-                results.append((traj, kwargs["env"]))
-                accumulated_steps[idx] += len(traj.steps)
-                if (
-                    max_steps is not None
-                    and (remaining_steps := max_steps - accumulated_steps[idx]) > 0
-                ):
-                    # submit another task if we haven't reached max_steps
-                    new_task = asyncio.create_task(
-                        rollout_with_args(
-                            idx,
-                            traj_id=uuid.uuid4().hex,
-                            env=environment_factory(),
-                            max_steps=remaining_steps,
-                            log_exceptions_immediately=log_exceptions_immediately,
+        with tqdm(
+            desc="Rollouts",
+            unit="rollout",
+            ncols=0,
+        ) as pbar:
+            while tasks:
+                done, pending = await asyncio.wait(
+                    tasks, return_when=asyncio.FIRST_COMPLETED
+                )
+                new_tasks = []
+                for task in done:
+                    idx, traj, kwargs = await task
+                    results.append((traj, kwargs["env"]))
+                    total_trajectories += 1
+                    pbar.update(1)
+
+                    steps_in_traj = len(traj.steps)
+                    accumulated_steps[idx] += steps_in_traj
+
+                    # Check for exceptions in this trajectory
+                    if traj.steps and traj.steps[-1].metadata.get("exception"):
+                        exc_str: str = str(traj.steps[-1].metadata["exception"])[
+                            :500
+                        ].replace('"', "'")
+                        exception_counter[exc_str] += 1
+                        num_exceptions = sum(exception_counter.values())
+                        pbar.set_postfix({"num_exceptions": num_exceptions})
+
+                    if (
+                        max_steps is not None
+                        and (remaining_steps := max_steps - accumulated_steps[idx]) > 0
+                    ):
+                        # submit another task if we haven't reached max_steps
+                        new_task = asyncio.create_task(
+                            rollout_with_args(
+                                idx,
+                                traj_id=uuid.uuid4().hex,
+                                env=environment_factory(),
+                                max_steps=remaining_steps,
+                                log_exceptions_immediately=log_exceptions_immediately,
+                            )
                         )
-                    )
-                    new_tasks.append(new_task)
+                        new_tasks.append(new_task)
 
-            tasks = list(pending) + new_tasks
+                tasks = list(pending) + new_tasks
+
+        # Final summary of exceptions (if any)
+        if exception_counter and not log_exceptions_immediately:
+            summary = ["Caught exceptions:", "Count  Exception"]
+            summary.extend(
+                f"{count:<6d} {exc:<50s}" for exc, count in exception_counter.items()
+            )
+            logger.info("\n".join(summary))
 
         return results
 
@@ -204,7 +234,7 @@ async def _sample_trajectories_from_envs(
         environments: Sequence[Environment],
         max_steps: int | None = None,
         *,
-        log_exceptions_immediately: bool = False,
+        log_exceptions_immediately: bool = True,
     ) -> list[Trajectory]:
         self.traj_buffer.clear()
         exception_counter: Counter = Counter()
@@ -261,7 +291,7 @@ async def _rollout(
         env: Environment,
         max_steps: int | None,
         *,
-        log_exceptions_immediately: bool = False,
+        log_exceptions_immediately: bool = True,
     ) -> Trajectory:
         trajectory = Trajectory(traj_id=traj_id)
 

From 3d7c20e7aef8ea320ad1b5eaef83aadd837a1ba2 Mon Sep 17 00:00:00 2001
From: Ori Kabeli <ori@futurehouse.org>
Date: Wed, 12 Mar 2025 03:48:39 -0500
Subject: [PATCH 13/21] Refactor Dask handling in transformer handler for
 improved exception management and memory efficiency

---
 src/ldp/alg/rollout.py                     |   2 +-
 src/ldp/nn/handlers/transformer_handler.py | 128 ++++++++++++++-------
 2 files changed, 90 insertions(+), 40 deletions(-)

diff --git a/src/ldp/alg/rollout.py b/src/ldp/alg/rollout.py
index 0ab2b7b6..da40f669 100644
--- a/src/ldp/alg/rollout.py
+++ b/src/ldp/alg/rollout.py
@@ -8,7 +8,7 @@
 from typing import Any, TypeVar
 
 from aviary.core import Environment, Message
-from tqdm import tqdm
+from tqdm.asyncio import tqdm
 
 from ldp.agent import Agent
 from ldp.data_structures import Trajectory, Transition
diff --git a/src/ldp/nn/handlers/transformer_handler.py b/src/ldp/nn/handlers/transformer_handler.py
index 537a9129..6ee5ae54 100644
--- a/src/ldp/nn/handlers/transformer_handler.py
+++ b/src/ldp/nn/handlers/transformer_handler.py
@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import asyncio
 import atexit
 import logging
 import os
@@ -10,14 +11,15 @@
 from enum import StrEnum, auto
 from functools import cache, partial, wraps
 from pathlib import Path
-from typing import Any, Concatenate, ParamSpec, Self, TypeVar, assert_never, cast
+from typing import Any, Concatenate, ParamSpec, Self, TypeVar, assert_never
 
 import accelerate
 import torch
 import torch.distributed as dist
 import tree
 from dask import config
-from dask.distributed import Client, as_completed
+from dask.distributed import Actor, ActorFuture, Client
+from distributed.utils import sync
 from pydantic import BaseModel, ConfigDict, Field, field_validator
 from torch import nn
 from torch.cuda import nccl
@@ -199,6 +201,7 @@ def model_generate(model: PreTrainedModel, *args, **kwargs):
                 kwargs["synced_gpus"] = True
             elif not synced_gpus:
                 raise ValueError("synced_gpus must be True when using FSDP.")
+        raise torch.OutOfMemoryError("yoyoyoyoyoyoyo test test test TODO")  # TODO remove
 
         # Summoning params per https://github.com/pytorch/pytorch/issues/100069
         # If model is not FSDP, this context manager is a no-op.
@@ -425,22 +428,29 @@ def _exec_func(
         args = tree.map_structure(to_device, args)
         kwargs = tree.map_structure(to_device, kwargs)
 
-        with torch.autocast(
-            device_type=self.module.device.type, dtype=self.module.dtype
-        ):
-            res = (
-                getattr(self, func)(*args, **kwargs)
-                if isinstance(func, str)
-                else func(self, *args, **kwargs)
-            )
+        try:
+            with torch.autocast(
+                device_type=self.module.device.type, dtype=self.module.dtype
+            ):
+                res = (
+                    getattr(self, func)(*args, **kwargs)
+                    if isinstance(func, str)
+                    else func(self, *args, **kwargs)
+                )
 
-        # Needed to prevent GPU memory leak to the main process scheduling the workers
-        if isinstance(res, GenerateDecoderOnlyOutput):
-            res.past_key_values = None
-            res["past_key_values"] = None
+            # Needed to prevent GPU memory leak to the main process scheduling the workers
+            if isinstance(res, GenerateDecoderOnlyOutput):
+                res.past_key_values = None
+                res["past_key_values"] = None
 
-        to_cpu = partial(_move_tensor, device=torch.device("cpu"))
-        return tree.map_structure(to_cpu, res)
+            to_cpu = partial(_move_tensor, device=torch.device("cpu"))
+            return tree.map_structure(to_cpu, res)
+        except Exception as e:
+            # Re-raise the exception with traceback preserved. For some exceptions, Dask
+            # modifies or loses the original traceback when crossing process boundaries.
+            # RuntimeError preserves the traceback when using with_traceback() of original
+            # exception.
+            raise RuntimeError(str(e)).with_traceback(e.__traceback__)  # noqa: B904
 
     def __del__(self) -> None:
         dist.destroy_process_group()
@@ -582,7 +592,7 @@ def get_cuda_visible_devices() -> int | None:
             futures.append(future_op)
             worker_ids.append(worker_id)
 
-        self.handlers = self.client_gather(futures)
+        self.actors: list[Actor] = self._client_gather(futures)
         self.worker_ids = worker_ids
 
     async def __call__(
@@ -633,28 +643,24 @@ def _submit_and_gather(
         """
         if split_data:
             chunker = TensorChunker(
-                num_chunks=len(self.handlers),
+                num_chunks=len(self.actors),
             )
             split_args, split_kwargs, dummy_flags = chunker.chunkify(*args, **kwargs)
         else:
-            split_args = [args] * len(self.handlers)
-            split_kwargs = [kwargs] * len(self.handlers)
+            split_args = [args] * len(self.actors)
+            split_kwargs = [kwargs] * len(self.actors)
 
         futures = [
-            self.client.submit(
-                handler._exec_func,
+            handler._exec_func(
                 func,
                 *args_i,
-                workers=[worker_id],
-                actor=True,
                 **kwargs_i,
             )
             for handler, worker_id, args_i, kwargs_i in zip(
-                self.handlers, self.worker_ids, split_args, split_kwargs, strict=True
+                self.actors, self.worker_ids, split_args, split_kwargs, strict=True
             )
         ]
-        results = self.client_gather(futures)
-        results = cast("list[TReturn]", [res.result().result() for res in results])
+        results: list[TReturn] = self._client_gather(futures)
 
         if split_data:
             return chunker.dechunkify(results, dummy_flags)
@@ -767,29 +773,73 @@ def teardown(self) -> None:
         if self._initialized:
             self.client.shutdown()
             self.cluster.close()
+            del self.client
+            del self.cluster
             self._initialized = False
 
     def __del__(self) -> None:
         self.teardown()
 
-    def client_gather(self, futures):
+    @staticmethod
+    def _wrap_dask_future(dask_future: ActorFuture):
+        """Converts a Dask ActorFuture into an awaitable asyncio.Future."""
+        loop = asyncio.get_running_loop()
+        return asyncio.ensure_future(loop.run_in_executor(None, dask_future.result))
+
+    def _client_gather(self, futures: list[ActorFuture]) -> list[Any]:
         """Gather results from futures, propagating exceptions as they arrive.
 
         Unlike client.gather() which waits for all futures to complete before raising
         any exceptions, this method processes futures as they complete and raises
         exceptions immediately. This is crucial when using FSDP where workers may
-        be stuck waiting for each other where one worker crashes, causing long hangs.
+        be stuck waiting for each other when one worker crashes, causing long hangs.
+
+        Note: Dask Actors currently have an issue where they're not working properly with
+        dask.gather() and can cause blocking issues or hide worker errors. This implementation
+        works around those limitations.
         """
-        # Initialize a list to hold results
-        results = [None] * len(futures)
-        for completed_future, result in as_completed(
-            futures, with_results=True, raise_errors=True
-        ):
-            # Find the index of the completed future
-            index = futures.index(completed_future)
-            # Store the result directly from as_completed
-            results[index] = result
-        return results
+
+        async def _gather_with_exception_handling(futures):
+            wrapped_futures = [self._wrap_dask_future(f) for f in futures]
+
+            try:
+                # Use asyncio.wait with FIRST_EXCEPTION instead of gather
+                done, pending = await asyncio.wait(
+                    wrapped_futures, timeout=120, return_when=asyncio.FIRST_EXCEPTION
+                )
+                                
+                exceptions = []
+                for future in done:
+                    exc = future.exception()
+                    if exc:
+                        exceptions.append(exc)
+                if exceptions:
+                    if len(exceptions) == 1:
+                        raise exceptions[0]
+                    raise ExceptionGroup("Multiple actor exceptions", exceptions)
+
+                if pending:
+                    pending_indices = sorted([
+                        wrapped_futures.index(p) for p in pending
+                    ])
+                    raise TimeoutError(
+                        f"Tasks didn't complete within timeout. {len(pending)} out of {len(wrapped_futures)} "
+                        f"still pending. Pending task indices: {pending_indices}"
+                    )
+
+                return await asyncio.gather(*wrapped_futures)
+            except Exception as e:
+                logger.exception("Error in dask workers")
+                for f in wrapped_futures:
+                    if not f.done():
+                        f.cancel()         
+                self.teardown()
+                # sys.exit(1) would wait for dask to finish, which can cause hanging
+                # when workers are in a deadlock. Use os._exit to force immediate termination
+                os._exit(1) 
+
+        # Use distributed.utils.sync to run the async function in the current thread
+        return sync(self.client.loop, _gather_with_exception_handling, futures)  # type: ignore[arg-type]
 
 
 # Helpers

From 15b936d1f12a211d67b98451c62e03b2bf3a9e81 Mon Sep 17 00:00:00 2001
From: Ori Kabeli <ori@futurehouse.org>
Date: Wed, 12 Mar 2025 05:49:12 -0500
Subject: [PATCH 14/21] Remove test OutOfMemoryError raise in
 AsyncTransformerInterface

---
 src/ldp/nn/handlers/transformer_handler.py | 46 ++++++++++++++++------
 1 file changed, 33 insertions(+), 13 deletions(-)

diff --git a/src/ldp/nn/handlers/transformer_handler.py b/src/ldp/nn/handlers/transformer_handler.py
index 6ee5ae54..ddf159f8 100644
--- a/src/ldp/nn/handlers/transformer_handler.py
+++ b/src/ldp/nn/handlers/transformer_handler.py
@@ -201,7 +201,6 @@ def model_generate(model: PreTrainedModel, *args, **kwargs):
                 kwargs["synced_gpus"] = True
             elif not synced_gpus:
                 raise ValueError("synced_gpus must be True when using FSDP.")
-        raise torch.OutOfMemoryError("yoyoyoyoyoyoyo test test test TODO")  # TODO remove
 
         # Summoning params per https://github.com/pytorch/pytorch/issues/100069
         # If model is not FSDP, this context manager is a no-op.
@@ -240,14 +239,22 @@ def __init__(self, config: TransformerHandlerConfig):
                 assert_never(config.lm_type)
         super().__init__(model)
         self.tokenizer = tokenizer
+        logger.info(
+            f"Initialized tokenizer: {type(tokenizer).__name__}, vocab size: {len(tokenizer)}"
+        )
+
         maybe_set_tokenizer_chat_template(
             self.tokenizer, self.config.lm_config.chat_template
         )
+        logger.info(f"Chat template: {self.config.lm_config.chat_template}")
 
         self._setup_accelerator()
+        logger.info(f"Accelerator set up with device: {self.accelerator.device}")
 
         if config.checkpoint is not None:
+            logger.info(f"Loading checkpoint from: {config.checkpoint}")
             self.load_checkpoint(config.checkpoint)
+        logger.info("Initialization complete")
 
     def _setup_accelerator(self):
         self.accelerator = accelerate.Accelerator(
@@ -391,6 +398,13 @@ def _setup_accelerator(self):
                 buffer_dtype=torch.bfloat16,
             )
 
+        logger.info(f"Setting up accelerator with bf16={bf16}")
+        logger.info(f"Worker config: offload_cpu={self.worker_config.offload_cpu}, "
+                   f"activation_checkpointing={self.worker_config.activation_checkpointing}, "
+                   f"cpu_ram_efficient_loading={self.worker_config.cpu_ram_efficient_loading}, "
+                   f"state_dict_type={self.worker_config.state_dict_type}, "
+                   f"backward_prefetch={self.worker_config.backward_prefetch}")
+        
         self.accelerator = accelerate.Accelerator(
             # See note in TransformerHandler._setup_accelerator() about this
             # mixed_precision=("bf16" if bf16 else "no"),
@@ -406,11 +420,16 @@ def _setup_accelerator(self):
                 backward_prefetch=self.worker_config.backward_prefetch,
             ),
         )
+        logger.info(f"Accelerator setup complete on rank {self.worker_config.rank}")
 
         if self.config.lm_config.device == "meta":
+            logger.info(f"Preparing model for FSDP with meta device on rank {self.worker_config.rank}")
             self.module = prepare_model_for_fsdp_with_meta_device(self.module)
+            logger.info(f"Meta device preparation complete on rank {self.worker_config.rank}")
 
+        logger.info(f"Preparing model with accelerator on rank {self.worker_config.rank}")
         self.module = self.accelerator.prepare(self.module)
+        logger.info(f"Model preparation complete on rank {self.worker_config.rank}, model device: {self.module.device}, dtype: {self.module.dtype}")
 
     def set_seed(self, seed: int) -> None:
         """Set the seed for the current worker."""
@@ -770,6 +789,7 @@ def save_checkpoint(self, ckpt: os.PathLike | str, **kwargs) -> None:
         self._submit_and_gather("save_checkpoint", ckpt, **kwargs)
 
     def teardown(self) -> None:
+        logger.info(f"Shutting down Dask cluster, is_initialized: {self._initialized}")
         if self._initialized:
             self.client.shutdown()
             self.cluster.close()
@@ -800,14 +820,14 @@ def _client_gather(self, futures: list[ActorFuture]) -> list[Any]:
         """
 
         async def _gather_with_exception_handling(futures):
-            wrapped_futures = [self._wrap_dask_future(f) for f in futures]
-
             try:
+                wrapped_futures = [self._wrap_dask_future(f) for f in futures]
+
                 # Use asyncio.wait with FIRST_EXCEPTION instead of gather
                 done, pending = await asyncio.wait(
-                    wrapped_futures, timeout=120, return_when=asyncio.FIRST_EXCEPTION
+                    wrapped_futures, timeout=1200, return_when=asyncio.FIRST_EXCEPTION
                 )
-                                
+
                 exceptions = []
                 for future in done:
                     exc = future.exception()
@@ -819,9 +839,7 @@ async def _gather_with_exception_handling(futures):
                     raise ExceptionGroup("Multiple actor exceptions", exceptions)
 
                 if pending:
-                    pending_indices = sorted([
-                        wrapped_futures.index(p) for p in pending
-                    ])
+                    pending_indices = sorted([wrapped_futures.index(p) for p in pending])
                     raise TimeoutError(
                         f"Tasks didn't complete within timeout. {len(pending)} out of {len(wrapped_futures)} "
                         f"still pending. Pending task indices: {pending_indices}"
@@ -829,19 +847,21 @@ async def _gather_with_exception_handling(futures):
 
                 return await asyncio.gather(*wrapped_futures)
             except Exception as e:
-                logger.exception("Error in dask workers")
-                for f in wrapped_futures:
-                    if not f.done():
-                        f.cancel()         
+                logger.exception("Error in dask workers: %s")
+                for future in wrapped_futures:
+                    future.cancel()
                 self.teardown()
                 # sys.exit(1) would wait for dask to finish, which can cause hanging
                 # when workers are in a deadlock. Use os._exit to force immediate termination
-                os._exit(1) 
+                # TODO: this is more of a hack, we should propagate special exception that is 
+                # not caught by the rollout manager.
+                os._exit(1)
 
         # Use distributed.utils.sync to run the async function in the current thread
         return sync(self.client.loop, _gather_with_exception_handling, futures)  # type: ignore[arg-type]
 
 
+
 # Helpers
 
 

From ad375011b868d17cddd77903e448c5d72cadc5c9 Mon Sep 17 00:00:00 2001
From: Ori Kabeli <ori@futurehouse.org>
Date: Wed, 12 Mar 2025 07:14:03 -0500
Subject: [PATCH 15/21] Refactor exception handling in ParallelAsyncTransformer
 for improved clarity and reliability

---
 src/ldp/nn/handlers/transformer_handler.py | 98 +++++++++-------------
 1 file changed, 40 insertions(+), 58 deletions(-)

diff --git a/src/ldp/nn/handlers/transformer_handler.py b/src/ldp/nn/handlers/transformer_handler.py
index ddf159f8..82e9d0e3 100644
--- a/src/ldp/nn/handlers/transformer_handler.py
+++ b/src/ldp/nn/handlers/transformer_handler.py
@@ -239,22 +239,15 @@ def __init__(self, config: TransformerHandlerConfig):
                 assert_never(config.lm_type)
         super().__init__(model)
         self.tokenizer = tokenizer
-        logger.info(
-            f"Initialized tokenizer: {type(tokenizer).__name__}, vocab size: {len(tokenizer)}"
-        )
 
         maybe_set_tokenizer_chat_template(
             self.tokenizer, self.config.lm_config.chat_template
         )
-        logger.info(f"Chat template: {self.config.lm_config.chat_template}")
 
         self._setup_accelerator()
-        logger.info(f"Accelerator set up with device: {self.accelerator.device}")
 
         if config.checkpoint is not None:
-            logger.info(f"Loading checkpoint from: {config.checkpoint}")
             self.load_checkpoint(config.checkpoint)
-        logger.info("Initialization complete")
 
     def _setup_accelerator(self):
         self.accelerator = accelerate.Accelerator(
@@ -398,13 +391,6 @@ def _setup_accelerator(self):
                 buffer_dtype=torch.bfloat16,
             )
 
-        logger.info(f"Setting up accelerator with bf16={bf16}")
-        logger.info(f"Worker config: offload_cpu={self.worker_config.offload_cpu}, "
-                   f"activation_checkpointing={self.worker_config.activation_checkpointing}, "
-                   f"cpu_ram_efficient_loading={self.worker_config.cpu_ram_efficient_loading}, "
-                   f"state_dict_type={self.worker_config.state_dict_type}, "
-                   f"backward_prefetch={self.worker_config.backward_prefetch}")
-        
         self.accelerator = accelerate.Accelerator(
             # See note in TransformerHandler._setup_accelerator() about this
             # mixed_precision=("bf16" if bf16 else "no"),
@@ -420,16 +406,11 @@ def _setup_accelerator(self):
                 backward_prefetch=self.worker_config.backward_prefetch,
             ),
         )
-        logger.info(f"Accelerator setup complete on rank {self.worker_config.rank}")
 
         if self.config.lm_config.device == "meta":
-            logger.info(f"Preparing model for FSDP with meta device on rank {self.worker_config.rank}")
             self.module = prepare_model_for_fsdp_with_meta_device(self.module)
-            logger.info(f"Meta device preparation complete on rank {self.worker_config.rank}")
 
-        logger.info(f"Preparing model with accelerator on rank {self.worker_config.rank}")
         self.module = self.accelerator.prepare(self.module)
-        logger.info(f"Model preparation complete on rank {self.worker_config.rank}, model device: {self.module.device}, dtype: {self.module.dtype}")
 
     def set_seed(self, seed: int) -> None:
         """Set the seed for the current worker."""
@@ -789,7 +770,6 @@ def save_checkpoint(self, ckpt: os.PathLike | str, **kwargs) -> None:
         self._submit_and_gather("save_checkpoint", ckpt, **kwargs)
 
     def teardown(self) -> None:
-        logger.info(f"Shutting down Dask cluster, is_initialized: {self._initialized}")
         if self._initialized:
             self.client.shutdown()
             self.cluster.close()
@@ -806,7 +786,26 @@ def _wrap_dask_future(dask_future: ActorFuture):
         loop = asyncio.get_running_loop()
         return asyncio.ensure_future(loop.run_in_executor(None, dask_future.result))
 
-    def _client_gather(self, futures: list[ActorFuture]) -> list[Any]:
+    @staticmethod
+    def _raise_exceptions(done, pending, wrapped_futures):
+        exceptions = []
+        for future in done:
+            exc = future.exception()
+            if exc:
+                exceptions.append(exc)
+        if exceptions:
+            if len(exceptions) == 1:
+                raise exceptions[0]
+            raise ExceptionGroup("Multiple actor exceptions", exceptions)
+
+        if pending:
+            pending_indices = sorted([wrapped_futures.index(p) for p in pending])
+            raise TimeoutError(
+                f"Tasks didn't complete within timeout. {len(pending)} out of {len(wrapped_futures)} "
+                f"still pending. Pending task indices: {pending_indices}"
+            )
+
+    async def _client_gather_async(self, futures):
         """Gather results from futures, propagating exceptions as they arrive.
 
         Unlike client.gather() which waits for all futures to complete before raising
@@ -818,48 +817,31 @@ def _client_gather(self, futures: list[ActorFuture]) -> list[Any]:
         dask.gather() and can cause blocking issues or hide worker errors. This implementation
         works around those limitations.
         """
+        try:
+            wrapped_futures = [self._wrap_dask_future(f) for f in futures]
 
-        async def _gather_with_exception_handling(futures):
-            try:
-                wrapped_futures = [self._wrap_dask_future(f) for f in futures]
+            # Use asyncio.wait with FIRST_EXCEPTION instead of gather
+            done, pending = await asyncio.wait(
+                wrapped_futures, timeout=1200, return_when=asyncio.FIRST_EXCEPTION
+            )
 
-                # Use asyncio.wait with FIRST_EXCEPTION instead of gather
-                done, pending = await asyncio.wait(
-                    wrapped_futures, timeout=1200, return_when=asyncio.FIRST_EXCEPTION
-                )
+            self._raise_exceptions(done, pending, wrapped_futures)
 
-                exceptions = []
-                for future in done:
-                    exc = future.exception()
-                    if exc:
-                        exceptions.append(exc)
-                if exceptions:
-                    if len(exceptions) == 1:
-                        raise exceptions[0]
-                    raise ExceptionGroup("Multiple actor exceptions", exceptions)
-
-                if pending:
-                    pending_indices = sorted([wrapped_futures.index(p) for p in pending])
-                    raise TimeoutError(
-                        f"Tasks didn't complete within timeout. {len(pending)} out of {len(wrapped_futures)} "
-                        f"still pending. Pending task indices: {pending_indices}"
-                    )
-
-                return await asyncio.gather(*wrapped_futures)
-            except Exception as e:
-                logger.exception("Error in dask workers: %s")
-                for future in wrapped_futures:
-                    future.cancel()
-                self.teardown()
-                # sys.exit(1) would wait for dask to finish, which can cause hanging
-                # when workers are in a deadlock. Use os._exit to force immediate termination
-                # TODO: this is more of a hack, we should propagate special exception that is 
-                # not caught by the rollout manager.
-                os._exit(1)
+            return await asyncio.gather(*wrapped_futures)
+        except Exception:
+            logger.exception("Error in dask workers: %s")
+            for future in wrapped_futures:
+                future.cancel()
+            self.teardown()
+            # sys.exit(1) would wait for dask to finish, which can cause hanging
+            # when workers are in a deadlock. Use os._exit to force immediate termination
+            # TODO: this is more of a hack, we should propagate special exception that is
+            # not caught by the rollout manager.
+            os._exit(1)
 
+    def _client_gather(self, futures: list[ActorFuture]) -> list[Any]:
         # Use distributed.utils.sync to run the async function in the current thread
-        return sync(self.client.loop, _gather_with_exception_handling, futures)  # type: ignore[arg-type]
-
+        return sync(self.client.loop, self._client_gather_async, futures)  # type: ignore[arg-type]
 
 
 # Helpers

From 62f42d108afa2aefce2035496f3413717a3fb7ca Mon Sep 17 00:00:00 2001
From: Ori Kabeli <ori@futurehouse.org>
Date: Wed, 12 Mar 2025 07:47:42 -0500
Subject: [PATCH 16/21] nits

---
 src/ldp/alg/rollout.py                 | 3 +--
 src/ldp/nn/agent/simple_local_agent.py | 1 +
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/ldp/alg/rollout.py b/src/ldp/alg/rollout.py
index da40f669..10dcf820 100644
--- a/src/ldp/alg/rollout.py
+++ b/src/ldp/alg/rollout.py
@@ -46,7 +46,6 @@ def reraise_exc_as(reraise: type[CaughtError], enabled: bool) -> Iterator[None]:
         yield
     except Exception as e:
         if enabled:
-            # Minimal logging instead of spamming. Detailed error stored in the trajectory's metadata.
             logger.debug(f"Reraising {reraise.exc_type} exception.")
             raise reraise(e) from None
         raise
@@ -81,7 +80,7 @@ async def sample_trajectories(
         max_steps: int | None = None,
         *,
         log_exceptions_immediately: bool = True,
-    ) -> list[tuple[Trajectory, Environment]] | list[Trajectory]:
+    ):
         """Sample trajectories from environments, either via factory or pre-created.
 
         There are two main ways to use this method:
diff --git a/src/ldp/nn/agent/simple_local_agent.py b/src/ldp/nn/agent/simple_local_agent.py
index 00957d0d..d7fbd86d 100644
--- a/src/ldp/nn/agent/simple_local_agent.py
+++ b/src/ldp/nn/agent/simple_local_agent.py
@@ -127,6 +127,7 @@ async def get_asv(
         return cast("OpResult[ToolRequestMessage]", result), next_state, 0.0
 
     def _validate_token_count(self, messages: list[Message], tools: list[Tool]):
+        """Asserts token count for the trajectory is within the limit."""
         if self.llm_model.max_traj_token_count is None:
             return
         messages_for_tokenizer = self._llm_call_op.prep_messages_for_tokenizer(messages)

From 1b486cfd29f73ed6c43e50bd93ddbb32814050af Mon Sep 17 00:00:00 2001
From: Ori Kabeli <ori@futurehouse.org>
Date: Wed, 19 Mar 2025 02:54:09 -0500
Subject: [PATCH 17/21] nits code review

---
 src/ldp/alg/rollout.py                     | 88 +++++++++++-----------
 src/ldp/graph/async_torch.py               | 14 +++-
 src/ldp/nn/handlers/transformer_handler.py |  2 +-
 3 files changed, 59 insertions(+), 45 deletions(-)

diff --git a/src/ldp/alg/rollout.py b/src/ldp/alg/rollout.py
index 10dcf820..74519936 100644
--- a/src/ldp/alg/rollout.py
+++ b/src/ldp/alg/rollout.py
@@ -5,13 +5,14 @@
 from collections import Counter
 from collections.abc import Callable, Iterator, Sequence
 from contextlib import contextmanager, nullcontext
-from typing import Any, TypeVar
+from typing import Any, TypeVar, overload
 
 from aviary.core import Environment, Message
 from tqdm.asyncio import tqdm
 
 from ldp.agent import Agent
 from ldp.data_structures import Trajectory, Transition
+from ldp.utils import format_error_details
 
 from .callbacks import Callback
 
@@ -72,67 +73,65 @@ def __init__(
         self.traj_buffer: dict[str, Trajectory] = {}
         self.callbacks = callbacks or []
 
-    async def sample_trajectories(
+    @overload
+    async def sample_trajectories(  # noqa: D418
         self,
-        environment_factory: Callable[[], TEnv] | None = None,
-        environments: Sequence[TEnv] | None = None,
+        environment_factory: Callable[[], TEnv],
         batch_size: int = 1,
         max_steps: int | None = None,
-        *,
-        log_exceptions_immediately: bool = True,
-    ):
-        """Sample trajectories from environments, either via factory or pre-created.
+    ) -> list[tuple[Trajectory, TEnv]]:
+        """Run rollouts in parallel, using a factory to construct environments.
 
-        There are two main ways to use this method:
+        We will construct `batch_size` environments and run rollouts on each of them.
+        If `max_steps` is set, rollouts will be truncated at this value. If a rollout
+        has fewer than `max_steps`, then a new environment will be constructed and another
+        rollout will be started until `max_steps` is reached.
 
-        1. Using an environment factory:
-            Run rollouts in parallel, using a factory to construct environments.
-            We will construct `batch_size` environments and run rollouts on each of them.
-            If `max_steps` is set, rollouts will be truncated at this value. If a rollout
-            has fewer than `max_steps`, then a new environment will be constructed and another
-            rollout will be started until `max_steps` is reached.
+        Args:
+            environment_factory: A no-argument callable that returns
+                an environment instance
+            batch_size (int, optional): Defaults to 1.
+            max_steps (int | None, optional): Max steps per rollout. Defaults to None (see above).
 
-            In this case, returns a list of (trajectory, environment) tuples.
+        Returns:
+            list[tuple[Trajectory, Environment]]: A list of (trajectory, environment) tuples: one per rollout.
+        """
 
-        2. Using a sequence of environments:
-            Run rollouts in parallel on a list of provided environments.
-            In this case, returns a list of trajectories.
+    @overload
+    async def sample_trajectories(  # noqa: D418
+        self,
+        environments: Sequence[Environment],
+        max_steps: int | None = None,
+    ) -> list[Trajectory]:
+        """Run rollouts in parallel on a list of provided environments.
 
         Args:
-            environment_factory: A no-argument callable that returns an environment instance
-            environments: A list of environments to run rollouts on
-            batch_size: Number of parallel environments to run when using environment_factory. Defaults to 1.
+            environments: A list of environments to run rollouts on.
             max_steps: Max steps per rollout. Defaults to None, in which case the rollouts are run
                 until environment returns done.
-            log_exceptions_immediately: Whether to log exceptions as they occur
-                or only at the end of the rollouts.
-
-        Returns:
-            Either list[tuple[Trajectory, Environment]] or list[Trajectory] depending on whether
-            environment_factory or environments is provided.
-
-        Raises:
-            TypeError: If neither environment_factory nor environments is provided.
         """
-        if environment_factory is not None:
-            assert environments is None, (
+
+    async def sample_trajectories(self, **kwargs):
+        if "environment_factory" in kwargs:
+            assert "environments" not in kwargs, (
                 "Cannot use environment_factory with environments"
             )
+
             return await self._sample_trajectories_from_env_factory(
-                environment_factory,
-                batch_size,
-                max_steps,
-                log_exceptions_immediately=log_exceptions_immediately,
+                kwargs["environment_factory"],
+                kwargs.get("batch_size", 1),
+                kwargs.get("max_steps"),
+                log_exceptions_immediately=kwargs.get("log_exceptions_immediately", True)
             )
 
-        if environments is not None:
-            assert environment_factory is None, (
+        if "environments" in kwargs:
+            assert "environment_factory" not in kwargs, (
                 "Cannot use environments with environment_factory"
             )
             return await self._sample_trajectories_from_envs(
-                environments,
-                max_steps,
-                log_exceptions_immediately=log_exceptions_immediately,
+                kwargs["environments"],
+                kwargs.get("max_steps"),
+                log_exceptions_immediately=kwargs.get("log_exceptions_immediately", True),
             )
 
         raise TypeError(
@@ -176,6 +175,7 @@ async def rollout_with_args(idx: int, **rollout_kwargs):
             desc="Rollouts",
             unit="rollout",
             ncols=0,
+            disable=log_exceptions_immediately,
         ) as pbar:
             while tasks:
                 done, pending = await asyncio.wait(
@@ -258,6 +258,7 @@ async def _sample_trajectories_from_envs(
             desc="Rollouts",
             unit="rollout",
             ncols=0,
+            disable=log_exceptions_immediately,
         ) as pbar:
             for task in asyncio.as_completed(tasks):
                 trajectory = await task
@@ -344,7 +345,8 @@ async def store_step(step: Transition):
             # NOTE: This trajectory should not be used for regular training.
             # We save the last transition here for debugging, etc.
             if log_exceptions_immediately:
-                logger.exception(f"Exception in rollout {traj_id}: {e.original_exc}")
+                error_details = format_error_details(e.original_exc)
+                logger.exception(f"Exception in rollout {traj_id}:\n{error_details}")
 
             await store_step(
                 Transition(
diff --git a/src/ldp/graph/async_torch.py b/src/ldp/graph/async_torch.py
index d0aa8274..2eb19e85 100644
--- a/src/ldp/graph/async_torch.py
+++ b/src/ldp/graph/async_torch.py
@@ -1,6 +1,7 @@
 __all__ = ["AsyncTorchModule", "async_protect_torch_call"]
 
 import asyncio
+import logging
 import operator
 import time
 from abc import ABC, abstractmethod
@@ -19,6 +20,9 @@
         "Please run `pip install ldp[nn]`."
     ) from None
 
+
+logger = logging.getLogger(__name__)
+
 _TORCH_LOCK = asyncio.Lock()
 
 # Supported devices here: https://pytorch.org/docs/stable/amp.html#torch.autocast
@@ -90,6 +94,7 @@ def __init__(
         self._work_buffer: list[tuple[float, UUID, dict[str, Any]]] = []
         self._result_buffer: dict[UUID, Any] = {}
         self._lock = asyncio.Lock()
+        self._exception_raised: Exception | None = None
 
     async def __call__(self, **kwargs):
         request_id = uuid4()
@@ -104,13 +109,20 @@ async def __call__(self, **kwargs):
                 # Only one coroutine allowed in here when:
                 # - modifying the result buffer
                 # - modifying the work buffer
+                if self._exception_raised is not None:
+                    logger.info("Exception raised in another coroutine")
+                    raise self._exception_raised
 
                 if request_id in self._result_buffer:
                     # Our request was fulfilled by this or another coroutine!
                     return self._result_buffer.pop(request_id)
 
                 # Try to run a batch.
-                await self._maybe_process_batch()
+                try:
+                    await self._maybe_process_batch()
+                except Exception as e:
+                    self._exception_raised = e
+                    raise
 
             # Sleep, to let another coroutine take over if it needs to
             await asyncio.sleep(0.0)
diff --git a/src/ldp/nn/handlers/transformer_handler.py b/src/ldp/nn/handlers/transformer_handler.py
index 82e9d0e3..7ca01b74 100644
--- a/src/ldp/nn/handlers/transformer_handler.py
+++ b/src/ldp/nn/handlers/transformer_handler.py
@@ -59,7 +59,7 @@
     # Gives us more time to debug a downed worker. TODO: see if there are negative consequences
     # of having this always enabled
     "distributed.comm.timeouts.connect": "300s",
-    "distributed.comm.timeouts.tcp": "300s",
+    "distributed.comm.timeouts.tcp": "1200s",
 })
 
 TReturn = TypeVar("TReturn")

From bddaa92f11cc69cd86ff974e474f73faa6967dde Mon Sep 17 00:00:00 2001
From: Ori Kabeli <ori@futurehouse.org>
Date: Wed, 19 Mar 2025 06:32:37 -0500
Subject: [PATCH 18/21] nits

---
 src/ldp/alg/rollout.py                 |  8 ++++++--
 src/ldp/nn/agent/simple_local_agent.py | 10 +++++-----
 2 files changed, 11 insertions(+), 7 deletions(-)

diff --git a/src/ldp/alg/rollout.py b/src/ldp/alg/rollout.py
index 74519936..28c58450 100644
--- a/src/ldp/alg/rollout.py
+++ b/src/ldp/alg/rollout.py
@@ -121,7 +121,9 @@ async def sample_trajectories(self, **kwargs):
                 kwargs["environment_factory"],
                 kwargs.get("batch_size", 1),
                 kwargs.get("max_steps"),
-                log_exceptions_immediately=kwargs.get("log_exceptions_immediately", True)
+                log_exceptions_immediately=kwargs.get(
+                    "log_exceptions_immediately", True
+                ),
             )
 
         if "environments" in kwargs:
@@ -131,7 +133,9 @@ async def sample_trajectories(self, **kwargs):
             return await self._sample_trajectories_from_envs(
                 kwargs["environments"],
                 kwargs.get("max_steps"),
-                log_exceptions_immediately=kwargs.get("log_exceptions_immediately", True),
+                log_exceptions_immediately=kwargs.get(
+                    "log_exceptions_immediately", True
+                ),
             )
 
         raise TypeError(
diff --git a/src/ldp/nn/agent/simple_local_agent.py b/src/ldp/nn/agent/simple_local_agent.py
index d7fbd86d..82d5f5ec 100644
--- a/src/ldp/nn/agent/simple_local_agent.py
+++ b/src/ldp/nn/agent/simple_local_agent.py
@@ -46,7 +46,7 @@ class AgentLMConfig(_LMConfig):
         ),
         validate_default=True,
     )
-    max_traj_token_count: int | None = Field(
+    max_messages_token_count: int | None = Field(
         default=None,
         description="If set, raise an error if the total tokens in the trajectory exceed this value.",
     )
@@ -128,7 +128,7 @@ async def get_asv(
 
     def _validate_token_count(self, messages: list[Message], tools: list[Tool]):
         """Asserts token count for the trajectory is within the limit."""
-        if self.llm_model.max_traj_token_count is None:
+        if self.llm_model.max_messages_token_count is None:
             return
         messages_for_tokenizer = self._llm_call_op.prep_messages_for_tokenizer(messages)
         tools_for_tokenizer = self._llm_call_op.prep_tools_for_tokenizer(tools)
@@ -138,12 +138,12 @@ def _validate_token_count(self, messages: list[Message], tools: list[Tool]):
             messages=messages_for_tokenizer,
             tools=tools_for_tokenizer,  # type: ignore[arg-type]
         )
-        if total_tokens > self.llm_model.max_traj_token_count:
+        if total_tokens > self.llm_model.max_messages_token_count:
             logger.error(
-                f"Token limit exceeded for trajectory: {total_tokens} > {self.llm_model.max_traj_token_count}"
+                f"Token limit exceeded for trajectory: {total_tokens} > {self.llm_model.max_messages_token_count}"
             )
             raise ValueError(
-                f"Token limit exceeded for trajectory: {total_tokens} > {self.llm_model.max_traj_token_count}"
+                f"Token limit exceeded for trajectory: {total_tokens} > {self.llm_model.max_messages_token_count}"
             )
 
     # TODO: maybe remove these recomputation methods. I added them to debug some things. But idk,

From 310131eab6a8bf461fec75f0672cdebc6c78fbe2 Mon Sep 17 00:00:00 2001
From: Ori Kabeli <ori@futurehouse.org>
Date: Wed, 19 Mar 2025 06:57:37 -0500
Subject: [PATCH 19/21] nits

---
 src/ldp/alg/rollout.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/src/ldp/alg/rollout.py b/src/ldp/alg/rollout.py
index 28c58450..e6645049 100644
--- a/src/ldp/alg/rollout.py
+++ b/src/ldp/alg/rollout.py
@@ -109,6 +109,9 @@ async def sample_trajectories(  # noqa: D418
             environments: A list of environments to run rollouts on.
             max_steps: Max steps per rollout. Defaults to None, in which case the rollouts are run
                 until environment returns done.
+            log_exceptions_immediately: Whether to log exceptions in the rollout immediately 
+                to the console. Defaults to True. If False, progress bar will show and a summary
+                will be logged after all rollouts are complete.
         """
 
     async def sample_trajectories(self, **kwargs):

From 4e4090890f70d52bbe2539bf6c5f7cc2d81bdd37 Mon Sep 17 00:00:00 2001
From: Ori Kabeli <ori@futurehouse.org>
Date: Wed, 19 Mar 2025 06:59:33 -0500
Subject: [PATCH 20/21] nits

---
 src/ldp/graph/async_torch.py | 4 ----
 1 file changed, 4 deletions(-)

diff --git a/src/ldp/graph/async_torch.py b/src/ldp/graph/async_torch.py
index 2eb19e85..f9f76626 100644
--- a/src/ldp/graph/async_torch.py
+++ b/src/ldp/graph/async_torch.py
@@ -109,10 +109,6 @@ async def __call__(self, **kwargs):
                 # Only one coroutine allowed in here when:
                 # - modifying the result buffer
                 # - modifying the work buffer
-                if self._exception_raised is not None:
-                    logger.info("Exception raised in another coroutine")
-                    raise self._exception_raised
-
                 if request_id in self._result_buffer:
                     # Our request was fulfilled by this or another coroutine!
                     return self._result_buffer.pop(request_id)

From b5323bc06eff1b5a7f7cdadd4351fd43d7eb5790 Mon Sep 17 00:00:00 2001
From: Ori Kabeli <ori@futurehouse.org>
Date: Wed, 19 Mar 2025 07:00:47 -0500
Subject: [PATCH 21/21] nits

---
 src/ldp/graph/async_torch.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/src/ldp/graph/async_torch.py b/src/ldp/graph/async_torch.py
index f9f76626..612e0dad 100644
--- a/src/ldp/graph/async_torch.py
+++ b/src/ldp/graph/async_torch.py
@@ -106,6 +106,10 @@ async def __call__(self, **kwargs):
 
         while True:
             async with self._lock:
+                if self._exception_raised is not None:
+                    logger.info("Exception raised in another coroutine")
+                    raise self._exception_raised
+
                 # Only one coroutine allowed in here when:
                 # - modifying the result buffer
                 # - modifying the work buffer