chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,229 @@
|
||||
import platform
|
||||
from typing import List
|
||||
|
||||
import tree # pip install dm_tree
|
||||
|
||||
import ray
|
||||
from ray.rllib.algorithms.algorithm_config import AlgorithmConfig
|
||||
from ray.rllib.policy.sample_batch import MultiAgentBatch, SampleBatch
|
||||
from ray.rllib.utils.actor_manager import FaultAwareApply
|
||||
from ray.rllib.utils.framework import try_import_torch
|
||||
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.typing import EpisodeType
|
||||
from ray.util.annotations import DeveloperAPI
|
||||
from ray.util.metrics import Counter, Histogram
|
||||
|
||||
torch, _ = try_import_torch()
|
||||
|
||||
|
||||
@DeveloperAPI(stability="alpha")
|
||||
class AggregatorActor(FaultAwareApply):
|
||||
"""Runs episode lists through ConnectorV2 pipeline and creates train batches.
|
||||
|
||||
The actor should be co-located with a Learner worker. Ideally, there should be one
|
||||
or two aggregator actors per Learner worker (having even more per Learner probably
|
||||
won't help. Then the main process driving the RL algo can perform the following
|
||||
execution logic:
|
||||
- query n EnvRunners to sample the environment and return n lists of episodes as
|
||||
Ray.ObjectRefs.
|
||||
- remote call the set of aggregator actors (in round-robin fashion) with these
|
||||
list[episodes] refs in async fashion.
|
||||
- gather the results asynchronously, as each actor returns refs pointing to
|
||||
ready-to-go train batches.
|
||||
- as soon as we have at least one train batch per Learner, call the LearnerGroup
|
||||
with the (already sharded) refs.
|
||||
- an aggregator actor - when receiving p refs to List[EpisodeType] - does:
|
||||
-- ray.get() the actual p lists and concatenate the p lists into one
|
||||
List[EpisodeType].
|
||||
-- pass the lists of episodes through its LearnerConnector pipeline
|
||||
-- buffer the output batches of this pipeline until enough batches have been
|
||||
collected for creating one train batch (matching the config's
|
||||
`train_batch_size_per_learner`).
|
||||
-- concatenate q batches into a train batch and return that train batch.
|
||||
- the algo main process then passes the ray.ObjectRef to the ready-to-go train batch
|
||||
to the LearnerGroup for calling each Learner with one train batch.
|
||||
"""
|
||||
|
||||
def __init__(self, config: AlgorithmConfig, rl_module_spec):
|
||||
self.config = config
|
||||
|
||||
# Set device and node.
|
||||
self._node = platform.node()
|
||||
self._device = torch.device("cpu")
|
||||
self.metrics: MetricsLogger = MetricsLogger(
|
||||
stats_cls_lookup=config.stats_cls_lookup,
|
||||
root=True,
|
||||
)
|
||||
|
||||
# Create the RLModule.
|
||||
# TODO (sven): For now, this RLModule (its weights) never gets updated.
|
||||
# The reason the module is needed is for the connector to know, which
|
||||
# sub-modules are stateful (and what their initial state tensors are), and
|
||||
# which IDs the submodules have (to figure out, whether its multi-agent or
|
||||
# not).
|
||||
self._module = rl_module_spec.build()
|
||||
self._module = self._module.as_multi_rl_module()
|
||||
|
||||
# Create the Learner connector pipeline.
|
||||
self._learner_connector = self.config.build_learner_connector(
|
||||
input_observation_space=None,
|
||||
input_action_space=None,
|
||||
device=self._device,
|
||||
)
|
||||
|
||||
# Ray metrics
|
||||
self._metrics_get_batch_time = Histogram(
|
||||
name="rllib_utils_aggregator_actor_get_batch_time",
|
||||
description="Time spent in AggregatorActor.get_batch()",
|
||||
boundaries=DEFAULT_HISTOGRAM_BOUNDARIES_SHORT_EVENTS,
|
||||
tag_keys=("rllib",),
|
||||
)
|
||||
self._metrics_get_batch_time.set_default_tags(
|
||||
{"rllib": self.__class__.__name__}
|
||||
)
|
||||
|
||||
self._metrics_episode_owner_died = Counter(
|
||||
name="rllib_utils_aggregator_actor_episode_owner_died_counter",
|
||||
description="N times ray.get() on an episode ref failed ",
|
||||
tag_keys=("rllib",),
|
||||
)
|
||||
self._metrics_episode_owner_died.set_default_tags(
|
||||
{"rllib": self.__class__.__name__}
|
||||
)
|
||||
|
||||
self._metrics_get_batch_input_episode_refs = Counter(
|
||||
name="rllib_utils_aggregator_actor_get_batch_input_episode_refs_counter",
|
||||
description="Number of episode refs received as input to get_batch()",
|
||||
tag_keys=("rllib",),
|
||||
)
|
||||
self._metrics_get_batch_input_episode_refs.set_default_tags(
|
||||
{"rllib": self.__class__.__name__}
|
||||
)
|
||||
|
||||
self._metrics_get_batch_output_batches = Counter(
|
||||
name="rllib_utils_aggregator_actor_get_batch_output_batches_counter",
|
||||
description="Number of policy batches output by get_batch()",
|
||||
tag_keys=("rllib",),
|
||||
)
|
||||
self._metrics_get_batch_output_batches.set_default_tags(
|
||||
{"rllib": self.__class__.__name__}
|
||||
)
|
||||
|
||||
def get_batch(self, episode_refs: List[ray.ObjectRef]):
|
||||
with TimerAndPrometheusLogger(self._metrics_get_batch_time):
|
||||
if len(episode_refs) > 0:
|
||||
self._metrics_get_batch_input_episode_refs.inc(value=len(episode_refs))
|
||||
|
||||
episodes: List[EpisodeType] = []
|
||||
# It's possible that individual refs are invalid due to the EnvRunner
|
||||
# that produced the ref has crashed or had its entire node go down.
|
||||
# In this case, try each ref individually and collect only valid results.
|
||||
try:
|
||||
episodes = tree.flatten(ray.get(episode_refs))
|
||||
except ray.exceptions.OwnerDiedError:
|
||||
for ref in episode_refs:
|
||||
try:
|
||||
episodes.extend(ray.get(ref))
|
||||
except ray.exceptions.OwnerDiedError:
|
||||
self._metrics_episode_owner_died.inc(value=1)
|
||||
|
||||
env_steps = sum(len(e) for e in episodes)
|
||||
|
||||
# If we have enough episodes collected to create a single train batch, pass
|
||||
# them at once through the connector to receive a single train batch.
|
||||
batch = self._learner_connector(
|
||||
episodes=episodes,
|
||||
rl_module=self._module,
|
||||
metrics=self.metrics,
|
||||
)
|
||||
# Convert to a dict into a `MultiAgentBatch`.
|
||||
# TODO (sven): Try to get rid of dependency on MultiAgentBatch (once our mini-
|
||||
# batch iterators support splitting over a dict).
|
||||
ma_batch = MultiAgentBatch(
|
||||
policy_batches={
|
||||
pid: SampleBatch(pol_batch) for pid, pol_batch in batch.items()
|
||||
},
|
||||
env_steps=env_steps,
|
||||
)
|
||||
self._metrics_get_batch_output_batches.inc(value=1)
|
||||
return ma_batch
|
||||
|
||||
def get_metrics(self):
|
||||
return self.metrics.reduce()
|
||||
|
||||
|
||||
def _get_env_runner_bundles(config):
|
||||
return [
|
||||
{
|
||||
"CPU": config.num_cpus_per_env_runner,
|
||||
"GPU": config.num_gpus_per_env_runner,
|
||||
**config.custom_resources_per_env_runner,
|
||||
}
|
||||
for _ in range(config.num_env_runners)
|
||||
]
|
||||
|
||||
|
||||
def _get_offline_eval_runner_bundles(config):
|
||||
return [
|
||||
{
|
||||
"CPU": config.num_cpus_per_offline_eval_runner,
|
||||
"GPU": config.num_gpus_per_offline_eval_runner,
|
||||
**config.custom_resources_per_offline_eval_runner,
|
||||
}
|
||||
for _ in range(config.num_offline_eval_runners)
|
||||
]
|
||||
|
||||
|
||||
def _get_learner_bundles(config):
|
||||
if config.num_learners == 0:
|
||||
if config.num_aggregator_actors_per_learner > 0:
|
||||
return [{"CPU": 1} for _ in range(config.num_aggregator_actors_per_learner)]
|
||||
else:
|
||||
return []
|
||||
|
||||
if config.num_cpus_per_learner != "auto":
|
||||
num_cpus_per_learner = config.num_cpus_per_learner
|
||||
elif config.num_gpus_per_learner == 0:
|
||||
num_cpus_per_learner = 1
|
||||
else:
|
||||
num_cpus_per_learner = 0
|
||||
|
||||
# aggregator actors are co-located with learners and use 1 CPU each
|
||||
bundles = [
|
||||
{
|
||||
"CPU": num_cpus_per_learner + config.num_aggregator_actors_per_learner,
|
||||
"GPU": config.num_gpus_per_learner,
|
||||
**(config.custom_resources_per_learner or {}),
|
||||
}
|
||||
for _ in range(config.num_learners)
|
||||
]
|
||||
|
||||
return bundles
|
||||
|
||||
|
||||
def _get_main_process_bundle(config):
|
||||
if config.num_learners == 0:
|
||||
if config.num_cpus_per_learner != "auto":
|
||||
num_cpus_per_learner = config.num_cpus_per_learner
|
||||
elif config.num_gpus_per_learner == 0:
|
||||
num_cpus_per_learner = 1
|
||||
else:
|
||||
num_cpus_per_learner = 0
|
||||
|
||||
bundle = {
|
||||
"CPU": max(num_cpus_per_learner, config.num_cpus_for_main_process),
|
||||
"GPU": config.num_gpus_per_learner,
|
||||
**config.custom_resources_for_main_process,
|
||||
}
|
||||
else:
|
||||
bundle = {
|
||||
"CPU": config.num_cpus_for_main_process,
|
||||
"GPU": 0,
|
||||
**config.custom_resources_for_main_process,
|
||||
}
|
||||
|
||||
return bundle
|
||||
Reference in New Issue
Block a user