wehub-resource-sync
ba4be087d5
CICD NeMo / cicd-main-unit-tests (push) Blocked by required conditions
CICD NeMo / cicd-main-speech (push) Blocked by required conditions
CICD NeMo / cicd-test-container-build (push) Blocked by required conditions
CICD NeMo / cicd-import-tests (push) Blocked by required conditions
CICD NeMo / L0_Setup_Test_Data_And_Models (push) Blocked by required conditions
CICD NeMo / Nemo_CICD_Test (push) Blocked by required conditions
CICD NeMo / Coverage (e2e) (push) Blocked by required conditions
CICD NeMo / Coverage (unit-test) (push) Blocked by required conditions
CodeQL / Analyze (python) (push) Waiting to run
Create PR to main with cherry-pick from release / cherry-pick (push) Failing after 0s
CICD NeMo / pre-flight (push) Failing after 0s
CICD NeMo / configure (push) Has been skipped
Build, validate, and release Neural Modules / pre-flight (push) Failing after 1s
CICD NeMo / code-linting (push) Has been skipped
CICD NeMo / cicd-wait-in-queue (push) Waiting to run
Build, validate, and release Neural Modules / release (push) Has been skipped
Build, validate, and release Neural Modules / release-summary (push) Has been cancelled
1440 lines
61 KiB
Python
1440 lines
61 KiB
Python
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
import filecmp
|
|
import json
|
|
import os
|
|
import shutil
|
|
import tempfile
|
|
from typing import Any, Callable, Dict, Optional, Set, Union
|
|
|
|
import pytest
|
|
import torch
|
|
from omegaconf import DictConfig, OmegaConf, open_dict
|
|
|
|
from nemo.collections.asr.models import EncDecCTCModel, EncDecCTCModelBPE
|
|
from nemo.core.classes import ModelPT
|
|
from nemo.core.connectors import save_restore_connector
|
|
from nemo.utils.app_state import AppState
|
|
from nemo.utils.exceptions import NeMoBaseException
|
|
|
|
|
|
@pytest.fixture(scope="module", autouse=True)
|
|
def set_env():
|
|
os.environ["HF_HOME"] = "/home/TestData/hf_home_test_save_restore"
|
|
yield
|
|
del os.environ["HF_HOME"]
|
|
|
|
|
|
def classpath(cls):
|
|
return f'{cls.__module__}.{cls.__name__}'
|
|
|
|
|
|
def get_dir_size(path='.'):
|
|
total = 0
|
|
with os.scandir(path) as it:
|
|
for entry in it:
|
|
if entry.is_file():
|
|
total += entry.stat().st_size
|
|
elif entry.is_dir():
|
|
total += get_dir_size(entry.path)
|
|
return total
|
|
|
|
|
|
def get_size(path='.'):
|
|
if os.path.isfile(path):
|
|
return os.path.getsize(path)
|
|
elif os.path.isdir(path):
|
|
return get_dir_size(path)
|
|
|
|
|
|
def getattr2(object, attr):
|
|
if not '.' in attr:
|
|
return getattr(object, attr)
|
|
else:
|
|
arr = attr.split('.')
|
|
return getattr2(getattr(object, arr[0]), '.'.join(arr[1:]))
|
|
|
|
|
|
def _is_json_serializable(value: Any) -> bool:
|
|
"""Test whether a variable can be encoded as json."""
|
|
if value is None or isinstance(value, (bool, int, float, str, list, dict)): # fast path
|
|
return True
|
|
try:
|
|
json.dumps(value)
|
|
return True
|
|
except (TypeError, OverflowError):
|
|
# OverflowError is raised if number is too large to encode
|
|
return False
|
|
|
|
|
|
class MockModel(ModelPT):
|
|
def __init__(self, cfg, trainer=None):
|
|
super(MockModel, self).__init__(cfg=cfg, trainer=trainer)
|
|
self.w = torch.nn.Linear(10, 1)
|
|
# mock temp file
|
|
if 'temp_file' in self.cfg and self.cfg.temp_file is not None:
|
|
self.setup_data_from_file(self.cfg.temp_file)
|
|
else:
|
|
self.temp_file = None
|
|
self.temp_data = None
|
|
|
|
def setup_data_from_file(self, temp_file):
|
|
"""
|
|
Load data from temp_file to `self.temp_data`
|
|
Allows to test changing resource after instantiation
|
|
"""
|
|
with open_dict(self.cfg):
|
|
self.cfg.temp_file = temp_file
|
|
self.temp_file = self.register_artifact('temp_file', self.cfg.temp_file)
|
|
with open(self.temp_file, 'r', encoding='utf-8') as f:
|
|
self.temp_data = f.readlines()
|
|
|
|
def change_stub_number(self, new_number: int):
|
|
"""
|
|
Change stub number in config, useful for testing nested models,
|
|
since child can mutate config independently
|
|
"""
|
|
self.cfg.stub_number = new_number
|
|
|
|
def forward(self, x):
|
|
y = self.w(x)
|
|
return y, self.cfg.temp_file
|
|
|
|
def setup_training_data(self, train_data_config: Union[DictConfig, Dict]):
|
|
self._train_dl = None
|
|
|
|
def setup_validation_data(self, val_data_config: Union[DictConfig, Dict]):
|
|
self._validation_dl = None
|
|
|
|
def setup_test_data(self, test_data_config: Union[DictConfig, Dict]):
|
|
self._test_dl = None
|
|
|
|
@classmethod
|
|
def list_available_models(cls):
|
|
return []
|
|
|
|
|
|
class MockModelWithChildren(MockModel):
|
|
"""
|
|
Mock Model, can contain 2 children (other NeMo models)
|
|
"""
|
|
|
|
def __init__(self, cfg, trainer=None):
|
|
super().__init__(cfg=cfg, trainer=trainer)
|
|
|
|
# variant 1 for creating nested NeMo model:
|
|
# load model directly from config
|
|
|
|
# variant 2 for creating nested NeMo model:
|
|
# - initialize child model from .nemo checkpoint, subconfig will be automatically saved
|
|
# - after saving model will be restored directly from subconfig (attribute `config_field` of self.cfg)
|
|
|
|
# child 1
|
|
self.child1_model: Optional[MockModel] # annotate type for IDE autocompletion and type checking
|
|
if cfg.get("child1_model") is not None:
|
|
self.register_nemo_submodule(
|
|
"child1_model",
|
|
config_field="child1_model",
|
|
model=MockModel(self.cfg.child1_model),
|
|
)
|
|
elif cfg.get("child1_model_path") is not None:
|
|
self.register_nemo_submodule(
|
|
"child1_model",
|
|
config_field="child1_model",
|
|
model=MockModel.restore_from(self.cfg.child1_model_path),
|
|
)
|
|
else:
|
|
self.child1_model = None
|
|
|
|
# child 2
|
|
# can have sub-children
|
|
self.child2_model: Optional[MockModelWithChildren] # annotate type for IDE autocompletion and type checking
|
|
if cfg.get("child2_model") is not None:
|
|
self.register_nemo_submodule(
|
|
"child2_model",
|
|
config_field="child2_model",
|
|
model=MockModelWithChildren(self.cfg.child2_model),
|
|
)
|
|
elif cfg.get("child2_model_path") is not None:
|
|
self.register_nemo_submodule(
|
|
"child2_model",
|
|
config_field="child2_model",
|
|
model=MockModelWithChildren.restore_from(self.cfg.child2_model_path),
|
|
)
|
|
else:
|
|
self.child2_model = None
|
|
|
|
|
|
class MockModelWithChildEncDecCTCBPE(MockModel):
|
|
"""
|
|
Mock Model, will contain EncDecCTC model as a child
|
|
Useful for testing nested models with children initialized from pretrained NeMo models
|
|
"""
|
|
|
|
def __init__(self, cfg, trainer=None):
|
|
super().__init__(cfg=cfg, trainer=trainer)
|
|
|
|
# variant 3 for creating nested NeMo model:
|
|
# - initialize child model from pretrained NeMo model, subconfig will be automatically saved
|
|
# - after saving model will be restored directly from subconfig (attribute `config_field` of self.cfg)
|
|
|
|
self.ctc_model: EncDecCTCModelBPE # annotate type for IDE autocompletion and type checking
|
|
|
|
if cfg.get("ctc_model", None) is not None:
|
|
self.register_nemo_submodule(
|
|
"ctc_model",
|
|
config_field="ctc_model",
|
|
model=EncDecCTCModelBPE(self.cfg.ctc_model),
|
|
)
|
|
else:
|
|
# model is mandatory
|
|
assert cfg.get("ctc_model_pretrained", None) is not None
|
|
self.register_nemo_submodule(
|
|
"ctc_model",
|
|
config_field="ctc_model",
|
|
model=EncDecCTCModelBPE.from_pretrained(self.cfg.ctc_model_pretrained),
|
|
)
|
|
|
|
|
|
class MockModelWithChildCustomConfigPath(MockModel):
|
|
"""
|
|
Mock Model, can contain 1 child
|
|
Path in config is not equal to name of the attribute
|
|
Config is stored in `child1_model_config`
|
|
Child model is stored in `child1_model` attribute
|
|
NB: This is not recommended if it's not necessary. But here we test that it works.
|
|
"""
|
|
|
|
def __init__(self, cfg, trainer=None):
|
|
super().__init__(cfg=cfg, trainer=trainer)
|
|
|
|
self.child1_model: Optional[MockModel] # annotate type for IDE autocompletion and type checking
|
|
if cfg.get("child1_model_config") is not None:
|
|
self.register_nemo_submodule(
|
|
"child1_model",
|
|
config_field="child1_model_config",
|
|
model=MockModel(self.cfg.child1_model_config),
|
|
)
|
|
else:
|
|
self.child1_model = None
|
|
|
|
|
|
class MockModelIncorrectWithNemoArtifact(MockModel):
|
|
"""
|
|
Incorrect model that tries to use .nemo model checkpoint as an artifact
|
|
Expected to fail, since it is not supported
|
|
"""
|
|
|
|
def __init__(self, cfg, trainer=None):
|
|
super().__init__(cfg=cfg, trainer=trainer)
|
|
|
|
assert cfg.get("child_model_path") is not None
|
|
# this will fail, since .nemo model checkpoint is not supported as an artifact
|
|
child_model_path = self.register_artifact("child_model_path", cfg.child_model_path)
|
|
self.child_model = ModelPT.restore_from(child_model_path)
|
|
|
|
|
|
def _mock_model_config():
|
|
conf = {'temp_file': None, 'target': classpath(MockModel), 'stub_number': 1}
|
|
conf = OmegaConf.create({'model': conf})
|
|
OmegaConf.set_struct(conf, True)
|
|
return conf
|
|
|
|
|
|
def _mock_model_with_children_config(
|
|
child1_model_path: Optional[str] = None,
|
|
child2_model_path: Optional[str] = None,
|
|
child2_model_cfg: Optional[DictConfig] = None,
|
|
) -> DictConfig:
|
|
"""
|
|
Child 1 always constructed from .nemo model checkpoint (optional)
|
|
Child 2 can be constructed directly from subconfig (optional) or from .nemo model checkpoint (optional)
|
|
"""
|
|
conf = {
|
|
'temp_file': None,
|
|
'target': classpath(MockModelWithChildren),
|
|
'child1_model': None,
|
|
'child1_model_path': child1_model_path,
|
|
'child2_model': child2_model_cfg,
|
|
'child2_model_path': child2_model_path,
|
|
'stub_number': 1,
|
|
}
|
|
conf = OmegaConf.create({'model': conf})
|
|
OmegaConf.set_struct(conf, True)
|
|
return conf
|
|
|
|
|
|
def _mock_model_with_child_encdecctcbpe_config(pretrained_model_name: str) -> DictConfig:
|
|
conf = {'temp_file': None, 'ctc_model_pretrained': pretrained_model_name, 'stub_number': 1}
|
|
conf = OmegaConf.create({'model': conf})
|
|
OmegaConf.set_struct(conf, True)
|
|
return conf
|
|
|
|
|
|
def _mock_model_with_child_custom_config_path_config():
|
|
conf = {
|
|
'temp_file': None,
|
|
'child1_model_config': _mock_model_config().model,
|
|
'target': classpath(MockModelWithChildCustomConfigPath),
|
|
'stub_number': 1,
|
|
}
|
|
conf = OmegaConf.create({'model': conf})
|
|
OmegaConf.set_struct(conf, True)
|
|
return conf
|
|
|
|
|
|
def _mock_model_incorrect_with_nemo_artifact_config(child_model_path: str):
|
|
conf = {'temp_file': None, 'child_model_path': child_model_path, 'stub_number': 1}
|
|
conf = OmegaConf.create({'model': conf})
|
|
OmegaConf.set_struct(conf, True)
|
|
return conf
|
|
|
|
|
|
class TestSaveRestore:
|
|
def __test_restore_elsewhere(
|
|
self,
|
|
model: ModelPT,
|
|
attr_for_eq_check: Set[str] = None,
|
|
override_config_path: Optional[Union[str, DictConfig]] = None,
|
|
map_location: Optional[Union[torch.device, str]] = None,
|
|
strict: bool = False,
|
|
return_config: bool = False,
|
|
):
|
|
"""Test's logic:
|
|
1. Save model into temporary folder (save_folder)
|
|
2. Copy .nemo file from save_folder to restore_folder
|
|
3. Delete save_folder
|
|
4. Attempt to restore from .nemo file in restore_folder and compare to original instance
|
|
"""
|
|
# Create a new temporary directory
|
|
with tempfile.TemporaryDirectory() as restore_folder:
|
|
with tempfile.TemporaryDirectory() as save_folder:
|
|
save_folder_path = save_folder
|
|
# Where model will be saved
|
|
model_save_path = os.path.join(save_folder, f"{model.__class__.__name__}.nemo")
|
|
model.save_to(save_path=model_save_path)
|
|
# Where model will be restored from
|
|
model_restore_path = os.path.join(restore_folder, f"{model.__class__.__name__}.nemo")
|
|
shutil.copy(model_save_path, model_restore_path)
|
|
# at this point save_folder should not exist
|
|
assert save_folder_path is not None and not os.path.exists(save_folder_path)
|
|
assert not os.path.exists(model_save_path)
|
|
assert os.path.exists(model_restore_path)
|
|
# attempt to restore
|
|
model_copy = model.__class__.restore_from(
|
|
restore_path=model_restore_path,
|
|
map_location=map_location,
|
|
strict=strict,
|
|
return_config=return_config,
|
|
override_config_path=override_config_path,
|
|
)
|
|
|
|
if return_config:
|
|
return model_copy
|
|
|
|
assert model.num_weights == model_copy.num_weights
|
|
if attr_for_eq_check is not None and len(attr_for_eq_check) > 0:
|
|
for attr in attr_for_eq_check:
|
|
assert getattr2(model, attr) == getattr2(model_copy, attr)
|
|
|
|
return model_copy
|
|
|
|
@pytest.mark.with_downloads()
|
|
@pytest.mark.unit
|
|
def test_EncDecCTCModelBPE_v2(self):
|
|
# TODO: Switch to using named configs because here we don't really care about weights
|
|
cn = EncDecCTCModelBPE.from_pretrained(model_name="stt_en_conformer_ctc_small")
|
|
self.__test_restore_elsewhere(model=cn, attr_for_eq_check=set(["decoder._feat_in", "decoder._num_classes"]))
|
|
|
|
@pytest.mark.unit
|
|
def test_mock_save_to_restore_from(self):
|
|
with tempfile.NamedTemporaryFile('w') as empty_file:
|
|
# Write some data
|
|
empty_file.writelines(["*****\n"])
|
|
empty_file.flush()
|
|
|
|
# Update config
|
|
cfg = _mock_model_config()
|
|
cfg.model.temp_file = empty_file.name
|
|
|
|
# Create model
|
|
model = MockModel(cfg=cfg.model, trainer=None)
|
|
model = model.to('cpu')
|
|
|
|
assert model.temp_file == empty_file.name
|
|
|
|
# Save test
|
|
model_copy = self.__test_restore_elsewhere(model, map_location='cpu')
|
|
|
|
# Restore test
|
|
diff = model.w.weight - model_copy.w.weight
|
|
# because of caching - cache gets prepended
|
|
assert os.path.basename(model_copy.temp_file).endswith(os.path.basename(model.temp_file))
|
|
assert diff.mean() <= 1e-9
|
|
# assert os.path.basename(model.temp_file) == model_copy.temp_file
|
|
assert model_copy.temp_data == ["*****\n"]
|
|
|
|
@pytest.mark.unit
|
|
def test_mock_restore_from_config_only(self):
|
|
with tempfile.NamedTemporaryFile('w') as empty_file:
|
|
# Write some data
|
|
empty_file.writelines(["*****\n"])
|
|
empty_file.flush()
|
|
|
|
# Update config
|
|
cfg = _mock_model_config()
|
|
cfg.model.temp_file = os.path.abspath(empty_file.name)
|
|
|
|
# Inject arbitrary config arguments (after creating model)
|
|
with open_dict(cfg.model):
|
|
cfg.model.xyz = "abc"
|
|
|
|
# Create model
|
|
model = MockModel(cfg=cfg.model, trainer=None)
|
|
model = model.to('cpu')
|
|
|
|
assert model.temp_file == empty_file.name
|
|
model_copy = self.__test_restore_elsewhere(model, map_location='cpu', return_config=False)
|
|
# because of caching - cache gets prepended
|
|
assert os.path.basename(model_copy.temp_file).endswith(os.path.basename(model.temp_file))
|
|
# assert filecmp.cmp(model.temp_file, model_copy._cfg.temp_file)
|
|
assert model.cfg.xyz == model_copy.cfg.xyz
|
|
|
|
@pytest.mark.unit
|
|
def test_mock_restore_from_config_override_with_OmegaConf(self):
|
|
with tempfile.NamedTemporaryFile('w') as empty_file:
|
|
# Write some data
|
|
empty_file.writelines(["*****\n"])
|
|
empty_file.flush()
|
|
|
|
# Update config
|
|
cfg = _mock_model_config()
|
|
cfg.model.temp_file = empty_file.name
|
|
|
|
# Create model
|
|
model = MockModel(cfg=cfg.model, trainer=None)
|
|
model = model.to('cpu')
|
|
|
|
assert model.temp_file == empty_file.name
|
|
|
|
# Inject arbitrary config arguments (after creating model)
|
|
with open_dict(cfg.model):
|
|
cfg.model.xyz = "abc"
|
|
|
|
# Save test (with overriden config as OmegaConf object)
|
|
model_copy = self.__test_restore_elsewhere(model, map_location='cpu', override_config_path=cfg)
|
|
|
|
# Restore test
|
|
diff = model.w.weight - model_copy.w.weight
|
|
assert diff.mean() <= 1e-9
|
|
assert model_copy.temp_data == ["*****\n"]
|
|
|
|
# Test that new config has arbitrary content
|
|
assert model_copy.cfg.xyz == "abc"
|
|
|
|
@pytest.mark.unit
|
|
def test_mock_restore_from_config_override_with_yaml(self):
|
|
with tempfile.NamedTemporaryFile('w') as empty_file, tempfile.NamedTemporaryFile('w') as config_file:
|
|
# Write some data
|
|
empty_file.writelines(["*****\n"])
|
|
empty_file.flush()
|
|
|
|
# Update config
|
|
cfg = _mock_model_config()
|
|
cfg.model.temp_file = empty_file.name
|
|
|
|
# Create model
|
|
model = MockModel(cfg=cfg.model, trainer=None)
|
|
model = model.to('cpu')
|
|
|
|
assert model.temp_file == empty_file.name
|
|
|
|
# Inject arbitrary config arguments (after creating model)
|
|
with open_dict(cfg.model):
|
|
cfg.model.xyz = "abc"
|
|
|
|
# Write new config into file
|
|
OmegaConf.save(cfg, config_file)
|
|
|
|
# Save test (with overriden config as OmegaConf object)
|
|
model_copy = self.__test_restore_elsewhere(
|
|
model, map_location='cpu', override_config_path=config_file.name
|
|
)
|
|
|
|
# Restore test
|
|
diff = model.w.weight - model_copy.w.weight
|
|
assert diff.mean() <= 1e-9
|
|
assert filecmp.cmp(model.temp_file, model_copy.temp_file)
|
|
assert model_copy.temp_data == ["*****\n"]
|
|
|
|
# Test that new config has arbitrary content
|
|
assert model_copy.cfg.xyz == "abc"
|
|
|
|
@pytest.mark.unit
|
|
def test_mock_save_to_restore_from_with_target_class(self):
|
|
with tempfile.NamedTemporaryFile('w') as empty_file:
|
|
# Write some data
|
|
empty_file.writelines(["*****\n"])
|
|
empty_file.flush()
|
|
|
|
# Update config
|
|
cfg = _mock_model_config()
|
|
cfg.model.temp_file = empty_file.name
|
|
|
|
# Create model
|
|
model = MockModel(cfg=cfg.model, trainer=None)
|
|
model = model.to('cpu') # type: MockModel
|
|
|
|
assert model.temp_file == empty_file.name
|
|
|
|
# Save file using MockModel
|
|
with tempfile.TemporaryDirectory() as save_folder:
|
|
save_path = os.path.join(save_folder, "temp.nemo")
|
|
model.save_to(save_path)
|
|
|
|
# Restore test (using ModelPT as restorer)
|
|
# This forces the target class = MockModel to be used as resolver
|
|
model_copy = ModelPT.restore_from(save_path, map_location='cpu')
|
|
# because of caching - cache gets prepended
|
|
assert os.path.basename(model_copy.temp_file).endswith(os.path.basename(model.temp_file))
|
|
# assert filecmp.cmp(model.temp_file, model_copy.temp_file)
|
|
# Restore test
|
|
diff = model.w.weight - model_copy.w.weight
|
|
assert diff.mean() <= 1e-9
|
|
assert isinstance(model_copy, MockModel)
|
|
assert model_copy.temp_data == ["*****\n"]
|
|
|
|
@pytest.mark.unit
|
|
def test_mock_save_to_restore_from_multiple_models(self):
|
|
with tempfile.NamedTemporaryFile('w') as empty_file, tempfile.NamedTemporaryFile('w') as empty_file2:
|
|
# Write some data
|
|
empty_file.writelines(["*****\n"])
|
|
empty_file.flush()
|
|
empty_file2.writelines(["+++++\n"])
|
|
empty_file2.flush()
|
|
|
|
# Update config + create ,pde;s
|
|
cfg = _mock_model_config()
|
|
cfg.model.temp_file = empty_file.name
|
|
cfg2 = _mock_model_config()
|
|
cfg2.model.temp_file = empty_file2.name
|
|
|
|
# Create models
|
|
model = MockModel(cfg=cfg.model, trainer=None)
|
|
model = model.to('cpu')
|
|
model2 = MockModel(cfg=cfg2.model, trainer=None)
|
|
model2 = model2.to('cpu')
|
|
|
|
assert model.temp_file == empty_file.name
|
|
assert model2.temp_file == empty_file2.name
|
|
|
|
# Save test
|
|
model_copy = self.__test_restore_elsewhere(model, map_location='cpu')
|
|
model2_copy = self.__test_restore_elsewhere(model2, map_location='cpu')
|
|
|
|
# Restore test
|
|
assert model_copy.temp_data == ["*****\n"]
|
|
assert model2_copy.temp_data == ["+++++\n"]
|
|
|
|
@pytest.mark.unit
|
|
def test_mock_save_to_restore_from_multiple_models_inverted_order(self):
|
|
with tempfile.NamedTemporaryFile('w') as empty_file, tempfile.NamedTemporaryFile('w') as empty_file2:
|
|
# Write some data
|
|
empty_file.writelines(["*****\n"])
|
|
empty_file.flush()
|
|
empty_file2.writelines(["+++++\n"])
|
|
empty_file2.flush()
|
|
|
|
# Update config + create ,pde;s
|
|
cfg = _mock_model_config()
|
|
cfg.model.temp_file = empty_file.name
|
|
cfg2 = _mock_model_config()
|
|
cfg2.model.temp_file = empty_file2.name
|
|
|
|
# Create models
|
|
model = MockModel(cfg=cfg.model, trainer=None)
|
|
model = model.to('cpu')
|
|
model2 = MockModel(cfg=cfg2.model, trainer=None)
|
|
model2 = model2.to('cpu')
|
|
|
|
assert model.temp_file == empty_file.name
|
|
assert model2.temp_file == empty_file2.name
|
|
|
|
# Save test (inverted order)
|
|
model2_copy = self.__test_restore_elsewhere(model2, map_location='cpu')
|
|
model_copy = self.__test_restore_elsewhere(model, map_location='cpu')
|
|
|
|
# Restore test
|
|
assert model_copy.temp_data == ["*****\n"]
|
|
assert model2_copy.temp_data == ["+++++\n"]
|
|
|
|
@pytest.mark.unit
|
|
def test_mock_save_to_restore_chained(self):
|
|
with tempfile.NamedTemporaryFile('w') as empty_file, tempfile.NamedTemporaryFile('w') as empty_file2:
|
|
# Write some data
|
|
empty_file.writelines(["*****\n"])
|
|
empty_file.flush()
|
|
|
|
# Update config + create ,pde;s
|
|
cfg = _mock_model_config()
|
|
cfg.model.temp_file = empty_file.name
|
|
|
|
# Create models
|
|
model = MockModel(cfg=cfg.model, trainer=None)
|
|
model = model.to('cpu')
|
|
|
|
assert model.temp_file == empty_file.name
|
|
|
|
def save_copy(model, save_folder, restore_folder):
|
|
# Where model will be saved
|
|
model_save_path = os.path.join(save_folder, f"{model.__class__.__name__}.nemo")
|
|
model.save_to(save_path=model_save_path)
|
|
# Where model will be restored from
|
|
model_restore_path = os.path.join(restore_folder, f"{model.__class__.__name__}.nemo")
|
|
shutil.copy(model_save_path, model_restore_path)
|
|
return model_restore_path
|
|
|
|
# Save test
|
|
with tempfile.TemporaryDirectory() as level4:
|
|
with tempfile.TemporaryDirectory() as level3:
|
|
with tempfile.TemporaryDirectory() as level2:
|
|
with tempfile.TemporaryDirectory() as level1:
|
|
path = save_copy(model, level1, level2)
|
|
model_copy2 = model.__class__.restore_from(path)
|
|
path = save_copy(model_copy2, level2, level3)
|
|
model_copy3 = model.__class__.restore_from(path)
|
|
path = save_copy(model_copy3, level3, level4)
|
|
model_copy = model.__class__.restore_from(path)
|
|
|
|
# Restore test
|
|
assert model_copy.temp_data == ["*****\n"]
|
|
|
|
# AppState test
|
|
appstate = AppState()
|
|
metadata = appstate.get_model_metadata_from_guid(model_copy.model_guid)
|
|
assert metadata.guid != model.model_guid
|
|
assert metadata.restoration_path == path
|
|
|
|
@pytest.mark.unit
|
|
def test_mock_save_to_multiple_times(self):
|
|
with tempfile.NamedTemporaryFile('w') as empty_file, tempfile.TemporaryDirectory() as tmpdir:
|
|
# Write some data
|
|
empty_file.writelines(["*****\n"])
|
|
empty_file.flush()
|
|
|
|
# Update config
|
|
cfg = _mock_model_config()
|
|
cfg.model.temp_file = empty_file.name
|
|
|
|
# Create model
|
|
model = MockModel(cfg=cfg.model, trainer=None) # type: MockModel
|
|
model = model.to('cpu')
|
|
|
|
assert model.temp_file == empty_file.name
|
|
|
|
# Save test
|
|
model.save_to(os.path.join(tmpdir, 'save_0.nemo'))
|
|
model.save_to(os.path.join(tmpdir, 'save_1.nemo'))
|
|
model.save_to(os.path.join(tmpdir, 'save_2.nemo'))
|
|
|
|
@pytest.mark.unit
|
|
def test_multiple_model_save_restore_connector(self):
|
|
class MySaveRestoreConnector(save_restore_connector.SaveRestoreConnector):
|
|
def save_to(self, model, save_path: str):
|
|
save_path = save_path.replace(".nemo", "_XYZ.nemo")
|
|
super(MySaveRestoreConnector, self).save_to(model, save_path)
|
|
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
# Update config
|
|
cfg = _mock_model_config()
|
|
# Create model
|
|
model = MockModel(cfg=cfg.model, trainer=None)
|
|
model_with_custom_connector = MockModel(cfg=cfg.model, trainer=None)
|
|
model_with_custom_connector._save_restore_connector = MySaveRestoreConnector()
|
|
model_with_custom_connector.save_to(os.path.join(tmpdir, 'save_custom.nemo'))
|
|
|
|
assert os.path.exists(os.path.join(tmpdir, 'save_custom_XYZ.nemo'))
|
|
assert isinstance(model._save_restore_connector, save_restore_connector.SaveRestoreConnector)
|
|
assert isinstance(model_with_custom_connector._save_restore_connector, MySaveRestoreConnector)
|
|
|
|
assert type(MockModel._save_restore_connector) == save_restore_connector.SaveRestoreConnector
|
|
|
|
@pytest.mark.unit
|
|
def test_restore_from_save_restore_connector(self):
|
|
class MySaveRestoreConnector(save_restore_connector.SaveRestoreConnector):
|
|
def save_to(self, model, save_path: str):
|
|
save_path = save_path.replace(".nemo", "_XYZ.nemo")
|
|
super().save_to(model, save_path)
|
|
|
|
class MockModelV2(MockModel):
|
|
pass
|
|
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
# Update config
|
|
cfg = _mock_model_config()
|
|
|
|
# Create model
|
|
save_path = os.path.join(tmpdir, 'save_custom.nemo')
|
|
model_with_custom_connector = MockModel(cfg=cfg.model, trainer=None)
|
|
model_with_custom_connector._save_restore_connector = MySaveRestoreConnector()
|
|
model_with_custom_connector.save_to(save_path)
|
|
|
|
assert os.path.exists(os.path.join(tmpdir, 'save_custom_XYZ.nemo'))
|
|
|
|
restored_model = MockModelV2.restore_from(
|
|
save_path.replace(".nemo", "_XYZ.nemo"), save_restore_connector=MySaveRestoreConnector()
|
|
)
|
|
assert type(restored_model) == MockModelV2
|
|
assert type(restored_model._save_restore_connector) == MySaveRestoreConnector
|
|
|
|
@pytest.mark.unit
|
|
def test_restore_from_save_restore_connector_return_config(self):
|
|
class MySaveRestoreConnector(save_restore_connector.SaveRestoreConnector):
|
|
def save_to(self, model, save_path: str):
|
|
save_path = save_path.replace(".nemo", "_XYZ.nemo")
|
|
super().save_to(model, save_path)
|
|
|
|
class MockModelV2(MockModel):
|
|
pass
|
|
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
# Update config
|
|
cfg = _mock_model_config()
|
|
|
|
# Create model
|
|
save_path = os.path.join(tmpdir, 'save_custom.nemo')
|
|
model_with_custom_connector = MockModel(cfg=cfg.model, trainer=None)
|
|
model_with_custom_connector._save_restore_connector = MySaveRestoreConnector()
|
|
model_with_custom_connector.save_to(save_path)
|
|
|
|
assert os.path.exists(os.path.join(tmpdir, 'save_custom_XYZ.nemo'))
|
|
|
|
restored_model_cfg = MockModelV2.restore_from(
|
|
save_path.replace(".nemo", "_XYZ.nemo"),
|
|
save_restore_connector=MySaveRestoreConnector(),
|
|
return_config=True,
|
|
)
|
|
assert isinstance(restored_model_cfg, DictConfig)
|
|
assert model_with_custom_connector.cfg == restored_model_cfg
|
|
|
|
@pytest.mark.unit
|
|
def test_restore_from_save_restore_connector_return_config_partial_tar_extraction(self):
|
|
class MySaveRestoreConnector(save_restore_connector.SaveRestoreConnector):
|
|
def save_to(self, model, save_path: str):
|
|
save_path = save_path.replace(".nemo", "_XYZ.nemo")
|
|
super().save_to(model, save_path)
|
|
|
|
class MockModelV2(MockModel):
|
|
pass
|
|
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
# Update config
|
|
cfg = _mock_model_config()
|
|
|
|
# Create model
|
|
save_path = os.path.join(tmpdir, 'save_custom.nemo')
|
|
model_with_custom_connector = MockModel(cfg=cfg.model, trainer=None)
|
|
model_with_custom_connector._save_restore_connector = MySaveRestoreConnector()
|
|
model_with_custom_connector.save_to(save_path)
|
|
|
|
true_save_path = os.path.join(tmpdir, 'save_custom_XYZ.nemo')
|
|
assert os.path.exists(true_save_path)
|
|
|
|
my_connector = MySaveRestoreConnector()
|
|
|
|
with tempfile.TemporaryDirectory() as config_tmpdir:
|
|
config_members = my_connector._filtered_tar_info(
|
|
true_save_path, filter_fn=lambda name: '.yaml' in name
|
|
)
|
|
my_connector._unpack_nemo_file(true_save_path, out_folder=config_tmpdir, members=config_members)
|
|
current_files = list(os.listdir(config_tmpdir))
|
|
|
|
assert len(current_files) == 1 # only config file should have been extracted, no pytorch params
|
|
config_filepath = current_files[0]
|
|
assert config_filepath.endswith(".yaml")
|
|
|
|
@pytest.mark.unit
|
|
def test_restore_from_save_restore_connector_unpacked_file(self):
|
|
class MySaveRestoreConnector(save_restore_connector.SaveRestoreConnector):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.pack_nemo_file = False
|
|
|
|
def save_to(self, model, save_path: str):
|
|
save_path = save_path.replace(".nemo", "_XYZ.nemo")
|
|
super().save_to(model, save_path)
|
|
|
|
class MockModelV2(MockModel):
|
|
pass
|
|
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
# Update config
|
|
cfg = _mock_model_config()
|
|
|
|
# Create model
|
|
save_path = os.path.join(tmpdir, 'temp_model')
|
|
os.makedirs(save_path, exist_ok=True)
|
|
model_with_custom_connector = MockModel(cfg=cfg.model, trainer=None)
|
|
model_with_custom_connector._save_restore_connector = MySaveRestoreConnector()
|
|
model_with_custom_connector.save_to(save_path + "/abc.nemo")
|
|
|
|
assert os.path.isdir(save_path)
|
|
assert len(os.listdir(save_path)) == 2 # config and pytorch params
|
|
|
|
restore_connector = MySaveRestoreConnector()
|
|
restore_connector.model_extracted_dir = save_path
|
|
restored_model = MockModelV2.restore_from(
|
|
save_path.replace(".nemo", "_XYZ.nemo"), save_restore_connector=restore_connector
|
|
)
|
|
assert type(restored_model) == MockModelV2
|
|
assert type(restored_model._save_restore_connector) == MySaveRestoreConnector
|
|
|
|
@pytest.mark.unit
|
|
def test_mock_model_model_collision(self):
|
|
# The usual pipeline is working just fine.
|
|
cfg = _mock_model_config()
|
|
model = MockModel(cfg=cfg.model, trainer=None) # type: MockModel
|
|
model = model.to('cpu')
|
|
|
|
# Let's create a custom config with a 'model.model' node.
|
|
cfg = _mock_model_config()
|
|
OmegaConf.set_struct(cfg, False)
|
|
cfg.model.model = 'aaa'
|
|
OmegaConf.set_struct(cfg, True)
|
|
|
|
# Failing due to collision.
|
|
with pytest.raises(ValueError, match="Creating model config node is forbidden"):
|
|
model = MockModel(cfg=cfg.model, trainer=None) # type: MockModel
|
|
model = model.to('cpu')
|
|
|
|
@pytest.mark.unit
|
|
@pytest.mark.parametrize("change_child_number", [False, True])
|
|
@pytest.mark.parametrize("child2_model_from_path", [False, True])
|
|
def test_mock_model_nested(self, change_child_number: bool, child2_model_from_path: bool):
|
|
"""
|
|
Test model with 2 children
|
|
Model and each child can be saved/restored separately
|
|
Model is constructed using saved child models (.nemo checkpoints)
|
|
|
|
Args:
|
|
change_child_number: if change_child_number is True, child model changes its config
|
|
without notifying parent model, and saved parent model should handle this correctly.
|
|
child2_model_from_path: if child2_model_from_path is True, child2 model is restored from .nemo checkpoint,
|
|
otherwise constructed directly from config. Child1 model always loaded from checkpoint.
|
|
"""
|
|
# children - models without sub-children
|
|
cfg_child1 = _mock_model_config()
|
|
cfg_child2 = _mock_model_with_children_config() # no children
|
|
|
|
# Create models
|
|
child1 = MockModel(cfg=cfg_child1.model, trainer=None)
|
|
child1 = child1.to('cpu')
|
|
with tempfile.TemporaryDirectory() as tmpdir_parent:
|
|
parent_path = os.path.join(tmpdir_parent, "parent.nemo")
|
|
with tempfile.TemporaryDirectory() as tmpdir_child:
|
|
# save children
|
|
child1_path = os.path.join(tmpdir_child, 'child1.nemo')
|
|
child1.save_to(child1_path)
|
|
if child2_model_from_path:
|
|
child2 = MockModelWithChildren(cfg=cfg_child2.model, trainer=None)
|
|
child2 = child2.to('cpu')
|
|
child2_path = os.path.join(tmpdir_child, 'child2.nemo')
|
|
child2.save_to(child2_path)
|
|
|
|
# create model with children using saved .nemo model checkpoints
|
|
cfg_parent = _mock_model_with_children_config(
|
|
child1_model_path=child1_path, child2_model_path=child2_path
|
|
)
|
|
else:
|
|
# child 2 model will be directly constructed from subconfig
|
|
cfg_parent = _mock_model_with_children_config(
|
|
child1_model_path=child1_path, child2_model_path=None, child2_model_cfg=cfg_child2.get("model")
|
|
)
|
|
|
|
parent = MockModelWithChildren(cfg_parent.model)
|
|
if change_child_number:
|
|
parent.child2_model.change_stub_number(10)
|
|
parent.save_to(parent_path)
|
|
# restore, separate children checkpoints are not available here (tmpdir_child destroyed)
|
|
parent = ModelPT.restore_from(parent_path)
|
|
# check model is transparent, child models can be accessed and can be saved/restored separately
|
|
_ = self.__test_restore_elsewhere(parent.child1_model, map_location='cpu')
|
|
child2 = self.__test_restore_elsewhere(parent.child2_model, map_location='cpu')
|
|
if change_child_number:
|
|
assert child2.cfg.stub_number == 10
|
|
|
|
# check model itself can be saved/restored
|
|
parent = self.__test_restore_elsewhere(parent, map_location='cpu')
|
|
if change_child_number:
|
|
assert parent.child2_model.cfg.stub_number == 10
|
|
|
|
@pytest.mark.unit
|
|
@pytest.mark.parametrize("change_child_resource", [False, True])
|
|
@pytest.mark.parametrize("child2_model_from_path", [False, True])
|
|
def test_mock_model_nested_with_resources(self, change_child_resource: bool, child2_model_from_path: bool):
|
|
"""
|
|
Test nested model with 2 children: model and each child can be saved/restored separately
|
|
child models and parent model itself contain resources
|
|
|
|
Args:
|
|
change_child_resource: if change_child_resource is True,
|
|
child model resources are changed after instantiation parent model.
|
|
child2_model_from_path: if child2_model_from_path is True, child2 model is restored from .nemo checkpoint,
|
|
otherwise constructed directly from config. Child1 model always loaded from checkpoint.
|
|
"""
|
|
with (
|
|
tempfile.NamedTemporaryFile('w') as file_child1,
|
|
tempfile.NamedTemporaryFile('w') as file_child2,
|
|
tempfile.NamedTemporaryFile('w') as file_child2_other,
|
|
tempfile.NamedTemporaryFile('w') as file_parent,
|
|
):
|
|
# write text data, use these files as resources
|
|
parent_data = ["*****\n"]
|
|
child1_data = ["+++++\n"]
|
|
child2_data = ["-----\n"]
|
|
child2_data_other = [".....\n"]
|
|
file_parent.writelines(parent_data)
|
|
file_parent.flush()
|
|
file_child1.writelines(child1_data)
|
|
file_child1.flush()
|
|
file_child2.writelines(child2_data)
|
|
file_child2.flush()
|
|
file_child2_other.writelines(child2_data_other)
|
|
file_child2_other.flush()
|
|
|
|
# construct child models with resources
|
|
# create configs
|
|
cfg_child1 = _mock_model_config()
|
|
cfg_child1.model.temp_file = file_child1.name
|
|
cfg_child2 = _mock_model_with_children_config() # no sub-children
|
|
cfg_child2.model.temp_file = file_child2.name
|
|
# create child models
|
|
child1 = MockModel(cfg=cfg_child1.model, trainer=None)
|
|
child1 = child1.to('cpu')
|
|
|
|
with tempfile.TemporaryDirectory() as tmpdir_parent:
|
|
parent_path = os.path.join(tmpdir_parent, "parent.nemo")
|
|
with tempfile.TemporaryDirectory() as tmpdir_child:
|
|
# save children
|
|
child1_path = os.path.join(tmpdir_child, 'child1.nemo')
|
|
child1.save_to(child1_path)
|
|
if child2_model_from_path:
|
|
child2 = MockModelWithChildren(cfg=cfg_child2.model, trainer=None)
|
|
child2 = child2.to('cpu')
|
|
child2_path = os.path.join(tmpdir_child, 'child2.nemo')
|
|
child2.save_to(child2_path)
|
|
|
|
# create model with children using saved .nemo model checkpoints
|
|
cfg_parent = _mock_model_with_children_config(
|
|
child1_model_path=child1_path, child2_model_path=child2_path
|
|
)
|
|
else:
|
|
# child 2 model will be directly constructed from subconfig
|
|
cfg_parent = _mock_model_with_children_config(
|
|
child1_model_path=child1_path,
|
|
child2_model_path=None,
|
|
child2_model_cfg=cfg_child2.get("model"),
|
|
)
|
|
|
|
cfg_parent.model.temp_file = file_parent.name # add resource
|
|
parent = MockModelWithChildren(cfg_parent.model)
|
|
if change_child_resource:
|
|
parent.child2_model.setup_data_from_file(file_child2_other.name)
|
|
parent.save_to(parent_path)
|
|
|
|
# restore, separate children checkpoints are not available here (tmpdir_child destroyed)
|
|
parent = ModelPT.restore_from(parent_path)
|
|
# check model is transparent, child models can be accessed and can be saved/restored separately
|
|
child1 = self.__test_restore_elsewhere(parent.child1_model, map_location='cpu')
|
|
child2 = self.__test_restore_elsewhere(parent.child2_model, map_location='cpu')
|
|
# test parent save/restore
|
|
parent = self.__test_restore_elsewhere(parent, map_location='cpu')
|
|
|
|
# test resources
|
|
# check separately restored child models
|
|
assert child1.temp_data == child1_data
|
|
if change_child_resource:
|
|
assert child2.temp_data == child2_data_other
|
|
else:
|
|
assert child2.temp_data == child2_data
|
|
# test parent model + child models
|
|
assert parent.temp_data == parent_data
|
|
assert parent.child1_model.temp_data == child1_data
|
|
if change_child_resource:
|
|
assert parent.child2_model.temp_data == child2_data_other
|
|
else:
|
|
assert parent.child2_model.temp_data == child2_data
|
|
|
|
@pytest.mark.unit
|
|
def test_mock_model_nested_with_resources_multiple_passes(self):
|
|
"""
|
|
Test nested model with 2 children: multiple save-restore passes
|
|
child models and parent model itself contain resources
|
|
"""
|
|
with (
|
|
tempfile.NamedTemporaryFile('w') as file_child1,
|
|
tempfile.NamedTemporaryFile('w') as file_child2,
|
|
tempfile.NamedTemporaryFile('w') as file_child2_other,
|
|
tempfile.NamedTemporaryFile('w') as file_parent,
|
|
):
|
|
# write text data, use these files as resources
|
|
parent_data = ["*****\n"]
|
|
child1_data = ["+++++\n"]
|
|
child2_data = ["-----\n"]
|
|
child2_data_other = [".....\n"]
|
|
file_parent.writelines(parent_data)
|
|
file_parent.flush()
|
|
file_child1.writelines(child1_data)
|
|
file_child1.flush()
|
|
file_child2.writelines(child2_data)
|
|
file_child2.flush()
|
|
file_child2_other.writelines(child2_data_other)
|
|
file_child2_other.flush()
|
|
|
|
# construct child models with resources
|
|
# create configs
|
|
cfg_child1 = _mock_model_config()
|
|
cfg_child1.model.temp_file = file_child1.name
|
|
cfg_child2 = _mock_model_with_children_config() # no sub-children
|
|
cfg_child2.model.temp_file = file_child2.name
|
|
# create child models
|
|
child1 = MockModel(cfg=cfg_child1.model, trainer=None)
|
|
child1 = child1.to('cpu')
|
|
child2 = MockModelWithChildren(cfg=cfg_child2.model, trainer=None)
|
|
child2 = child2.to('cpu')
|
|
|
|
with (
|
|
tempfile.TemporaryDirectory() as tmpdir_parent1,
|
|
tempfile.TemporaryDirectory() as tmpdir_parent2,
|
|
tempfile.TemporaryDirectory() as tmpdir_parent3,
|
|
tempfile.TemporaryDirectory() as tmpdir_parent4,
|
|
):
|
|
parent_path1 = os.path.join(tmpdir_parent1, "parent.nemo")
|
|
parent_path2 = os.path.join(tmpdir_parent2, "parent.nemo")
|
|
with tempfile.TemporaryDirectory() as tmpdir_child:
|
|
# save children
|
|
child1_path = os.path.join(tmpdir_child, 'child1.nemo')
|
|
child1.save_to(child1_path)
|
|
child2_path = os.path.join(tmpdir_child, 'child2.nemo')
|
|
child2.save_to(child2_path)
|
|
|
|
# create model with children using saved "nemo" checkpoints
|
|
cfg_parent = _mock_model_with_children_config(
|
|
child1_model_path=child1_path, child2_model_path=child2_path
|
|
)
|
|
cfg_parent.model.temp_file = file_parent.name # add resource
|
|
parent = MockModelWithChildren(cfg_parent.model)
|
|
|
|
# save-restore first pass
|
|
# save to different locations
|
|
parent.save_to(parent_path1)
|
|
parent.save_to(parent_path2)
|
|
|
|
# restore, separate children checkpoints are not available here (tmpdir_child destroyed)
|
|
parent1 = ModelPT.restore_from(parent_path1)
|
|
parent2 = ModelPT.restore_from(parent_path2)
|
|
|
|
# check resources
|
|
for parent in (parent1, parent2):
|
|
assert parent.temp_data == parent_data
|
|
assert parent.child1_model.temp_data == child1_data
|
|
assert parent.child2_model.temp_data == child2_data
|
|
|
|
del parent2 # use parent1 for second pass
|
|
|
|
# save-restore second pass
|
|
parent_path3 = os.path.join(tmpdir_parent3, "parent.nemo")
|
|
parent_path4 = os.path.join(tmpdir_parent4, "parent.nemo")
|
|
parent1.save_to(parent_path3)
|
|
parent1.save_to(parent_path4)
|
|
|
|
parent3 = ModelPT.restore_from(parent_path3)
|
|
parent4 = ModelPT.restore_from(parent_path4)
|
|
|
|
# check resources
|
|
for parent in (parent3, parent4):
|
|
assert parent.temp_data == parent_data
|
|
assert parent.child1_model.temp_data == child1_data
|
|
assert parent.child2_model.temp_data == child2_data
|
|
|
|
@pytest.mark.unit
|
|
def test_mock_model_nested_double_with_resources(self):
|
|
"""
|
|
test nested model: parent -> child_with_child -> child; model and each child can be saved/restored separately
|
|
all models can contain resources
|
|
"""
|
|
with (
|
|
tempfile.NamedTemporaryFile('w') as file_child,
|
|
tempfile.NamedTemporaryFile('w') as file_child_with_child,
|
|
tempfile.NamedTemporaryFile('w') as file_parent,
|
|
):
|
|
# write text data, use these files as resources
|
|
parent_data = ["*****\n"]
|
|
child_with_child_data = ["+++++\n"]
|
|
child_data = ["-----\n"]
|
|
file_parent.writelines(parent_data)
|
|
file_parent.flush()
|
|
file_child_with_child.writelines(child_with_child_data)
|
|
file_child_with_child.flush()
|
|
file_child.writelines(child_data)
|
|
file_child.flush()
|
|
|
|
# construct child model (leaf) with resource
|
|
cfg_child = _mock_model_config()
|
|
cfg_child.model.temp_file = file_child.name
|
|
child = MockModel(cfg=cfg_child.model, trainer=None)
|
|
child = child.to('cpu')
|
|
|
|
with tempfile.TemporaryDirectory() as tmpdir_parent:
|
|
parent_path = os.path.join(tmpdir_parent, "parent.nemo")
|
|
with tempfile.TemporaryDirectory() as tmpdir_child_with_child:
|
|
child_with_child_path = os.path.join(tmpdir_child_with_child, 'child_with_child.nemo')
|
|
with tempfile.TemporaryDirectory() as tmpdir_child:
|
|
# save child
|
|
child_path = os.path.join(tmpdir_child, 'child.nemo')
|
|
child.save_to(child_path)
|
|
|
|
# create child model with child
|
|
cfg_child_with_child = _mock_model_with_children_config(
|
|
child1_model_path=None, child2_model_path=child_path
|
|
)
|
|
cfg_child_with_child.model.temp_file = file_child_with_child.name
|
|
child_with_child = MockModelWithChildren(cfg_child_with_child.model)
|
|
child_with_child.save_to(child_with_child_path)
|
|
# create parent model with child-with-child, leaf checkpoint is not available here
|
|
cfg_parent = _mock_model_with_children_config(
|
|
child1_model_path=None, child2_model_path=child_with_child_path
|
|
)
|
|
cfg_parent.model.temp_file = file_parent.name
|
|
parent = MockModelWithChildren(cfg_parent.model)
|
|
parent.save_to(parent_path)
|
|
|
|
# restore, separate children checkpoints are not available here
|
|
# tmpdir_child, tmpdir_child_with_child are destroyed
|
|
parent = ModelPT.restore_from(parent_path)
|
|
# model is transparent, children and model itself can be saved/restored
|
|
child = self.__test_restore_elsewhere(parent.child2_model.child2_model, map_location='cpu')
|
|
child_with_child = self.__test_restore_elsewhere(parent.child2_model, map_location='cpu')
|
|
parent = self.__test_restore_elsewhere(parent, map_location='cpu')
|
|
|
|
# test resources for all restored models
|
|
# leaf model
|
|
assert child.temp_data == child_data
|
|
# child with child
|
|
assert child_with_child.temp_data == child_with_child_data
|
|
assert child_with_child.child2_model.temp_data == child_data
|
|
# parent
|
|
assert parent.temp_data == parent_data
|
|
assert parent.child2_model.temp_data == child_with_child_data
|
|
assert parent.child2_model.child2_model.temp_data == child_data
|
|
|
|
# check named_nemo_modules: parent -> child2 -> child2.child2,
|
|
# tuples of (attribute_path, cfg_path, module)
|
|
named_nemo_modules = list(parent.named_nemo_modules())
|
|
etalon_nemo_modules = [
|
|
("", "", parent),
|
|
("child2_model", "child2_model", parent.child2_model),
|
|
("child2_model.child2_model", "child2_model.child2_model", parent.child2_model.child2_model),
|
|
]
|
|
assert len(named_nemo_modules) == len(etalon_nemo_modules)
|
|
for etalon, actual in zip(etalon_nemo_modules, named_nemo_modules):
|
|
assert etalon[0] == actual[0]
|
|
assert etalon[1] == actual[1]
|
|
assert etalon[2] is actual[2]
|
|
|
|
@pytest.mark.unit
|
|
@pytest.mark.with_downloads
|
|
def test_mock_model_nested_child_from_pretrained(self):
|
|
"""
|
|
Test nested model with child initialized from pretrained model
|
|
"""
|
|
cfg = _mock_model_with_child_encdecctcbpe_config("stt_en_conformer_ctc_small")
|
|
parent = MockModelWithChildEncDecCTCBPE(cfg=cfg.model, trainer=None)
|
|
with tempfile.TemporaryDirectory() as tmpdir_parent:
|
|
parent_path = os.path.join(tmpdir_parent, "parent.nemo")
|
|
|
|
# save, then restore
|
|
parent.save_to(parent_path)
|
|
parent = ModelPT.restore_from(parent_path)
|
|
|
|
# test child can be saved/restored
|
|
_ = self.__test_restore_elsewhere(parent.ctc_model, map_location='cpu')
|
|
# test parent can be saved/restored
|
|
parent = self.__test_restore_elsewhere(parent, map_location='cpu')
|
|
assert isinstance(parent.ctc_model, EncDecCTCModel)
|
|
|
|
assert _is_json_serializable(parent.ctc_model.hparams_initial)
|
|
|
|
@pytest.mark.unit
|
|
def test_mock_model_nested_custom_config_field(self):
|
|
"""
|
|
Test nested model with custom config field not equal to attribute name
|
|
Config is stored in `child1_model_config`
|
|
Child model is stored in `child1_model` attribute
|
|
"""
|
|
with tempfile.NamedTemporaryFile('w') as file_child1, tempfile.NamedTemporaryFile('w') as file_parent:
|
|
# write text data, use these files as resources
|
|
parent_data = ["*****\n"]
|
|
child1_data = ["+++++\n"]
|
|
file_parent.writelines(parent_data)
|
|
file_parent.flush()
|
|
file_child1.writelines(child1_data)
|
|
file_child1.flush()
|
|
|
|
cfg = _mock_model_with_child_custom_config_path_config()
|
|
cfg.model.temp_file = file_parent.name
|
|
cfg.model.child1_model_config.temp_file = file_child1.name
|
|
|
|
# construct parent model
|
|
parent = MockModelWithChildCustomConfigPath(cfg=cfg.model, trainer=None)
|
|
with tempfile.TemporaryDirectory() as tmpdir_parent:
|
|
parent_path = os.path.join(tmpdir_parent, "parent.nemo")
|
|
|
|
# save, then restore
|
|
parent.save_to(parent_path)
|
|
parent = ModelPT.restore_from(parent_path)
|
|
# test child can be saved/restored
|
|
_ = self.__test_restore_elsewhere(parent.child1_model, map_location='cpu')
|
|
# test parent can be saved/restored
|
|
parent = self.__test_restore_elsewhere(parent, map_location='cpu')
|
|
|
|
# check data
|
|
assert parent.temp_data == parent_data
|
|
assert parent.child1_model.temp_data == child1_data
|
|
|
|
# check named_nemo_modules: parent -> child, tuples of (attribute_path, cfg_path, module)
|
|
named_nemo_modules = list(parent.named_nemo_modules())
|
|
etalon_nemo_modules = [("", "", parent), ("child1_model", "child1_model_config", parent.child1_model)]
|
|
assert len(named_nemo_modules) == len(etalon_nemo_modules)
|
|
for etalon, actual in zip(etalon_nemo_modules, named_nemo_modules):
|
|
assert etalon[0] == actual[0]
|
|
assert etalon[1] == actual[1]
|
|
assert etalon[2] is actual[2]
|
|
|
|
@pytest.mark.unit
|
|
def test_using_nemo_checkpoint_as_artifact_disallowed(self):
|
|
"""
|
|
Test that using nemo checkpoint as artifact is disallowed
|
|
"""
|
|
cfg_child = _mock_model_config()
|
|
child = MockModel(cfg=cfg_child.model, trainer=None).to("cpu")
|
|
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
child_path = os.path.join(tmpdir, "child.nemo")
|
|
child.save_to(child_path)
|
|
|
|
cfg_parent = _mock_model_incorrect_with_nemo_artifact_config(child_path)
|
|
with pytest.raises(NeMoBaseException):
|
|
# registering .nemo checkpoint as an artifact is not allowed
|
|
_ = MockModelIncorrectWithNemoArtifact(cfg=cfg_parent.model, trainer=None)
|
|
|
|
@pytest.mark.unit
|
|
def test_restore_from_save_restore_connector_extracted_dir(self):
|
|
class MySaveRestoreConnector(save_restore_connector.SaveRestoreConnector):
|
|
def save_to(self, model, save_path: str):
|
|
save_path = save_path.replace(".nemo", "_XYZ.nemo")
|
|
super().save_to(model, save_path)
|
|
|
|
class MockModelV2(MockModel):
|
|
pass
|
|
|
|
with tempfile.TemporaryDirectory() as extracted_tempdir:
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
# Update config
|
|
cfg = _mock_model_config()
|
|
|
|
# Create model
|
|
save_path = os.path.join(tmpdir, 'save_custom.nemo')
|
|
model_with_custom_connector = MockModel(cfg=cfg.model, trainer=None)
|
|
model_with_custom_connector._save_restore_connector = MySaveRestoreConnector()
|
|
model_with_custom_connector.save_to(save_path)
|
|
|
|
nemo_filepath = os.path.join(tmpdir, 'save_custom_XYZ.nemo')
|
|
assert os.path.exists(nemo_filepath)
|
|
|
|
# extract the contents to this dir apriori
|
|
# simulate by extracting now before calling restore_from
|
|
connector = MySaveRestoreConnector()
|
|
MySaveRestoreConnector._unpack_nemo_file(nemo_filepath, extracted_tempdir)
|
|
assert get_size(extracted_tempdir) > 0
|
|
|
|
# delete the old directory and preserve only the new extracted directory (escape scope of old dir)
|
|
|
|
# next, set the model's extracted directory path
|
|
connector.model_extracted_dir = extracted_tempdir
|
|
|
|
# note, we pass in the "old" nemo_filepath, stored somewhere other than the extracted directory
|
|
# this nemo_filepath is no longer valid, and has been deleted.
|
|
restored_model = MockModelV2.restore_from(nemo_filepath, save_restore_connector=connector)
|
|
assert type(restored_model) == MockModelV2
|
|
assert type(restored_model._save_restore_connector) == MySaveRestoreConnector
|
|
|
|
# assert models have correct restoration information and paths
|
|
appstate = AppState()
|
|
original_metadata = appstate.get_model_metadata_from_guid(model_with_custom_connector.model_guid)
|
|
assert original_metadata.restoration_path is None
|
|
|
|
restored_metadata = appstate.get_model_metadata_from_guid(restored_model.model_guid)
|
|
assert restored_metadata.restoration_path is not None
|
|
|
|
# assert that the restore path was the path of the pre-extracted directory
|
|
# irrespective of whether an old `nemo_filepath` (which doesnt exist anymore) was passed to restore_from.
|
|
assert extracted_tempdir in restored_metadata.restoration_path
|
|
assert extracted_tempdir not in nemo_filepath
|
|
assert not os.path.exists(nemo_filepath)
|
|
|
|
# test for parameter equality
|
|
model_with_custom_connector = model_with_custom_connector.to('cpu')
|
|
restored_model = restored_model.to('cpu')
|
|
|
|
original_state_dict = model_with_custom_connector.state_dict()
|
|
restored_state_dict = restored_model.state_dict()
|
|
for orig, restored in zip(original_state_dict.keys(), restored_state_dict.keys()):
|
|
assert (original_state_dict[orig] - restored_state_dict[restored]).abs().mean() < 1e-6
|
|
|
|
@pytest.mark.unit
|
|
def test_hf_model_filter(self):
|
|
filt = ModelPT.get_hf_model_filter()
|
|
assert isinstance(filt, dict)
|
|
assert 'nemo' in filt['filter']
|
|
|
|
@pytest.mark.with_downloads()
|
|
@pytest.mark.unit
|
|
def test_hf_model_info(self):
|
|
filt = ModelPT.get_hf_model_filter()
|
|
|
|
# check no override results
|
|
model_infos = ModelPT.search_huggingface_models(model_filter=None)
|
|
model_infos = [next(model_infos) for _ in range(5)]
|
|
assert len(model_infos) > 0
|
|
|
|
# check with default override results (should match above)
|
|
default_model_infos = ModelPT.search_huggingface_models(model_filter=filt)
|
|
default_model_infos = [next(default_model_infos) for _ in range(5)]
|
|
assert len(model_infos) == len(default_model_infos)
|
|
|
|
@pytest.mark.pleasefixme()
|
|
@pytest.mark.with_downloads()
|
|
@pytest.mark.unit
|
|
def test_hf_model_info_with_card_data(self):
|
|
filt = ModelPT.get_hf_model_filter()
|
|
|
|
# check no override results
|
|
model_infos = ModelPT.search_huggingface_models(model_filter=filt)
|
|
model_infos = [next(model_infos) for _ in range(5)]
|
|
assert len(model_infos) > 0
|
|
|
|
# check overriden defaults
|
|
filt['cardData'] = True
|
|
model_infos = ModelPT.search_huggingface_models(model_filter=filt)
|
|
|
|
for info in model_infos:
|
|
if hasattr(info, 'cardData'):
|
|
assert info.cardData is not None
|
|
break
|
|
|
|
@pytest.mark.with_downloads()
|
|
@pytest.mark.unit
|
|
def test_hf_model_info_with_limited_results(self):
|
|
filt = ModelPT.get_hf_model_filter()
|
|
|
|
# check no override results
|
|
model_infos = ModelPT.search_huggingface_models(model_filter=filt)
|
|
model_infos = [next(model_infos) for _ in range(6)]
|
|
assert len(model_infos) > 0
|
|
|
|
# check overriden defaults
|
|
filt['limit'] = 5
|
|
new_model_infos = ModelPT.search_huggingface_models(model_filter=filt)
|
|
new_model_infos = list(new_model_infos)
|
|
assert len(new_model_infos) <= 5
|
|
assert len(new_model_infos) < len(model_infos)
|
|
|
|
@pytest.mark.unit
|
|
@pytest.mark.parametrize(
|
|
"filter_method,tar_input",
|
|
[
|
|
(save_restore_connector.SaveRestoreConnector._filtered_recursive_walk, False),
|
|
(save_restore_connector.SaveRestoreConnector._filtered_tar_info, True),
|
|
],
|
|
)
|
|
def test_filtering_methods(self, filter_method: Callable, tar_input: bool):
|
|
def touch(path):
|
|
with open(path, 'a'):
|
|
os.utime(path, None)
|
|
|
|
def filter_even_children(path: str):
|
|
if not path[-1].isdigit():
|
|
return False
|
|
return int(path[-1]) % 2 == 0
|
|
|
|
cwd = os.getcwd()
|
|
# Since we os.chdir to a temp directory. Tests can fail if we don't jump back to the CWD.
|
|
# This try:finally block ensures we don't get left in an ephemeral directory
|
|
try:
|
|
with tempfile.TemporaryDirectory() as output_dir, tempfile.TemporaryDirectory() as nemo_base_dir:
|
|
os.chdir(output_dir)
|
|
os.makedirs('grand/parent', exist_ok=True)
|
|
os.makedirs('grand/aunt', exist_ok=True)
|
|
for i in range(3):
|
|
touch(f'grand/parent/child_{i}')
|
|
touch(f'grand/aunt/child_{i}')
|
|
|
|
if tar_input:
|
|
path = f'{nemo_base_dir}/model.nemo'
|
|
save_restore_connector.SaveRestoreConnector._make_nemo_file_from_folder(
|
|
filename=path, source_dir=output_dir
|
|
)
|
|
else:
|
|
path = '.'
|
|
|
|
expected_paths = set(
|
|
(
|
|
'./grand/aunt/child_0',
|
|
'./grand/aunt/child_2',
|
|
'./grand/parent/child_0',
|
|
'./grand/parent/child_2',
|
|
)
|
|
)
|
|
|
|
observed_paths = filter_method(path, filter_fn=filter_even_children)
|
|
if tar_input:
|
|
observed_paths = set((p.name for p in observed_paths))
|
|
else:
|
|
observed_paths = set(observed_paths)
|
|
|
|
assert expected_paths == observed_paths
|
|
finally:
|
|
os.chdir(cwd)
|
|
|
|
|
|
class _NonTensorPayload:
|
|
"""A non-tensor, picklable object that torch's weights_only unpickler rejects."""
|
|
|
|
def __init__(self, value):
|
|
self.value = value
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_load_state_dict_honors_force_no_weights_only(tmp_path, monkeypatch):
|
|
"""TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD must let a checkpoint with non-tensor objects load."""
|
|
ckpt = tmp_path / "weights.ckpt"
|
|
torch.save({"w": torch.zeros(2), "meta": _NonTensorPayload(11)}, ckpt)
|
|
monkeypatch.setenv("TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD", "1")
|
|
|
|
state = save_restore_connector.SaveRestoreConnector._load_state_dict_from_disk(str(ckpt))
|
|
|
|
assert torch.equal(state["w"], torch.zeros(2))
|
|
assert state["meta"].value == 11
|