[BUGFIX] Stabilize the default test suite and bugfixes

tests/hparams_optim/test_mlflow_with_optuna.py → examples/example_mlflow_optuna_cifar10_demo.py

@@ -1,5 +1,11 @@
-''' Script created by PeterC to test mlflowand optuna integration library for model monitoring and tracking - 02-07-2024 '''
+"""MLflow and Optuna CIFAR10 demo script.
 
+This is intentionally an example, not a pytest module: it downloads CIFAR10,
+starts long optimization runs, and may start external tracking services.
+"""
+
+import numpy as np
+import pyTorchAutoForge  # Custom torch tools
 import optuna
 import mlflow
 import mlflow.pytorch
@@ -11,41 +17,48 @@
 import matplotlib.pyplot as plt
 
 # Import modules
-import os, subprocess, time, logging
+import os
+import subprocess
+import time
+import logging
 
 # Set up logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger()
 logger.setLevel(logging.INFO)
 
-import pyTorchAutoForge # Custom torch tools
-import numpy as np
 
-def StartMLflowUI(port:int=5000):
+def StartMLflowUI(port: int = 5000):
 
     # Start MLflow UI
     os.system('mlflow ui --port ' + str(port))
-    process = subprocess.Popen(['mlflow', 'ui', '--port ' + f'{port}', '&'], 
+    process = subprocess.Popen(['mlflow', 'ui', '--port ' + f'{port}', '&'],
                                stdout=subprocess.PIPE, stderr=subprocess.PIPE)
     print(f'MLflow UI started with PID: {process.pid}, on port: {port}')
-    time.sleep(1) # Ensure the server has started
+    time.sleep(1)  # Ensure the server has started
     if process.poll() is None:
         print('MLflow UI is running OK.')
     else:
         raise RuntimeError('MLflow UI failed to start. Run stopped.')
 
     return process
 
+
 # %% Datasets loading (global)
 # Load CIFAR-10 dataset from torchvision
-transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
-train_dataset = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
-test_dataset = datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)
+transform = transforms.Compose(
+    [transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
+train_dataset = datasets.CIFAR10(
+    root='./data', train=True, download=True, transform=transform)
+test_dataset = datasets.CIFAR10(
+    root='./data', train=False, download=True, transform=transform)
 train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
 test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False)
 
 # TEST EXAMPLE BY GPT
 # %% Optuna model for trial optimization
+
+
 class SimpleCNN(nn.Module):
     def __init__(self, trial):
         super(SimpleCNN, self).__init__()
@@ -56,33 +69,37 @@ def __init__(self, trial):
 
         # NOTE: What are the second inputs to suggest_int? Are they the lower and upper bounds?
 
-        in_channels = 3 # Number of channels in the input image (RGB)
+        in_channels = 3  # Number of channels in the input image (RGB)
 
         # Add convolutional blocks to the model up to "num_conv_layers" --> optimization parameter
         for i in range(num_conv_layers):
-            
+
             # Number of output channels for the ith convolutional block
             out_channels = trial.suggest_int(f'filters_{i}', 16, 64)
             # Kernel size for the ith convolutional block
             kernel_size = trial.suggest_int(f'kernel_size_{i}', 3, 5)
 
             # Add convolutional block to the model. layers.add_module() takes input:(name, module)
             # NOTE: this is easily repurposed for the Model AutoBuilder
-            self.layers.add_module(f'conv{i}', nn.Conv2d(in_channels, out_channels, kernel_size))
+            self.layers.add_module(f'conv{i}', nn.Conv2d(
+                in_channels, out_channels, kernel_size))
             self.layers.add_module(f'relu{i}', nn.ReLU())
             self.layers.add_module(f'maxpool{i}', nn.MaxPool2d(2))
 
-            in_channels = out_channels # Number of input channels to conv2d is the number of output channels from the previous conv2d
+            # Number of input channels to conv2d is the number of output channels from the previous conv2d
+            in_channels = out_channels
 
-        self.layers.add_module('flatten', nn.Flatten()) # Add flatten before NN
+        self.layers.add_module('flatten', nn.Flatten()
+                               )  # Add flatten before NN
 
         # Dense layers
         num_dense_layers = trial.suggest_int('num_dense_layers', 1, 3)
         in_features = self._get_conv_output((3, 32, 32))
 
         for i in range(num_dense_layers):
             out_features = trial.suggest_int(f'units_{i}', 32, 256)
-            self.layers.add_module(f'fc{i}', nn.Linear(in_features, out_features))
+            self.layers.add_module(
+                f'fc{i}', nn.Linear(in_features, out_features))
             self.layers.add_module(f'relu_fc{i}', nn.ReLU())
             dropout_rate = trial.suggest_float(f'dropout_rate_{i}', 0.1, 0.5)
             self.layers.add_module(f'dropout{i}', nn.Dropout(dropout_rate))
@@ -101,6 +118,8 @@ def _get_conv_output(self, shape):
             return output.size(1)
 
 # %% Optuna objective function
+
+
 def objective(trial):
     # Create new run_IN within the current session
     with mlflow.start_run():
@@ -111,11 +130,13 @@ def objective(trial):
         model = SimpleCNN(trial).to(device)
 
         # Select the optimizer
-        optimizer_name = trial.suggest_categorical('optimizer', ['Adam', 'RMSprop'])
-        lr = trial.suggest_float('lr', 1e-6, 1e-2, log=True) # NOTE: What are the inputs to suggest_float?
+        optimizer_name = trial.suggest_categorical(
+            'optimizer', ['Adam', 'RMSprop'])
+        # NOTE: What are the inputs to suggest_float?
+        lr = trial.suggest_float('lr', 1e-6, 1e-2, log=True)
         optimizer = getattr(optim, optimizer_name)(model.parameters(), lr=lr)
         # NOTE: What is "getattr" function? It should be a pytorch function: likely a method of "model", TBC
-    
+
         # Loss function definition
         criterion = nn.CrossEntropyLoss()
 
@@ -137,12 +158,12 @@ def objective(trial):
             mlflow.log_param(f'units_{i}', trial.params.get(f'units_{i}'))
             mlflow.log_param(
                 f'dropout_rate_{i}', trial.params.get(f'dropout_rate_{i}'))
-            
+
         # Training the current trial model
         for epoch in range(10):
             model.train()
             for batch in train_loader:
-                inputs, targets = batch # Unpack tuples
+                inputs, targets = batch  # Unpack tuples
                 inputs, targets = inputs.to(device), targets.to(device)
                 optimizer.zero_grad()
 
@@ -168,19 +189,21 @@ def objective(trial):
                     correct += (predicted == targets).sum().item()
 
             accuracy = correct / total
-            
+
             # Log the accuracy of the model using mlflow
             mlflow.log_metric('Accuracy value', accuracy, step=epoch)
 
-            trial.report(accuracy, epoch) # Report the accuracy of the model to optuna pruner
+            # Report the accuracy of the model to optuna pruner
+            trial.report(accuracy, epoch)
 
             if trial.should_prune():
                 # Mark the run as killed in mlflow
                 mlflow.end_run(status='KILLED')
-                raise optuna.TrialPruned() # Raise an optuna exception to stop the trial due to pruning
-
+                # Raise an optuna exception to stop the trial due to pruning
+                raise optuna.TrialPruned()
+
         # MLflow: log trial data (it will have a unique ID)
-        # NOTE: this could be matched with the AutoBuilder configuration such that the parameters 
+        # NOTE: this could be matched with the AutoBuilder configuration such that the parameters
         # of config file are automatically logged by mlflow with their names
         mlflow.end_run(status='FINISHED')
 
@@ -191,42 +214,45 @@ def main():
     print('---------------------------- TEST SCRIPT: MLflow and Optuna functionalities ----------------------------\n')
     # %% MLflow tracking initialization
     port = 7000
-    #StartMLflowUI(port) # Start MLflow UI
+    # StartMLflowUI(port) # Start MLflow UI
 
     # TODO: add change of script execution working directory to the current directory
-    
+
     # %% Optuna study configuration
     if not (os.path.exists('testdata/optuna_db')):
         os.makedirs('testdata/optuna_db')
 
     studyName = 'CIFAR10_CNN_OptimizationExample'
     optunaStudyObj = optuna.create_study(study_name=studyName,
-                                         storage='sqlite:///testdata/{studyName}.db'.format(studyName=os.path.join('optuna_db', studyName)),
+                                         storage='sqlite:///testdata/{studyName}.db'.format(
+                                             studyName=os.path.join('optuna_db', studyName)),
                                          load_if_exists=True,
                                          direction='maximize',
                                          sampler=optuna.samplers.TPESampler(),
                                          pruner=optuna.pruners.SuccessiveHalvingPruner(min_resource=1, reduction_factor=2,
                                                                                        min_early_stopping_rate=1))
-    
+
     # Mlflow experiment name
     mlflow.set_experiment(studyName)
 
     # %% Optuna optimization
     optunaStudyObj.optimize(objective, n_trials=100, timeout=1800)
-    
+
     # Print the best trial
-    print('Number of finished trials:', len(optunaStudyObj.trials)) # Get number of finished trials
+    # Get number of finished trials
+    print('Number of finished trials:', len(optunaStudyObj.trials))
     print('Best trial:')
 
-    trial = optunaStudyObj.best_trial # Get the best trial from study object
-    print('  Value: {:.4f}'.format(trial.value)) # Loss function value for the best trial
+    trial = optunaStudyObj.best_trial  # Get the best trial from study object
+    # Loss function value for the best trial
+    print('  Value: {:.4f}'.format(trial.value))
 
     # Print parameters of the best trial
     print('  Params: ')
     for key, value in trial.params.items():
         print('    {}: {}'.format(key, value))
-    
+
 
 # %% Main call
 if __name__ == "__main__":
-    main()
+    main()

pyTorchAutoForge/api/torch/torchModulesIO.py

@@ -36,7 +36,7 @@ def SaveModel(model: torch.nn.Module,
               model_filename: str | pathlib.Path, 
               save_mode : AutoForgeModuleSaveMode | str = AutoForgeModuleSaveMode.MODEL_ARCH_STATE, 
               example_input: torch.Tensor | None = None, 
-              target_device: str = 'cpu', 
+              target_device: str | torch.device = 'cpu',
               model_base_name : str | None = None) -> None:
     """
     Saves a PyTorch model to a file.
@@ -57,7 +57,7 @@ def SaveModel(model: torch.nn.Module,
             Defaults to AutoForgeModuleSaveMode.MODEL_ARCH_STATE.
         example_input (torch.Tensor | None, optional): A sample input tensor for tracing or scripting.
             Defaults to None.
-        target_device (str, optional): The target device (e.g., 'cpu' or 'cuda:0') to save the model.
+        target_device (str | torch.device, optional): The target device (e.g., 'cpu' or 'cuda:0') to save the model.
             Defaults to 'cpu'.
         model_base_name (str | None, optional): An optional base name for the model.
             Defaults to None.
@@ -94,9 +94,10 @@ def SaveModel(model: torch.nn.Module,
     else: 
         extension = '.pth'
 
+    target_device = torch.device(target_device)
+
     # Format target device string to remove ':' from name
-    target_device_name = target_device
-    target_device_name = target_device_name.replace(':', '')
+    target_device_name = str(target_device).replace(':', '')
 
     # Form filename for saving
     # Check if device is in model name and remove it
@@ -283,5 +284,3 @@ def ValidateDictLoading(model: torch.nn.Module | torch.nn.ModuleDict | torch.nn.
 
     else:
         print("All model parameters are correctly loaded.")
-
-

tests/__init__.py

@@ -0,0 +1 @@
+"""Shared pytest support package for local tests."""

tests/api/tcp/test_start_server.py

@@ -1,50 +1,64 @@
-import numpy as np
+from __future__ import annotations
 
-# Custom imports
-from pyTorchAutoForge.api.tcp import DataProcessor, pytcp_server, pytcp_requestHandler, ProcessingMode
 import threading
 
-# MAIN SCRIPT
-def main():
-    print('\n\n----------------------------------- RUNNING: test_start_server.py -----------------------------------\n')
-
-    # %% TCP SERVER INITIALIZATION
-    HOST1, PORT1 = "127.0.0.1", 50000 # Define host and port for the first server
-    HOST2, PORT2 = "127.0.0.1", 50001 # Define host and port for the second server
-
-    def dummy_function(inputData):
-        return inputData
-
-    # Define DataProcessor object for RequestHandler
-    dataProcessorObj = DataProcessor(dummy_function, np.float32, 1024,
-                                        ENDIANNESS='little', DYNAMIC_BUFFER_MODE=True,
-                                        PRE_PROCESSING_MODE=ProcessingMode.TENSOR)
-
-    dataProcessorObj_multi = DataProcessor(dummy_function, np.float32, 1024, 
-                                                 ENDIANNESS='little', DYNAMIC_BUFFER_MODE=True, 
-                                                 PRE_PROCESSING_MODE=ProcessingMode.MULTI_TENSOR)
-
-    def start_server(host, port, dataProcessorObj):
-        with pytcp_server((host, port), pytcp_requestHandler, dataProcessorObj, bindAndActivate=True) as server:
-            try:
-                print(f'\nServer initialized correctly on {host}:{port}. Set in "serve_forever" mode.')
-                server.serve_forever()
-            except KeyboardInterrupt:
-                print(f"\nServer on {host}:{port} is gracefully shutting down =D.")
-                server.shutdown()
-                server.server_close()
-
-    # Start two servers on separate threads
-    thread1 = threading.Thread(target=start_server, args=( HOST1, PORT1, dataProcessorObj) )
-    thread2 = threading.Thread(target=start_server, args=( HOST2, PORT2, dataProcessorObj_multi) )
-
-    thread1.start()
-    thread2.start()
-
-    # Wait for the threads to finish processing
-    thread1.join()
-    thread2.join()
-
-if __name__ == "__main__":
-    main()
+import numpy as np
+
+from pyTorchAutoForge.api.tcp import (
+    DataProcessor,
+    ProcessingMode,
+    pytcp_requestHandler,
+    pytcp_server,
+)
+
+
+def _Identity(input_data_: np.ndarray) -> np.ndarray:
+    return input_data_
+
+
+def test_data_processor_tensor_roundtrip() -> None:
+    processor_ = DataProcessor(
+        _Identity,
+        np.float32,
+        1024,
+        ENDIANNESS="little",
+        DYNAMIC_BUFFER_MODE=True,
+        PRE_PROCESSING_MODE=ProcessingMode.TENSOR,
+    )
+    input_array_ = np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32)
+
+    output_buffer_ = processor_.process(
+        processor_.TensorToBytesBuffer(input_array_))
+    output_array_, output_shape_ = processor_.BytesBufferToTensor(output_buffer_[
+                                                                  4:])
+
+    assert output_shape_ == input_array_.shape
+    assert np.array_equal(output_array_, input_array_)
+
+
+def test_tcp_server_starts_and_shuts_down_without_hanging() -> None:
+    processor_ = DataProcessor(
+        _Identity,
+        np.float32,
+        1024,
+        ENDIANNESS="little",
+        DYNAMIC_BUFFER_MODE=True,
+        PRE_PROCESSING_MODE=ProcessingMode.TENSOR,
+    )
+
+    with pytcp_server(
+        ("127.0.0.1", 0),
+        pytcp_requestHandler,
+        processor_,
+        bindAndActivate=True,
+    ) as server_:
+        thread_ = threading.Thread(target=server_.serve_forever, daemon=True)
+        thread_.start()
+
+        assert server_.server_address[0] == "127.0.0.1"
+        assert server_.server_address[1] > 0
+
+        server_.shutdown()
+        thread_.join(timeout=2.0)
 
+    assert not thread_.is_alive()