diff --git a/envlogger/backends/riegeli_backend_writer.py b/envlogger/backends/riegeli_backend_writer.py index 42824a9..c8108ef 100644 --- a/envlogger/backends/riegeli_backend_writer.py +++ b/envlogger/backends/riegeli_backend_writer.py @@ -53,14 +53,16 @@ def __init__( writer_options: Comma-separated list of options that are passed to Riegeli RecordWriter as is. flush_scheduler: This controls when data is flushed to permanent storage. - If `None`, it defaults to a step-wise Bernoulli scheduler with 1/5000 - chances of flushing. + If `None`, it defaults to a TimeScheduler with 600 seconds (10 minutes) + interval. **base_kwargs: arguments for the base class. """ super().__init__(**base_kwargs) self._data_directory = data_directory if flush_scheduler is None: - self._flush_scheduler = schedulers.BernoulliStepScheduler(1.0 / 5000) + self._flush_scheduler = schedulers.TimeScheduler( + time_interval_seconds=600.0 + ) else: self._flush_scheduler = flush_scheduler self._data_writer = riegeli_dataset_writer.RiegeliDatasetWriter() diff --git a/envlogger/backends/riegeli_backend_writer_test.py b/envlogger/backends/riegeli_backend_writer_test.py index 716d5a8..735c103 100644 --- a/envlogger/backends/riegeli_backend_writer_test.py +++ b/envlogger/backends/riegeli_backend_writer_test.py @@ -205,6 +205,18 @@ def test_episode_step_out_of_bounds_indices(self): -len(episode) - 1) + def test_default_scheduler(self): + """RiegeliBackendWriter should use TimeScheduler(600) by default.""" + temp_dir = self.create_tempdir() + backend = riegeli_backend_writer.RiegeliBackendWriter( + data_directory=temp_dir.full_path + ) + scheduler = backend._flush_scheduler + self.assertIsInstance(scheduler, schedulers.TimeScheduler) + assert isinstance(scheduler, schedulers.TimeScheduler) + self.assertEqual(scheduler._time_interval_seconds, 600.0) + backend.close() + if __name__ == '__main__': absltest.main() diff --git a/envlogger/backends/schedulers.py b/envlogger/backends/schedulers.py index 05241d7..6658f88 100644 --- a/envlogger/backends/schedulers.py +++ b/envlogger/backends/schedulers.py @@ -16,6 +16,7 @@ """Common logging scheduling strategies.""" from collections.abc import Callable +import time from envlogger import step_data import numpy as np @@ -155,3 +156,32 @@ def __call__(self, data: step_data.StepData): if data.timestep.last(): self._episode_counter += 1 return should_log + + +class TimeScheduler: + """Returns `True` if at least `time_interval_seconds` have passed. + + This is useful for time-based flushing. + """ + + def __init__( + self, + time_interval_seconds: float, + clock_fn: Callable[[], float] = time.time, + ): + if time_interval_seconds <= 0: + raise ValueError( + f'time_interval_seconds must be positive, got {time_interval_seconds}' + ) + + self._time_interval_seconds = time_interval_seconds + self._clock_fn = clock_fn + self._last_flush_time = self._clock_fn() + + def __call__(self, unused_data: step_data.StepData | None = None) -> bool: + """Returns `True` if at least `time_interval_seconds` have passed.""" + current_time = self._clock_fn() + if current_time - self._last_flush_time >= self._time_interval_seconds: + self._last_flush_time = current_time + return True + return False diff --git a/envlogger/backends/schedulers_test.py b/envlogger/backends/schedulers_test.py index 538728c..f110c22 100644 --- a/envlogger/backends/schedulers_test.py +++ b/envlogger/backends/schedulers_test.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Tests for schedulers.""" +"""Tests for EnvLogger backend flush schedulers (e.g. step, episode, time).""" from absl.testing import absltest from absl.testing import parameterized @@ -364,5 +364,46 @@ def test_list_episodes_logspace(self): self.assertFalse(scheduler(timestep)) +class TimeSchedulerTest(parameterized.TestCase): + + def test_invalid_args(self): + self.assertRaises( + ValueError, schedulers.TimeScheduler, time_interval_seconds=-1 + ) + self.assertRaises( + ValueError, schedulers.TimeScheduler, time_interval_seconds=0 + ) + + def test_time_scheduler(self): + """TimeScheduler should return True only after the time interval has passed.""" + current_time = 0.0 + + def mock_clock(): + return current_time + + # 10 seconds interval + scheduler = schedulers.TimeScheduler( + time_interval_seconds=10.0, clock_fn=mock_clock + ) + + # First call at t=0 should be False + self.assertFalse(scheduler(None)) + + # Advance time by 5 seconds + current_time = 5.0 + self.assertFalse(scheduler(None)) + + # Advance time by another 5 seconds (total 10) + current_time = 10.0 + self.assertTrue(scheduler(None)) + + # Next call immediately after should be False + self.assertFalse(scheduler(None)) + + # Advance time by 11 seconds (total 21) + current_time = 21.0 + self.assertTrue(scheduler(None)) + + if __name__ == '__main__': absltest.main()