[low-prority] remove unused imports

pipelinerl/async_llm.py

@@ -1,7 +1,7 @@
 import logging
 import aiohttp
 
-from tapeagents.core import LLMCall, LLMOutput, Prompt, TokenLogprob
+from tapeagents.core import LLMCall, LLMOutput, Prompt
 from tapeagents.llms.trainable import TrainableLLM
 
 

pipelinerl/load_datasets.py

@@ -8,7 +8,6 @@
 import random
 from datasets import load_dataset
 from omegaconf import DictConfig
-from tqdm import tqdm
 
 """
 math_verify expects the following LaTeX format for the gold answer (with $ or \\boxed).

pipelinerl/run_actor.py

@@ -23,7 +23,6 @@
 from pipelinerl.math_rollouts import RolloutResult, generate_math_rollout
 from pipelinerl.state import TrainerState
 import asyncio
-from collections import defaultdict
 from pipelinerl.streams import (
     SingleStreamSpec,
     StreamSpec,

pipelinerl/run_llm.py

@@ -1,9 +1,6 @@
 import asyncio
-import json
 import logging
-import os
 import signal
-from pydantic import TypeAdapter
 import torch
 import uvloop
 from vllm import AsyncLLMEngine
@@ -21,7 +18,6 @@
 )
 from vllm.engine.arg_utils import AsyncEngineArgs
 from vllm.entrypoints.openai.tool_parsers import ToolParserManager
-from vllm.logger import init_logger
 from vllm._version import version
 from vllm.worker.worker import Worker
 from vllm.executor.multiproc_worker_utils import ProcessWorkerWrapper
@@ -33,8 +29,7 @@
 from vllm.worker.multi_step_model_runner import MultiStepModelRunner
 
 
-import torch.distributed as dist
-from pipelinerl.run_finetune import TrainerMessage, WeightUpdateRequest
+from pipelinerl.run_finetune import WeightUpdateRequest
 import pipelinerl.torch_utils
 
 logger = logging.getLogger(__name__)
@@ -157,7 +152,7 @@ async def receive_weight_update(self, message: WeightUpdateRequest):
             self.driver_worker.receive_weight_update(message)
             for future in futures:
                 future.get()
-            logger.info(f"All workers received weight updates")
+            logger.info("All workers received weight updates")
 
 
 async def run_server(args, **uvicorn_kwargs) -> None:

pipelinerl/utils.py

@@ -3,17 +3,13 @@
 import logging
 import os
 import shutil
-import subprocess
 import time
 from pathlib import Path
 import traceback
-from typing import Dict, Mapping, Optional, TextIO, Union, List
-import threading
+from typing import Dict, Mapping, List
 import numpy as np
 import psutil
 import requests
-import torch
-from tenacity import retry, stop_after_attempt, wait_exponential
 from transformers import PreTrainedTokenizer
 from collections import defaultdict