2022-02-02 01:20:37 -08:00
|
|
|
from ray.rllib.offline.dataset_reader import (
|
2022-05-24 22:14:25 -07:00
|
|
|
_get_resource_bundles as dataset_reader_get_resource_bundles,
|
2022-02-02 01:20:37 -08:00
|
|
|
)
|
2022-05-24 22:14:25 -07:00
|
|
|
from ray.rllib.utils.annotations import PublicAPI
|
2022-02-02 01:20:37 -08:00
|
|
|
from ray.rllib.utils.typing import PartialTrainerConfigDict
|
|
|
|
from typing import Dict, List
|
|
|
|
|
|
|
|
|
2022-05-24 22:14:25 -07:00
|
|
|
@PublicAPI
|
2022-02-02 01:20:37 -08:00
|
|
|
def get_offline_io_resource_bundles(
|
|
|
|
config: PartialTrainerConfigDict,
|
|
|
|
) -> List[Dict[str, float]]:
|
|
|
|
# DatasetReader is the only offline I/O component today that
|
|
|
|
# requires compute resources.
|
|
|
|
if config["input"] == "dataset":
|
|
|
|
return dataset_reader_get_resource_bundles(config["input_config"])
|
|
|
|
else:
|
|
|
|
return []
|