Feat: Make atexit flush registration optional For _HTTPBackgroundLogger

py/src/braintrust/logger.py

@@ -1035,7 +1035,14 @@ def __init__(self, api_conn: LazyValue[HTTPConnection]):
         # Counter for tracking overflow uploads (useful for testing)
         self._overflow_upload_count = 0
 
-        atexit.register(self._finalize)
+        disable_atexit_flush = False
+        try:
+            disable_atexit_flush = os.environ["BRAINTRUST_DISABLE_ATEXIT_FLUSH"] in ("True", "1")
+        except:
+            pass
+
+        if not disable_atexit_flush:
+            atexit.register(self._finalize)
 
     def enforce_queue_size_limit(self, enforce: bool) -> None:
         """

py/src/braintrust/test_logger.py

@@ -125,6 +125,36 @@ def test_init_with_repo_info_does_not_raise(self):
         assert metadata.project.id == "test-project-id"
         assert metadata.experiment.name == "test-exp"
 
+    def test_init_enable_atexit_flush(self):
+        from braintrust.logger import _HTTPBackgroundLogger
+
+        api_con_response = lambda : {
+            "project": {"id": "test-project-id", "name": "test-project"},
+            "experiment": {"id": "test-exp-id", "name": "test-exp"},
+        }
+
+        with patch("atexit.register") as mock_register:
+            _HTTPBackgroundLogger(LazyValue(api_con_response, use_mutex=False))  # type: ignore
+            mock_register.assert_called()
+
+    def test_init_disable_atexit_flush(self):
+        from braintrust.logger import _HTTPBackgroundLogger
+
+        api_con_response = lambda : {
+            "project": {"id": "test-project-id", "name": "test-project"},
+            "experiment": {"id": "test-exp-id", "name": "test-exp"},
+        }
+
+        with patch.dict(os.environ, {"BRAINTRUST_DISABLE_ATEXIT_FLUSH": "True"}):
+            with patch("atexit.register") as mock_register:
+                _HTTPBackgroundLogger(LazyValue(api_con_response, use_mutex=False))  # type: ignore
+                mock_register.assert_not_called()
+
+        with patch.dict(os.environ, {"BRAINTRUST_DISABLE_ATEXIT_FLUSH": "1"}):
+            with patch("atexit.register") as mock_register:
+                _HTTPBackgroundLogger(LazyValue(api_con_response, use_mutex=False))  # type: ignore
+                mock_register.assert_not_called()
+
 class TestLogger(TestCase):
     def test_extract_attachments_no_op(self):
         attachments: List[BaseAttachment] = []