wehub-resource-sync
633 lines
26 KiB
Python
633 lines
26 KiB
Python
import atexit
|
||
import logging
|
||
import queue
|
||
import threading
|
||
import weakref
|
||
from queue import Queue
|
||
from typing import Any, Dict, List
|
||
|
||
import ray
|
||
from ray.rllib.algorithms.impala.impala import LEARNER_RESULTS_CURR_ENTROPY_COEFF_KEY
|
||
from ray.rllib.core import COMPONENT_RL_MODULE
|
||
from ray.rllib.core.learner.learner import Learner
|
||
from ray.rllib.core.learner.training_data import TrainingData
|
||
from ray.rllib.core.rl_module.apis import ValueFunctionAPI
|
||
from ray.rllib.utils.annotations import (
|
||
OverrideToImplementCustomLogic_CallToSuperRecommended,
|
||
override,
|
||
)
|
||
from ray.rllib.utils.framework import try_import_torch
|
||
from ray.rllib.utils.lambda_defaultdict import LambdaDefaultDict
|
||
from ray.rllib.utils.metrics import (
|
||
ALL_MODULES,
|
||
NUM_ENV_STEPS_SAMPLED_LIFETIME,
|
||
)
|
||
from ray.rllib.utils.metrics.metrics_logger import MetricsLogger
|
||
from ray.rllib.utils.metrics.ray_metrics import (
|
||
DEFAULT_HISTOGRAM_BOUNDARIES_SHORT_EVENTS,
|
||
TimerAndPrometheusLogger,
|
||
)
|
||
from ray.rllib.utils.schedules.scheduler import Scheduler
|
||
from ray.rllib.utils.typing import ModuleID, ResultDict
|
||
from ray.util.metrics import Gauge, Histogram
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
torch, _ = try_import_torch()
|
||
|
||
GPU_LOADER_QUEUE_WAIT_TIMER = "gpu_loader_queue_wait_timer"
|
||
GPU_LOADER_LOAD_TO_GPU_TIMER = "gpu_loader_load_to_gpu_timer"
|
||
LEARNER_THREAD_IN_QUEUE_WAIT_TIMER = "learner_thread_in_queue_wait_timer"
|
||
LEARNER_THREAD_ENV_STEPS_DROPPED = "learner_thread_env_steps_dropped"
|
||
LEARNER_THREAD_UPDATE_TIMER = "learner_thread_update_timer"
|
||
RAY_GET_EPISODES_TIMER = "ray_get_episodes_timer"
|
||
|
||
QUEUE_SIZE_GPU_LOADER_QUEUE = "queue_size_gpu_loader_queue"
|
||
QUEUE_SIZE_LEARNER_THREAD_QUEUE = "queue_size_learner_thread_queue"
|
||
QUEUE_SIZE_RESULTS_QUEUE = "queue_size_results_queue"
|
||
|
||
# Aggregation cycle size.
|
||
BATCHES_PER_AGGREGATION = 10
|
||
|
||
# Stop sentinel for the `_LearnerThread`
|
||
_STOP_SENTINEL = object()
|
||
|
||
|
||
class IMPALALearner(Learner):
|
||
@override(Learner)
|
||
def __init__(self, *args, **kwargs):
|
||
super().__init__(*args, **kwargs)
|
||
|
||
# Ray metrics
|
||
self._metrics_learner_impala_update = Histogram(
|
||
name="rllib_learner_impala_update_time",
|
||
description="Time spent in the 'IMPALALearner.update()' method.",
|
||
boundaries=DEFAULT_HISTOGRAM_BOUNDARIES_SHORT_EVENTS,
|
||
tag_keys=("rllib",),
|
||
)
|
||
self._metrics_learner_impala_update.set_default_tags(
|
||
{"rllib": self.__class__.__name__}
|
||
)
|
||
|
||
self._metrics_learner_impala_update_solve_refs = Histogram(
|
||
name="rllib_learner_impala_update_solve_refs_time",
|
||
description="Time spent on resolving refs in the 'Learner.update()'",
|
||
boundaries=DEFAULT_HISTOGRAM_BOUNDARIES_SHORT_EVENTS,
|
||
tag_keys=("rllib",),
|
||
)
|
||
self._metrics_learner_impala_update_solve_refs.set_default_tags(
|
||
{"rllib": self.__class__.__name__}
|
||
)
|
||
|
||
self._metrics_learner_impala_update_make_batch_if_necessary = Histogram(
|
||
name="rllib_learner_impala_update_make_batch_if_necessary_time",
|
||
description="Time spent on making a batch in the 'Learner.update()'.",
|
||
boundaries=DEFAULT_HISTOGRAM_BOUNDARIES_SHORT_EVENTS,
|
||
tag_keys=("rllib",),
|
||
)
|
||
self._metrics_learner_impala_update_make_batch_if_necessary.set_default_tags(
|
||
{"rllib": self.__class__.__name__}
|
||
)
|
||
|
||
self._metrics_learner_impala_get_learner_state_time = Histogram(
|
||
name="rllib_learner_impala_get_learner_state_time",
|
||
description="Time spent on get_state() in IMPALALearner.update().",
|
||
boundaries=DEFAULT_HISTOGRAM_BOUNDARIES_SHORT_EVENTS,
|
||
tag_keys=("rllib",),
|
||
)
|
||
self._metrics_learner_impala_get_learner_state_time.set_default_tags(
|
||
{"rllib": self.__class__.__name__}
|
||
)
|
||
|
||
# Set the aggregation threshold to the broadcast interval. We return
|
||
# a state at the same time the metrics are aggregated.
|
||
global BATCHES_PER_AGGREGATION
|
||
BATCHES_PER_AGGREGATION = self.config.broadcast_interval
|
||
|
||
@override(Learner)
|
||
def build(self) -> None:
|
||
super().build()
|
||
|
||
# APPO/IMPALA require RLock for thread safety around metrics.
|
||
self.metrics._threading_lock = threading.RLock()
|
||
|
||
# Aggregation signaling (replaces condition-variable contention) ---
|
||
self._agg_event = threading.Event()
|
||
self._submitted_updates = 0 # producer-side counter (update thread(s))
|
||
self._num_updates = 0 # learner-side counter
|
||
self._num_updates_lock = threading.Lock()
|
||
# Set the update kwargs passed in the main thread for use in the learner thread.
|
||
self._update_kwargs = {}
|
||
|
||
self._model_io_lock = threading.RLock()
|
||
|
||
self._learner_state_queue = Queue(maxsize=1)
|
||
self._learner_state_lock = threading.Lock()
|
||
self._learner_state = None
|
||
|
||
# Dict mapping module IDs to entropy Scheduler instances.
|
||
self.entropy_coeff_schedulers_per_module: Dict[
|
||
ModuleID, Scheduler
|
||
] = LambdaDefaultDict(
|
||
lambda module_id: Scheduler(
|
||
fixed_value_or_schedule=(
|
||
self.config.get_config_for_module(module_id).entropy_coeff
|
||
),
|
||
framework=self.framework,
|
||
device=self._device,
|
||
)
|
||
)
|
||
|
||
# Create queues as bounded queues to create real back-pressure & stabilize
|
||
# GPU memory usage.
|
||
|
||
# Small loader in-queue to keep threads busy without flooding.
|
||
# TODO (simon): Do extensive testing to find an optimal queue size.
|
||
loader_qsize = max(2, 10 * self.config.num_gpu_loader_threads)
|
||
# Note, we are passing now the timesteps dictionary through the queue.
|
||
self._gpu_loader_in_queue: "Queue[tuple[TrainingData, Dict[str, Any]]]" = Queue(
|
||
maxsize=loader_qsize
|
||
)
|
||
|
||
# Learner in-queue must be tiny. 1 strictly serializes GPU-resident batches.
|
||
# TODO (simon): Add a parameter to define queue size.
|
||
if not hasattr(self, "_learner_thread_in_queue"):
|
||
self._learner_thread_in_queue: "Queue[tuple[Any, Dict[str, Any]]]" = Queue(
|
||
maxsize=self.config.learner_queue_size
|
||
)
|
||
|
||
# Get the rank of this learner, if necessary.
|
||
self._rank: int = (
|
||
torch.distributed.get_rank() if torch.distributed.is_initialized() else 0
|
||
)
|
||
|
||
# Define the out-queue for the metrics from the `_LearnerThread`.
|
||
# TODO (simon): Add types for items.
|
||
self._learner_thread_out_queue: "Queue[Dict[str, Any]]" = Queue()
|
||
|
||
# Create and start `_GPULoaderThread`(s).
|
||
if self.config.num_gpus_per_learner > 0:
|
||
self._gpu_loader_threads: List[threading.Thread] = [
|
||
_GPULoaderThread(
|
||
in_queue=self._gpu_loader_in_queue,
|
||
out_queue=self._learner_thread_in_queue,
|
||
device=self._device,
|
||
metrics_logger=self.metrics,
|
||
)
|
||
for _ in range(self.config.num_gpu_loader_threads)
|
||
]
|
||
for t in self._gpu_loader_threads:
|
||
t.start()
|
||
|
||
# Create and start the `_LearnerThread`.
|
||
self._learner_thread: threading.Thread = _LearnerThread(
|
||
update_method=Learner.update,
|
||
in_queue=self._learner_thread_in_queue,
|
||
out_queue=self._learner_thread_out_queue,
|
||
learner=self,
|
||
)
|
||
self._learner_thread.start()
|
||
|
||
@override(Learner)
|
||
def update(
|
||
self,
|
||
training_data: TrainingData,
|
||
*,
|
||
timesteps: Dict[str, Any],
|
||
return_state: bool = False,
|
||
**kwargs,
|
||
) -> ResultDict:
|
||
"""
|
||
|
||
Args:
|
||
batch:
|
||
timesteps:
|
||
return_state: Whether to include one of the Learner worker's state from
|
||
after the update step in the returned results dict (under the
|
||
`_rl_module_state_after_update` key). Note that after an update, all
|
||
Learner workers' states should be identical, so we use the first
|
||
Learner's state here. Useful for avoiding an extra `get_weights()` call,
|
||
e.g. for synchronizing EnvRunner weights.
|
||
**kwargs:
|
||
|
||
Returns:
|
||
|
||
"""
|
||
# Set the update kwargs passed in the main thread for use in the learner thread.
|
||
self._update_kwargs = kwargs
|
||
|
||
with TimerAndPrometheusLogger(self._metrics_learner_impala_update):
|
||
# Get the train batch from the object store.
|
||
with TimerAndPrometheusLogger(
|
||
self._metrics_learner_impala_update_solve_refs
|
||
):
|
||
# Resolve object refs and ensure we have a proper batch object.
|
||
# TODO (simon): Check, if we can resolve the object references and
|
||
# run the pipeline on the GPULoaderThreads.
|
||
training_data.solve_refs()
|
||
|
||
with TimerAndPrometheusLogger(
|
||
self._metrics_learner_impala_update_make_batch_if_necessary
|
||
):
|
||
batch = self._make_batch_if_necessary(training_data=training_data)
|
||
assert batch is not None
|
||
|
||
# Enqeue the batch (bounded backpressure).
|
||
if self.config.num_gpus_per_learner > 0:
|
||
# Pass timesteps alongside batch (no globals).
|
||
self._gpu_loader_in_queue.put((batch, timesteps))
|
||
# Only occasionally log loader queue size.
|
||
if (self._submitted_updates & 0xFF) == 0:
|
||
self.metrics.log_value(
|
||
(ALL_MODULES, QUEUE_SIZE_GPU_LOADER_QUEUE),
|
||
self._gpu_loader_in_queue.qsize(),
|
||
window=1,
|
||
)
|
||
# TODO (simon): Check, if we want to get here stats from the
|
||
# RingBuffer.
|
||
else:
|
||
# No GPU loader: directly enqueue to learner queue.
|
||
_LearnerThread.enqueue(
|
||
self._learner_thread_in_queue, (batch, timesteps), self.metrics
|
||
)
|
||
|
||
# Return the module state, if requested and available.
|
||
if return_state:
|
||
try:
|
||
with self._learner_state_lock:
|
||
self._learner_state = self._learner_state_queue.get_nowait()
|
||
except queue.Empty:
|
||
logger.debug("No learner state available in the queue yet.")
|
||
|
||
# Every 20th block call we submit results. Otherwise we keep the
|
||
# thread running without interruption to avoid thread contention.
|
||
self._submitted_updates += 1
|
||
if (self._submitted_updates % BATCHES_PER_AGGREGATION) != 0:
|
||
result = {}
|
||
if return_state and self._learner_state:
|
||
result["_rl_module_state_after_update"] = self._learner_state
|
||
return result
|
||
|
||
# Result submission: wait until learner finished BATCHES_PER_AGGREGATION updates (blocking).
|
||
self._agg_event.wait()
|
||
# Reset the aggregation event to keep the `_LearnerThread` running.
|
||
self._agg_event.clear()
|
||
|
||
if self._learner_thread_out_queue:
|
||
try:
|
||
result = self._learner_thread_out_queue.get(timeout=0.001)
|
||
except queue.Empty:
|
||
result = {}
|
||
|
||
# Return the module state, if requested and existent.
|
||
if return_state and self._learner_state:
|
||
result["_rl_module_state_after_update"] = self._learner_state
|
||
|
||
return result
|
||
|
||
@OverrideToImplementCustomLogic_CallToSuperRecommended
|
||
def before_gradient_based_update(self, *, timesteps: Dict[str, Any]) -> None:
|
||
super().before_gradient_based_update(timesteps=timesteps)
|
||
|
||
for module_id in self.module.keys():
|
||
# Update entropy coefficient via our Scheduler.
|
||
new_entropy_coeff = self.entropy_coeff_schedulers_per_module[
|
||
module_id
|
||
].update(timestep=timesteps.get(NUM_ENV_STEPS_SAMPLED_LIFETIME, 0))
|
||
self.metrics.log_value(
|
||
(module_id, LEARNER_RESULTS_CURR_ENTROPY_COEFF_KEY),
|
||
new_entropy_coeff,
|
||
window=1,
|
||
)
|
||
|
||
@override(Learner)
|
||
def remove_module(self, module_id: str):
|
||
super().remove_module(module_id)
|
||
self.entropy_coeff_schedulers_per_module.pop(module_id)
|
||
|
||
@override(Learner)
|
||
def shutdown(self) -> None:
|
||
# Stop the learner thread deterministically: setting the stop event
|
||
# and enqueuing a sentinel wakes the consumer if it's blocked on
|
||
# `_in_queue.get()`. Then `join` ensures it has fully exited before
|
||
# we return, so any subsequent `ray.shutdown()`/interpreter teardown
|
||
# can't race with the daemon thread.
|
||
thread = getattr(self, "_learner_thread", None)
|
||
if thread is not None and thread.is_alive():
|
||
thread.request_stop()
|
||
thread.join(timeout=5.0)
|
||
|
||
@classmethod
|
||
@override(Learner)
|
||
def rl_module_required_apis(cls) -> list[type]:
|
||
# In order for a PPOLearner to update an RLModule, it must implement the
|
||
# following APIs:
|
||
return [ValueFunctionAPI]
|
||
|
||
|
||
ImpalaLearner = IMPALALearner
|
||
|
||
|
||
class _GPULoaderThread(threading.Thread):
|
||
def __init__(
|
||
self,
|
||
*,
|
||
in_queue: "Queue[tuple[TrainingData, Dict[str, Any]]]",
|
||
out_queue: "Queue[tuple[Any, Dict[str, Any]]]",
|
||
device: "torch.device",
|
||
metrics_logger: MetricsLogger,
|
||
):
|
||
super().__init__(name="_GPULoaderThread")
|
||
self.daemon = True
|
||
|
||
self._in_queue = in_queue
|
||
self._out_queue = out_queue
|
||
self._device = device
|
||
self.metrics = metrics_logger
|
||
|
||
# Use a single CUDA stream for each loader thread.
|
||
self._use_cuda_stream = (
|
||
torch is not None
|
||
and hasattr(torch, "cuda")
|
||
and device is not None
|
||
and getattr(device, "type", None) == "cuda"
|
||
)
|
||
self._stream = (
|
||
torch.cuda.Stream(device=self._device) if self._use_cuda_stream else None
|
||
)
|
||
|
||
self._metrics_impala_gpu_loader_thread_step_time = Histogram(
|
||
name="rllib_learner_impala_gpu_loader_thread_step_time",
|
||
description="Time taken in seconds for gpu loader thread _step.",
|
||
boundaries=DEFAULT_HISTOGRAM_BOUNDARIES_SHORT_EVENTS,
|
||
tag_keys=("rllib",),
|
||
)
|
||
self._metrics_impala_gpu_loader_thread_step_time.set_default_tags(
|
||
{"rllib": "IMPALA/GPULoaderThread"}
|
||
)
|
||
|
||
self._metrics_impala_gpu_loader_thread_step_in_queue_get_time = Histogram(
|
||
name="rllib_learner_impala_gpu_loader_thread_step_get_time",
|
||
description="Time taken in seconds for gpu loader thread _step _in_queue.get().",
|
||
boundaries=DEFAULT_HISTOGRAM_BOUNDARIES_SHORT_EVENTS,
|
||
tag_keys=("rllib",),
|
||
)
|
||
self._metrics_impala_gpu_loader_thread_step_in_queue_get_time.set_default_tags(
|
||
{"rllib": "IMPALA/GPULoaderThread"}
|
||
)
|
||
|
||
self._metrics_impala_gpu_loader_thread_step_load_to_gpu_time = Histogram(
|
||
name="rllib_learner_impala_gpu_loader_thread_step_load_to_gpu_time",
|
||
description="Time taken in seconds for GPU loader thread _step to load batch to GPU.",
|
||
boundaries=DEFAULT_HISTOGRAM_BOUNDARIES_SHORT_EVENTS,
|
||
tag_keys=("rllib",),
|
||
)
|
||
self._metrics_impala_gpu_loader_thread_step_load_to_gpu_time.set_default_tags(
|
||
{"rllib": "IMPALA/GPULoaderThread"}
|
||
)
|
||
|
||
self._metrics_impala_gpu_loader_thread_in_qsize_beginning_of_step = Gauge(
|
||
name="rllib_impala_gpu_loader_thread_in_qsize_beginning_of_step",
|
||
description="Size of the _GPULoaderThread in-queue size, at the beginning of the step.",
|
||
tag_keys=("rllib",),
|
||
)
|
||
self._metrics_impala_gpu_loader_thread_in_qsize_beginning_of_step.set_default_tags(
|
||
{"rllib": "IMPALA/GPULoaderThread"}
|
||
)
|
||
|
||
# Robust pinned-memory copy: fall back if batch contains CUDA tensors already.
|
||
# TODO (simon): Find a more compliant solution.
|
||
def _to_device_safe(self, batch):
|
||
try:
|
||
return batch.to_device(self._device, pin_memory=True)
|
||
except RuntimeError as e:
|
||
msg = str(e)
|
||
if "only dense CPU tensors can be pinned" in msg or "pin_memory" in msg:
|
||
return batch.to_device(self._device, pin_memory=False)
|
||
raise
|
||
|
||
def run(self) -> None:
|
||
while True:
|
||
with TimerAndPrometheusLogger(
|
||
self._metrics_impala_gpu_loader_thread_step_time
|
||
):
|
||
self._step()
|
||
|
||
def _step(self) -> None:
|
||
self._metrics_impala_gpu_loader_thread_in_qsize_beginning_of_step.set(
|
||
value=self._in_queue.qsize()
|
||
)
|
||
# Get a new batch (CPU) and the global timesteps from the loader in--queue (blocking).
|
||
with self.metrics.log_time((ALL_MODULES, GPU_LOADER_QUEUE_WAIT_TIMER)):
|
||
with TimerAndPrometheusLogger(
|
||
self._metrics_impala_gpu_loader_thread_step_in_queue_get_time
|
||
):
|
||
ma_batch_on_cpu, timesteps = self._in_queue.get()
|
||
|
||
# Load the batch onto the GPU device; enable pinned memory for async copies.
|
||
with self.metrics.log_time((ALL_MODULES, GPU_LOADER_LOAD_TO_GPU_TIMER)):
|
||
if self._use_cuda_stream and self._stream is not None:
|
||
# Issue copies on a non-default stream so they can overlap with compute.
|
||
with torch.cuda.stream(self._stream):
|
||
ma_batch_on_gpu = self._to_device_safe(ma_batch_on_cpu)
|
||
# TODO (simon): Maybe use the `use_stream` in `convert_to_tensor`.
|
||
# No explicit synching here. Consumer will naturally serialize when needed.
|
||
else:
|
||
ma_batch_on_gpu = self._to_device_safe(ma_batch_on_cpu)
|
||
|
||
# Enqueue to Learner thread’s in-queue (GPU-resident batch and global timesteps).
|
||
_LearnerThread.enqueue(
|
||
self._out_queue, (ma_batch_on_gpu, timesteps), self.metrics
|
||
)
|
||
|
||
|
||
class _LearnerThread(threading.Thread):
|
||
def __init__(
|
||
self,
|
||
*,
|
||
update_method,
|
||
in_queue: "Queue[tuple[Any, Dict[str, Any]]]",
|
||
out_queue: "Queue[Dict[str, Any]]",
|
||
learner: IMPALALearner,
|
||
):
|
||
super().__init__(name="_LearnerThread")
|
||
self.daemon = True
|
||
self.learner = learner
|
||
self._update_method = update_method
|
||
# Note, we pass now the timesteps dictionary through the queue.
|
||
self._in_queue: "queue.Queue[tuple[Any, Dict[str, Any]]]" = in_queue
|
||
# TODO (simon): Type hints.
|
||
self._out_queue = out_queue
|
||
self._stop_event = threading.Event()
|
||
|
||
# Ray metrics
|
||
self._metrics_learner_impala_thread_step = Histogram(
|
||
name="rllib_learner_impala_learner_thread_step_time",
|
||
description="Time taken in seconds for learner thread _step.",
|
||
boundaries=DEFAULT_HISTOGRAM_BOUNDARIES_SHORT_EVENTS,
|
||
tag_keys=("rllib",),
|
||
)
|
||
self._metrics_learner_impala_thread_step.set_default_tags(
|
||
{"rllib": "IMPALA/LearnerThread"}
|
||
)
|
||
|
||
self._metrics_learner_impala_thread_step_update = Histogram(
|
||
name="rllib_learner_impala_learner_thread_step_update_time",
|
||
description="Time taken in seconds for learner thread _step update.",
|
||
boundaries=DEFAULT_HISTOGRAM_BOUNDARIES_SHORT_EVENTS,
|
||
tag_keys=("rllib",),
|
||
)
|
||
self._metrics_learner_impala_thread_step_update.set_default_tags(
|
||
{"rllib": "IMPALA/LearnerThread"}
|
||
)
|
||
|
||
# Stop cleanly at interpreter shutdown so the daemon thread doesn't
|
||
# get killed mid-call inside an auto_init-wrapped Ray API (which
|
||
# would otherwise trigger e.g. `start_reaper` -> `preexec_fn not
|
||
# supported at interpreter shutdown`). Use a weakref so this hook
|
||
# doesn't pin the thread (and therefore the Learner) alive.
|
||
weak_self = weakref.ref(self)
|
||
|
||
def _request_stop_at_exit():
|
||
t = weak_self()
|
||
if t is not None:
|
||
t.request_stop()
|
||
|
||
atexit.register(_request_stop_at_exit)
|
||
|
||
# Keeps compatibility, but thread-safe.
|
||
@property
|
||
def stopped(self) -> bool:
|
||
return self._stop_event.is_set()
|
||
|
||
# Call this to stop the thread and wake it if it's blocked on .get()
|
||
def request_stop(self) -> None:
|
||
self._stop_event.set()
|
||
# Wake the consumer if it's blocked on an empty queue
|
||
try:
|
||
self._in_queue.put_nowait(_STOP_SENTINEL)
|
||
except queue.Full:
|
||
# If the queue is full, the consumer will wake soon anyway.
|
||
logger.warning(
|
||
"_LearnerThread.request_stop(): in_queue is full; cannot enqueue stop sentinel."
|
||
)
|
||
|
||
def run(self) -> None:
|
||
while True:
|
||
# Returns always `True` until stop-signal/sentinel is sent.
|
||
if not self.step():
|
||
break
|
||
|
||
def step(self) -> bool:
|
||
# Get a batch and wait, if the input queue is empty (blocking; no polling).
|
||
with self.learner.metrics.log_time(
|
||
(ALL_MODULES, LEARNER_THREAD_IN_QUEUE_WAIT_TIMER)
|
||
):
|
||
item = self._in_queue.get()
|
||
|
||
# Handle the stop/sentinel signal(s).
|
||
# TODO (simon): Check, if we need `None` for belt-and-suspenders/comp.
|
||
if item is _STOP_SENTINEL or self.stopped:
|
||
try:
|
||
self._in_queue.task_done()
|
||
except Exception:
|
||
logger.warning(
|
||
"_LearnerThread._in_queue.task_done() failed during stop handling."
|
||
)
|
||
# Signal `run` to exit.
|
||
return False
|
||
|
||
# Extract the multi-agent batch and the timesteps dictionary.
|
||
ma_batch_on_gpu, timesteps = item
|
||
|
||
# Update the `RLModule`, but do not reduce metrics.
|
||
with self.learner.metrics.log_time((ALL_MODULES, LEARNER_THREAD_UPDATE_TIMER)):
|
||
with TimerAndPrometheusLogger(
|
||
self._metrics_learner_impala_thread_step_update
|
||
):
|
||
self._update_method(
|
||
self=self.learner,
|
||
training_data=TrainingData(batch=ma_batch_on_gpu),
|
||
timesteps=timesteps,
|
||
_no_metrics_reduce=True,
|
||
# Include the learner update kwargs set in the main thread.
|
||
**self.learner._update_kwargs,
|
||
)
|
||
|
||
# Signal queue done (unblocks producer’s put when bounded)
|
||
try:
|
||
self._in_queue.task_done()
|
||
finally:
|
||
# Set the Aggregation counter and signal this event (atomic).
|
||
with self.learner._num_updates_lock:
|
||
self.learner._num_updates += 1
|
||
# Check, if we need to aggregate.
|
||
do_agg = self.learner._num_updates == BATCHES_PER_AGGREGATION
|
||
if do_agg:
|
||
# Reset the update counter inside the lock.
|
||
self.learner._num_updates = 0
|
||
|
||
# If we need to aggregate, reduce metrics and queue them.
|
||
if do_agg:
|
||
# If in multi-learner setup, safeguard state retrieval within barriers.
|
||
if torch.distributed.is_initialized():
|
||
torch.distributed.barrier()
|
||
# Only the first rank retrieves the state.
|
||
if self.learner._rank == 0:
|
||
with self.learner._model_io_lock, torch.inference_mode():
|
||
learner_state = self.learner.get_state(
|
||
# Only return the state of those RLModules that are trainable.
|
||
components=[
|
||
COMPONENT_RL_MODULE + "/" + mid
|
||
for mid in self.learner.module.keys()
|
||
if self.learner.should_module_be_updated(mid)
|
||
],
|
||
inference_only=True,
|
||
)
|
||
learner_state[COMPONENT_RL_MODULE] = ray.put(
|
||
learner_state[COMPONENT_RL_MODULE]
|
||
)
|
||
try:
|
||
if (self.learner._submitted_updates & ~0xFF) != (
|
||
(self.learner._submitted_updates - BATCHES_PER_AGGREGATION)
|
||
& ~0xFF
|
||
):
|
||
with self.learner._learner_state_lock:
|
||
self.learner.metrics.log_value(
|
||
(ALL_MODULES, "learner_thread_state_queue_size"),
|
||
self.learner._learner_state_queue.qsize(),
|
||
window=1,
|
||
)
|
||
# Remove any old learner state in the queue.
|
||
self.learner._learner_state_queue.get_nowait()
|
||
except queue.Empty:
|
||
logger.debug("No old learner state to remove from the queue.")
|
||
|
||
# Pass the learner state into the queue to the main process.
|
||
self.learner._learner_state_queue.put_nowait(learner_state)
|
||
self.learner.metrics.log_value(
|
||
(ALL_MODULES, "learner_thread_out_queue_size"),
|
||
self._out_queue.qsize(),
|
||
window=1,
|
||
)
|
||
|
||
# Reduce metrics and pass them into the queue for the main process.
|
||
self._out_queue.put(self.learner.metrics.reduce())
|
||
# Notify all listeners that aggregation is done and results can be
|
||
# retrieved.
|
||
self.learner._agg_event.set()
|
||
if torch.distributed.is_initialized():
|
||
torch.distributed.barrier()
|
||
|
||
# Keep running (see `run` method).
|
||
return True
|
||
|
||
@staticmethod
|
||
def enqueue(
|
||
learner_queue: "queue.Queue[tuple[Any, Dict[str, Any]]]",
|
||
batch_with_ts,
|
||
metrics: MetricsLogger,
|
||
):
|
||
# Put the batch into the queue (blocking if thread is updating).
|
||
learner_queue.put(batch_with_ts, block=True)
|