From b8000e8e44cb9d31f72891bd8bbf89b90e4d104e Mon Sep 17 00:00:00 2001
From: Weishan Li <wli51@jh.edu>
Date: Thu, 6 Aug 2026 11:56:08 -0600
Subject: [PATCH 1/3] Add optimizer state saving functionality as new trainer
 contract, update existing properties by best effort inferring some. Remove
 some confusing and unused properties. Adjust model saving defaults.

---
 .../trainers/AbstractTrainer.py               | 56 ++++++++++++++++---
 .../trainers/logging_gan_trainer.py           |  1 -
 .../trainers/trainer_protocol.py              | 26 ++++++++-
 .../trainers/trainer_utils/__init__.py        |  2 +
 .../trainers/trainer_utils/save_model.py      |  5 +-
 .../trainers/trainer_utils/save_optimizer.py  | 39 +++++++++++++
 tests/trainers/test_abstract_trainer.py       |  5 +-
 7 files changed, 119 insertions(+), 15 deletions(-)
 create mode 100644 src/virtual_stain_flow/trainers/trainer_utils/save_optimizer.py

diff --git a/src/virtual_stain_flow/trainers/AbstractTrainer.py b/src/virtual_stain_flow/trainers/AbstractTrainer.py
index 932faf3..82bf209 100644
--- a/src/virtual_stain_flow/trainers/AbstractTrainer.py
+++ b/src/virtual_stain_flow/trainers/AbstractTrainer.py
@@ -13,7 +13,7 @@
 from torch.utils.data import DataLoader
 
 from .trainer_protocol import TrainerProtocol
-from .trainer_utils import EarlyStopHelper, save_model
+from .trainer_utils import EarlyStopHelper, save_model, save_optimizer_state
 from ..metrics.AbstractMetrics import AbstractMetrics
 from ..engine.progress import Progress
 from ..datasets.data_split import default_random_split
@@ -42,6 +42,7 @@ def __init__(
         test_ratio: Optional[float] = 0.15,
         metrics: Dict[str, AbstractMetrics] = None,
         device: Optional[torch.device] = None,
+        epoch: Optional[int] = 0,
         early_termination_metric: Optional[str] = None,
         early_termination_mode: Literal['min', 'max'] = "min",
         **kwargs,
@@ -102,17 +103,19 @@ def __init__(
             **kwargs
         )
         self._init_state(
+            epoch,
             early_termination_metric, early_termination_mode, **kwargs)
 
     def _init_state(
         self, 
+        epoch,
         early_termination_metric: Optional[str] = None,
         early_termination_mode: Literal['min', 'max'] = "min",
         **kwargs
     ):
 
         # Epoch state
-        self._epoch = 0
+        self._epoch = epoch
         
         # Progress tracking for loss weight scheduling
         self._progress = Progress(epoch=0, step=0)
@@ -183,7 +186,6 @@ def _init_data(
                 **kwargs
             )
 
-            self._batch_size = batch_size
             self._train_ratio, self._val_ratio, self._test_ratio = (
                 train_ratio, val_ratio, test_ratio
             )
@@ -194,6 +196,11 @@ def _init_data(
                 "or provide at least train_loader."
             )
 
+        self._batch_size = self._train_loader.batch_size if hasattr(self._train_loader, 'batch_size') else None
+        self._train_n = len(self._train_loader.dataset) if hasattr(self._train_loader, 'dataset') else None
+        self._val_n = len(self._val_loader.dataset) if hasattr(self._val_loader, 'dataset') else None
+        self._test_n = len(self._test_loader.dataset) if hasattr(self._test_loader, 'dataset') else None
+
         return None
 
     @abstractmethod
@@ -323,7 +330,6 @@ def train(
         if hasattr(logger, "on_train_start"):
             logger.on_train_start()
 
-        self._epochs = epochs
         self._epoch_pbar: Optional[tqdm] = tqdm(
             range(epochs), desc="Training", unit="epoch") if verbose else None
         iterable = self._epoch_pbar if self._epoch_pbar else range(epochs)
@@ -420,6 +426,30 @@ def save_model(
             save_best_model=best_model
         )
 
+    def save_optimizer_state(
+        self, 
+        save_path: pathlib.Path, 
+        file_name_prefix: Optional[str] = None, 
+        file_name_suffix: Optional[str] = None, 
+        file_ext: str = '.pth',
+        recent: bool = True
+    ) -> Optional[List[pathlib.Path]]:
+        """
+        Save the optimizer state to the specified path.
+        """
+        if not recent:
+            raise NotImplementedError(
+                "Saving non-recent optimizer states is not implemented yet."
+            )
+        file_name_suffix = file_name_suffix or 'recent'
+        return save_optimizer_state(
+            trainer=self,
+            save_path=save_path,
+            file_name_prefix=file_name_prefix,
+            file_name_suffix=file_name_suffix,
+            file_ext=file_ext
+        )
+
     """
     Log property
     """
@@ -452,6 +482,18 @@ def val_ratio(self):
     @property
     def test_ratio(self):
         return self._test_ratio
+
+    @property
+    def train_n(self):
+        return self._train_n
+
+    @property
+    def val_n(self):
+        return self._val_n
+
+    @property
+    def test_n(self):
+        return self._test_n
     
     @property
     def model(self):
@@ -468,11 +510,7 @@ def device(self):
     @property
     def batch_size(self):
         return self._batch_size
-    
-    @property
-    def epochs(self):
-        return self._epochs
-    
+        
     @property
     def patience(self):
         return self._patience
diff --git a/src/virtual_stain_flow/trainers/logging_gan_trainer.py b/src/virtual_stain_flow/trainers/logging_gan_trainer.py
index e51cc14..ca9a25e 100644
--- a/src/virtual_stain_flow/trainers/logging_gan_trainer.py
+++ b/src/virtual_stain_flow/trainers/logging_gan_trainer.py
@@ -84,7 +84,6 @@ def __init__(
         super().__init__(
             model=generator, # register generator as main model for early stopping
             optimizer=generator_optimizer,
-            losses=generator_loss_group,
             device=device,
             **kwargs
         )
diff --git a/src/virtual_stain_flow/trainers/trainer_protocol.py b/src/virtual_stain_flow/trainers/trainer_protocol.py
index 5fc3cb0..e0de924 100644
--- a/src/virtual_stain_flow/trainers/trainer_protocol.py
+++ b/src/virtual_stain_flow/trainers/trainer_protocol.py
@@ -44,6 +44,18 @@ def train(self, *args: Any, **kwargs: Any) -> None: ...
     @property
     def epoch(self) -> int: ...
 
+    @property
+    def batch_size(self) -> int: ...
+
+    @property
+    def train_n(self) -> int: ...
+
+    @property
+    def val_n(self) -> int: ...
+
+    @property
+    def test_n(self) -> int: ...
+
     @property
     def device(self) -> torch.device: ...
 
@@ -59,9 +71,19 @@ def best_model(self) -> Optional[torch.nn.Module]: ...
     def save_model(
         self, 
         save_path: pathlib.Path, 
-        file_name_prefix: Optional[str], 
-        file_name_suffix: Optional[str], 
+        file_name_prefix: Optional[str] = None, 
+        file_name_suffix: Optional[str] = None, 
         file_ext: str = '.pth',
         best_model: bool = True,
     ) -> Optional[List[pathlib.Path]]:
         ...
+
+    def save_optimizer_state(
+        self, 
+        save_path: pathlib.Path, 
+        file_name_prefix: Optional[str] = None, 
+        file_name_suffix: Optional[str] = None, 
+        file_ext: str = '.pth',
+        recent: bool = True,
+    ) -> Optional[List[pathlib.Path]]:
+        ...
diff --git a/src/virtual_stain_flow/trainers/trainer_utils/__init__.py b/src/virtual_stain_flow/trainers/trainer_utils/__init__.py
index f0a09fa..a706015 100644
--- a/src/virtual_stain_flow/trainers/trainer_utils/__init__.py
+++ b/src/virtual_stain_flow/trainers/trainer_utils/__init__.py
@@ -7,10 +7,12 @@
     _get_latest_metric_value,
 )
 from .save_model import save_model
+from .save_optimizer import save_optimizer_state
 
 
 __all__ = [
     "EarlyStopHelper",
     "_get_latest_metric_value",
     "save_model",
+    "save_optimizer_state",
 ]
diff --git a/src/virtual_stain_flow/trainers/trainer_utils/save_model.py b/src/virtual_stain_flow/trainers/trainer_utils/save_model.py
index 528d81b..8063199 100644
--- a/src/virtual_stain_flow/trainers/trainer_utils/save_model.py
+++ b/src/virtual_stain_flow/trainers/trainer_utils/save_model.py
@@ -6,12 +6,15 @@
 def save_model(
     trainer: 'TrainerProtocol',
     save_path: Path,
-    file_name_prefix: str = 'generator',
+    file_name_prefix: Optional[str] = None,
     file_name_suffix: Optional[str] = None,
     file_ext: str = '.pth',
     save_best_model: bool = True,
 ) -> List[Path]:
 
+    if file_name_prefix is None:
+        file_name_prefix = 'generator'
+
     if file_name_suffix is None:
         file_name_suffix = 'weights_' + (
             'best' if save_best_model else str(trainer.epoch)
diff --git a/src/virtual_stain_flow/trainers/trainer_utils/save_optimizer.py b/src/virtual_stain_flow/trainers/trainer_utils/save_optimizer.py
new file mode 100644
index 0000000..97e899c
--- /dev/null
+++ b/src/virtual_stain_flow/trainers/trainer_utils/save_optimizer.py
@@ -0,0 +1,39 @@
+from pathlib import Path
+from typing import Optional, List
+
+import torch
+
+from ..trainer_protocol import TrainerProtocol
+
+
+def save_optimizer_state(
+    trainer: 'TrainerProtocol',
+    save_path: Path,
+    file_name_prefix: Optional[str] = None,
+    file_name_suffix: Optional[str] = None,
+    file_ext: str = '.pth',
+) -> List[Path]:
+
+    if file_name_prefix is None:
+        file_name_prefix = 'optimizer'
+
+    if file_name_suffix is None:
+        file_name_suffix = f"{trainer.epoch}"
+
+    optimizer = trainer.optimizer
+
+    if optimizer is None:
+        return []
+
+    save_file = save_path / f"{file_name_prefix}_{file_name_suffix}{file_ext}"
+
+    torch.save(
+        optimizer.state_dict(), 
+        save_file
+    )
+
+    if save_file.exists():
+        return [save_file]
+
+    return []
+ 
diff --git a/tests/trainers/test_abstract_trainer.py b/tests/trainers/test_abstract_trainer.py
index 2cb9099..0d22901 100644
--- a/tests/trainers/test_abstract_trainer.py
+++ b/tests/trainers/test_abstract_trainer.py
@@ -697,8 +697,9 @@ def test_batch_size_property_with_loader_init(self, minimal_model, minimal_optim
             device=torch.device('cpu')
         )
         
-        # When providing loaders, batch_size is set to None
-        assert trainer.batch_size is None
+        # When providing loaders, batch_size is inferred from the train loader
+        # should therefore match
+        assert trainer.batch_size is train_dataloader.batch_size
     
     def test_batch_size_property_default_value(self, minimal_model, minimal_optimizer, dataset_for_splitting):
         """Verify that batch_size property uses default value when not specified."""

From 3c199222918c29d3ff4e5de9d8d6be8161dcc2fa Mon Sep 17 00:00:00 2001
From: Weishan Li <wli51@jh.edu>
Date: Thu, 6 Aug 2026 12:01:52 -0600
Subject: [PATCH 2/3] Implement AutoTrainerLogger for adding auto logging of
 certain useful trainer metadata and couple logging of model weights by
 logging of optimizer state to facilitate resuming of complete trainings in
 whole epoch increments.

---
 .../vsf_logging/MlflowLogger.py               | 123 +++---------------
 .../vsf_logging/auto_loggers/__init__.py      |   2 +
 .../auto_loggers/optimizer_config_logger.py   |  12 +-
 .../auto_loggers/trainer_config_logger.py     |  35 +++++
 .../vsf_logging/logger_utils/__init__.py      |   6 +
 .../vsf_logging/logger_utils/log_artifacts.py |  83 ++++++++++++
 6 files changed, 155 insertions(+), 106 deletions(-)
 create mode 100644 src/virtual_stain_flow/vsf_logging/auto_loggers/trainer_config_logger.py
 create mode 100644 src/virtual_stain_flow/vsf_logging/logger_utils/__init__.py
 create mode 100644 src/virtual_stain_flow/vsf_logging/logger_utils/log_artifacts.py

diff --git a/src/virtual_stain_flow/vsf_logging/MlflowLogger.py b/src/virtual_stain_flow/vsf_logging/MlflowLogger.py
index 8933283..c534378 100644
--- a/src/virtual_stain_flow/vsf_logging/MlflowLogger.py
+++ b/src/virtual_stain_flow/vsf_logging/MlflowLogger.py
@@ -15,7 +15,9 @@
     AutoLossGroupConfigLogger,
     AutoModelConfigLogger,
     AutoOptimizerConfigLogger,
+    AutoTrainerLogger
 )
+from .logger_utils import _log_artifact, _log_trainer_artifact
 from .callbacks.LoggerCallback import (
     AbstractLoggerCallback,
     log_type
@@ -145,6 +147,7 @@ def __init__(
         self._model_config_logger = AutoModelConfigLogger(self)
         self._optimizer_config_logger = AutoOptimizerConfigLogger(self)
         self._loss_group_config_logger = AutoLossGroupConfigLogger(self)
+        self._trainer_logger = AutoTrainerLogger(self)
 
         return None
     
@@ -208,7 +211,7 @@ def on_train_start(self):
         self._model_config_logger.log_model_configs(self.trainer)
         self._optimizer_config_logger.log_optimizer_configs(self.trainer)
         self._loss_group_config_logger.log_loss_group_configs(self.trainer)
-
+        self._trainer_logger.log_trainer_config(self.trainer)
         
         for callback in self.callbacks:
             # TODO consider if we want hasattr checks
@@ -248,16 +251,10 @@ def on_epoch_end(self):
 
         if self._save_model_every_n_epochs is not None:
             if self.trainer.epoch % self._save_model_every_n_epochs == 0:
-                self._save_model_weights(
-                    artifact_path='weights',
-                    best_model=False
-                )
+                _log_trainer_artifact(self.trainer, best_model=False)
 
         if self._save_best_model:
-            self._save_model_weights(
-                artifact_path='weights',
-                best_model=True
-            )
+            _log_trainer_artifact(self.trainer, best_model=True)
 
         # Call on_epoch_end for all registered callbacks
         for callback in self.callbacks:
@@ -323,14 +320,14 @@ def end_run(self):
             print("No active MLflow run to end.")
 
     """
-    Exposed? logging methods
+    Exposed logging methods
     """    
     def log_artifact(
-            self,
-            tag: str,
-            file_path: pathlib.Path,
-            stage: Optional[str] = None
-        ):
+        self,
+        tag: str,
+        file_path: pathlib.Path,
+        stage: Optional[str] = None
+    ):
         """
         Log an artifact to MLflow.
 
@@ -339,29 +336,11 @@ def log_artifact(
         :param stage: Optional stage to categorize the artifact, defaults to None.
         :raises TypeError: If file_path is not a pathlib.Path instance.
         """
-
-        if not isinstance(file_path, pathlib.Path):
-            raise TypeError("file_path must be a pathlib.Path instance.")
-
-        artifact_path = ''
-        artifact_ext = file_path.suffix.lower()
-        if artifact_ext in ['.png', '.jpg', '.jpeg', '.pdf', '.svg']:
-            # log as plot artifact
-            artifact_path += 'plots/'
-        elif artifact_ext in ['.pth', '.pt']:
-            # log as model artifact
-            artifact_path += 'weights/'
-        else:
-            # log as generic artifact
-            artifact_path += 'artifacts/'
-        
-        if stage is not None:
-            artifact_path += f"{stage}/"
-        artifact_path += f"{tag}"
-
-        mlflow.log_artifact(
-            str(file_path), 
-            artifact_path=artifact_path
+        log_subdirs = [stage] if stage is not None else []
+        log_subdirs = log_subdirs + [tag] if tag is not None else log_subdirs
+        _log_artifact(
+            file_path=file_path,
+            artifact_subdirs=log_subdirs
         )
 
     def log_metric(
@@ -464,31 +443,6 @@ def _log_callback_output(
                 continue
                 # raise TypeError("Unsupported callback return type for logging.")
 
-    def _save_model_weights(
-        self,
-        prefix: Optional[str] = None,
-        suffix: Optional[str] = None,
-        artifact_path: str = "weights",
-        best_model: bool = True
-    ):
-        with tempfile.TemporaryDirectory() as tmpdirname:
-            
-            tmpdirpath = pathlib.Path(tmpdirname)
-            
-            saved_file_paths = self.trainer.save_model(
-                save_path=tmpdirpath,
-                file_name_prefix=prefix,
-                file_name_suffix=suffix,
-                file_ext='.pth',
-                best_model=best_model
-            )
-
-            for saved_file_path in (saved_file_paths or []):
-                mlflow.log_artifact(
-                    str(saved_file_path), 
-                    artifact_path=artifact_path
-                )
-
     def log_config(
         self,
         tag: str,
@@ -550,7 +504,6 @@ def __check_trainer_bound(
         if self.trainer is None:
             raise RuntimeError("No trainer bound to logger. Cannot access trainer attributes.")
 
-
     def get_epoch(
             self
         ) -> int:
@@ -595,48 +548,8 @@ def get_model(
     @property
     def run_id(self):
         return self._run_id
-    
-    """
-    Unimplemented helper that might be useful
-    """
-    def _log_dict_as_yaml(
-        self,
-        dict: Dict[str, Any],
-    ):
-        """
-        Log a dictionary as a YAML file in MLflow.
-        
-        :param dict: The dictionary to log.
-        """
-        
-        # TODO implement this method to convert dict to YAML and log it
-        raise NotImplementedError(
-            "log_dict_as_yaml method is not implemented. "
-            "Please implement this method to log dictionary as YAML."
-        )
-    
-    def _log_dict_as_param(
-        self
-    ):
-        """
-        Log a dictionary as parameters in MLflow.
-        
-        :return: None
-        :raises NotImplementedError: If the method is not implemented.
-        
-        This method is intended to log a dictionary as parameters in MLflow.
-        It is currently not implemented and raises a NotImplementedError.
-        """
-
-        # 1. DO flatten dict
-
-        # 2. DO mlflow log
-        
-        raise NotImplementedError(
-            "log_dict_as_param method is not implemented. "
-            "Please implement this method to log dictionary as parameters."
-        )
 
+    
     """
     Overridden destructor method to ensure MLflow run is ended
     """
diff --git a/src/virtual_stain_flow/vsf_logging/auto_loggers/__init__.py b/src/virtual_stain_flow/vsf_logging/auto_loggers/__init__.py
index 2860937..2625d41 100644
--- a/src/virtual_stain_flow/vsf_logging/auto_loggers/__init__.py
+++ b/src/virtual_stain_flow/vsf_logging/auto_loggers/__init__.py
@@ -1,9 +1,11 @@
 from .loss_group_config_logger import AutoLossGroupConfigLogger
 from .model_config_logger import AutoModelConfigLogger
 from .optimizer_config_logger import AutoOptimizerConfigLogger
+from .trainer_config_logger import AutoTrainerLogger
 
 __all__ = [
     "AutoModelConfigLogger",
     "AutoOptimizerConfigLogger",
     "AutoLossGroupConfigLogger",
+    "AutoTrainerLogger",
 ]
diff --git a/src/virtual_stain_flow/vsf_logging/auto_loggers/optimizer_config_logger.py b/src/virtual_stain_flow/vsf_logging/auto_loggers/optimizer_config_logger.py
index fa3f344..c95537e 100644
--- a/src/virtual_stain_flow/vsf_logging/auto_loggers/optimizer_config_logger.py
+++ b/src/virtual_stain_flow/vsf_logging/auto_loggers/optimizer_config_logger.py
@@ -1,4 +1,5 @@
 from typing import Any, Dict, List, Optional
+import inspect
 
 import mlflow
 from torch.optim import Optimizer
@@ -44,13 +45,22 @@ def log_optimizer_configs(
                 continue
 
             try:
+                
+                defaults = dict(optimizer.defaults)
+                init_signature = inspect.signature(optimizer.__class__.__init__)
+                valid_params = init_signature.parameters.keys()
+                
                 opt_config: Optional[Dict[str, Any]] = {
                     "class_path": (
                         f"{optimizer.__class__.__module__}."
                         f"{optimizer.__class__.__name__}"
                     ),
-                    "defaults": dict(optimizer.defaults),
+                    "defaults": defaults,
+                    "init": {
+                        k: v for k, v in defaults.items() if k in valid_params
+                    }
                 }
+
             except Exception as e:
                 print(f"Could not get optimizer config for logging: {e}")
                 opt_config = None
diff --git a/src/virtual_stain_flow/vsf_logging/auto_loggers/trainer_config_logger.py b/src/virtual_stain_flow/vsf_logging/auto_loggers/trainer_config_logger.py
new file mode 100644
index 0000000..cf6c292
--- /dev/null
+++ b/src/virtual_stain_flow/vsf_logging/auto_loggers/trainer_config_logger.py
@@ -0,0 +1,35 @@
+from typing import Any, Optional
+
+from ...trainers.trainer_protocol import TrainerProtocol
+
+
+class AutoTrainerLogger:
+    """
+    Auto-log trainer metadata to MLflow.
+    """
+
+    def __init__(self, logger: Any) -> None:
+        self._logger = logger
+
+    def log_trainer_config(self, trainer: Optional[TrainerProtocol]) -> None:
+        
+        if trainer is None:
+            return
+
+        config = {
+            "class_path": f"{trainer.__class__.__module__}.{trainer.__class__.__name__}",
+            "device": str(trainer.device), # device used for training
+            "batch_size": trainer.batch_size, # batch size used for training
+            "train_n": trainer.train_n,
+            "val_n": trainer.val_n,
+            "test_n": trainer.test_n,
+        }
+
+        try:
+            self._logger.log_config(
+                tag="trainer",
+                config=config,
+                stage=None,
+            )
+        except Exception as e:
+            print(f"Could not log trainer config: {e}")
diff --git a/src/virtual_stain_flow/vsf_logging/logger_utils/__init__.py b/src/virtual_stain_flow/vsf_logging/logger_utils/__init__.py
new file mode 100644
index 0000000..04ad9c8
--- /dev/null
+++ b/src/virtual_stain_flow/vsf_logging/logger_utils/__init__.py
@@ -0,0 +1,6 @@
+from .log_artifacts import _log_artifact, _log_trainer_artifact
+
+__all__ = [
+    "_log_artifact",
+    "_log_trainer_artifact",
+]
diff --git a/src/virtual_stain_flow/vsf_logging/logger_utils/log_artifacts.py b/src/virtual_stain_flow/vsf_logging/logger_utils/log_artifacts.py
new file mode 100644
index 0000000..199fe7c
--- /dev/null
+++ b/src/virtual_stain_flow/vsf_logging/logger_utils/log_artifacts.py
@@ -0,0 +1,83 @@
+from typing import Optional, List
+import pathlib
+import tempfile
+
+import mlflow
+
+from virtual_stain_flow.trainers.trainer_protocol import TrainerProtocol
+
+
+def _log_artifact(
+    file_path: pathlib.Path,
+    artifact_path: Optional[str] = None,
+    artifact_subdirs: Optional[List[str]] = None
+) -> None:
+    """
+    Logs a single artifact to MLflow.
+
+    :param file_path: The path to the file to log as an artifact.
+    :param artifact_path: Optional artifact path within the MLflow run, defaults to None.
+    :param artifact_subdirs: Optional list of subdirectories to include in the artifact path, defaults to None.
+    :raises TypeError: If file_path is not a pathlib.Path instance.
+    """
+
+    if not isinstance(file_path, pathlib.Path):
+        raise TypeError("file_path must be a pathlib.Path instance.")
+
+    if artifact_path is None:
+        artifact_ext = file_path.suffix.lower()
+        if artifact_ext in ['.png', '.jpg', '.jpeg', '.pdf', '.svg']:
+            # log as plot artifact
+            artifact_path = 'plots'
+        elif artifact_ext in ['.pth', '.pt']:
+            # log as model artifact
+            artifact_path = 'weights'
+        else:
+            # log as generic artifact
+            artifact_path = 'artifacts'
+
+    artifact_subdirs = [] if artifact_subdirs is None else artifact_subdirs
+
+    path_parts = [artifact_path, *artifact_subdirs]
+    clean_parts = [
+        part.replace('\\', '/').strip('/')
+        for part in path_parts
+        if part and part.strip('/\\')
+    ]
+
+    # Build a lexical POSIX path for MLflow without resolving it on the host OS.
+    artifact_path = str(pathlib.PurePosixPath(*clean_parts)) if clean_parts else ''
+
+    mlflow.log_artifact(str(file_path), artifact_path=artifact_path)
+
+
+def _log_trainer_artifact(
+    trainer: "TrainerProtocol",
+    best_model: bool = True,
+) -> None:
+    """
+    Logs the model and optimizer state artifacts from a trainer to MLflow.
+    The most recent optimizer state is saved and logged regardless of the best_model flag.
+
+    :param trainer: The trainer instance adhering to TrainerProtocol.
+    :param best_model: Whether to log only the best model, defaults to True.
+    :raises TypeError: If the provided trainer does not adhere to TrainerProtocol.
+    """
+
+    if not isinstance(trainer, TrainerProtocol):
+        raise TypeError("The provided trainer must adhere to the TrainerProtocol.")
+    
+    with tempfile.TemporaryDirectory() as tmpdirname:
+        
+        tmpdirpath = pathlib.Path(tmpdirname)
+        saved_model_paths = trainer.save_model(
+            save_path=tmpdirpath, best_model=best_model
+        )
+        for saved_model_path in (saved_model_paths or []):
+            _log_artifact(saved_model_path, artifact_path='weights')
+
+        saved_optimizer_paths = trainer.save_optimizer_state(
+            save_path=tmpdirpath, recent=True
+        )
+        for saved_optimizer_path in (saved_optimizer_paths or []):
+            _log_artifact(saved_optimizer_path, artifact_path='optimizer')

From bd3036b49eabf703a4c0a06ecf9a06d7e4839c7c Mon Sep 17 00:00:00 2001
From: Weishan Li <wli51@jh.edu>
Date: Thu, 6 Aug 2026 12:12:53 -0600
Subject: [PATCH 3/3] Add unit tests for model and optimizer state saving
 functionality

---
 .../trainers/trainer_utils/test_save_model.py | 23 +++++++++++++++++++
 .../trainer_utils/test_save_optimizer.py      | 23 +++++++++++++++++++
 2 files changed, 46 insertions(+)
 create mode 100644 tests/trainers/trainer_utils/test_save_model.py
 create mode 100644 tests/trainers/trainer_utils/test_save_optimizer.py

diff --git a/tests/trainers/trainer_utils/test_save_model.py b/tests/trainers/trainer_utils/test_save_model.py
new file mode 100644
index 0000000..be65008
--- /dev/null
+++ b/tests/trainers/trainer_utils/test_save_model.py
@@ -0,0 +1,23 @@
+"""Standalone tests for save_model helper."""
+
+from types import SimpleNamespace
+
+from virtual_stain_flow.trainers.trainer_utils.save_model import save_model
+
+
+def test_save_model_returns_empty_when_target_model_is_none(tmp_path):
+    trainer = SimpleNamespace(model=None, best_model=None, epoch=2)
+
+    saved_paths = save_model(trainer=trainer, save_path=tmp_path, save_best_model=True)
+
+    assert saved_paths == []
+
+
+def test_save_model_saves_current_model_with_default_name(mock_model_with_save, tmp_path):
+    trainer = SimpleNamespace(model=mock_model_with_save, best_model=None, epoch=3)
+
+    saved_paths = save_model(trainer=trainer, save_path=tmp_path, save_best_model=False)
+
+    assert len(saved_paths) == 1
+    assert saved_paths[0].exists()
+    assert saved_paths[0].name == "generator_weights_3.pth"
diff --git a/tests/trainers/trainer_utils/test_save_optimizer.py b/tests/trainers/trainer_utils/test_save_optimizer.py
new file mode 100644
index 0000000..dac7d52
--- /dev/null
+++ b/tests/trainers/trainer_utils/test_save_optimizer.py
@@ -0,0 +1,23 @@
+"""Standalone tests for save_optimizer_state helper."""
+
+from types import SimpleNamespace
+
+from virtual_stain_flow.trainers.trainer_utils.save_optimizer import save_optimizer_state
+
+
+def test_save_optimizer_state_returns_empty_for_missing_optimizer(tmp_path):
+    trainer = SimpleNamespace(optimizer=None, epoch=7)
+
+    saved_paths = save_optimizer_state(trainer=trainer, save_path=tmp_path)
+
+    assert saved_paths == []
+
+
+def test_save_optimizer_state_saves_with_default_name(minimal_optimizer, tmp_path):
+    trainer = SimpleNamespace(optimizer=minimal_optimizer, epoch=7)
+
+    saved_paths = save_optimizer_state(trainer=trainer, save_path=tmp_path)
+
+    assert len(saved_paths) == 1
+    assert saved_paths[0].exists()
+    assert saved_paths[0].name == "optimizer_7.pth"