update

Author: deependujha

tests/tests_pytorch/deprecated_api/test_no_removal_version.py

@@ -2,7 +2,7 @@
 from unittest.mock import Mock
 
 import pytest
-import pytorch_lightning_enterprise.utils.imports
+import pytorch_lightning_enterprise.plugins.precision.bitsandbytes
 import torch.nn
 
 import lightning.fabric
@@ -63,7 +63,7 @@ def test_fsdp_precision_plugin():
 
 def test_bitsandbytes_precision_plugin(monkeypatch):
     monkeypatch.setattr(lightning.fabric.plugins.precision.bitsandbytes, "_BITSANDBYTES_AVAILABLE", True)
-    monkeypatch.setattr(pytorch_lightning_enterprise.utils.imports, "_BITSANDBYTES_AVAILABLE", True)
+    monkeypatch.setattr(pytorch_lightning_enterprise.plugins.precision.bitsandbytes, "_BITSANDBYTES_AVAILABLE", True)
     bitsandbytes_mock = Mock()
     monkeypatch.setitem(sys.modules, "bitsandbytes", bitsandbytes_mock)
 

tests/tests_pytorch/graveyard/test_tpu.py

@@ -1,3 +1,4 @@
+import os
 from importlib import import_module
 
 import pytest
@@ -39,3 +40,9 @@ def test_graveyard_no_device(import_path, name):
     cls = getattr(module, name)
     with pytest.deprecated_call(match="is deprecated"), pytest.raises(ModuleNotFoundError, match="torch_xla"):
         cls()
+
+    # teardown
+    # ideally, we should call the plugin's teardown method, but since the class
+    # instantiation itself fails, we directly manipulate the env vars here
+    os.environ.pop("XLA_USE_BF16", None)
+    os.environ.pop("XLA_USE_F16", None)

tests/tests_pytorch/loggers/test_mlflow.py

@@ -304,7 +304,7 @@ def test_mlflow_logger_experiment_calls_with_synchronous(mlflow_mock, tmp_path,
 
 
 @mock.patch("pytorch_lightning_enterprise.loggers.mlflow._get_resolve_tags", Mock())
-@mock.patch("pytorch_lightning_enterprise.utils.imports._MLFLOW_SYNCHRONOUS_AVAILABLE", False)
+@mock.patch("pytorch_lightning_enterprise.loggers.mlflow._MLFLOW_SYNCHRONOUS_AVAILABLE", False)
 def test_mlflow_logger_no_synchronous_support(mlflow_mock, tmp_path):
     """Test that the logger does not support synchronous flag."""
     time = mlflow_mock.entities.time