mirror of
https://github.com/gryf/coach.git
synced 2025-12-17 19:20:19 +01:00
Implement frame-work agnostic rollout and training workers (#137)
* Added checkpoint state file to coach checkpointing. * Removed TF specific code from rollout_worker, training_worker, and s3_data_store
This commit is contained in:
committed by
Balaji Subramaniam
parent
4a6c404070
commit
5332013bd1
24
rl_coach/tests/test_checkpoint.py
Normal file
24
rl_coach/tests/test_checkpoint.py
Normal file
@@ -0,0 +1,24 @@
|
||||
import os
|
||||
import pytest
|
||||
import tempfile
|
||||
|
||||
from rl_coach import checkpoint
|
||||
|
||||
|
||||
@pytest.mark.unit_test
|
||||
def test_get_checkpoint_state():
|
||||
files = ['4.test.ckpt.ext', '2.test.ckpt.ext', '3.test.ckpt.ext', '1.test.ckpt.ext', 'prefix.10.test.ckpt.ext']
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
[open(os.path.join(temp_dir, fn), 'a').close() for fn in files]
|
||||
checkpoint_state = checkpoint.get_checkpoint_state(temp_dir, all_checkpoints=True)
|
||||
assert checkpoint_state.model_checkpoint_path == os.path.join(temp_dir, '4.test.ckpt')
|
||||
assert checkpoint_state.all_model_checkpoint_paths == \
|
||||
[os.path.join(temp_dir, f[:-4]) for f in sorted(files[:-1])]
|
||||
|
||||
reader = checkpoint.CheckpointStateReader(temp_dir, checkpoint_state_optional=False)
|
||||
assert reader.get_latest() is None
|
||||
assert len(reader.get_all()) == 0
|
||||
|
||||
reader = checkpoint.CheckpointStateReader(temp_dir)
|
||||
assert reader.get_latest().num == 4
|
||||
assert [ckp.num for ckp in reader.get_all()] == [1, 2, 3, 4]
|
||||
Reference in New Issue
Block a user