hotfix: Batch size finder calling wrong super class functions (#85)

* fix: Fix batch size finder callback not calling proper super functions

* fix: Fix hardcoded value

* build: Bump version 1.3.6 -> 1.3.7

* docs: Update changelog

Approved By: @rcmalli

CHANGELOG.md

@@ -2,6 +2,13 @@
 # Changelog
 All notable changes to this project will be documented in this file.
 
+### [1.3.7]
+
+#### Fixed
+
+- Fix BatchSizeFinder calling wrong super functions
+- Fix ModelManager get_latest_version calling an hardcoded model
+
 ### [1.3.6]
 
 #### Fixed

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "quadra"
-version = "1.3.6"
+version = "1.3.7"
 description = "Deep Learning experiment orchestration library"
 authors = [
   { name = "Alessandro Polidori", email = "[email protected]" },
@@ -121,7 +121,7 @@ repository = "https://github.com/orobix/quadra"
 
 # Adapted from https://realpython.com/pypi-publish-python-package/#version-your-package
 [tool.bumpver]
-current_version = "1.3.6"
+current_version = "1.3.7"
 version_pattern = "MAJOR.MINOR.PATCH"
 commit_message = "build: Bump version {old_version} -> {new_version}"
 commit = true

quadra/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "1.3.6"
+__version__ = "1.3.7"
 
 
 def get_version():

quadra/callbacks/lightning.py

@@ -76,14 +76,19 @@ def __init__(
         self.find_test_batch_size = find_test_batch_size
         self.find_predict_batch_size = find_predict_batch_size
 
-    def on_train_start(self, trainer: pl.Trainer, pl_module: pl.LightningModule) -> None:
-        if not self.find_train_batch_size:
+    def on_fit_start(self, trainer: pl.Trainer, pl_module: pl.LightningModule) -> None:
+        if not self.find_train_batch_size or trainer.state.stage is None:
+            # If called during validation skip it as it will be triggered during on_validation_start
+            return None
+
+        if trainer.state.stage.value != "train":
             return None
 
         if not isinstance(pl_module.model, nn.Module):
             raise ValueError("The model must be a nn.Module")
         pl_module.model.train()
-        return super().on_train_epoch_start(trainer, pl_module)
+
+        return super().on_fit_start(trainer, pl_module)
 
     def on_validation_start(self, trainer: pl.Trainer, pl_module: pl.LightningModule) -> None:
         if not self.find_validation_batch_size:
@@ -92,7 +97,8 @@ def on_validation_start(self, trainer: pl.Trainer, pl_module: pl.LightningModule
         if not isinstance(pl_module.model, nn.Module):
             raise ValueError("The model must be a nn.Module")
         pl_module.model.eval()
-        return super().on_validation_epoch_start(trainer, pl_module)
+
+        return super().on_validation_start(trainer, pl_module)
 
     def on_test_start(self, trainer: pl.Trainer, pl_module: pl.LightningModule) -> None:
         if not self.find_test_batch_size:
@@ -101,7 +107,8 @@ def on_test_start(self, trainer: pl.Trainer, pl_module: pl.LightningModule) -> N
         if not isinstance(pl_module.model, nn.Module):
             raise ValueError("The model must be a nn.Module")
         pl_module.model.eval()
-        return super().on_test_epoch_start(trainer, pl_module)
+
+        return super().on_test_start(trainer, pl_module)
 
     def on_predict_start(self, trainer: pl.Trainer, pl_module: pl.LightningModule) -> None:
         if not self.find_predict_batch_size:
@@ -110,4 +117,5 @@ def on_predict_start(self, trainer: pl.Trainer, pl_module: pl.LightningModule) -
         if not isinstance(pl_module.model, nn.Module):
             raise ValueError("The model must be a nn.Module")
         pl_module.model.eval()
-        return super().on_predict_epoch_start(trainer, pl_module)
+
+        return super().on_predict_start(trainer, pl_module)

quadra/utils/model_manager.py

@@ -122,7 +122,7 @@ def get_latest_version(self, model_name: str) -> ModelVersion:
         Returns:
             The model version
         """
-        latest_version = max(int(x.version) for x in self.client.get_latest_versions("manager_model"))
+        latest_version = max(int(x.version) for x in self.client.get_latest_versions(model_name))
         model_version = self.client.get_model_version(model_name, latest_version)
 
         return model_version