[pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Author: pre-commit-ci[bot]

applications/ColossalChat/coati/distributed/agent/agentic_producer.py

@@ -1,5 +1,4 @@
 import copy
-import random
 import re
 from typing import Any, Dict
 from uuid import uuid4

applications/ColossalChat/coati/distributed/launch.py

@@ -87,7 +87,9 @@ def launch_distributed(
     num_samples = get_jsonl_size_fast(dataset_path)
     global_inference_batch_size = inference_batch_size * num_producers
     num_update_per_episode = num_samples // global_inference_batch_size
-    num_recv_per_update = inference_batch_size // inference_microbatch_size if "async-agentic" not in inference_backend else 1
+    num_recv_per_update = (
+        inference_batch_size // inference_microbatch_size if "async-agentic" not in inference_backend else 1
+    )
 
     run_name = f"{inference_backend}_bs_{train_batch_size * train_dp_size}_temp_{generate_config['temperature']:.01f}_top_p_{generate_config['top_p']:.02f}"
     wandb_group_name = str(uuid.uuid4())

applications/ColossalChat/coati/distributed/producer.py

@@ -847,6 +847,7 @@ class AsyncSimpleProducer(BaseAsyncProducer):
     Asyncronous version of the producer that uses vLLM for generation.
     This class is designed to handle multiple producer actors and distribute tasks among them.
     """
+
     @torch.no_grad()
     async def rollout(self, input_ids, attention_mask, **kwargs):
         # naive rollout strategy without load balancing

applications/ColossalChat/coati/distributed/utils.py

@@ -1,12 +1,14 @@
 import json
 import os
+import random
 from typing import Any, Dict, List
-import asyncio
+
+import ray
 import torch
 from filelock import FileLock
-import random
+
 from colossalai.shardformer.layer.loss import dist_log_prob
-import ray
+
 
 def unbind_batch(batch: Dict[str, torch.Tensor]) -> List[Dict[str, torch.Tensor]]:
     batches = []
@@ -166,6 +168,7 @@ def safe_append_to_jsonl_file(file_path, data):
                 json_line = json.dumps(entry, ensure_ascii=False)
                 f.write(json_line + "\n")
 
+
 @ray.remote
 class LoadBalancer:
     def __init__(self, worker_counts):
@@ -180,10 +183,9 @@ def get_next_worker(self, worker_type, amount=1):
         chosen = random.choice(candidates)
         self.load[worker_type][chosen] += amount
         return chosen, self.load[worker_type]
-    
+
     def increase_load(self, worker_type, worker_id, amount=1):
         self.load[worker_type][worker_id] += amount
-        
+
     def decrease_load(self, worker_type, worker_id, amount=1):
         self.load[worker_type][worker_id] -= amount
-