chore: import upstream snapshot with attribution

tests/utils/test_enum.py

@@ -0,0 +1,41 @@
+# Copyright (c) 2023, 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.
+
+
+from nemo.utils.enum import PrettyStrEnum
+
+
+class ASRModelType(PrettyStrEnum):
+    CTC = "ctc"
+    RNNT = "rnnt"
+
+
+class TestPrettyStrEnum:
+    def test_incorrect_value(self):
+        """Test pretty error message for invalid value"""
+        try:
+            ASRModelType("incorrect")
+        except ValueError as e:
+            assert str(e) == "incorrect is not a valid ASRModelType. Possible choices: ctc, rnnt"
+
+    def test_correct_value(self):
+        """Test that correct value is accepted"""
+        assert ASRModelType("ctc") == ASRModelType.CTC
+
+    def test_str(self):
+        """
+        Test that str() returns the source value,
+        useful for serialization/deserialization and user-friendly logging
+        """
+        assert str(ASRModelType("ctc")) == "ctc"

tests/utils/test_env_var_parsing.py

@@ -0,0 +1,332 @@
+# Copyright (c) 2022, 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 json
+import os
+from datetime import date, datetime
+from decimal import Decimal
+from unittest import mock
+
+import pytest
+
+from nemo.utils.env_var_parsing import (
+    CoercionError,
+    RequiredSettingMissingError,
+    get_envbool,
+    get_envdate,
+    get_envdatetime,
+    get_envdecimal,
+    get_envdict,
+    get_envfloat,
+    get_envint,
+    get_envlist,
+)
+
+
+class TestEnvironmentVariableParsing:
+
+    def test_get_envint_returns_int_value(self):
+        """Test that get_envint returns the integer value from environment variable."""
+        with mock.patch.dict(os.environ, {'TEST_INT': '42'}):
+            assert get_envint('TEST_INT') == 42
+
+    def test_get_envint_with_default(self):
+        """Test that get_envint returns the default value when env var is missing."""
+        assert get_envint('NONEXISTENT_INT', 123) == 123
+
+    def test_get_envint_required_missing(self):
+        """Test that get_envint raises an exception when required env var is missing."""
+        with pytest.raises(RequiredSettingMissingError):
+            get_envint('NONEXISTENT_INT')
+
+    def test_get_envint_coercion_error(self):
+        """Test that get_envint raises a CoercionError for non-integer values."""
+        with mock.patch.dict(os.environ, {'TEST_INT': 'not-an-int'}):
+            with pytest.raises(CoercionError):
+                get_envint('TEST_INT')
+
+    def test_get_envint_negative_value(self):
+        """Test that get_envint correctly handles negative integers."""
+        with mock.patch.dict(os.environ, {'TEST_INT': '-42'}):
+            assert get_envint('TEST_INT') == -42
+
+    def test_get_envfloat_returns_float_value(self):
+        """Test that get_envfloat returns the float value from environment variable."""
+        with mock.patch.dict(os.environ, {'TEST_FLOAT': '3.14'}):
+            assert get_envfloat('TEST_FLOAT') == 3.14
+
+    def test_get_envfloat_with_integer_string(self):
+        """Test that get_envfloat correctly converts integer strings to floats."""
+        with mock.patch.dict(os.environ, {'TEST_FLOAT': '42'}):
+            value = get_envfloat('TEST_FLOAT')
+            assert value == 42.0
+            assert isinstance(value, float)
+
+    def test_get_envfloat_with_default(self):
+        """Test that get_envfloat returns the default value when env var is missing."""
+        assert get_envfloat('NONEXISTENT_FLOAT', 3.14) == 3.14
+
+    def test_get_envfloat_required_missing(self):
+        """Test that get_envfloat raises an exception when required env var is missing."""
+        with pytest.raises(RequiredSettingMissingError):
+            get_envfloat('NONEXISTENT_FLOAT')
+
+    def test_get_envfloat_coercion_error(self):
+        """Test that get_envfloat raises a CoercionError for non-float values."""
+        with mock.patch.dict(os.environ, {'TEST_FLOAT': 'not-a-float'}):
+            with pytest.raises(CoercionError):
+                get_envfloat('TEST_FLOAT')
+
+    def test_get_envfloat_scientific_notation(self):
+        """Test that get_envfloat correctly handles scientific notation."""
+        with mock.patch.dict(os.environ, {'TEST_FLOAT': '1.23e-4'}):
+            assert get_envfloat('TEST_FLOAT') == 1.23e-4
+
+    def test_get_envfloat_negative_value(self):
+        """Test that get_envfloat correctly handles negative values."""
+        with mock.patch.dict(os.environ, {'TEST_FLOAT': '-3.14'}):
+            assert get_envfloat('TEST_FLOAT') == -3.14
+
+    # Tests for get_envbool
+    def test_get_envbool_true_values(self):
+        """Test that get_envbool returns True for various truthy values."""
+        true_values = ['true', 'True', 'TRUE', '1', 'yes', 'Yes', 'YES', 'y', 'Y', 't', 'T']
+        for val in true_values:
+            with mock.patch.dict(os.environ, {'TEST_BOOL': val}):
+                assert get_envbool('TEST_BOOL') is True
+
+    def test_get_envbool_false_values(self):
+        """Test that get_envbool returns False for various falsy values."""
+        false_values = ['false', 'False', 'FALSE', '0', 'no', 'No', 'NO', 'n', 'N', 'f', 'F', 'none', 'None', 'NONE']
+        for val in false_values:
+            with mock.patch.dict(os.environ, {'TEST_BOOL': val}):
+                assert get_envbool('TEST_BOOL') is False
+
+    def test_get_envbool_with_default(self):
+        """Test that get_envbool returns the default value when env var is missing."""
+        assert get_envbool('NONEXISTENT_BOOL', True) is True
+        assert get_envbool('NONEXISTENT_BOOL', False) is False
+
+    def test_get_envbool_required_missing(self):
+        """Test that get_envbool raises an exception when required env var is missing."""
+        with pytest.raises(RequiredSettingMissingError):
+            get_envbool('NONEXISTENT_BOOL')
+
+    def test_get_envbool_non_boolean_value(self):
+        """Test that get_envbool interprets non-standard values as True."""
+        with mock.patch.dict(os.environ, {'TEST_BOOL': 'something-else'}):
+            assert get_envbool('TEST_BOOL') is True
+
+    # Tests for get_envdecimal
+    def test_get_envdecimal_returns_decimal_value(self):
+        """Test that get_envdecimal returns the Decimal value from environment variable."""
+        with mock.patch.dict(os.environ, {'TEST_DECIMAL': '3.14'}):
+            value = get_envdecimal('TEST_DECIMAL')
+            assert value == Decimal('3.14')
+            assert isinstance(value, Decimal)
+
+    def test_get_envdecimal_with_integer_string(self):
+        """Test that get_envdecimal correctly converts integer strings to Decimals."""
+        with mock.patch.dict(os.environ, {'TEST_DECIMAL': '42'}):
+            value = get_envdecimal('TEST_DECIMAL')
+            assert value == Decimal('42')
+            assert isinstance(value, Decimal)
+
+    def test_get_envdecimal_with_default(self):
+        """Test that get_envdecimal returns the default value when env var is missing."""
+        default_value = Decimal('3.14')
+        assert get_envdecimal('NONEXISTENT_DECIMAL', default_value) == default_value
+
+    def test_get_envdecimal_required_missing(self):
+        """Test that get_envdecimal raises an exception when required env var is missing."""
+        with pytest.raises(RequiredSettingMissingError):
+            get_envdecimal('NONEXISTENT_DECIMAL')
+
+    def test_get_envdecimal_coercion_error(self):
+        """Test that get_envdecimal raises a CoercionError for non-decimal values."""
+        with mock.patch.dict(os.environ, {'TEST_DECIMAL': 'not-a-decimal'}):
+            with pytest.raises(CoercionError):
+                get_envdecimal('TEST_DECIMAL')
+
+    def test_get_envdecimal_negative_value(self):
+        """Test that get_envdecimal correctly handles negative values."""
+        with mock.patch.dict(os.environ, {'TEST_DECIMAL': '-3.14'}):
+            assert get_envdecimal('TEST_DECIMAL') == Decimal('-3.14')
+
+    def test_get_envdecimal_high_precision(self):
+        """Test that get_envdecimal preserves high precision values."""
+        with mock.patch.dict(os.environ, {'TEST_DECIMAL': '3.1415926535897932384626433832795028841971'}):
+            value = get_envdecimal('TEST_DECIMAL')
+            assert value == Decimal('3.1415926535897932384626433832795028841971')
+
+    # Tests for get_envdate
+    def test_get_envdate_returns_date_value(self):
+        """Test that get_envdate returns the date value from environment variable."""
+        with mock.patch.dict(os.environ, {'TEST_DATE': '2023-05-15'}):
+            value = get_envdate('TEST_DATE')
+            assert value == date(2023, 5, 15)
+            assert isinstance(value, date)
+
+    def test_get_envdate_with_different_formats(self):
+        """Test that get_envdate handles different date formats."""
+        date_formats = {
+            '2023-05-15': date(2023, 5, 15),
+            '15-05-2023': date(2023, 5, 15),
+            '05/15/2023': date(2023, 5, 15),
+            '15 May 2023': date(2023, 5, 15),
+            'May 15, 2023': date(2023, 5, 15),
+        }
+
+        for date_str, expected_date in date_formats.items():
+            with mock.patch.dict(os.environ, {'TEST_DATE': date_str}):
+                assert get_envdate('TEST_DATE') == expected_date
+
+    def test_get_envdate_with_default(self):
+        """Test that get_envdate returns the default value when env var is missing."""
+        default_value = date(2023, 5, 15)
+        assert get_envdate('NONEXISTENT_DATE', default_value) == default_value
+
+    def test_get_envdate_required_missing(self):
+        """Test that get_envdate raises an exception when required env var is missing."""
+        with pytest.raises(RequiredSettingMissingError):
+            get_envdate('NONEXISTENT_DATE')
+
+    def test_get_envdate_coercion_error(self):
+        """Test that get_envdate raises a CoercionError for invalid date values."""
+        with mock.patch.dict(os.environ, {'TEST_DATE': 'not-a-date'}):
+            with pytest.raises(CoercionError):
+                get_envdate('TEST_DATE')
+
+    # Tests for get_envdatetime
+    def test_get_envdatetime_returns_datetime_value(self):
+        """Test that get_envdatetime returns the datetime value from environment variable."""
+        with mock.patch.dict(os.environ, {'TEST_DATETIME': '2023-05-15T14:30:45'}):
+            value = get_envdatetime('TEST_DATETIME')
+            assert value == datetime(2023, 5, 15, 14, 30, 45)
+            assert isinstance(value, datetime)
+
+    def test_get_envdatetime_with_different_formats(self):
+        """Test that get_envdatetime handles different datetime formats."""
+        datetime_formats = {
+            '2023-05-15T14:30:45': datetime(2023, 5, 15, 14, 30, 45),
+            '2023-05-15 14:30:45': datetime(2023, 5, 15, 14, 30, 45),
+            '15-05-2023 14:30:45': datetime(2023, 5, 15, 14, 30, 45),
+            '05/15/2023 14:30:45': datetime(2023, 5, 15, 14, 30, 45),
+            '15 May 2023 14:30:45': datetime(2023, 5, 15, 14, 30, 45),
+        }
+
+        for datetime_str, expected_datetime in datetime_formats.items():
+            with mock.patch.dict(os.environ, {'TEST_DATETIME': datetime_str}):
+                assert get_envdatetime('TEST_DATETIME') == expected_datetime
+
+    def test_get_envdatetime_with_default(self):
+        """Test that get_envdatetime returns the default value when env var is missing."""
+        default_value = datetime(2023, 5, 15, 14, 30, 45)
+        assert get_envdatetime('NONEXISTENT_DATETIME', default_value) == default_value
+
+    def test_get_envdatetime_required_missing(self):
+        """Test that get_envdatetime raises an exception when required env var is missing."""
+        with pytest.raises(RequiredSettingMissingError):
+            get_envdatetime('NONEXISTENT_DATETIME')
+
+    def test_get_envdatetime_coercion_error(self):
+        """Test that get_envdatetime raises a CoercionError for invalid datetime values."""
+        with mock.patch.dict(os.environ, {'TEST_DATETIME': 'not-a-datetime'}):
+            with pytest.raises(CoercionError):
+                get_envdatetime('TEST_DATETIME')
+
+    def test_get_envdatetime_with_timezone(self):
+        """Test that get_envdatetime handles timezone information."""
+        with mock.patch.dict(os.environ, {'TEST_DATETIME': '2023-05-15T14:30:45+0200'}):
+            dt = get_envdatetime('TEST_DATETIME')
+            assert dt.year == 2023
+            assert dt.month == 5
+            assert dt.day == 15
+            assert dt.hour == 14
+            assert dt.minute == 30
+            assert dt.second == 45
+
+    # Tests for get_envlist
+    def test_get_envlist_returns_list_value(self):
+        """Test that get_envlist returns the list value from environment variable."""
+        with mock.patch.dict(os.environ, {'TEST_LIST': 'item1 item2 item3'}):
+            value = get_envlist('TEST_LIST')
+            assert value == ['item1', 'item2', 'item3']
+            assert isinstance(value, list)
+
+    def test_get_envlist_with_custom_separator(self):
+        """Test that get_envlist handles custom separators."""
+        with mock.patch.dict(os.environ, {'TEST_LIST': 'item1,item2,item3'}):
+            value = get_envlist('TEST_LIST', separator=',')
+            assert value == ['item1', 'item2', 'item3']
+
+    def test_get_envlist_with_default(self):
+        """Test that get_envlist returns the default value when env var is missing."""
+        default_value = ['default1', 'default2']
+        assert get_envlist('NONEXISTENT_LIST', default_value) == default_value
+
+    def test_get_envlist_required_missing(self):
+        """Test that get_envlist raises an exception when required env var is missing."""
+        with pytest.raises(RequiredSettingMissingError):
+            get_envlist('NONEXISTENT_LIST')
+
+    def test_get_envlist_empty_string(self):
+        """Test that get_envlist handles empty strings."""
+        with mock.patch.dict(os.environ, {'TEST_LIST': ''}):
+            value = get_envlist('TEST_LIST')
+            assert value == ['']
+
+    def test_get_envlist_multiple_words(self):
+        """Test that get_envlist correctly splits words with spaces."""
+        with mock.patch.dict(os.environ, {'TEST_LIST': 'word1 "phrase with spaces" word3'}):
+            value = get_envlist('TEST_LIST')
+            assert value == ['word1', '"phrase', 'with', 'spaces"', 'word3']
+
+    # Tests for get_envdict
+    def test_get_envdict_returns_dict_value(self):
+        """Test that get_envdict returns the dict value from environment variable."""
+        test_dict = {'key1': 'value1', 'key2': 42, 'key3': True}
+        with mock.patch.dict(os.environ, {'TEST_DICT': json.dumps(test_dict)}):
+            value = get_envdict('TEST_DICT')
+            assert value == test_dict
+            assert isinstance(value, dict)
+
+    def test_get_envdict_with_default(self):
+        """Test that get_envdict returns the default value when env var is missing."""
+        default_value = {'default_key': 'default_value'}
+        assert get_envdict('NONEXISTENT_DICT', default_value) == default_value
+
+    def test_get_envdict_required_missing(self):
+        """Test that get_envdict raises an exception when required env var is missing."""
+        with pytest.raises(RequiredSettingMissingError):
+            get_envdict('NONEXISTENT_DICT')
+
+    def test_get_envdict_coercion_error(self):
+        """Test that get_envdict raises a CoercionError for invalid JSON."""
+        with mock.patch.dict(os.environ, {'TEST_DICT': 'not-valid-json'}):
+            with pytest.raises(CoercionError):
+                get_envdict('TEST_DICT')
+
+    def test_get_envdict_complex_dict(self):
+        """Test that get_envdict handles complex nested dictionaries."""
+        complex_dict = {
+            'string': 'value',
+            'number': 42,
+            'boolean': True,
+            'list': [1, 2, 3],
+            'nested': {'key1': 'value1', 'key2': [4, 5, 6]},
+        }
+        with mock.patch.dict(os.environ, {'TEST_DICT': json.dumps(complex_dict)}):
+            value = get_envdict('TEST_DICT')
+            assert value == complex_dict

tests/utils/test_file_utils.py

@@ -0,0 +1,51 @@
+# Copyright (c) 2025, 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.
+
+from unittest.mock import patch
+
+import pytest
+
+from nemo.utils.file_utils import robust_copy
+
+
+class TestRobustCopy:
+    @pytest.mark.unit
+    def test_robust_copy_success(self, tmp_path):
+        """Tests that robust_copy uses shutil.copy2 and does not fall back."""
+        src_file = tmp_path / "source.txt"
+        src_file.write_text("test content")
+        dest_file = tmp_path / "dest.txt"
+
+        with (
+            patch('nemo.utils.file_utils.shutil.copy2') as mock_copy2,
+            patch('nemo.utils.file_utils.shutil.copy') as mock_copy,
+        ):
+            robust_copy(src_file, dest_file)
+            mock_copy2.assert_called_once_with(src_file, dest_file)
+            mock_copy.assert_not_called()
+
+    @pytest.mark.unit
+    def test_robust_copy_fallback(self, tmp_path):
+        """Tests that robust_copy falls back to shutil.copy if shutil.copy2 fails."""
+        src_file = tmp_path / "source.txt"
+        src_file.write_text("test content")
+        dest_file = tmp_path / "dest.txt"
+
+        with (
+            patch('nemo.utils.file_utils.shutil.copy2', side_effect=PermissionError("copy2 fails")) as mock_copy2,
+            patch('nemo.utils.file_utils.shutil.copy') as mock_copy,
+        ):
+            robust_copy(src_file, dest_file)
+            mock_copy2.assert_called_once_with(src_file, dest_file)
+            mock_copy.assert_called_once_with(src_file, dest_file)

tests/utils/test_flash_precision.py

@@ -0,0 +1,135 @@
+# Copyright (c) 2026, 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 torch
+import torch.nn as nn
+from lightning.pytorch.plugins import HalfPrecision
+from omegaconf import DictConfig
+
+from nemo.utils.trainer_utils import FlashPrecision, HalfPrecisionForAudio, resolve_trainer_cfg
+
+
+class TestForwardContext:
+    def test_default_dtype_remains_fp32_during_forward(self):
+        plugin = FlashPrecision("bf16-flash")
+        with plugin.forward_context():
+            assert torch.get_default_dtype() == torch.float32
+
+    def test_implicit_tensor_creation_is_fp32(self):
+        plugin = FlashPrecision("bf16-flash")
+        with plugin.forward_context():
+            assert torch.zeros(10).dtype == torch.float32
+            assert torch.ones(10).dtype == torch.float32
+            assert torch.empty(10).dtype == torch.float32
+
+
+class TestConvertModule:
+    def test_convert_module_casts_plain_fp32_module(self):
+        model = nn.Sequential(nn.Linear(10, 10), nn.Linear(10, 5))
+        plugin = FlashPrecision("bf16-flash")
+        plugin.convert_module(model)
+        assert model[0].weight.dtype == torch.bfloat16
+        assert model[0].bias.dtype == torch.bfloat16
+        assert model[1].weight.dtype == torch.bfloat16
+
+    def test_convert_module_skips_models_with_existing_dtype_policy(self):
+        model = nn.Sequential(nn.Linear(10, 10), nn.Linear(10, 5))
+        model[0].to(dtype=torch.bfloat16)
+        plugin = FlashPrecision("bf16-flash")
+        plugin.convert_module(model)
+        assert model[0].weight.dtype == torch.bfloat16
+        assert model[1].weight.dtype == torch.float32
+
+
+class TestConvertInput:
+    def test_preserves_audio_tensors(self):
+        plugin = FlashPrecision("bf16-flash")
+        batch = {"audio": torch.randn(1, 16000), "tokens": torch.randn(1, 10)}
+        converted = plugin.convert_input(batch)
+        assert converted["audio"].dtype == torch.float32
+        assert converted["tokens"].dtype == torch.bfloat16
+
+    def test_handles_nested_dicts(self):
+        plugin = FlashPrecision("bf16-flash")
+        batch = {
+            "inputs": {"audio_signal": torch.randn(1, 16000), "text_ids": torch.randn(1, 10)},
+            "labels": torch.randn(1, 5),
+        }
+        converted = plugin.convert_input(batch)
+        assert converted["inputs"]["audio_signal"].dtype == torch.float32
+        assert converted["inputs"]["text_ids"].dtype == torch.bfloat16
+        assert converted["labels"].dtype == torch.bfloat16
+
+    def test_non_dict_input_converted(self):
+        plugin = FlashPrecision("bf16-flash")
+        t = torch.randn(4, 8)
+        converted = plugin.convert_input(t)
+        assert converted.dtype == torch.bfloat16
+
+    def test_non_float_tensors_unchanged(self):
+        plugin = FlashPrecision("bf16-flash")
+        batch = {"ids": torch.tensor([1, 2, 3], dtype=torch.long), "values": torch.randn(3)}
+        converted = plugin.convert_input(batch)
+        assert converted["ids"].dtype == torch.long
+        assert converted["values"].dtype == torch.bfloat16
+
+
+class TestHalfPrecisionRegression:
+    def test_half_precision_does_change_default_dtype(self):
+        plugin = HalfPrecision("bf16-true")
+        with plugin.forward_context():
+            assert torch.get_default_dtype() == torch.bfloat16
+            assert torch.zeros(10).dtype == torch.bfloat16
+
+        assert torch.get_default_dtype() == torch.float32
+
+
+class TestResolveTrainerCfg:
+    def test_bf16_flash_creates_flash_precision(self):
+        cfg = DictConfig({"precision": "bf16-flash"})
+        resolved = resolve_trainer_cfg(cfg)
+        plugins = resolved["plugins"]
+        assert "precision" not in resolved
+        assert len(plugins) == 1
+        assert isinstance(plugins[0], FlashPrecision)
+        assert plugins[0].precision == "bf16-flash"
+        assert plugins[0]._desired_input_dtype == torch.bfloat16
+
+    def test_fp16_flash_creates_flash_precision(self):
+        cfg = DictConfig({"precision": "fp16-flash"})
+        resolved = resolve_trainer_cfg(cfg)
+        plugins = resolved["plugins"]
+        assert isinstance(plugins[0], FlashPrecision)
+        assert plugins[0].precision == "fp16-flash"
+        assert plugins[0]._desired_input_dtype == torch.float16
+
+    def test_legacy_automodel_aliases_resolve_to_flash_precision(self):
+        cfg = DictConfig({"precision": "bf16-automodel"})
+        resolved = resolve_trainer_cfg(cfg)
+        plugins = resolved["plugins"]
+        assert isinstance(plugins[0], FlashPrecision)
+        assert plugins[0].precision == "bf16-flash"
+
+        cfg = DictConfig({"precision": "fp16-automodel"})
+        resolved = resolve_trainer_cfg(cfg)
+        plugins = resolved["plugins"]
+        assert isinstance(plugins[0], FlashPrecision)
+        assert plugins[0].precision == "fp16-flash"
+
+    def test_bf16_true_still_creates_half_precision_for_audio(self):
+        cfg = DictConfig({"precision": "bf16-true"})
+        resolved = resolve_trainer_cfg(cfg)
+        plugins = resolved["plugins"]
+        assert len(plugins) == 1
+        assert isinstance(plugins[0], HalfPrecisionForAudio)

tests/utils/test_get_rank.py

@@ -0,0 +1,170 @@
+# Copyright (c) 2023, 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 os
+from unittest import mock
+
+import pytest
+import torch
+
+from nemo.utils.get_rank import get_last_rank, get_rank, is_global_rank_zero
+
+
+class TestIsGlobalRankZero:
+    """Test the is_global_rank_zero function with various environment variable settings."""
+
+    @pytest.fixture(autouse=True)
+    def setup_method(self):
+        """Clear all relevant environment variables before each test."""
+        for var in ["RANK", "SLURM_PROCID", "OMPI_COMM_WORLD_RANK", "NODE_RANK", "GROUP_RANK", "LOCAL_RANK"]:
+            if var in os.environ:
+                del os.environ[var]
+
+    def test_default_behavior(self):
+        """Test the default behavior when no environment variables are set."""
+        assert is_global_rank_zero() is True
+
+    def test_with_pytorch_rank_0(self):
+        """Test when RANK=0 (pytorch environment)."""
+        os.environ["RANK"] = "0"
+        assert is_global_rank_zero() is True
+
+    def test_with_pytorch_rank_nonzero(self):
+        """Test when RANK is not 0 (pytorch environment)."""
+        os.environ["RANK"] = "1"
+        assert is_global_rank_zero() is False
+
+    def test_with_slurm_rank_0(self):
+        """Test when SLURM_PROCID=0 (SLURM environment)."""
+        os.environ["SLURM_PROCID"] = "0"
+        assert is_global_rank_zero() is True
+
+    def test_with_slurm_rank_nonzero(self):
+        """Test when SLURM_PROCID is not 0 (SLURM environment)."""
+        os.environ["SLURM_PROCID"] = "1"
+        assert is_global_rank_zero() is False
+
+    def test_with_mpi_rank_0(self):
+        """Test when OMPI_COMM_WORLD_RANK=0 (MPI environment)."""
+        os.environ["OMPI_COMM_WORLD_RANK"] = "0"
+        assert is_global_rank_zero() is True
+
+    def test_with_mpi_rank_nonzero(self):
+        """Test when OMPI_COMM_WORLD_RANK is not 0 (MPI environment)."""
+        os.environ["OMPI_COMM_WORLD_RANK"] = "1"
+        assert is_global_rank_zero() is False
+
+    def test_with_node_rank_0_local_rank_0(self):
+        """Test when NODE_RANK=0 and LOCAL_RANK=0."""
+        os.environ["NODE_RANK"] = "0"
+        os.environ["LOCAL_RANK"] = "0"
+        assert is_global_rank_zero() is True
+
+    def test_with_node_rank_0_local_rank_nonzero(self):
+        """Test when NODE_RANK=0 but LOCAL_RANK is not 0."""
+        os.environ["NODE_RANK"] = "0"
+        os.environ["LOCAL_RANK"] = "1"
+        assert is_global_rank_zero() is False
+
+    def test_with_node_rank_nonzero(self):
+        """Test when NODE_RANK is not 0."""
+        os.environ["NODE_RANK"] = "1"
+        os.environ["LOCAL_RANK"] = "0"
+        assert is_global_rank_zero() is False
+
+    def test_with_group_rank_fallback(self):
+        """Test using GROUP_RANK as fallback for NODE_RANK."""
+        os.environ["GROUP_RANK"] = "0"
+        os.environ["LOCAL_RANK"] = "0"
+        assert is_global_rank_zero() is True
+
+        os.environ["GROUP_RANK"] = "1"
+        assert is_global_rank_zero() is False
+
+    def test_env_var_precedence(self):
+        """Test that environment variables are checked in the expected order of precedence."""
+        # RANK has highest precedence
+        os.environ["RANK"] = "0"
+        os.environ["SLURM_PROCID"] = "1"
+        os.environ["OMPI_COMM_WORLD_RANK"] = "1"
+        assert is_global_rank_zero() is True
+
+        os.environ["RANK"] = "1"
+        os.environ["SLURM_PROCID"] = "0"
+        assert is_global_rank_zero() is False
+
+        # Without RANK, SLURM_PROCID has next precedence
+        del os.environ["RANK"]
+        assert is_global_rank_zero() is True
+
+        os.environ["SLURM_PROCID"] = "1"
+        os.environ["OMPI_COMM_WORLD_RANK"] = "0"
+        assert is_global_rank_zero() is False
+
+        # Without RANK and SLURM_PROCID, OMPI_COMM_WORLD_RANK has next precedence
+        del os.environ["SLURM_PROCID"]
+        assert is_global_rank_zero() is True
+
+
+class TestGetRank:
+    """Test the get_rank function."""
+
+    @pytest.fixture(autouse=True)
+    def setup_method(self):
+        """Clear all relevant environment variables before each test."""
+        for var in ["RANK", "SLURM_PROCID", "OMPI_COMM_WORLD_RANK", "NODE_RANK", "GROUP_RANK", "LOCAL_RANK"]:
+            if var in os.environ:
+                del os.environ[var]
+
+    @mock.patch("torch.distributed.is_initialized", return_value=False)
+    def test_not_distributed(self, mock_is_initialized):
+        """Test when not in a distributed environment."""
+        assert get_rank() == 0
+
+    @mock.patch("torch.distributed.is_initialized", return_value=True)
+    @mock.patch("torch.distributed.get_rank", return_value=2)
+    def test_distributed_not_global_rank_zero(self, mock_dist_get_rank, mock_is_initialized):
+        """Test when in a distributed environment and not global rank zero."""
+        # Make sure is_global_rank_zero() returns False
+        os.environ["RANK"] = "1"
+        assert get_rank() == 2
+        mock_dist_get_rank.assert_called_once()
+
+    @mock.patch("torch.distributed.is_initialized", return_value=True)
+    @mock.patch("torch.distributed.get_rank", return_value=0)
+    def test_distributed_global_rank_zero(self, mock_dist_get_rank, mock_is_initialized):
+        """Test when in a distributed environment and is global rank zero."""
+        # Global rank is zero
+        os.environ["RANK"] = "0"
+        assert get_rank() == 0
+        # Should not call torch.distributed.get_rank() when is_global_rank_zero() is True
+        mock_dist_get_rank.assert_not_called()
+
+
+class TestGetLastRank:
+    """Test the get_last_rank function."""
+
+    @mock.patch("torch.distributed.is_initialized", return_value=False)
+    def test_not_distributed(self, mock_is_initialized):
+        """Test when not in a distributed environment."""
+        assert get_last_rank() == 0
+        mock_is_initialized.assert_called_once()
+
+    @mock.patch("torch.distributed.is_initialized", return_value=True)
+    @mock.patch("torch.distributed.get_world_size", return_value=4)
+    def test_distributed(self, mock_get_world_size, mock_is_initialized):
+        """Test when in a distributed environment."""
+        assert get_last_rank() == 3  # world_size - 1
+        mock_is_initialized.assert_called_once()
+        mock_get_world_size.assert_called_once()

tests/utils/test_import_utils.py

@@ -0,0 +1,267 @@
+# Copyright (c) 2025, 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 sys
+import types
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from nemo.utils.import_utils import UnavailableError, UnavailableMeta, is_unavailable, safe_import, safe_import_from
+
+
+class TestUnavailableMeta:
+    """Test suite for the UnavailableMeta metaclass."""
+
+    def test_metaclass_creation(self):
+        """Test that UnavailableMeta creates a class with the expected properties."""
+        # Create a class using UnavailableMeta
+        TestClass = UnavailableMeta("TestClass", (), {})
+
+        # The class name should be prefixed with "MISSING"
+        assert TestClass.__name__ == "MISSINGTestClass"
+
+        # The default error message should be set
+        assert TestClass._msg == "TestClass could not be imported"
+
+    def test_custom_error_message(self):
+        """Test that a custom error message can be provided."""
+        custom_msg = "Custom error message"
+        TestClass = UnavailableMeta("TestClass", (), {"_msg": custom_msg})
+
+        assert TestClass._msg == custom_msg
+
+        # Verify the message is used in exceptions
+        with pytest.raises(UnavailableError, match=custom_msg):
+            TestClass()
+
+    def test_call_raises_error(self):
+        """Test that attempting to instantiate the class raises UnavailableError."""
+        TestClass = UnavailableMeta("TestClass", (), {})
+
+        with pytest.raises(UnavailableError):
+            TestClass()
+
+        with pytest.raises(UnavailableError):
+            TestClass(1, 2, 3, key="value")
+
+    def test_attribute_access_raises_error(self):
+        """Test that accessing attributes raises UnavailableError."""
+        TestClass = UnavailableMeta("TestClass", (), {})
+
+        with pytest.raises(UnavailableError):
+            TestClass.some_attribute
+
+    def test_arithmetic_operations_raise_error(self):
+        """Test that arithmetic operations raise UnavailableError."""
+        TestClass = UnavailableMeta("TestClass", (), {})
+
+        operations = [
+            lambda c: c + 1,
+            lambda c: 1 + c,  # __radd__
+            lambda c: c - 1,
+            lambda c: 1 - c,  # __rsub__
+            lambda c: c * 2,
+            lambda c: 2 * c,  # __rmul__
+            lambda c: c / 2,
+            lambda c: 2 / c,  # __rtruediv__
+            lambda c: c // 2,
+            lambda c: 2 // c,  # __rfloordiv__
+            lambda c: c**2,
+            lambda c: 2**c,  # __rpow__
+            lambda c: -c,  # __neg__
+            lambda c: abs(c),  # __abs__
+        ]
+
+        for op in operations:
+            with pytest.raises(UnavailableError):
+                op(TestClass)
+
+    def test_comparison_operations_raise_error(self):
+        """Test that comparison operations raise UnavailableError."""
+        TestClass = UnavailableMeta("TestClass", (), {})
+        another_class = UnavailableMeta("AnotherClass", (), {})
+
+        comparisons = [
+            lambda c: c == another_class,
+            lambda c: c != another_class,
+            lambda c: c < another_class,
+            lambda c: c <= another_class,
+            lambda c: c > another_class,
+            lambda c: c >= another_class,
+        ]
+
+        for comp in comparisons:
+            with pytest.raises(UnavailableError):
+                comp(TestClass)
+
+    def test_container_operations_raise_error(self):
+        """Test that container operations raise UnavailableError."""
+        TestClass = UnavailableMeta("TestClass", (), {})
+
+        with pytest.raises(UnavailableError):
+            len(TestClass)
+
+        with pytest.raises(UnavailableError):
+            TestClass[0]
+
+        with pytest.raises(UnavailableError):
+            TestClass[0] = 1
+
+        with pytest.raises(UnavailableError):
+            del TestClass[0]
+
+        with pytest.raises(UnavailableError):
+            iter(TestClass)
+
+    def test_descriptor_operations_raise_error(self):
+        """Test that descriptor operations raise UnavailableError."""
+        TestClass = UnavailableMeta("TestClass", (), {})
+
+        class DummyClass:
+            prop = TestClass
+
+        dummy = DummyClass()
+
+        with pytest.raises(UnavailableError):
+            TestClass.__get__(None, None)
+
+        with pytest.raises(UnavailableError):
+            TestClass.__delete__(None)
+
+
+class TestSafeImport:
+    def test_successful_import(self):
+        """Test safe_import with a module that exists."""
+        module, success = safe_import("os")
+        assert success is True
+        assert isinstance(module, types.ModuleType)
+        assert module.__name__ == "os"
+
+    def test_failed_import(self):
+        """Test safe_import with a module that doesn't exist."""
+        module, success = safe_import("nonexistent_module")
+        assert success is False
+        assert is_unavailable(module)
+        assert type(module) is UnavailableMeta
+
+    def test_import_with_custom_message(self):
+        """Test safe_import with a custom error message."""
+        custom_msg = "Custom error message"
+        module, success = safe_import("nonexistent_module", msg=custom_msg)
+
+        assert success is False
+        assert is_unavailable(module)
+
+        # Verify the custom message is used when trying to use the module
+        with pytest.raises(UnavailableError, match=custom_msg):
+            module()
+
+    def test_import_with_alternative(self):
+        """Test safe_import with an alternative module."""
+        alt_module = object()
+        module, success = safe_import("nonexistent_module", alt=alt_module)
+
+        assert success is False
+        assert module is alt_module
+
+    def test_unavailable_module_raises_error_when_used(self):
+        """Test that using a UnavailableMeta placeholder raises UnavailableError."""
+        module, success = safe_import("nonexistent_module")
+
+        assert success is False
+
+        # Test various operations that should raise UnavailableError
+        with pytest.raises(UnavailableError):
+            module()
+
+        with pytest.raises(UnavailableError):
+            module.attribute
+
+        with pytest.raises(UnavailableError):
+            module + 1
+
+        with pytest.raises(UnavailableError):
+            module == 1
+
+
+class TestSafeImportFrom:
+    def test_successful_import_from(self):
+        """Test safe_import_from with a symbol that exists."""
+        symbol, success = safe_import_from("os", "path")
+        assert success is True
+
+        import os
+
+        assert symbol is os.path
+
+    def test_failed_import_from_nonexistent_module(self):
+        """Test safe_import_from with a module that doesn't exist."""
+        symbol, success = safe_import_from("nonexistent_module", "nonexistent_symbol")
+        assert success is False
+        assert is_unavailable(symbol)
+
+    def test_failed_import_from_nonexistent_symbol(self):
+        """Test safe_import_from with a symbol that doesn't exist in an existing module."""
+        symbol, success = safe_import_from("os", "nonexistent_symbol")
+        assert success is False
+        assert is_unavailable(symbol)
+
+    def test_import_from_with_custom_message(self):
+        """Test safe_import_from with a custom error message."""
+        custom_msg = "Custom error message for symbol"
+        symbol, success = safe_import_from("os", "nonexistent_symbol", msg=custom_msg)
+
+        assert success is False
+
+        # Verify the custom message is used when trying to use the symbol
+        with pytest.raises(UnavailableError, match=custom_msg):
+            symbol()
+
+    def test_import_from_with_alternative(self):
+        """Test safe_import_from with an alternative symbol."""
+        alt_symbol = object()
+        symbol, success = safe_import_from("os", "nonexistent_symbol", alt=alt_symbol)
+
+        assert success is False
+        assert symbol is alt_symbol
+
+    def test_fallback_module(self):
+        """Test safe_import_from with a fallback module."""
+        # First import fails, but fallback succeeds
+        with patch('importlib.import_module') as mock_import:
+            # Mock the first import to fail as AttributeError
+            def side_effect(name):
+                if name == "primary_module":
+                    raise AttributeError("Symbol not found")
+                elif name == "fallback_module":
+                    mock_module = MagicMock()
+                    mock_module.symbol = "fallback_symbol"
+                    return mock_module
+                else:
+                    raise ImportError(f"Unexpected module: {name}")
+
+            mock_import.side_effect = side_effect
+
+            symbol, success = safe_import_from("primary_module", "symbol", fallback_module="fallback_module")
+
+            assert success is True
+            assert symbol == "fallback_symbol"
+
+    def test_fallback_module_both_fail(self):
+        """Test safe_import_from when both primary and fallback modules fail."""
+        symbol, success = safe_import_from("nonexistent_primary", "symbol", fallback_module="nonexistent_fallback")
+
+        assert success is False
+        assert is_unavailable(symbol)

tests/utils/test_msc_utils.py

@@ -0,0 +1,35 @@
+# Copyright (c) 2025, 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.
+
+from pathlib import Path
+from unittest import mock
+
+import pytest
+
+from nemo.utils.msc_utils import is_multistorageclient_url
+
+
+def test_is_multistorageclient_url_with_msc_not_installed():
+    with mock.patch('nemo.utils.msc_utils.HAVE_MSC', False):
+        assert not is_multistorageclient_url('/tmp/path/to/data.bin')
+        assert not is_multistorageclient_url(Path('/tmp/path/to/data.bin'))
+
+        with pytest.raises(ValueError):
+            is_multistorageclient_url('msc://profile/path/to/data.bin')
+
+
+def test_is_multistorageclient_url_with_msc_installed():
+    with mock.patch('nemo.utils.msc_utils.HAVE_MSC', True):
+        assert is_multistorageclient_url('msc://profile/path/to/data.bin')
+        assert not is_multistorageclient_url('/tmp/path/to/data.bin')

tests/utils/test_tar_utils.py

@@ -0,0 +1,154 @@
+# Copyright (c) 2026, 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 ast
+import io
+import subprocess
+import tarfile
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+from omegaconf import OmegaConf
+
+from nemo.collections.asr.parts.mixins.mixins import ASRBPEMixin
+from nemo.utils.app_state import AppState
+from nemo.utils.model_utils import load_config, save_artifacts
+from nemo.utils.notebook_utils import download_an4
+from nemo.utils.tar_utils import TarPathTraversalError
+
+REPO_ROOT = Path(__file__).parents[2]
+RAW_TAR_EXTRACTION_ALLOWLIST = {REPO_ROOT / "nemo/utils/tar_utils.py"}
+
+
+def _is_tarfile_open_call(node):
+    return (
+        isinstance(node, ast.Call)
+        and isinstance(node.func, ast.Attribute)
+        and node.func.attr == "open"
+        and isinstance(node.func.value, ast.Name)
+        and node.func.value.id == "tarfile"
+    )
+
+
+def _add_tar_file(tar, name, data=b""):
+    info = tarfile.TarInfo(name=name)
+    info.size = len(data)
+    tar.addfile(info, io.BytesIO(data))
+
+
+def test_load_config_rejects_traversal_model_config(tmp_path):
+    nemo_path = tmp_path / "malicious.nemo"
+    with tarfile.open(nemo_path, "w:") as tar:
+        _add_tar_file(tar, "../model_config.yaml", b"model:\n  value: 1\n")
+
+    with pytest.raises(TarPathTraversalError):
+        load_config(str(nemo_path))
+
+
+def test_save_artifacts_rejects_traversal_artifact_path(tmp_path):
+    nemo_path = tmp_path / "malicious.nemo"
+    with tarfile.open(nemo_path, "w:") as tar:
+        _add_tar_file(tar, "model_config.yaml", b"model:\n  value: 1\n")
+        _add_tar_file(tar, "../tokenizer.model", b"payload")
+
+    model = SimpleNamespace(
+        cfg=OmegaConf.create({"tokenizer": {"library": "sentencepiece"}}),
+        artifacts={"tokenizer.model": SimpleNamespace(path="nemo:../tokenizer.model")},
+    )
+    AppState().model_restore_path = str(nemo_path)
+
+    with pytest.raises(TarPathTraversalError):
+        save_artifacts(model, str(tmp_path / "output"))
+
+
+def test_save_artifacts_rejects_absolute_artifact_path_before_rename(tmp_path):
+    secret_path = tmp_path / "secret.txt"
+    secret_path.write_text("sensitive")
+    nemo_path = tmp_path / "malicious.nemo"
+    with tarfile.open(nemo_path, "w:") as tar:
+        _add_tar_file(tar, "nested/model_config.yaml", b"model:\n  value: 1\n")
+        _add_tar_file(tar, f"nested/{secret_path.as_posix()}", b"payload")
+
+    model = SimpleNamespace(
+        cfg=OmegaConf.create({"tokenizer": {"library": "sentencepiece"}}),
+        artifacts={"tokenizer.model": SimpleNamespace(path=f"nemo:{secret_path.as_posix()}")},
+    )
+    output_dir = tmp_path / "output"
+    AppState().model_restore_path = str(nemo_path)
+
+    with pytest.raises(TarPathTraversalError):
+        save_artifacts(model, str(output_dir))
+
+    assert secret_path.read_text() == "sensitive"
+    assert not (output_dir / "tokenizer.model").exists()
+
+
+def test_asr_tokenizer_rename_uses_member_basename():
+    assert ASRBPEMixin._get_extracted_tokenizer_name("prefix_/tmp/tokenizer_vocab") == "vocab"
+
+
+def test_download_an4_rejects_traversal_archive(tmp_path):
+    data_dir = tmp_path / "data"
+    data_dir.mkdir()
+    an4_path = data_dir / "an4_sphere.tar.gz"
+    with tarfile.open(an4_path, "w:gz") as tar:
+        _add_tar_file(tar, "an4/etc/an4_train.transcription")
+        _add_tar_file(tar, "an4/etc/an4_test.transcription")
+        _add_tar_file(tar, "../nemo_poc_traversal_test", b"payload")
+
+    with pytest.raises(TarPathTraversalError):
+        download_an4(str(data_dir))
+
+    assert not (tmp_path / "nemo_poc_traversal_test").exists()
+
+
+def test_no_raw_tar_extraction_outside_safe_helper():
+    raw_extractions = []
+    python_files = subprocess.run(
+        ["git", "ls-files", "-z", "--cached", "--others", "--exclude-standard", "--", "*.py"],
+        cwd=REPO_ROOT,
+        check=True,
+        capture_output=True,
+    ).stdout.split(b"\0")
+    for python_file in python_files:
+        if not python_file:
+            continue
+        path = REPO_ROOT / python_file.decode()
+        if path in RAW_TAR_EXTRACTION_ALLOWLIST:
+            continue
+
+        tar_vars = set()
+        tree = ast.parse(path.read_text(), filename=str(path))
+        for node in ast.walk(tree):
+            if isinstance(node, ast.Assign) and _is_tarfile_open_call(node.value):
+                for target in node.targets:
+                    if isinstance(target, ast.Name):
+                        tar_vars.add(target.id)
+            elif isinstance(node, ast.With):
+                for item in node.items:
+                    if _is_tarfile_open_call(item.context_expr) and isinstance(item.optional_vars, ast.Name):
+                        tar_vars.add(item.optional_vars.id)
+
+        for node in ast.walk(tree):
+            if not isinstance(node, ast.Call) or not isinstance(node.func, ast.Attribute):
+                continue
+            if (
+                node.func.attr in {"extract", "extractall"}
+                and isinstance(node.func.value, ast.Name)
+                and node.func.value.id in tar_vars
+            ):
+                raw_extractions.append(f"{path.relative_to(REPO_ROOT)}:{node.lineno}")
+
+    assert raw_extractions == []

tests/utils/test_trainer_utils.py

@@ -0,0 +1,34 @@
+# Copyright (c) 2025, 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.
+
+from lightning.pytorch.strategies import DDPStrategy
+from omegaconf import OmegaConf
+
+from nemo.utils.trainer_utils import resolve_trainer_cfg
+
+
+def test_resolve_trainer_cfg_strategy():
+    cfg = OmegaConf.create({"strategy": "ddp"})
+    ans = resolve_trainer_cfg(cfg)
+    assert isinstance(ans, dict)
+    assert ans["strategy"] == "ddp"
+
+    cfg = OmegaConf.create(
+        {"strategy": {"_target_": "lightning.pytorch.strategies.DDPStrategy", "gradient_as_bucket_view": True}}
+    )
+    ans = resolve_trainer_cfg(cfg)
+    assert isinstance(ans, dict)
+    assert isinstance(ans["strategy"], DDPStrategy)
+    assert "gradient_as_bucket_view" in ans["strategy"]._ddp_kwargs
+    assert ans["strategy"]._ddp_kwargs["gradient_as_bucket_view"] == True

tests/utils/test_utils.py

@@ -0,0 +1,92 @@
+# Copyright (c) 2022, 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 os
+from unittest import mock
+
+import pytest
+
+from nemo import __version__ as NEMO_VERSION
+from nemo.utils.data_utils import (
+    ais_binary,
+    ais_endpoint_to_dir,
+    bucket_and_object_from_uri,
+    is_datastore_path,
+    resolve_cache_dir,
+)
+
+
+class TestDataUtils:
+    @pytest.mark.unit
+    def test_resolve_cache_dir(self):
+        """Test cache dir path."""
+        TEST_NEMO_ENV_CACHE_DIR = 'TEST_NEMO_ENV_CACHE_DIR'
+        with mock.patch('nemo.constants.NEMO_ENV_CACHE_DIR', TEST_NEMO_ENV_CACHE_DIR):
+
+            envar_to_resolved_path = {
+                '/path/to/cache': '/path/to/cache',
+                'relative/path': os.path.join(os.getcwd(), 'relative/path'),
+                '': os.path.expanduser(f'~/.cache/torch/NeMo/NeMo_{NEMO_VERSION}'),
+            }
+
+            for envar, expected_path in envar_to_resolved_path.items():
+                # Set envar
+                os.environ[TEST_NEMO_ENV_CACHE_DIR] = envar
+                # Check path
+                uut_path = resolve_cache_dir().as_posix()
+                assert uut_path == expected_path, f'Expected: {expected_path}, got {uut_path}'
+
+    @pytest.mark.unit
+    def test_is_datastore_path(self):
+        """Test checking for datastore path."""
+        # Positive examples
+        assert is_datastore_path('ais://positive/example')
+        # Negative examples
+        assert not is_datastore_path('ais/negative/example')
+        assert not is_datastore_path('/negative/example')
+        assert not is_datastore_path('negative/example')
+
+    @pytest.mark.unit
+    def test_bucket_and_object_from_uri(self):
+        """Test getting bucket and object from URI."""
+        # Positive examples
+        assert bucket_and_object_from_uri('ais://bucket/object') == ('bucket', 'object')
+        assert bucket_and_object_from_uri('ais://bucket_2/object/is/here') == ('bucket_2', 'object/is/here')
+
+        # Negative examples: invalid URI
+        with pytest.raises(ValueError):
+            bucket_and_object_from_uri('/local/file')
+
+        with pytest.raises(ValueError):
+            bucket_and_object_from_uri('local/file')
+
+    @pytest.mark.unit
+    def test_ais_endpoint_to_dir(self):
+        """Test converting an AIS endpoint to dir."""
+        assert ais_endpoint_to_dir('http://local:123') == os.path.join('local', '123')
+        assert ais_endpoint_to_dir('http://1.2.3.4:567') == os.path.join('1.2.3.4', '567')
+
+        with pytest.raises(ValueError):
+            ais_endpoint_to_dir('local:123')
+
+    @pytest.mark.unit
+    def test_ais_binary(self):
+        """Test cache dir path."""
+        with mock.patch('shutil.which', lambda x: '/test/path/ais'):
+            assert ais_binary() == '/test/path/ais'
+
+        # Negative example: AIS binary cannot be found
+        with mock.patch('shutil.which', lambda x: None), mock.patch('os.path.isfile', lambda x: None):
+            ais_binary.cache_clear()
+            assert ais_binary() is None