| import os |
| import unittest |
| import json |
| from enum import Enum |
| import time |
| from mlagents.trainers.training_status import ( |
| StatusType, |
| StatusMetaData, |
| GlobalTrainingStatus, |
| ) |
| from mlagents.trainers.policy.checkpoint_manager import ( |
| ModelCheckpointManager, |
| ModelCheckpoint, |
| ) |
|
|
|
|
| def test_globaltrainingstatus(tmpdir): |
| path_dir = os.path.join(tmpdir, "test.json") |
|
|
| GlobalTrainingStatus.set_parameter_state("Category1", StatusType.LESSON_NUM, 3) |
| GlobalTrainingStatus.save_state(path_dir) |
|
|
| with open(path_dir) as fp: |
| test_json = json.load(fp) |
|
|
| assert "Category1" in test_json |
| assert StatusType.LESSON_NUM.value in test_json["Category1"] |
| assert test_json["Category1"][StatusType.LESSON_NUM.value] == 3 |
| assert "metadata" in test_json |
|
|
| GlobalTrainingStatus.load_state(path_dir) |
| restored_val = GlobalTrainingStatus.get_parameter_state( |
| "Category1", StatusType.LESSON_NUM |
| ) |
| assert restored_val == 3 |
|
|
| |
| unknown_category = GlobalTrainingStatus.get_parameter_state( |
| "Category3", StatusType.LESSON_NUM |
| ) |
|
|
| class FakeStatusType(Enum): |
| NOTAREALKEY = "notarealkey" |
|
|
| unknown_key = GlobalTrainingStatus.get_parameter_state( |
| "Category1", FakeStatusType.NOTAREALKEY |
| ) |
| assert unknown_category is None |
| assert unknown_key is None |
|
|
|
|
| def test_model_management(tmpdir): |
|
|
| results_path = os.path.join(tmpdir, "results") |
| brain_name = "Mock_brain" |
| final_model_path = os.path.join(results_path, brain_name) |
| test_checkpoint_list = [ |
| { |
| "steps": 1, |
| "file_path": os.path.join(final_model_path, f"{brain_name}-1.nn"), |
| "reward": 1.312, |
| "creation_time": time.time(), |
| "auxillary_file_paths": [], |
| }, |
| { |
| "steps": 2, |
| "file_path": os.path.join(final_model_path, f"{brain_name}-2.nn"), |
| "reward": 1.912, |
| "creation_time": time.time(), |
| "auxillary_file_paths": [], |
| }, |
| { |
| "steps": 3, |
| "file_path": os.path.join(final_model_path, f"{brain_name}-3.nn"), |
| "reward": 2.312, |
| "creation_time": time.time(), |
| "auxillary_file_paths": [], |
| }, |
| ] |
| GlobalTrainingStatus.set_parameter_state( |
| brain_name, StatusType.CHECKPOINTS, test_checkpoint_list |
| ) |
|
|
| new_checkpoint_4 = ModelCheckpoint( |
| 4, os.path.join(final_model_path, f"{brain_name}-4.nn"), 2.678, time.time() |
| ) |
| ModelCheckpointManager.add_checkpoint(brain_name, new_checkpoint_4, 4) |
| assert len(ModelCheckpointManager.get_checkpoints(brain_name)) == 4 |
|
|
| new_checkpoint_5 = ModelCheckpoint( |
| 5, os.path.join(final_model_path, f"{brain_name}-5.nn"), 3.122, time.time() |
| ) |
| ModelCheckpointManager.add_checkpoint(brain_name, new_checkpoint_5, 4) |
| assert len(ModelCheckpointManager.get_checkpoints(brain_name)) == 4 |
|
|
| final_model_path = f"{final_model_path}.nn" |
| final_model_time = time.time() |
| current_step = 6 |
| final_model = ModelCheckpoint( |
| current_step, final_model_path, 3.294, final_model_time |
| ) |
|
|
| ModelCheckpointManager.track_final_checkpoint(brain_name, final_model) |
| assert len(ModelCheckpointManager.get_checkpoints(brain_name)) == 4 |
|
|
| check_checkpoints = GlobalTrainingStatus.saved_state[brain_name][ |
| StatusType.CHECKPOINTS.value |
| ] |
| assert check_checkpoints is not None |
|
|
| final_model = GlobalTrainingStatus.saved_state[StatusType.FINAL_CHECKPOINT.value] |
| assert final_model is not None |
|
|
|
|
| class StatsMetaDataTest(unittest.TestCase): |
| def test_metadata_compare(self): |
| |
| with self.assertLogs("mlagents.trainers", level="WARNING") as cm: |
| default_metadata = StatusMetaData() |
| version_statsmetadata = StatusMetaData(mlagents_version="test") |
| default_metadata.check_compatibility(version_statsmetadata) |
|
|
| torch_version_statsmetadata = StatusMetaData(torch_version="test") |
| default_metadata.check_compatibility(torch_version_statsmetadata) |
|
|
| |
| assert len(cm.output) == 2 |
|
|
