Merge pull request #3090 from bghira/agent/dataset-batch-cache-versioning

Keep dataset batch size current across cache reloads

Author: bghira

simpletuner/helpers/data_backend/factory.py

@@ -3376,6 +3376,7 @@ def _handle_bucket_operations(
 
     def _handle_config_versioning(self, backend: Dict[str, Any], init_backend: Dict[str, Any]) -> None:
         """Handle configuration versioning and validation."""
+        runtime_mutable_keys = ("train_batch_size",)
         excluded_keys = [
             "probability",
             "repeats",
@@ -3393,6 +3394,7 @@ def _handle_config_versioning(self, backend: Dict[str, Any], init_backend: Dict[
             "start_epoch",
             "hash_filenames",  # always enabled, not user-configurable
             "_s2v_audio_autoinjected",  # runtime flag, not user-configurable
+            *runtime_mutable_keys,
         ]
         _latest_config_version = latest_config_version()
         current_config_version = _latest_config_version
@@ -3427,6 +3429,14 @@ def _handle_config_versioning(self, backend: Dict[str, Any], init_backend: Dict[
                             )
                         init_backend["config"][key] = prev_config[key]
 
+        metadata_backend_config = getattr(init_backend["metadata_backend"], "config", None)
+        if isinstance(metadata_backend_config, dict):
+            for key in runtime_mutable_keys:
+                if key in init_backend["config"]:
+                    metadata_backend_config[key] = init_backend["config"][key]
+                else:
+                    metadata_backend_config.pop(key, None)
+
         runtime_linkage_keys = (
             "conditioning_data",
             "conditioning",

tests/test_factory_edge_cases.py

@@ -537,6 +537,72 @@ def test_eval_backend_config_forces_train_batch_size_one(self):
 
         self.assertEqual(result["config"]["train_batch_size"], 1)
 
+    def _apply_cached_train_batch_size(self, backend, cached_train_batch_size):
+        from simpletuner.helpers.data_backend.factory import FactoryRegistry, init_backend_config
+
+        init_backend = init_backend_config(backend, self.args, self.accelerator)
+        effective_train_batch_size = init_backend["config"]["train_batch_size"]
+        metadata_backend = MagicMock()
+        metadata_backend.config = {"train_batch_size": cached_train_batch_size}
+        metadata_backend.batch_size = effective_train_batch_size
+        metadata_backend.__len__.return_value = 1
+        init_backend["metadata_backend"] = metadata_backend
+        init_backend["data_backend"] = MagicMock()
+
+        factory = FactoryRegistry(
+            args=self.args,
+            accelerator=self.accelerator,
+            text_encoders=self.text_encoders,
+            tokenizers=self.tokenizers,
+            model=self.model,
+        )
+        sampler = MagicMock(caption_strategy="filename")
+        with (
+            patch("simpletuner.helpers.data_backend.factory.StateTracker.set_data_backend_config") as set_config,
+            patch("simpletuner.helpers.data_backend.factory.print_bucket_info"),
+            patch("simpletuner.helpers.data_backend.factory.MultiAspectDataset"),
+            patch("simpletuner.helpers.data_backend.factory.MultiAspectSampler", return_value=sampler) as sampler_cls,
+            patch("simpletuner.helpers.data_backend.factory.torch.utils.data.DataLoader"),
+        ):
+            factory._handle_config_versioning(backend, init_backend)
+            factory._create_dataset_and_sampler(backend, init_backend, conditioning_type=None)
+
+        set_config.assert_called_once_with(init_backend["id"], init_backend["config"])
+        self.assertEqual(init_backend["config"]["train_batch_size"], effective_train_batch_size)
+        self.assertEqual(metadata_backend.batch_size, effective_train_batch_size)
+        self.assertEqual(metadata_backend.config["train_batch_size"], effective_train_batch_size)
+        self.assertEqual(sampler_cls.call_args.kwargs["batch_size"], effective_train_batch_size)
+        return init_backend
+
+    def test_cached_batch_size_does_not_override_changed_global_default(self):
+        self.args.train_batch_size = 4
+        backend = {
+            "id": "image-global-batch",
+            "type": "local",
+            "dataset_type": "image",
+            "instance_data_dir": self.temp_dir,
+        }
+
+        result = self._apply_cached_train_batch_size(backend, cached_train_batch_size=2)
+
+        self.assertNotIn("train_batch_size", backend)
+        self.assertEqual(result["config"]["train_batch_size"], 4)
+
+    def test_cached_batch_size_does_not_override_changed_dataset_value(self):
+        self.args.train_batch_size = 8
+        backend = {
+            "id": "image-dataset-batch",
+            "type": "local",
+            "dataset_type": "image",
+            "train_batch_size": 4,
+            "instance_data_dir": self.temp_dir,
+        }
+
+        result = self._apply_cached_train_batch_size(backend, cached_train_batch_size=2)
+
+        self.assertEqual(backend["train_batch_size"], 4)
+        self.assertEqual(result["config"]["train_batch_size"], 4)
+
     def test_inline_conditioning_auto_generation_for_image_dataset(self):
         """Inline conditioning blocks on image datasets should spawn auto-generated conditioning datasets."""
         from simpletuner.helpers.data_backend.factory import FactoryRegistry