partial

Author: vmoens

docs/source/reference/collectors_weightsync.rst

@@ -5,7 +5,6 @@
 from __future__ import annotations
 
 import time
-
 from typing import Any, Literal
 
 import torch
@@ -612,7 +611,6 @@ def get_wandb_run_id(wandb_logger):
     """
     try:
         # Wait a bit for wandb to initialize
-        import time
 
         max_attempts = 10
         for attempt in range(max_attempts):

sota-implementations/expert-iteration/ei_utils.py

@@ -7,8 +7,10 @@
 import argparse
 import gc
 import importlib.util
+import threading
 
 import time
+from concurrent.futures import ThreadPoolExecutor, wait
 from functools import partial
 
 import pytest
@@ -412,8 +414,6 @@ def slow_forward(self, td_input, **kwargs):
 @pytest.fixture
 def monkey_patch_forward_for_instrumentation():
     """Fixture to monkey patch the forward method to add detailed processing event tracking."""
-    import threading
-    import time
 
     # Track processing events
     processing_events = []
@@ -2706,8 +2706,6 @@ def test_batching_min_batch_size_one_immediate_processing(
         monkey_patch_forward_for_timing,
     ):
         """Test that with min_batch_size=1, first request is processed immediately and subsequent ones are grouped."""
-        import time
-        from concurrent.futures import ThreadPoolExecutor, wait
 
         # Create wrapper using helper function
         wrapper = create_batching_test_wrapper(

test/llm/test_wrapper.py

@@ -4096,6 +4096,7 @@ def test_start_multi(self, total_frames, cls):
         "weight_sync_scheme",
         [None, MultiProcessWeightSyncScheme, SharedMemWeightSyncScheme],
     )
+    @pytest.mark.flaky(reruns=3, reruns_delay=0.5)
     def test_start_update_policy(self, total_frames, cls, weight_sync_scheme):
         rb = ReplayBuffer(storage=LazyMemmapStorage(max_size=1000))
         env = CountingEnv()

test/test_collector.py

@@ -18,7 +18,6 @@
 from copy import copy
 from functools import partial
 from sys import platform
-from torchrl import logger as torchrl_logger
 
 import numpy as np
 
@@ -39,6 +38,7 @@
 from tensordict.nn import TensorDictModule, TensorDictSequential, WrapModule
 from tensordict.utils import _unravel_key_to_tuple, assert_allclose_td
 from torch import multiprocessing as mp, nn, Tensor
+from torchrl import logger as torchrl_logger
 from torchrl._utils import _replace_last, prod, set_auto_unwrap_transformed_env
 
 from torchrl.collectors import MultiSyncDataCollector, SyncDataCollector
@@ -57,7 +57,6 @@
     Unbounded,
     UnboundedContinuous,
 )
-from torchrl.envs.transforms import TransformedEnv
 from torchrl.envs import (
     ActionMask,
     BinarizeReward,
@@ -139,7 +138,14 @@
 from torchrl.envs.transforms.vc1 import _has_vc
 from torchrl.envs.transforms.vip import _VIPNet, VIPRewardTransform
 from torchrl.envs.utils import check_env_specs, MarlGroupMapType, step_mdp
-from torchrl.modules import GRUModule, LSTMModule, MLP, ProbabilisticActor, TanhNormal, RandomPolicy
+from torchrl.modules import (
+    GRUModule,
+    LSTMModule,
+    MLP,
+    ProbabilisticActor,
+    RandomPolicy,
+    TanhNormal,
+)
 from torchrl.modules.utils import get_primers_from_module
 from torchrl.record.recorder import VideoRecorder
 from torchrl.testing.modules import BiasModule