minerva-opt 1.0.0__py3-none-any.whl

minerva_opt/__init__.py

@@ -0,0 +1,3 @@
+from minerva_opt.pipelines.hyperparameter_search import RayHyperParameterSearch
+
+__all__ = ["RayHyperParameterSearch"]

minerva_opt/callbacks/__init__.py

@@ -0,0 +1,9 @@
+from minerva_opt.callbacks.ray_callbacks import (
+    TrainerReportKeepOnlyLastCallback,
+    TrainerReportOnIntervalCallback,
+)
+
+__all__ = [
+    "TrainerReportKeepOnlyLastCallback",
+    "TrainerReportOnIntervalCallback",
+]

minerva_opt/callbacks/ray_callbacks.py

@@ -0,0 +1,89 @@
+import os
+import shutil
+import tempfile
+from pathlib import Path
+
+import lightning.pytorch as L
+from ray import train
+from ray._common.usage.usage_lib import TagKey, record_extra_usage_tag
+from ray.train import Checkpoint
+
+
+class TrainerReportOnIntervalCallback(L.Callback):
+
+    CHECKPOINT_NAME = "checkpoint.ckpt"
+
+    def __init__(self, interval: int = 1) -> None:
+        super().__init__()
+        self.trial_name = train.get_context().get_trial_name()
+        self.local_rank = train.get_context().get_local_rank()
+        self.tmpdir_prefix = Path(tempfile.gettempdir(), self.trial_name).as_posix()
+        self.interval = interval
+        self.step = 0
+        if os.path.isdir(self.tmpdir_prefix) and self.local_rank == 0:
+            shutil.rmtree(self.tmpdir_prefix)
+
+        record_extra_usage_tag(TagKey.TRAIN_LIGHTNING_RAYTRAINREPORTCALLBACK, "1")
+
+    def on_train_epoch_end(self, trainer: L.Trainer, pl_module: L.LightningModule) -> None:
+        metrics = trainer.callback_metrics
+        metrics = {k: v.item() for k, v in metrics.items()}
+        metrics["epoch"] = trainer.current_epoch
+        metrics["step"] = trainer.global_step
+
+        tmpdir = Path(self.tmpdir_prefix, str(trainer.current_epoch)).as_posix()
+        os.makedirs(tmpdir, exist_ok=True)
+
+        if self.step % self.interval == 0:
+            ckpt_path = Path(tmpdir, self.CHECKPOINT_NAME).as_posix()
+            trainer.save_checkpoint(ckpt_path, weights_only=False)
+            checkpoint = Checkpoint.from_directory(tmpdir)
+            train.report(metrics=metrics, checkpoint=checkpoint)
+        else:
+            train.report(metrics=metrics)
+
+        trainer.strategy.barrier()
+
+        if self.local_rank == 0:
+            shutil.rmtree(tmpdir)
+
+        self.step += 1
+
+
+class TrainerReportKeepOnlyLastCallback(L.Callback):
+
+    CHECKPOINT_NAME = "checkpoint.ckpt"
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.trial_name = train.get_context().get_trial_name()
+        self.local_rank = train.get_context().get_local_rank()
+        self.tmpdir_prefix = Path(tempfile.gettempdir(), self.trial_name).as_posix()
+        if os.path.isdir(self.tmpdir_prefix) and self.local_rank == 0:
+            shutil.rmtree(self.tmpdir_prefix)
+
+        record_extra_usage_tag(TagKey.TRAIN_LIGHTNING_RAYTRAINREPORTCALLBACK, "1")
+
+    def on_train_epoch_end(self, trainer: L.Trainer, pl_module: L.LightningModule) -> None:
+        metrics = trainer.callback_metrics
+        metrics = {k: v.item() for k, v in metrics.items()}
+        metrics["epoch"] = trainer.current_epoch
+        metrics["step"] = trainer.global_step
+
+        tmpdir = Path(self.tmpdir_prefix, "last").as_posix()
+
+        # Delete previous epoch's checkpoint before writing the new one
+        if os.path.isdir(tmpdir):
+            shutil.rmtree(tmpdir)
+        os.makedirs(tmpdir, exist_ok=True)
+
+        ckpt_path = Path(tmpdir, self.CHECKPOINT_NAME).as_posix()
+        trainer.save_checkpoint(ckpt_path, weights_only=False)
+
+        checkpoint = Checkpoint.from_directory(tmpdir)
+        train.report(metrics=metrics, checkpoint=checkpoint)
+
+        trainer.strategy.barrier()
+
+        if self.local_rank == 0:
+            shutil.rmtree(tmpdir)

minerva_opt/pipelines/__init__.py

@@ -0,0 +1,3 @@
+from minerva_opt.pipelines.hyperparameter_search import RayHyperParameterSearch
+
+__all__ = ["RayHyperParameterSearch"]

minerva_opt/pipelines/hyperparameter_search.py

@@ -0,0 +1,154 @@
+from copy import deepcopy
+from typing import Any, Dict, List, Literal, Optional
+
+import lightning.pytorch as L
+from lightning.pytorch.strategies import Strategy
+from ray import tune
+from ray.train import CheckpointConfig, RunConfig, ScalingConfig
+from ray.train.lightning import RayDDPStrategy, RayLightningEnvironment, prepare_trainer
+from ray.train.torch import TorchTrainer
+from ray.tune.result_grid import ResultGrid
+from ray.tune.schedulers import ASHAScheduler, TrialScheduler
+from ray.tune.search import ConcurrencyLimiter
+from ray.tune.search.searcher import Searcher
+
+from minerva.pipelines.base import Pipeline
+from minerva.utils.typing import PathLike
+from minerva_opt.callbacks.ray_callbacks import TrainerReportOnIntervalCallback
+
+
+class RayHyperParameterSearch(Pipeline):
+    """Hyperparameter search pipeline using Ray Tune and PyTorch Lightning.
+
+    Supports any Ray Tune search algorithm via the `search_alg` parameter in
+    `_search`. When no algorithm is provided, Ray's default random/grid search
+    is used. To use Bayesian optimization, pass a `HyperOptSearch` instance.
+    """
+
+    def __init__(
+        self,
+        model: type,
+        search_space: Dict[str, Any],
+        log_dir: Optional[PathLike] = None,
+        save_run_status: bool = True,
+        seed: Optional[int] = None,
+    ):
+        super().__init__(log_dir=log_dir, save_run_status=save_run_status, seed=seed)
+        self.model = model
+        self.search_space = search_space
+
+    def _search(
+        self,
+        data: L.LightningDataModule,
+        ckpt_path: Optional[PathLike],
+        devices: str = "auto",
+        accelerator: str = "auto",
+        strategy: Optional[Strategy] = None,
+        callbacks: Optional[List[Any]] = None,
+        plugins: Optional[List[Any]] = None,
+        num_nodes: int = 1,
+        debug_mode: bool = False,
+        scaling_config: Optional[ScalingConfig] = None,
+        run_config: Optional[RunConfig] = None,
+        tuner_metric: str = "val_loss",
+        tuner_mode: str = "min",
+        num_samples: int = 10,
+        scheduler: Optional[TrialScheduler] = None,
+        search_alg: Optional[Searcher] = None,
+        max_concurrent: int = 4,
+        max_epochs: int = 100,
+        checkpoint_interval: int = 1,
+    ) -> ResultGrid:
+
+        def _train_func(config):
+            dm = deepcopy(data)
+            model = self.model(**config)
+            trainer = L.Trainer(
+                max_epochs=max_epochs,
+                devices=devices,
+                accelerator=accelerator,
+                strategy=strategy or RayDDPStrategy(find_unused_parameters=True),
+                callbacks=callbacks or [TrainerReportOnIntervalCallback(checkpoint_interval)],
+                plugins=plugins or [RayLightningEnvironment()],
+                enable_progress_bar=False,
+                num_nodes=num_nodes,
+                enable_checkpointing=False if debug_mode else None,
+            )
+            trainer = prepare_trainer(trainer)
+            trainer.fit(model, dm, ckpt_path=ckpt_path)
+
+        scheduler = scheduler or ASHAScheduler(
+            time_attr="training_iteration",
+            metric=tuner_metric,
+            mode=tuner_mode,
+            max_t=max_epochs,
+            grace_period=max(1, max_epochs // 10),
+        )
+
+        scaling_config = scaling_config or ScalingConfig(
+            num_workers=1, use_gpu=True, resources_per_worker={"GPU": 1}
+        )
+
+        run_config = run_config or RunConfig(
+            checkpoint_config=CheckpointConfig(
+                num_to_keep=1,
+                checkpoint_score_attribute=tuner_metric,
+                checkpoint_score_order=tuner_mode,
+            )
+        )
+
+        if search_alg is not None:
+            search_alg = ConcurrencyLimiter(search_alg, max_concurrent=max_concurrent)
+
+        ray_trainer = TorchTrainer(
+            _train_func,
+            scaling_config=scaling_config,
+            run_config=run_config,
+        )
+
+        tuner = tune.Tuner(
+            ray_trainer,
+            param_space={"train_loop_config": self.search_space},
+            tune_config=tune.TuneConfig(
+                metric=tuner_metric,
+                mode=tuner_mode,
+                num_samples=num_samples,
+                scheduler=scheduler,
+                search_alg=search_alg,
+            ),
+        )
+
+        results = tuner.fit()
+        best = results.get_best_result()
+        print(f"Best config: {best.config}")
+        print(f"Best {tuner_metric}: {best.metrics.get(tuner_metric)}")
+        return results
+
+    def _test(self, data: L.LightningDataModule, ckpt_path: Optional[PathLike]) -> Any:
+        raise NotImplementedError(
+            "Load the best checkpoint from the ResultGrid returned by _search "
+            "and call trainer.test() directly."
+        )
+
+    def _run(
+        self,
+        data: L.LightningDataModule,
+        task: Optional[Literal["search", "test"]] = None,
+        ckpt_path: Optional[PathLike] = None,
+        **kwargs,
+    ) -> Any:
+        if task == "search" or task is None:
+            return self._search(data, ckpt_path, **kwargs)
+        elif task == "test":
+            return self._test(data, ckpt_path)
+
+
+def main():
+    from jsonargparse import CLI
+
+    print("Hyper Searching")
+    CLI(RayHyperParameterSearch, as_positional=False)
+
+
+if __name__ == "__main__":
+    main()

minerva_opt-1.0.0.dist-info/METADATA

@@ -0,0 +1,141 @@
+Metadata-Version: 2.4
+Name: minerva-opt
+Version: 1.0.0
+Summary: Hyperparameter optimization extensions for Minerva.
+Author-email: Gabriel Gutierrez <gabriel.bgs00@gmail.com>
+Project-URL: Bug Tracker, https://github.com/gabrielbg0/Minerva-OPT/issues
+Project-URL: Homepage, https://github.com/gabrielbg0/Minerva-OPT
+Keywords: Deep Learning,Hyperparameter Optimization,Machine Learning,Pytorch,Ray Tune,Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python
+Classifier: Programming Language :: Python :: 3
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: Operating System :: OS Independent
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+Requires-Dist: minerva==0.3.10-beta
+Requires-Dist: ray[tune]>=2.55
+Requires-Dist: hyperopt>=0.2.7
+Provides-Extra: dev
+Requires-Dist: black; extra == "dev"
+Requires-Dist: flake8; extra == "dev"
+Requires-Dist: pytest; extra == "dev"
+Requires-Dist: pytest-coverage; extra == "dev"
+
+# Minerva-OPT
+
+[![Auto Release](https://github.com/gabrielbg0/Minerva-OPT/actions/workflows/auto-release.yml/badge.svg)](https://github.com/gabrielbg0/Minerva-OPT/actions/workflows/auto-release.yml)
+
+Hyperparameter optimization extensions for [Minerva](https://github.com/discovery-unicamp/Minerva), powered by [Ray Tune](https://docs.ray.io/en/latest/tune/index.html).
+
+## Description
+
+Minerva-OPT provides a `RayHyperParameterSearch` pipeline that wraps Ray Tune and PyTorch Lightning to run distributed hyperparameter searches on top of any Minerva-compatible model. It supports random search, grid search, and Bayesian optimization (via HyperOpt), with early stopping through the ASHA scheduler.
+
+### Features
+
+- **Drop-in Minerva pipeline**: inherits from `minerva.pipelines.base.Pipeline`, so it integrates with Minerva's logging, reproducibility, and run-status tracking out of the box.
+- **Flexible search algorithms**: use Ray Tune's default random/grid search or pass any `ray.tune.search.Searcher` (e.g. `HyperOptSearch` for Bayesian optimization).
+- **ASHA early stopping**: trials are stopped early based on intermediate results; `grace_period` and `max_t` are derived automatically from `max_epochs`.
+- **Distributed training**: uses `RayDDPStrategy` and `RayLightningEnvironment` for multi-worker trials.
+- **Checkpointing**: keeps only the best checkpoint per trial, scored on the target metric.
+
+## Installation
+
+Requires Python 3.10+.
+
+### With uv (recommended)
+
+```bash
+uv pip install minerva-opt
+```
+
+### With pip
+
+```bash
+pip install minerva-opt
+```
+
+## Usage
+
+### Random / grid search
+
+```python
+from ray import tune
+from minerva_opt.pipelines.hyperparameter_search import RayHyperParameterSearch
+
+search_space = {
+    "learning_rate": tune.loguniform(1e-4, 1e-1),
+    "hidden_size": tune.choice([64, 128, 256]),
+}
+
+pipeline = RayHyperParameterSearch(
+    model=MyLightningModel,  # class, not instance — instantiated per trial as MyLightningModel(**config)
+    search_space=search_space,
+    log_dir="logs/",
+)
+
+results = pipeline.run(data=my_data_module, num_samples=20, max_epochs=50)
+best = results.get_best_result()
+print(best.config)
+```
+
+### Bayesian optimization with HyperOpt
+
+```python
+from ray.tune.search.hyperopt import HyperOptSearch
+from hyperopt import hp
+
+search_space = {
+    "learning_rate": tune.loguniform(1e-4, 1e-1),
+    "dropout": tune.uniform(0.1, 0.5),
+}
+
+pipeline = RayHyperParameterSearch(
+    model=MyLightningModel,
+    search_space=search_space,
+)
+
+results = pipeline.run(
+    data=my_data_module,
+    search_alg=HyperOptSearch(),
+    num_samples=30,
+    max_epochs=100,
+    tuner_metric="val_loss",
+    tuner_mode="min",
+)
+```
+
+### Key `run()` parameters
+
+| Parameter             | Default      | Description                                            |
+| --------------------- | ------------ | ------------------------------------------------------ |
+| `data`                | —            | `LightningDataModule` to train on                      |
+| `task`                | `"search"`   | `"search"` to run the sweep                            |
+| `ckpt_path`           | `None`       | Resume training from a checkpoint                      |
+| `num_samples`         | `10`         | Number of trials to run                                |
+| `max_epochs`          | `100`        | Max epochs per trial                                   |
+| `tuner_metric`        | `"val_loss"` | Metric to optimize                                     |
+| `tuner_mode`          | `"min"`      | `"min"` or `"max"`                                     |
+| `search_alg`          | `None`       | Any `ray.tune.search.Searcher`; `None` = random search |
+| `max_concurrent`      | `4`          | Max concurrent trials (when using a `search_alg`)      |
+| `scheduler`           | ASHA         | Override the trial scheduler                           |
+| `scaling_config`      | 1 GPU/worker | Override Ray `ScalingConfig`                           |
+| `checkpoint_interval` | `1`          | Save a checkpoint every N epochs                       |
+| `debug_mode`          | `False`      | Disable checkpointing for fast iteration               |
+
+## Requirements
+
+- `minerva >= 0.3.10b0`
+- `ray[tune] >= 2.55`
+- `hyperopt >= 0.2.7`
+
+## License
+
+MIT License. See [LICENSE](LICENSE) for details.
+
+## Contact
+
+For questions or bug reports, open an issue on the [GitHub issue tracker](https://github.com/gabrielbg0/Minerva-OPT/issues).

minerva_opt-1.0.0.dist-info/RECORD

@@ -0,0 +1,9 @@
+minerva_opt/__init__.py,sha256=uHOTGjXAEfB6Tvy9JztZTEJqqLJ7FmmMfLI8bN7OZK4,119
+minerva_opt/callbacks/__init__.py,sha256=OKAaTaN3_sbwAuWx0ufZaOiyDsXlq1t7_hfjCDYo4GQ,223
+minerva_opt/callbacks/ray_callbacks.py,sha256=bKWu2Z8tDMlsoz32kZoc_VrqfjxMpT4G9pJeYZSHeyo,3223
+minerva_opt/pipelines/__init__.py,sha256=uHOTGjXAEfB6Tvy9JztZTEJqqLJ7FmmMfLI8bN7OZK4,119
+minerva_opt/pipelines/hyperparameter_search.py,sha256=-0Ghvo-gXY3Z57ctJbvx12nqmT0QlgLwhK9TX6EKkUk,5356
+minerva_opt-1.0.0.dist-info/METADATA,sha256=PrRCunlNIbP6t3UZgRMWtodmXYdLVB-42vps7wshoXo,5576
+minerva_opt-1.0.0.dist-info/WHEEL,sha256=aeYiig01lYGDzBgS8HxWXOg3uV61G9ijOsup-k9o1sk,91
+minerva_opt-1.0.0.dist-info/top_level.txt,sha256=yWSjtBACKoMzc6a44LuLa5AaDmCKBGeWg27-hrVs1DQ,12
+minerva_opt-1.0.0.dist-info/RECORD,,

minerva_opt-1.0.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (82.0.1)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

minerva_opt-1.0.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+minerva_opt