From 8b6f0d3224055e5e028569e31cfd56316f7ce29e Mon Sep 17 00:00:00 2001 From: Eric Liang Date: Sat, 10 Aug 2019 13:53:27 -0700 Subject: [PATCH] [rllib] Fix output API when lz4 not installed (#5421) --- rllib/offline/json_writer.py | 4 ++-- rllib/utils/compression.py | 5 +++++ 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/rllib/offline/json_writer.py b/rllib/offline/json_writer.py index 679b00158..cf2009e3d 100644 --- a/rllib/offline/json_writer.py +++ b/rllib/offline/json_writer.py @@ -19,7 +19,7 @@ from ray.rllib.policy.sample_batch import MultiAgentBatch from ray.rllib.offline.io_context import IOContext from ray.rllib.offline.output_writer import OutputWriter from ray.rllib.utils.annotations import override, PublicAPI -from ray.rllib.utils.compression import pack +from ray.rllib.utils.compression import pack, compression_supported logger = logging.getLogger(__name__) @@ -99,7 +99,7 @@ class JsonWriter(OutputWriter): def _to_jsonable(v, compress): - if compress: + if compress and compression_supported(): return str(pack(v)) elif isinstance(v, np.ndarray): return v.tolist() diff --git a/rllib/utils/compression.py b/rllib/utils/compression.py index de7bef390..20a19ead6 100644 --- a/rllib/utils/compression.py +++ b/rllib/utils/compression.py @@ -23,6 +23,11 @@ except ImportError: LZ4_ENABLED = False +@DeveloperAPI +def compression_supported(): + return LZ4_ENABLED + + @DeveloperAPI def pack(data): if LZ4_ENABLED: