Make DeepCompile failure cleanup idempotent

Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>

Author: tohtana

deepspeed/compile/init_z3.py

@@ -78,8 +78,11 @@ def backend_with_failure_cleanup(*args, **kwargs):
         try:
             return backend_fn(*args, **kwargs)
         except Exception:
-            get_deepcompile_handle().cleanup()
-            engine._set_deepcompile_active(False)
+            if engine.is_deepcompile_active():
+                try:
+                    get_deepcompile_handle().cleanup()
+                finally:
+                    engine._set_deepcompile_active(False)
             raise
 
     return backend_with_failure_cleanup

tests/unit/compile/test_zero3_grad_dtype.py

@@ -204,14 +204,12 @@ def failing_backend():
     fake_engine._release_deepcompile_dynamo_config = (
         lambda: DeepSpeedEngine._release_deepcompile_dynamo_config(fake_engine))
     fake_engine._set_deepcompile_active = lambda active: DeepSpeedEngine._set_deepcompile_active(fake_engine, active)
+    fake_engine.is_deepcompile_active = lambda: fake_engine._deepcompile_active
 
     try:
-        try:
-            backend()
-        except RuntimeError as exc:
-            assert str(exc) == "compile failed"
-        else:
-            raise AssertionError("failing backend did not raise")
+        for _ in range(2):
+            with pytest.raises(RuntimeError, match="compile failed"):
+                backend()
 
         assert FakeDynamo.config.force_parameter_static_shapes is True
         assert FakeDynamo.config.force_nn_module_property_static_shapes is False
@@ -223,7 +221,6 @@ def failing_backend():
 
         fake_engine.optimizer = None
         fake_engine.checkpoint_engine = None
-        fake_engine.is_deepcompile_active = lambda: fake_engine._deepcompile_active
         DeepSpeedEngine.destroy(fake_engine)
         assert cleanup_states == [True]
     finally: