updates

src/axolotl/telemetry/__init__.py

@@ -1,5 +1,15 @@
 """Init for axolotl.telemetry module."""
 
-from .manager import ModelConfig, TelemetryConfig, TelemetryManager, init_telemetry_manager
+from .manager import (
+    ModelConfig,
+    TelemetryConfig,
+    TelemetryManager,
+    init_telemetry_manager,
+)
 
-__all__ = ["TelemetryConfig", "TelemetryManager", "ModelConfig", "init_telemetry_manager"]
+__all__ = [
+    "TelemetryConfig",
+    "TelemetryManager",
+    "ModelConfig",
+    "init_telemetry_manager",
+]

src/axolotl/telemetry/manager.py

@@ -126,8 +126,7 @@ class TelemetryManager:
 
         base_model = base_model.lower()
         return any(
-            org.lower() in base_model 
+            org.lower() in base_model for org in self.whitelist.get("organizations", [])
-            for org in self.whitelist.get("organizations", [])
         )
 
     def _init_posthog(self):
@@ -192,7 +191,7 @@ class TelemetryManager:
                         "run_id": self.run_id,
                         "system_info": system_info,
                         **properties,
-                    }
+                    },
                 )
             except Exception as e:
                 logger.warning(f"Failed to send telemetry event: {e}")

tests/telemetry/test_manager.py

@@ -4,7 +4,7 @@ from unittest.mock import patch
 import pytest
 import yaml
 
-from axolotl.telemetry import TelemetryConfig, TelemetryManager, ModelConfig
+from axolotl.telemetry import ModelConfig, TelemetryConfig, TelemetryManager
 
 
 @pytest.fixture
@@ -12,7 +12,7 @@ def mock_whitelist(tmp_path):
     """Create a temporary whitelist file for testing"""
     whitelist_content = {
         "organizations": ["meta", "mistral"],
-        "models": ["llama", "mistral-7b"]
+        "models": ["llama", "mistral-7b"],
     }
     whitelist_file = tmp_path / "whitelist.yaml"
     with open(whitelist_file, "w") as f:
@@ -24,8 +24,7 @@ def mock_whitelist(tmp_path):
 def config(mock_whitelist):
     """Create a TelemetryConfig with test settings"""
     return TelemetryConfig(
-        host="https://test.posthog.com",
+        host="https://test.posthog.com", whitelist_path=mock_whitelist
-        whitelist_path=mock_whitelist
     )
 
 
@@ -51,10 +50,7 @@ def test_telemetry_opt_in():
 
 def test_do_not_track_override():
     """Test that DO_NOT_TRACK overrides AXOLOTL_TELEMETRY"""
-    with patch.dict(os.environ, {
+    with patch.dict(os.environ, {"AXOLOTL_TELEMETRY": "1", "DO_NOT_TRACK": "1"}):
-        "AXOLOTL_TELEMETRY": "1",
-        "DO_NOT_TRACK": "1"
-    }):
         manager = TelemetryManager(TelemetryConfig())
         assert not manager.enabled
 
@@ -97,7 +93,7 @@ def test_model_tracking(manager):
             tokenizer_config={},
             flash_attention=True,
             quantization_config=None,
-            training_approach="lora"
+            training_approach="lora",
         )
 
         with patch("posthog.capture") as mock_capture:
@@ -106,7 +102,10 @@ def test_model_tracking(manager):
 
             assert mock_capture.called
             assert mock_capture.call_args[1]["event"] == "model_load"
-            assert mock_capture.call_args[1]["properties"]["model_config"] == model_config.to_dict()
+            assert (
+                mock_capture.call_args[1]["properties"]["model_config"]
+                == model_config.to_dict()
+            )
 
 
 def test_training_context(manager):