Merge pull request #3089 from bghira/agent/dataset-batch-strict-validation

Validate dataset train batch sizes strictly

Author: bghira

simpletuner/helpers/data_backend/config/image.py

@@ -4,7 +4,7 @@
 from dataclasses import dataclass, field
 from typing import Any, Dict, List, Optional, Union
 
-from simpletuner.helpers.data_backend.dataset_types import DatasetType
+from simpletuner.helpers.data_backend.dataset_types import DatasetType, parse_positive_train_batch_size
 from simpletuner.helpers.training.state_tracker import StateTracker
 
 from . import validators
@@ -213,10 +213,7 @@ def _get_arg(key: str, default: Any = None) -> Any:
         config.disable_validation = backend_dict.get("disable_validation", False)
         train_batch_size = backend_dict.get("train_batch_size")
         if train_batch_size not in (None, ""):
-            try:
-                config.train_batch_size = int(train_batch_size)
-            except (TypeError, ValueError) as exc:
-                raise ValueError(f"(id={config.id}) train_batch_size must be a positive integer.") from exc
+            config.train_batch_size = parse_positive_train_batch_size(train_batch_size, config.id)
         if "hash_filenames" in backend_dict and config.backend_type != "csv":
             config.hash_filenames = backend_dict.get("hash_filenames")
         config.source_dataset_id = backend_dict.get("source_dataset_id")
@@ -372,8 +369,8 @@ def validate(self, args: Dict[str, Any]) -> None:
             self._validate_video_settings(args)
 
         validators.check_for_caption_filter_list_misuse(self.dataset_type, False, self.id)
-        if self.train_batch_size is not None and self.train_batch_size < 1:
-            raise ValueError(f"(id={self.id}) train_batch_size must be a positive integer.")
+        if self.train_batch_size is not None:
+            self.train_batch_size = parse_positive_train_batch_size(self.train_batch_size, self.id)
 
     def _validate_controlnet_requirements(self, args: Dict[str, Any]) -> None:
         def _get_controlnet_flag(source: Any) -> Optional[bool]:

simpletuner/helpers/data_backend/dataset_types.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 from enum import Enum
+from numbers import Integral
 from typing import Any, Iterable, Mapping, Optional, Sequence
 
 
@@ -72,6 +73,26 @@ def get_arg_value(args: Any, key: str, default: Any = None) -> Any:
     return getattr(args, key, default)
 
 
+def parse_positive_train_batch_size(raw_value: Any, backend_id: Optional[str] = None) -> int:
+    """Parse a positive integer dataset training batch size."""
+    error_message = f"(id={backend_id}) train_batch_size must be a positive integer."
+
+    if isinstance(raw_value, bool):
+        raise ValueError(error_message)
+    if isinstance(raw_value, Integral):
+        batch_size = int(raw_value)
+    elif isinstance(raw_value, str) and raw_value.isascii() and raw_value.isdecimal():
+        batch_size = int(raw_value)
+        if raw_value != str(batch_size):
+            raise ValueError(error_message)
+    else:
+        raise ValueError(error_message)
+
+    if batch_size < 1:
+        raise ValueError(error_message)
+    return batch_size
+
+
 def resolve_dataset_train_batch_size(
     backend: Mapping[str, Any],
     args: Any,
@@ -88,12 +109,4 @@ def resolve_dataset_train_batch_size(
         raw_value = get_arg_value(args, "train_batch_size", 1)
 
     resolved_backend_id = backend_id if backend_id is not None else backend.get("id")
-    try:
-        batch_size = int(raw_value)
-    except (TypeError, ValueError) as exc:
-        raise ValueError(f"(id={resolved_backend_id}) train_batch_size must be a positive integer.") from exc
-
-    if batch_size < 1:
-        raise ValueError(f"(id={resolved_backend_id}) train_batch_size must be a positive integer.")
-
-    return batch_size
+    return parse_positive_train_batch_size(raw_value, resolved_backend_id)

tests/test_backend_config.py

@@ -8,6 +8,8 @@
 from typing import Any, Dict
 from unittest.mock import Mock, patch
 
+import numpy as np
+
 from simpletuner.helpers.data_backend.config import (
     BaseBackendConfig,
     ImageBackendConfig,
@@ -335,18 +337,40 @@ def test_train_batch_size_round_trip(self):
         self.assertEqual(config.train_batch_size, 3)
         self.assertEqual(output["config"]["train_batch_size"], 3)
 
-    def test_train_batch_size_must_be_positive(self):
-        backend_dict = {
-            "id": "image_test",
-            "type": "local",
-            "dataset_type": "image",
-            "train_batch_size": 0,
-        }
+    def test_train_batch_size_must_be_a_positive_integer(self):
+        invalid_values = (True, False, 2.5, 3.0, "2.5", 0, "0", -1, "-1", "03", "+3", "three")
+
+        for value in invalid_values:
+            with self.subTest(value=value):
+                backend_dict = {
+                    "id": "image_test",
+                    "type": "local",
+                    "dataset_type": "image",
+                    "train_batch_size": value,
+                }
+
+                with self.assertRaisesRegex(
+                    ValueError,
+                    r"\(id=image_test\) train_batch_size must be a positive integer\.",
+                ):
+                    ImageBackendConfig.from_dict(backend_dict, self.args)
+
+    def test_validate_rejects_invalid_direct_train_batch_size(self):
+        config = ImageBackendConfig(id="image_test", train_batch_size=True)
+
+        with self.assertRaisesRegex(
+            ValueError,
+            r"\(id=image_test\) train_batch_size must be a positive integer\.",
+        ):
+            config.validate(self.args)
 
-        config = ImageBackendConfig.from_dict(backend_dict, self.args)
+    def test_validate_normalizes_integral_train_batch_size(self):
+        config = ImageBackendConfig(id="image_test", train_batch_size=np.int64(3))
 
-        with self.assertRaisesRegex(ValueError, "train_batch_size"):
-            config.validate(self.args)
+        config.validate(self.args)
+
+        self.assertIs(type(config.train_batch_size), int)
+        self.assertIs(type(config.to_dict()["config"]["train_batch_size"]), int)
 
     def test_vae_cache_ondemand_round_trip(self):
         backend_dict = {

tests/test_dataset_types.py

@@ -0,0 +1,54 @@
+"""Tests for dataset type helpers."""
+
+import unittest
+
+from simpletuner.helpers.data_backend.dataset_types import (
+    DatasetType,
+    parse_positive_train_batch_size,
+    resolve_dataset_train_batch_size,
+)
+
+
+class DatasetTrainBatchSizeTestCase(unittest.TestCase):
+    def test_parser_accepts_positive_integers_and_canonical_strings(self):
+        for value in (1, 3, "1", "3"):
+            with self.subTest(value=value):
+                self.assertEqual(parse_positive_train_batch_size(value, "dataset"), int(value))
+
+    def test_parser_rejects_non_positive_or_non_integer_values(self):
+        invalid_values = (True, False, 2.5, 3.0, "2.5", 0, "0", -1, "-1", "03", "+3", "three", None)
+
+        for value in invalid_values:
+            with self.subTest(value=value):
+                with self.assertRaisesRegex(
+                    ValueError,
+                    r"\(id=dataset\) train_batch_size must be a positive integer\.",
+                ):
+                    parse_positive_train_batch_size(value, "dataset")
+
+    def test_resolver_validates_dataset_and_global_values(self):
+        with self.assertRaisesRegex(ValueError, "train_batch_size"):
+            resolve_dataset_train_batch_size(
+                {"id": "dataset", "dataset_type": "image", "train_batch_size": 2.5},
+                {"train_batch_size": 1},
+            )
+
+        with self.assertRaisesRegex(ValueError, "train_batch_size"):
+            resolve_dataset_train_batch_size(
+                {"id": "dataset", "dataset_type": "image"},
+                {"train_batch_size": True},
+            )
+
+    def test_eval_resolver_forces_batch_size_one_before_validation(self):
+        self.assertEqual(
+            resolve_dataset_train_batch_size(
+                {"id": "eval", "dataset_type": "eval", "train_batch_size": 2.5},
+                {"train_batch_size": True},
+                dataset_type=DatasetType.EVAL,
+            ),
+            1,
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()