From aecebab23ce18745d337bd07747dd917461760c1 Mon Sep 17 00:00:00 2001 From: PeterC Date: Wed, 20 May 2026 15:56:15 +0200 Subject: [PATCH 1/7] Add utils for tests setup (pytest) --- tests/__init__.py | 1 + tests/conftest.py | 37 +++++++++++++++++++++++++++++++++++++ tests/helpers.py | 27 +++++++++++++++++++++++++++ 3 files changed, 65 insertions(+) create mode 100644 tests/__init__.py create mode 100644 tests/helpers.py diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..a18217b --- /dev/null +++ b/tests/__init__.py @@ -0,0 +1 @@ +"""Shared pytest support package for local tests.""" diff --git a/tests/conftest.py b/tests/conftest.py index 6fa9f72..063936e 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -19,6 +19,43 @@ ) +def pytest_addoption(parser: pytest.Parser) -> None: + parser.addoption( + "--run-slow", + action="store_true", + default=False, + help="Run tests marked slow.", + ) + parser.addoption( + "--run-gpu", + action="store_true", + default=False, + help="Run tests marked gpu.", + ) + parser.addoption( + "--run-visual", + action="store_true", + default=False, + help="Run tests marked visual.", + ) + + +def pytest_collection_modifyitems(config: pytest.Config, + items: list[pytest.Item], + ) -> None: + marker_options_ = { + "slow": ("--run-slow", "slow test skipped by default"), + "gpu": ("--run-gpu", "GPU test skipped by default"), + "visual": ("--run-visual", "visual test skipped by default"), + } + + for item_ in items: + for marker_name_, (option_name_, reason_) in marker_options_.items(): + if marker_name_ in item_.keywords and not config.getoption(option_name_): + item_.add_marker(pytest.mark.skip(reason=reason_)) + break + + def _cleanup_legacy_output_dirs() -> None: repo_root_resolved_ = REPO_ROOT.resolve() for output_dir_ in LEGACY_TEST_OUTPUT_DIRS: diff --git a/tests/helpers.py b/tests/helpers.py new file mode 100644 index 0000000..183b671 --- /dev/null +++ b/tests/helpers.py @@ -0,0 +1,27 @@ +from __future__ import annotations + +from functools import lru_cache + +import pytest +import torch + + +@lru_cache(maxsize=1) +def CudaIsUsable() -> bool: + """Return True only when CUDA kernels can execute, not just be discovered.""" + if not torch.cuda.is_available(): + return False + + try: + device_ = torch.device("cuda") + tensor_ = torch.ones(1, device=device_) + result_ = (tensor_ + 1.0).item() + torch.cuda.synchronize(device_) + return result_ == 2.0 + except Exception: + return False + + +def RequireCudaUsable() -> None: + if not CudaIsUsable(): + pytest.skip("CUDA runtime is not usable in this environment.") From dfe2f8423bd42e31e02e48908d4fa4c86f70756a Mon Sep 17 00:00:00 2001 From: PeterC Date: Wed, 20 May 2026 15:56:44 +0200 Subject: [PATCH 2/7] Cleanup old tests --- tests/api/tcp/test_torch_model_over_tcp.py | 8 - tests/evaluation/test_ModelExplainer.py | 1 - .../hparams_optim/test_mlflow_with_optuna.py | 232 ------------------ tests/model_building/test_TorchModel.py | 13 - 4 files changed, 254 deletions(-) delete mode 100644 tests/api/tcp/test_torch_model_over_tcp.py delete mode 100644 tests/evaluation/test_ModelExplainer.py delete mode 100644 tests/hparams_optim/test_mlflow_with_optuna.py delete mode 100644 tests/model_building/test_TorchModel.py diff --git a/tests/api/tcp/test_torch_model_over_tcp.py b/tests/api/tcp/test_torch_model_over_tcp.py deleted file mode 100644 index 57639ba..0000000 --- a/tests/api/tcp/test_torch_model_over_tcp.py +++ /dev/null @@ -1,8 +0,0 @@ -import pyTorchAutoForge - -def main(): - pass - - -if __name__ == "__main__": - main() \ No newline at end of file diff --git a/tests/evaluation/test_ModelExplainer.py b/tests/evaluation/test_ModelExplainer.py deleted file mode 100644 index f87f5c1..0000000 --- a/tests/evaluation/test_ModelExplainer.py +++ /dev/null @@ -1 +0,0 @@ -# TODO \ No newline at end of file diff --git a/tests/hparams_optim/test_mlflow_with_optuna.py b/tests/hparams_optim/test_mlflow_with_optuna.py deleted file mode 100644 index 317378f..0000000 --- a/tests/hparams_optim/test_mlflow_with_optuna.py +++ /dev/null @@ -1,232 +0,0 @@ -''' Script created by PeterC to test mlflowand optuna integration library for model monitoring and tracking - 02-07-2024 ''' - -import optuna -import mlflow -import mlflow.pytorch -import torch -import torch.nn as nn -import torch.optim as optim -from torch.utils.data import DataLoader -from torchvision import datasets, transforms -import matplotlib.pyplot as plt - -# Import modules -import os, subprocess, time, 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): - - # Start MLflow UI - os.system('mlflow ui --port ' + str(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 - 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) -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__() - self.layers = nn.Sequential() - - # Convolutional layers - num_conv_layers = trial.suggest_int('num_conv_layers', 1, 2) - - # 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) - - # 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'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 - - 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'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)) - in_features = out_features - - self.layers.add_module('output', nn.Linear(in_features, 10)) - - def forward(self, x): - return self.layers(x) - - # Function to compute the shape of the output of the a convolutional layer - def _get_conv_output(self, shape): - with torch.no_grad(): - input_data = torch.rand(1, *shape) - output = self.layers(input_data) - return output.size(1) - -# %% Optuna objective function -def objective(trial): - # Create new run_IN within the current session - with mlflow.start_run(): - - device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') - - # Create the trial model - 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 = 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() - - # MLflow: log trial data - mlflow.log_param('optimizer', optimizer_name) - mlflow.log_param('learning_rate', lr) - mlflow.log_param('num_conv_layers', - trial.params.get('num_conv_layers')) - # NOTE: integration with optuna --> parameters are got directly from the ith trial object under evaluation - mlflow.log_param('num_dense_layers', - trial.params.get('num_dense_layers')) - - for i in range(trial.params.get('num_conv_layers')): - mlflow.log_param(f'filters_{i}', trial.params.get(f'filters_{i}')) - mlflow.log_param( - f'kernel_size_{i}', trial.params.get(f'kernel_size_{i}')) - - for i in range(trial.params.get('num_dense_layers')): - 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 = inputs.to(device), targets.to(device) - optimizer.zero_grad() - - # Evaluate model - outputs = model(inputs) - # Compute loss - loss = criterion(outputs, targets) - # Perform backpropagation - loss.backward() - optimizer.step() - - # Model validation - model.eval() - correct = 0 - total = 0 - with torch.no_grad(): - for batch in test_loader: - inputs, targets = batch - inputs, targets = inputs.to(device), targets.to(device) - outputs = model(inputs) - _, predicted = torch.max(outputs.data, 1) - total += targets.size(0) - 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 - - 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 - - # MLflow: log trial data (it will have a unique ID) - # 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') - - return accuracy - - -def main(): - print('---------------------------- TEST SCRIPT: MLflow and Optuna functionalities ----------------------------\n') - # %% MLflow tracking initialization - port = 7000 - #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)), - 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 - 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 - - # 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() \ No newline at end of file diff --git a/tests/model_building/test_TorchModel.py b/tests/model_building/test_TorchModel.py deleted file mode 100644 index c5fe3d3..0000000 --- a/tests/model_building/test_TorchModel.py +++ /dev/null @@ -1,13 +0,0 @@ -import torch -from pyTorchAutoForge.utils import GetDevice -from pyTorchAutoForge.api.torch.torchModulesIO import * -from pyTorchAutoForge.model_building import AutoForgeModule - - - -def main(): - pass - -if __name__ == '__main__': - main() - From 19955be059c43f16e5ae4df4737ce62469220b0d Mon Sep 17 00:00:00 2001 From: PeterC Date: Wed, 20 May 2026 15:57:31 +0200 Subject: [PATCH 3/7] Improve markers and gates for exec --- tests/datasets/test_ImagesAugmentation.py | 27 +++++++++++++------ tests/datasets/test_auxiliary_functions.py | 5 ++++ .../test_image_processing_operators.py | 3 +++ .../test_model_training_manager.py | 7 ++++- 4 files changed, 33 insertions(+), 9 deletions(-) diff --git a/tests/datasets/test_ImagesAugmentation.py b/tests/datasets/test_ImagesAugmentation.py index c57c04b..40741ac 100644 --- a/tests/datasets/test_ImagesAugmentation.py +++ b/tests/datasets/test_ImagesAugmentation.py @@ -10,6 +10,7 @@ import matplotlib.pyplot as plt from pyTorchAutoForge.datasets import ImagesLabelsCachedDataset, AugmentationConfig, ImageAugmentationsHelper from pyTorchAutoForge.utils.conversion_utils import torch_to_numpy, numpy_to_torch +from tests.helpers import CudaIsUsable, RequireCudaUsable import numpy as np from time import perf_counter import PIL @@ -194,6 +195,9 @@ def test_detect_border_crossing_white_masks(): # %% Integrated tests +@pytest.mark.integration +@pytest.mark.slow +@pytest.mark.visual def test_synthetic_mask_augmentation(): augs_datakey = [DataKey.IMAGE, DataKey.KEYPOINTS] @@ -219,7 +223,7 @@ def test_synthetic_mask_augmentation(): contrast_aug_prob=1.0, input_normalization_factor=255.0, enable_auto_input_normalization=True, - device='cuda' if torch.cuda.is_available() else 'cpu' + device='cuda' if CudaIsUsable() else 'cpu' ) augs_helper = ImageAugmentationsHelper(cfg) @@ -281,6 +285,9 @@ def test_synthetic_mask_augmentation(): plt.close() +@pytest.mark.integration +@pytest.mark.slow +@pytest.mark.visual def test_sample_images_augmentation(): # Load sample images imgs = _load_sample_images() @@ -330,7 +337,7 @@ def test_sample_images_augmentation(): input_normalization_factor=255.0, enable_auto_input_normalization=True, enable_batch_validation_check=True, - device='cuda' if torch.cuda.is_available() else 'cpu' + device='cuda' if CudaIsUsable() else 'cpu' ) augs_helper = ImageAugmentationsHelper(cfg) @@ -380,6 +387,7 @@ def test_sample_images_augmentation(): plt.close() +@pytest.mark.visual def test_random_softbinarize_only_sample_images_with_viz(): # Load sample images np.random.seed(0) @@ -455,6 +463,8 @@ def test_random_softbinarize_only_sample_images_with_viz(): plt.close() +@pytest.mark.integration +@pytest.mark.visual def test_AugmentationSequential(): from kornia.augmentation import AugmentationSequential, RandomAffine, RandomHorizontalFlip from kornia.constants import DataKey @@ -542,7 +552,7 @@ def augment_data_batch(*inputs: torch.Tensor) -> tuple[torch.Tensor, ...]: # Chain parametrize to test combinations -@pytest.mark.parametrize("device", ["cpu", "cuda"]) +@pytest.mark.parametrize("device", ["cpu", pytest.param("cuda", marks=pytest.mark.gpu)]) @pytest.mark.parametrize("shift_aug_prob", [0, 1]) @pytest.mark.parametrize("rotation_aug_prob", [0, 1]) @pytest.mark.parametrize("gaussian_noise_aug_prob", [0, 1]) @@ -557,8 +567,8 @@ def test_augmentation_helper_preserves_device(device, brightness_aug_prob, contrast_aug_prob) -> None: - if device == "cuda" and not torch.cuda.is_available(): - pytest.skip("CUDA not available.") + if device == "cuda": + RequireCudaUsable() # Create a minimal configuration for the augmentations helper. cfg = AugmentationConfig( @@ -601,10 +611,11 @@ def test_augmentation_helper_preserves_device(device, _assert_helper_modules_on_device(augs_helper, device_t) +@pytest.mark.gpu +@pytest.mark.slow def test_augmentation_helper_timing_cpu_vs_cuda(): """Test execution time cpu vs cuda for augmentations helper""" - if not torch.cuda.is_available(): - pytest.skip("CUDA not available.") + RequireCudaUsable() num_iters_ = 500 cfg_kwargs = dict( max_shift_img_fraction=(0.2, 0.2), @@ -1106,7 +1117,7 @@ def test_helper_metadata_supports_external_sun_direction_label_update(): # %% MANUAL TEST CALLS if __name__ == '__main__': - device = "cuda" if torch.cuda.is_available() else "cpu" + device = "cuda" if CudaIsUsable() else "cpu" shift_aug_prob = 0.0 rotation_aug_prob = 0.0 gaussian_noise_aug_prob = 0.0 diff --git a/tests/datasets/test_auxiliary_functions.py b/tests/datasets/test_auxiliary_functions.py index 837eda8..1290d6e 100644 --- a/tests/datasets/test_auxiliary_functions.py +++ b/tests/datasets/test_auxiliary_functions.py @@ -3,9 +3,14 @@ import os import numpy as np +import pytest from pyTorchAutoForge.setup.AutoForgeInit import Is_session_headless + +@pytest.mark.integration +@pytest.mark.slow +@pytest.mark.visual def test_LoadDatasetToMem(): import matplotlib diff --git a/tests/model_building/test_image_processing_operators.py b/tests/model_building/test_image_processing_operators.py index f42ef8f..2dd856f 100644 --- a/tests/model_building/test_image_processing_operators.py +++ b/tests/model_building/test_image_processing_operators.py @@ -5,6 +5,7 @@ from pyTorchAutoForge.datasets.LabelsClasses import PTAF_Datakey from kornia.constants import DataKey import torch +import pytest from pyTorchAutoForge.model_building.backbones.image_processing_operators import Compute_threshold_mask, Apply_sobel_gradient, Apply_laplacian_of_gaussian, Compute_distance_transform_map, Compute_local_variance_map from pyTorchAutoForge.utils import torch_to_numpy, numpy_to_torch from PIL import Image @@ -192,6 +193,8 @@ def _run_all_batched_images_(image_names, apply_augs, augmentation_module: Image plt.close() +@pytest.mark.integration +@pytest.mark.visual def test_all_operators(): # ---- Configuration ---- this_file_path = os.path.dirname(os.path.abspath(__file__)) diff --git a/tests/optimization/test_model_training_manager.py b/tests/optimization/test_model_training_manager.py index 6f31e7f..431b401 100644 --- a/tests/optimization/test_model_training_manager.py +++ b/tests/optimization/test_model_training_manager.py @@ -7,8 +7,13 @@ from torchvision import models import optuna, mlflow +from tests.helpers import CudaIsUsable -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA GPU required for this test.") + +@pytest.mark.gpu +@pytest.mark.integration +@pytest.mark.slow +@pytest.mark.skipif(not CudaIsUsable(), reason="Usable CUDA runtime required for this test.") def test_ModelTrainingManager(): from torchvision import transforms From c6d2135c70e11d17378a955b1d399b0ce139ec6b Mon Sep 17 00:00:00 2001 From: PeterC Date: Wed, 20 May 2026 15:58:15 +0200 Subject: [PATCH 4/7] Add misc tests for utils and tcp api (legacy) --- tests/api/tcp/test_start_server.py | 104 ++++++++++++++++------------- tests/utils/test_timing_utils.py | 81 ++++++++++++++-------- 2 files changed, 110 insertions(+), 75 deletions(-) diff --git a/tests/api/tcp/test_start_server.py b/tests/api/tcp/test_start_server.py index 00756bf..10e35f3 100644 --- a/tests/api/tcp/test_start_server.py +++ b/tests/api/tcp/test_start_server.py @@ -1,50 +1,62 @@ -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() diff --git a/tests/utils/test_timing_utils.py b/tests/utils/test_timing_utils.py index 854edc3..3af9c98 100644 --- a/tests/utils/test_timing_utils.py +++ b/tests/utils/test_timing_utils.py @@ -1,40 +1,63 @@ +from __future__ import annotations + +from itertools import count + +import pytest + +import pyTorchAutoForge.utils.timing_utils as timing_utils from pyTorchAutoForge.utils import timeit_averaged, timeit_averaged_ -import time -@timeit_averaged(2) -def dummy_function(): - print("Dummy function called") - time.sleep(1) -def test_timeit_averaged(): - print("Testing timeit_averaged as decorator...") - dummy_function() +def _PatchPerfCounter(monkeypatch: pytest.MonkeyPatch, + step: float = 0.01, + ) -> None: + counter_ = count() + + def FakePerfCounter() -> float: + return next(counter_) * step + + monkeypatch.setattr(timing_utils.time, "perf_counter", FakePerfCounter) + + +def test_timeit_averaged_returns_wrapped_result(monkeypatch: pytest.MonkeyPatch) -> None: + _PatchPerfCounter(monkeypatch) + calls_: list[int] = [] + + @timeit_averaged(2) + def SampleFunction(value_: int) -> int: + calls_.append(value_) + return value_ + 1 + + assert SampleFunction(4) == 5 + assert calls_ == [4, 4] + + +def test_timeit_averaged_preserves_function_metadata(monkeypatch: pytest.MonkeyPatch) -> None: + _PatchPerfCounter(monkeypatch) + + @timeit_averaged(1) + def SampleFunction() -> str: + return "ok" + + assert SampleFunction.__name__ == "SampleFunction" + assert SampleFunction() == "ok" -def test_timeit_averaged_wrapped(): - print("Testing timeit_averaged wrapped...") - dummy_function() -def test_timeit_averaged_function(): - def sample_function(x, y): - time.sleep(0.5) - return x + y +def test_timeit_averaged_function() -> None: + def SampleFunction(x_: int, y_: int) -> int: + return x_ + y_ - num_trials = 3 - average_time = timeit_averaged_(sample_function, num_trials, 2, 3) + average_time_ = timeit_averaged_(SampleFunction, 3, 2, 3) - assert isinstance(average_time, float), "Average time should be a float" - assert average_time > 0, "Average time should be greater than 0" - assert average_time < 1, "Average time should be less than 1 for this test case" + assert isinstance(average_time_, float) + assert average_time_ >= 0.0 -def test_timeit_averaged_function_with_kwargs(): - def sample_function(x, y, delay=0.5): - time.sleep(delay) - return x * y - num_trials = 2 - average_time = timeit_averaged_(sample_function, num_trials, 3, 4, delay=0.3) +def test_timeit_averaged_function_with_kwargs() -> None: + def SampleFunction(x_: int, y_: int, scale_: int = 1) -> int: + return scale_ * x_ * y_ - assert isinstance(average_time, float), "Average time should be a float" - assert average_time > 0, "Average time should be greater than 0" - assert average_time < 1, "Average time should be less than 1 for this test case" + average_time_ = timeit_averaged_(SampleFunction, 2, 3, 4, scale_=2) + assert isinstance(average_time_, float) + assert average_time_ >= 0.0 From fb81cd391342cfa0360e0ebf9fa271ba16f13585 Mon Sep 17 00:00:00 2001 From: PeterC Date: Wed, 20 May 2026 15:58:52 +0200 Subject: [PATCH 5/7] Minor bugfix and tests --- pyTorchAutoForge/api/torch/torchModulesIO.py | 11 +++++------ tests/api/tcp/test_start_server.py | 6 ++++-- tests/api/torch/test_torch_modules_io.py | 19 +++++++++++++++++++ 3 files changed, 28 insertions(+), 8 deletions(-) create mode 100644 tests/api/torch/test_torch_modules_io.py diff --git a/pyTorchAutoForge/api/torch/torchModulesIO.py b/pyTorchAutoForge/api/torch/torchModulesIO.py index a5ab6ac..8cbe886 100644 --- a/pyTorchAutoForge/api/torch/torchModulesIO.py +++ b/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.") - - diff --git a/tests/api/tcp/test_start_server.py b/tests/api/tcp/test_start_server.py index 10e35f3..b59da74 100644 --- a/tests/api/tcp/test_start_server.py +++ b/tests/api/tcp/test_start_server.py @@ -27,8 +27,10 @@ def test_data_processor_tensor_roundtrip() -> None: ) 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:]) + 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_) diff --git a/tests/api/torch/test_torch_modules_io.py b/tests/api/torch/test_torch_modules_io.py new file mode 100644 index 0000000..2cb125c --- /dev/null +++ b/tests/api/torch/test_torch_modules_io.py @@ -0,0 +1,19 @@ +from pathlib import Path + +import torch + +from pyTorchAutoForge.api.torch import AutoForgeModuleSaveMode, SaveModel + + +def test_SaveModel_accepts_torch_device(tmp_path: Path) -> None: + model = torch.nn.Linear(2, 1) + model_path = tmp_path / "linear_model" + + SaveModel( + model=model, + model_filename=model_path, + save_mode=AutoForgeModuleSaveMode.MODEL_STATE_DICT, + target_device=torch.device("cpu"), + ) + + assert (tmp_path / "linear_model_statedict.pth").is_file() From 5a5ccac516a389d80bdc18cfe64a2a952180e08b Mon Sep 17 00:00:00 2001 From: PeterC Date: Wed, 20 May 2026 15:59:47 +0200 Subject: [PATCH 6/7] Add pytest marker --- tests/model_building/test_models_export.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/model_building/test_models_export.py b/tests/model_building/test_models_export.py index bd18df6..87ec769 100644 --- a/tests/model_building/test_models_export.py +++ b/tests/model_building/test_models_export.py @@ -1,8 +1,10 @@ import torch +import pytest from pathlib import Path from pyTorchAutoForge.model_building.backbones.efficient_net import EfficientNetConfig, FeatureExtractorFactory # Export efficient net to ONNX +@pytest.mark.export def test_efficientnet_basic_backbone_export(tmp_path: Path): # Create configuration cfg = EfficientNetConfig( From b307d8491eb41c777401b862831776553fd9c665 Mon Sep 17 00:00:00 2001 From: PeterC Date: Wed, 20 May 2026 16:00:21 +0200 Subject: [PATCH 7/7] Add demo of optuna with mlflow logging --- .../example_mlflow_optuna_cifar10_demo.py | 258 ++++++++++++++++++ 1 file changed, 258 insertions(+) create mode 100644 examples/example_mlflow_optuna_cifar10_demo.py diff --git a/examples/example_mlflow_optuna_cifar10_demo.py b/examples/example_mlflow_optuna_cifar10_demo.py new file mode 100644 index 0000000..ebdbcc1 --- /dev/null +++ b/examples/example_mlflow_optuna_cifar10_demo.py @@ -0,0 +1,258 @@ +"""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 +import torch +import torch.nn as nn +import torch.optim as optim +from torch.utils.data import DataLoader +from torchvision import datasets, transforms +import matplotlib.pyplot as plt + +# Import modules +import os +import subprocess +import time +import logging + +# Set up logging +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger() +logger.setLevel(logging.INFO) + + +def StartMLflowUI(port: int = 5000): + + # Start MLflow UI + os.system('mlflow ui --port ' + str(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 + 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) +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__() + self.layers = nn.Sequential() + + # Convolutional layers + num_conv_layers = trial.suggest_int('num_conv_layers', 1, 2) + + # 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) + + # 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'relu{i}', nn.ReLU()) + self.layers.add_module(f'maxpool{i}', nn.MaxPool2d(2)) + + # 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 + + # 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'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)) + in_features = out_features + + self.layers.add_module('output', nn.Linear(in_features, 10)) + + def forward(self, x): + return self.layers(x) + + # Function to compute the shape of the output of the a convolutional layer + def _get_conv_output(self, shape): + with torch.no_grad(): + input_data = torch.rand(1, *shape) + output = self.layers(input_data) + return output.size(1) + +# %% Optuna objective function + + +def objective(trial): + # Create new run_IN within the current session + with mlflow.start_run(): + + device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') + + # Create the trial model + model = SimpleCNN(trial).to(device) + + # Select the optimizer + 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() + + # MLflow: log trial data + mlflow.log_param('optimizer', optimizer_name) + mlflow.log_param('learning_rate', lr) + mlflow.log_param('num_conv_layers', + trial.params.get('num_conv_layers')) + # NOTE: integration with optuna --> parameters are got directly from the ith trial object under evaluation + mlflow.log_param('num_dense_layers', + trial.params.get('num_dense_layers')) + + for i in range(trial.params.get('num_conv_layers')): + mlflow.log_param(f'filters_{i}', trial.params.get(f'filters_{i}')) + mlflow.log_param( + f'kernel_size_{i}', trial.params.get(f'kernel_size_{i}')) + + for i in range(trial.params.get('num_dense_layers')): + 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 = inputs.to(device), targets.to(device) + optimizer.zero_grad() + + # Evaluate model + outputs = model(inputs) + # Compute loss + loss = criterion(outputs, targets) + # Perform backpropagation + loss.backward() + optimizer.step() + + # Model validation + model.eval() + correct = 0 + total = 0 + with torch.no_grad(): + for batch in test_loader: + inputs, targets = batch + inputs, targets = inputs.to(device), targets.to(device) + outputs = model(inputs) + _, predicted = torch.max(outputs.data, 1) + total += targets.size(0) + correct += (predicted == targets).sum().item() + + accuracy = correct / total + + # Log the accuracy of the model using mlflow + mlflow.log_metric('Accuracy value', accuracy, step=epoch) + + # 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 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 + # of config file are automatically logged by mlflow with their names + mlflow.end_run(status='FINISHED') + + return accuracy + + +def main(): + print('---------------------------- TEST SCRIPT: MLflow and Optuna functionalities ----------------------------\n') + # %% MLflow tracking initialization + port = 7000 + # 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)), + 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 + # 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 + # 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()