fix: TelemetryCallback __bases__ crash + 12 TRL integration tests (#231)

The dynamic __bases__ assignment to inject TrainerCallback as a base
class fails in Python: "deallocator differs from object". Fixed by
creating a proper subclass at definition time instead.

12 new tests:
- Mock rollout_func: correct keys, count, reward variance
- Config separation: TrainingConfig has no TRL fields, wrapper accepts trl_config
- Wrapper construction: all callback combinations, trl_config passthrough
- TelemetryCallback: importable, fires events

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: abrichr

openadapt_evals/integrations/trl_callbacks.py

@@ -162,16 +162,17 @@ def on_train_end(
 
 
 # Register as a TrainerCallback subclass at import time so TRL recognizes it.
-# If transformers is not installed, the class still works as a plain object
-# (the callback methods are called by name, not by inheritance check in recent
-# TRL versions).
+# If transformers is installed, wrap with proper inheritance.
+# We can't patch __bases__ after the fact (Python doesn't allow it when
+# deallocators differ), so we create a subclass instead.
 try:
     from transformers import TrainerCallback as _TrainerCallback
 
-    # Dynamically add TrainerCallback as a base class
-    TelemetryCallback.__bases__ = (_TrainerCallback,) + TelemetryCallback.__bases__
+    class _TelemetryCallbackWithBase(_TrainerCallback, TelemetryCallback):
+        """TelemetryCallback with proper TrainerCallback inheritance."""
+        pass
+
+    # Replace the module-level name so imports get the subclass
+    TelemetryCallback = _TelemetryCallbackWithBase  # type: ignore[misc]
 except ImportError:
-    logger.debug(
-        "transformers not installed; TelemetryCallback will work as a "
-        "duck-typed callback but won't inherit from TrainerCallback"
-    )
+    pass  # TelemetryCallback works as duck-typed callback without inheritance

tests/test_trl_integration.py

@@ -0,0 +1,241 @@
+"""Tests for TRL GRPOTrainer integration.
+
+Validates the rollout_func, mock mode, config separation, and wrapper
+without requiring a GPU, real model, or WAA server.
+"""
+
+from __future__ import annotations
+
+from unittest.mock import MagicMock, patch
+import pytest
+
+
+# ---------------------------------------------------------------------------
+# Mock rollout_func tests
+# ---------------------------------------------------------------------------
+
+
+class TestMockRolloutFunc:
+    """Test the mock rollout function from train_trl_grpo.py."""
+
+    def _make_task_configs(self, n=3):
+        """Create simple task configs."""
+        from openadapt_evals.task_config import TaskConfig
+
+        configs = []
+        for i in range(n):
+            tc = MagicMock(spec=TaskConfig)
+            tc.name = f"Task {i}"
+            tc.id = f"task-{i}"
+            tc.milestones = [MagicMock() for _ in range(2)]
+            tc.max_steps = 10
+            configs.append(tc)
+        return configs
+
+    def test_mock_returns_correct_keys(self):
+        """Mock rollout returns prompt_ids, completion_ids, logprobs, env_reward."""
+        # Import the mock creator from the training script
+        import importlib.util
+        spec = importlib.util.spec_from_file_location(
+            "train_trl_grpo", "scripts/train_trl_grpo.py",
+        )
+        mod = importlib.util.module_from_spec(spec)
+        spec.loader.exec_module(mod)
+
+        configs = self._make_task_configs()
+        rollout_func = mod.create_mock_rollout_func(configs)
+
+        mock_trainer = MagicMock()
+        mock_trainer.args.num_generations = 4
+
+        result = rollout_func(["Task 0", "Task 1"], mock_trainer)
+
+        assert "prompt_ids" in result
+        assert "completion_ids" in result
+        assert "logprobs" in result
+        assert "env_reward" in result
+
+    def test_mock_returns_correct_count(self):
+        """Mock returns num_prompts * num_generations entries."""
+        import importlib.util
+        spec = importlib.util.spec_from_file_location(
+            "train_trl_grpo", "scripts/train_trl_grpo.py",
+        )
+        mod = importlib.util.module_from_spec(spec)
+        spec.loader.exec_module(mod)
+
+        configs = self._make_task_configs()
+        rollout_func = mod.create_mock_rollout_func(configs)
+
+        mock_trainer = MagicMock()
+        mock_trainer.args.num_generations = 4
+
+        result = rollout_func(["Task 0", "Task 1"], mock_trainer)
+
+        expected = 2 * 4  # 2 prompts * 4 generations
+        assert len(result["env_reward"]) == expected
+        assert len(result["prompt_ids"]) == expected
+
+    def test_mock_has_reward_variance(self):
+        """Mock produces different reward values (needed for GRPO)."""
+        import importlib.util
+        spec = importlib.util.spec_from_file_location(
+            "train_trl_grpo", "scripts/train_trl_grpo.py",
+        )
+        mod = importlib.util.module_from_spec(spec)
+        spec.loader.exec_module(mod)
+
+        configs = self._make_task_configs()
+        rollout_func = mod.create_mock_rollout_func(configs)
+
+        mock_trainer = MagicMock()
+        mock_trainer.args.num_generations = 8
+
+        # Run multiple times to get reward variance (randomized)
+        all_rewards = []
+        for _ in range(5):
+            result = rollout_func(["Task 0"], mock_trainer)
+            all_rewards.extend(result["env_reward"])
+
+        unique_rewards = set(all_rewards)
+        assert len(unique_rewards) > 1, (
+            f"Mock should produce reward variance, got {unique_rewards}"
+        )
+
+
+# ---------------------------------------------------------------------------
+# Config separation tests
+# ---------------------------------------------------------------------------
+
+
+class TestConfigSeparation:
+    """Verify TrainingConfig and TRL GRPOConfig have clean separation."""
+
+    def test_training_config_has_no_trl_fields(self):
+        """TrainingConfig should NOT have loss_type, gradient_accumulation, etc."""
+        from openadapt_evals.training.standalone.config import TrainingConfig
+
+        tc = TrainingConfig()
+        # These belong to TRL's GRPOConfig, not ours
+        assert not hasattr(tc, "loss_type"), "loss_type belongs in GRPOConfig"
+        assert not hasattr(tc, "gradient_accumulation_steps"), "belongs in GRPOConfig"
+        assert not hasattr(tc, "per_device_train_batch_size"), "belongs in GRPOConfig"
+        assert not hasattr(tc, "bf16"), "belongs in GRPOConfig"
+        assert not hasattr(tc, "report_to"), "belongs in GRPOConfig"
+        assert not hasattr(tc, "use_vllm"), "belongs in GRPOConfig"
+
+    def test_training_config_has_our_fields(self):
+        """TrainingConfig should have OpenAdapt-specific fields."""
+        from openadapt_evals.training.standalone.config import TrainingConfig
+
+        tc = TrainingConfig()
+        assert hasattr(tc, "server_url")
+        assert hasattr(tc, "task_dir")
+        assert hasattr(tc, "constrained_decoding")
+        assert hasattr(tc, "max_new_tokens")
+        assert hasattr(tc, "vision_loss_mode")
+        assert hasattr(tc, "model_name")
+        assert hasattr(tc, "use_unsloth")
+        assert hasattr(tc, "weave_project")
+
+    def test_wrapper_accepts_trl_config(self):
+        """The TRL wrapper accepts a trl_config kwarg."""
+        from openadapt_evals.training.trl_wrapper import GRPOTrainer
+        from openadapt_evals.training.standalone.config import TrainingConfig
+
+        tc = TrainingConfig(task_dir="tasks/")
+
+        # Should not crash — trl_config is stored, not used until train()
+        trainer = GRPOTrainer(tc, trl_config="mock_grpo_config")
+        assert trainer._trl_config == "mock_grpo_config"
+
+    def test_wrapper_defaults_without_trl_config(self):
+        """Without trl_config, wrapper builds defaults from TrainingConfig."""
+        from openadapt_evals.training.trl_wrapper import GRPOTrainer
+        from openadapt_evals.training.standalone.config import TrainingConfig
+
+        tc = TrainingConfig(task_dir="tasks/")
+        trainer = GRPOTrainer(tc)
+        assert trainer._trl_config is None  # will build defaults in train()
+
+
+# ---------------------------------------------------------------------------
+# TRL wrapper construction tests
+# ---------------------------------------------------------------------------
+
+
+class TestTRLWrapperConstruction:
+    """Test the wrapper can be constructed with all callback combinations."""
+
+    def test_no_callbacks(self):
+        from openadapt_evals.training.trl_wrapper import GRPOTrainer
+        from openadapt_evals.training.standalone.config import TrainingConfig
+
+        trainer = GRPOTrainer(TrainingConfig())
+        assert trainer._on_model_loaded is None
+        assert trainer._on_step_complete is None
+
+    def test_all_callbacks(self):
+        from openadapt_evals.training.trl_wrapper import GRPOTrainer
+        from openadapt_evals.training.standalone.config import TrainingConfig
+
+        fn = lambda *a, **kw: None
+        trainer = GRPOTrainer(
+            TrainingConfig(),
+            on_model_loaded=fn,
+            on_before_collect=fn,
+            on_rollout_complete=fn,
+            on_step_complete=fn,
+        )
+        assert trainer._on_model_loaded is fn
+        assert trainer._on_before_collect is fn
+        assert trainer._on_rollout_complete is fn
+        assert trainer._on_step_complete is fn
+
+    def test_trl_config_passthrough(self):
+        """TRL config is stored as-is, not translated."""
+        from openadapt_evals.training.trl_wrapper import GRPOTrainer
+        from openadapt_evals.training.standalone.config import TrainingConfig
+
+        mock_trl = MagicMock()
+        mock_trl.loss_type = "dapo"
+        mock_trl.output_dir = "/tmp/test"
+
+        trainer = GRPOTrainer(TrainingConfig(), trl_config=mock_trl)
+        assert trainer._trl_config.loss_type == "dapo"
+        assert trainer._trl_config.output_dir == "/tmp/test"
+
+
+# ---------------------------------------------------------------------------
+# TelemetryCallback tests
+# ---------------------------------------------------------------------------
+
+
+class TestTelemetryCallback:
+    """Test the TRL TelemetryCallback."""
+
+    def test_callback_importable(self):
+        try:
+            from openadapt_evals.integrations.trl_callbacks import TelemetryCallback
+            cb = TelemetryCallback()
+            assert cb is not None
+        except ImportError:
+            pytest.skip("trl_callbacks not available")
+
+    def test_callback_fires_events(self):
+        try:
+            from openadapt_evals.integrations.trl_callbacks import TelemetryCallback
+        except ImportError:
+            pytest.skip("trl_callbacks not available")
+
+        cb = TelemetryCallback()
+        # These should not crash even without a real trainer
+        args = MagicMock()
+        state = MagicMock()
+        state.global_step = 5
+        state.log_history = [{"loss": 0.5, "reward_mean": 0.7}]
+        control = MagicMock()
+
+        with patch("openadapt_evals.telemetry.capture_event"):
+            cb.on_train_begin(args, state, control)
+            cb.on_step_end(args, state, control)