wehub-resource-sync
32 lines
1.2 KiB
Python
32 lines
1.2 KiB
Python
from typing import TYPE_CHECKING, Dict, List
|
|
|
|
from ray.rllib.utils.annotations import PublicAPI
|
|
|
|
if TYPE_CHECKING:
|
|
from ray.rllib.algorithms.algorithm_config import AlgorithmConfig
|
|
|
|
DEFAULT_NUM_CPUS_PER_TASK = 0.5
|
|
|
|
|
|
@PublicAPI
|
|
def get_offline_io_resource_bundles(
|
|
config: "AlgorithmConfig",
|
|
) -> List[Dict[str, float]]:
|
|
# DatasetReader is the only offline I/O component today that
|
|
# requires compute resources.
|
|
if config.input_ == "dataset":
|
|
input_config = config.input_config
|
|
# TODO (Kourosh): parallelism is use for reading the dataset, which defaults to
|
|
# num_workers. This logic here relies on the information that dataset reader
|
|
# will have the same logic. So to remove the information leakage, inside
|
|
# Algorithm config, we should set parallelism to num_workers if not specified
|
|
# and only deal with parallelism here or in dataset_reader.py. same thing is
|
|
# true with cpus_per_task.
|
|
parallelism = input_config.get("parallelism", config.get("num_env_runners", 1))
|
|
cpus_per_task = input_config.get(
|
|
"num_cpus_per_read_task", DEFAULT_NUM_CPUS_PER_TASK
|
|
)
|
|
return [{"CPU": cpus_per_task} for _ in range(parallelism)]
|
|
else:
|
|
return []
|