pre commit fix

Author: chenyushuo

tests/buffer/sample_strategy_test.py

@@ -1,12 +1,15 @@
 import asyncio
-from collections import deque
 import shutil
+from collections import deque
+
 import torch
+
 from tests.tools import RayUnittestBaseAysnc, get_template_config
-from trinity.algorithm.sample_strategy.sample_strategy import SAMPLE_STRATEGY, SampleStrategy
+from trinity.algorithm.sample_strategy.sample_strategy import (
+    SAMPLE_STRATEGY,
+    SampleStrategy,
+)
 from trinity.buffer.buffer import get_buffer_writer
-from trinity.common.config import ExperienceBufferConfig
-from trinity.common.constants import StorageType
 from trinity.common.experience import Experience
 
 
@@ -29,16 +32,30 @@ def _default_exp_list(self):
             ]
             for i in range(self.num_steps)
         ]
-    
+
     def _default_steps(self):
         return [0, 5, 10, 15]
 
     async def _verify_model_version(self, step, expected_versions):
         batch, metrics, _ = await self.sample_strategy.sample(step=step)
-        self.assertEqual(batch.rewards.tolist(), expected_versions, f"Model versions mismatch at step {step}")
-        self.assertEqual(metrics['sample/model_version/min'], min(expected_versions), f"Min model version mismatch at step {step}")
-        self.assertEqual(metrics['sample/model_version/max'], max(expected_versions), f"Max model version mismatch at step {step}")
-        self.assertEqual(metrics['sample/model_version/mean'], sum(expected_versions) / len(expected_versions), f"Mean model version mismatch at step {step}")
+        self.assertEqual(
+            batch.rewards.tolist(), expected_versions, f"Model versions mismatch at step {step}"
+        )
+        self.assertEqual(
+            metrics["sample/model_version/min"],
+            min(expected_versions),
+            f"Min model version mismatch at step {step}",
+        )
+        self.assertEqual(
+            metrics["sample/model_version/max"],
+            max(expected_versions),
+            f"Max model version mismatch at step {step}",
+        )
+        self.assertEqual(
+            metrics["sample/model_version/mean"],
+            sum(expected_versions) / len(expected_versions),
+            f"Mean model version mismatch at step {step}",
+        )
 
     async def _verify_sampling_model_versions(self, exps_list, expected_model_versions_map):
         # Initialize buffer writer and sample strategy
@@ -51,7 +68,7 @@ async def _verify_sampling_model_versions(self, exps_list, expected_model_versio
             buffer_config=self.config.buffer,
             **self.config.algorithm.sample_strategy_args,
         )
-        
+
         # Write experiences to buffer, while sample and validate model versions
         current_task = None
         for step, exps in enumerate(exps_list):
@@ -106,7 +123,7 @@ async def test_default_queue_staleness_control_sample_strategy(self):
 
         await self._verify_sampling_model_versions(exps_list, expected_model_versions_map)
 
-    def _simulate_priority_queue(self, steps, staleness_limit = float('inf')):
+    def _simulate_priority_queue(self, steps, staleness_limit=float("inf")):
         expected_model_versions_map = {}
         buffer = deque()
         exp_pool = deque()