tests/core/test_save_restore.py

# 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_EncDecCTCModel(self):
        # TODO: Switch to using named configs because here we don't really care about weights
        qn = EncDecCTCModel.from_pretrained(model_name="QuartzNet15x5Base-En")
        self.__test_restore_elsewhere(model=qn, attr_for_eq_check=set(["decoder._feat_in", "decoder._num_classes"]))

    @pytest.mark.with_downloads()
    @pytest.mark.unit
    def test_EncDecCTCModelBPE(self):
        # TODO: Switch to using named configs because here we don't really care about weights
        cn = EncDecCTCModelBPE.from_pretrained(model_name="stt_en_citrinet_256")
        self.__test_restore_elsewhere(model=cn, attr_for_eq_check=set(["decoder._feat_in", "decoder._num_classes"]))

    @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.with_downloads()
    @pytest.mark.unit
    def test_EncDecCTCModelBPE_v3(self):
        # TODO: Switch to using named configs because here we don't really care about weights
        cn = EncDecCTCModelBPE.from_pretrained(model_name="stt_en_squeezeformer_ctc_xsmall_ls")
        self.__test_restore_elsewhere(model=cn, attr_for_eq_check=set(["decoder._feat_in", "decoder._num_classes"]))

    @pytest.mark.with_downloads()
    @pytest.mark.unit
    def test_EncDecCTCModelBPE_HF(self):
        # TODO: Switch to using named configs because here we don't really care about weights
        # Specifically use ModelPT instead of EncDecCTCModelBPE in order to test target class resolution.
        cn = ModelPT.from_pretrained(model_name="nvidia/stt_en_citrinet_256_ls")
        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 filt['library'] == 'nemo'

    @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)