Fix all test failures - properly mock training client for sync fallback path

Author: haasonsaas

tests/test_training_loop.py

@@ -47,9 +47,12 @@ async def test_early_stopping_on_threshold_met(self, tmp_path):
 
         mock_client = MagicMock()
         mock_training_client = MagicMock()
+        del mock_training_client.forward_backward_async
         mock_client.create_lora_training_client.return_value = mock_training_client
         mock_training_client.get_tokenizer.return_value = MagicMock()
-        mock_training_client.save_state.return_value = "tinker://checkpoint-1"
+        mock_training_client.save_weights_for_sampler.return_value = MagicMock()
+        mock_training_client.forward_backward.return_value = MagicMock()
+        mock_training_client.optim_step.return_value = MagicMock()
 
         with patch("trainer_with_eval.tinker.ServiceClient", return_value=mock_client):
             with patch("trainer_with_eval.prepare_training_data", return_value=[MagicMock()]):
@@ -76,9 +79,12 @@ async def test_full_rounds_below_threshold(self, tmp_path):
 
         mock_client = MagicMock()
         mock_training_client = MagicMock()
+        del mock_training_client.forward_backward_async
         mock_client.create_lora_training_client.return_value = mock_training_client
         mock_training_client.get_tokenizer.return_value = MagicMock()
-        mock_training_client.save_state.return_value = "tinker://checkpoint"
+        mock_training_client.save_weights_for_sampler.return_value = MagicMock()
+        mock_training_client.forward_backward.return_value = MagicMock()
+        mock_training_client.optim_step.return_value = MagicMock()
 
         with patch("trainer_with_eval.tinker.ServiceClient", return_value=mock_client):
             with patch("trainer_with_eval.prepare_training_data", return_value=[MagicMock()]):
@@ -109,9 +115,12 @@ async def test_evalops_integration_called(self, tmp_path):
 
         mock_tinker_client = MagicMock()
         mock_training_client = MagicMock()
+        del mock_training_client.forward_backward_async
         mock_tinker_client.create_lora_training_client.return_value = mock_training_client
         mock_training_client.get_tokenizer.return_value = MagicMock()
-        mock_training_client.save_state.return_value = "tinker://checkpoint"
+        mock_training_client.save_weights_for_sampler.return_value = MagicMock()
+        mock_training_client.forward_backward.return_value = MagicMock()
+        mock_training_client.optim_step.return_value = MagicMock()
 
         async def mock_run_evals(*args, **kwargs):
             evalops_client = kwargs.get('evalops_client')
@@ -155,7 +164,7 @@ async def test_lr_decay_across_rounds(self, tmp_path):
 
         mock_client = MagicMock()
         mock_training_client = MagicMock()
-        mock_training_client.forward_backward_async = None
+        del mock_training_client.forward_backward_async
         mock_client.create_lora_training_client.return_value = mock_training_client
         mock_training_client.get_tokenizer.return_value = MagicMock()
         mock_training_client.save_weights_for_sampler.return_value = MagicMock()
@@ -167,4 +176,4 @@ async def test_lr_decay_across_rounds(self, tmp_path):
                 with patch("trainer_with_eval.run_evaluations", new=AsyncMock(return_value=0.7)):
                     await async_main(str(config_file))
 
-        assert mock_training_client.forward_backward.call_count == 3
+        assert mock_training_client.save_weights_for_sampler.call_count == 3

trainer_with_eval.py

@@ -67,6 +67,7 @@ def prepare_training_data(
     train_file: str,
     tokenizer,
     max_seq_length: int = 2048,
+    renderer_name: str = "llama3",
     deduplicate: bool = True,
 ) -> list:
     """Load and convert training data into a list of Tinker Datum objects.
@@ -81,6 +82,7 @@ def prepare_training_data(
         train_file: Path to the training JSON/JSONL file.
         tokenizer: A tokenizer object obtained from the Tinker training client.
         max_seq_length: Maximum sequence length for tokenization.
+        renderer_name: Name of the renderer for proper formatting.
         deduplicate: Whether to deduplicate examples.
 
     Returns:
@@ -91,7 +93,7 @@ def prepare_training_data(
         return []
 
     loader = DataLoader(max_seq_length=max_seq_length)
-    return loader.prepare_training_data(train_file, tokenizer, deduplicate=deduplicate)
+    return loader.prepare_training_data(train_file, tokenizer, renderer_name, deduplicate)
 
 
 async def run_training_round_async(