Update

[ghstack-poisoned]

Author: vmoens

.github/unittest/llm/scripts_llm/setup_env.sh

@@ -6,10 +6,13 @@
 # Do not install PyTorch and torchvision here, otherwise they also get cached.
 
 set -e
-apt-get update && apt-get upgrade -y && apt-get install -y git cmake
+export DEBIAN_FRONTEND=noninteractive
+export TZ=UTC
+apt-get update
+apt-get install -yq --no-install-recommends git cmake
 # Avoid error: "fatal: unsafe repository"
 git config --global --add safe.directory '*'
-apt-get install -y wget \
+apt-get install -yq --no-install-recommends wget \
     gcc \
     g++ \
     unzip \
@@ -27,7 +30,10 @@ apt-get install -y wget \
     libgles2
 
 # Upgrade specific package
-apt-get upgrade -y libstdc++6
+apt-get install -yq --no-install-recommends --only-upgrade libstdc++6
+
+apt-get clean
+rm -rf /var/lib/apt/lists/*
 
 this_dir="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )"
 root_dir="$(git rev-parse --show-toplevel)"

torchrl/modules/llm/backends/vllm/vllm_async.py

@@ -20,12 +20,9 @@
 from concurrent.futures import ThreadPoolExecutor, wait
 from typing import Any, Literal, TYPE_CHECKING
 
-import ray
 
 import torch
 
-from ray.util.placement_group import placement_group, remove_placement_group
-from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
 from torchrl._utils import logger as torchrl_logger
 
 # Import RLvLLMEngine and shared utilities
@@ -43,6 +40,24 @@
 TIMEOUT_SECONDS = os.getenv("TORCHRL_VLLM_TIMEOUT_SECONDS", 300)
 
 
+def _get_ray():
+    """Import Ray on demand to avoid global import side-effects.
+
+    Returns:
+        ModuleType: The imported Ray module.
+
+    Raises:
+        ImportError: If Ray is not installed.
+    """
+    try:
+        import ray  # type: ignore
+
+        return ray
+    except Exception as e:  # pragma: no cover - surfaced to callers
+        raise ImportError(
+            "ray is not installed. Please install it with `pip install ray`."
+        ) from e
+
 class _AsyncvLLMWorker:
     """Async vLLM worker for Ray with weight update capabilities.
 
@@ -267,7 +282,7 @@ async def generate(
                 "vllm is not installed. Please install it with `pip install vllm`."
             )
 
-        from vllm import RequestOutput, SamplingParams, TokensPrompt
+        from vllm import SamplingParams, TokensPrompt
 
         # Track whether input was originally a single prompt
         single_prompt_input = False
@@ -474,11 +489,7 @@ def _gpus_per_replica(engine_args: AsyncEngineArgs) -> int:
     )
 
 
-# Create Ray remote versions
-if ray is not None and _has_vllm:
-    _AsyncLLMEngineActor = ray.remote(num_cpus=0, num_gpus=0)(_AsyncLLMEngine)
-else:
-    _AsyncLLMEngineActor = None
+# Ray actor wrapper is created lazily in __init__ to avoid global Ray import.
 
 
 class AsyncVLLM(RLvLLMEngine):
@@ -583,17 +594,18 @@ def __init__(
             raise ImportError(
                 "vllm is not installed. Please install it with `pip install vllm`."
             )
-        if ray is None:
-            raise ImportError(
-                "ray is not installed. Please install it with `pip install ray`."
-            )
+        # Lazily import ray only when constructing the actor class to avoid global import
 
         # Enable prefix caching by default for better performance
         engine_args.enable_prefix_caching = enable_prefix_caching
 
         self.engine_args = engine_args
         self.num_replicas = num_replicas
-        self.actor_class = actor_class or _AsyncLLMEngineActor
+        if actor_class is None:
+            ray = _get_ray()
+            self.actor_class = ray.remote(num_cpus=0, num_gpus=0)(_AsyncLLMEngine)
+        else:
+            self.actor_class = actor_class
         self.actors: list = []
         self._launched = False
         self._service_id = uuid.uuid4().hex[
@@ -608,6 +620,11 @@ def _launch(self):
             torchrl_logger.warning("AsyncVLLMEngineService already launched")
             return
 
+        # Local imports to avoid global Ray dependency
+        ray = _get_ray()
+        from ray.util.placement_group import placement_group
+        from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
+
         torchrl_logger.info(
             f"Launching {self.num_replicas} async vLLM engine actors..."
         )
@@ -938,6 +955,7 @@ def generate(
         Returns:
             RequestOutput | list[RequestOutput]: Generated outputs from vLLM.
         """
+        ray = _get_ray()
         # Check if this is a batch request
         if self._is_batch(prompts, prompt_token_ids):
             # Handle batched input by unbinding and sending individual requests
@@ -1062,6 +1080,9 @@ def shutdown(self):
             f"Shutting down {len(self.actors)} async vLLM engine actors..."
         )
 
+        ray = _get_ray()
+        from ray.util.placement_group import remove_placement_group
+
         # Kill all actors
         for i, actor in enumerate(self.actors):
             try:
@@ -1254,6 +1275,7 @@ def _update_weights_with_nccl_broadcast_simple(
         )
 
         updated_weights = 0
+        ray = _get_ray()
         with torch.cuda.device(0):  # Ensure we're on the correct CUDA device
             for name, weight in gpu_weights.items():
                 # Convert dtype to string name (like periodic-mono)
@@ -1330,6 +1352,7 @@ def get_num_unfinished_requests(
                 "AsyncVLLM service must be launched before getting request counts"
             )
 
+        ray = _get_ray()
         if actor_index is not None:
             if not (0 <= actor_index < len(self.actors)):
                 raise IndexError(
@@ -1360,6 +1383,7 @@ def get_cache_usage(self, actor_index: int | None = None) -> float | list[float]
                 "AsyncVLLM service must be launched before getting cache usage"
             )
 
+        ray = _get_ray()
         if actor_index is not None:
             if not (0 <= actor_index < len(self.actors)):
                 raise IndexError(
@@ -1672,6 +1696,7 @@ def _select_by_requests(self) -> int:
             futures = [
                 actor.get_num_unfinished_requests.remote() for actor in self.actors
             ]
+            ray = _get_ray()
             request_counts = ray.get(futures)
 
         # Find the actor with minimum pending requests
@@ -1699,6 +1724,7 @@ def _select_by_cache_usage(self) -> int:
         else:
             # Query actors directly
             futures = [actor.get_cache_usage.remote() for actor in self.actors]
+            ray = _get_ray()
             cache_usages = ray.get(futures)
 
         # Find the actor with minimum cache usage
@@ -1838,7 +1864,8 @@ def _is_actor_overloaded(self, actor_index: int) -> bool:
                 futures = [
                     actor.get_num_unfinished_requests.remote() for actor in self.actors
                 ]
-                request_counts = ray.get(futures)
+            ray = _get_ray()
+            request_counts = ray.get(futures)
 
             if not request_counts:
                 return False
@@ -1887,8 +1914,9 @@ def get_stats(self) -> dict[str, Any]:
                 cache_futures = [
                     actor.get_cache_usage.remote() for actor in self.actors
                 ]
-                request_counts = ray.get(request_futures)
-                cache_usages = ray.get(cache_futures)
+            ray = _get_ray()
+            request_counts = ray.get(request_futures)
+            cache_usages = ray.get(cache_futures)
 
             for i, (requests, cache_usage) in enumerate(
                 zip(request_counts, cache_usages)