lightray 0.1.0__py3-none-any.whl

lightray/__init__.py

@@ -0,0 +1 @@
+from .tune import run

lightray/callbacks.py

@@ -0,0 +1,63 @@
+import os
+import shutil
+import time
+from pathlib import Path
+
+from botocore.exceptions import ClientError, ConnectTimeoutError
+from ray import train
+from ray.train import Checkpoint
+from ray.train.lightning import RayTrainReportCallback
+
+BOTO_RETRY_EXCEPTIONS = (ClientError, ConnectTimeoutError)
+
+
+def report_with_retries(metrics, checkpoint, retries: int = 10):
+    """
+    Call `train.report`, which will persist checkpoints to s3,
+    retrying after any possible errors
+    """
+    for _ in range(retries):
+        try:
+            train.report(metrics=metrics, checkpoint=checkpoint)
+            break
+        except BOTO_RETRY_EXCEPTIONS:
+            time.sleep(5)
+            continue
+
+
+class TrainReportCallback(RayTrainReportCallback):
+    """
+    Equivalent of the RayTrainReportCallback
+    (https://docs.ray.io/en/latest/train/api/doc/ray.train.lightning.RayTrainReportCallback.html)
+    except adds a retry mechanism to the `train.report` call
+    """
+
+    def on_train_epoch_end(self, trainer, pl_module) -> None:
+        # Creates a checkpoint dir with fixed name
+        tmpdir = Path(
+            self.tmpdir_prefix, str(trainer.current_epoch)
+        ).as_posix()
+        os.makedirs(tmpdir, exist_ok=True)
+
+        # Fetch metrics
+        metrics = trainer.callback_metrics
+        metrics = {k: v.item() for k, v in metrics.items()}
+
+        # (Optional) Add customized metrics
+        metrics["epoch"] = trainer.current_epoch
+        metrics["step"] = trainer.global_step
+
+        # Save checkpoint to local
+        ckpt_path = Path(tmpdir, self.CHECKPOINT_NAME).as_posix()
+        trainer.save_checkpoint(ckpt_path, weights_only=False)
+
+        # Report to train session
+        checkpoint = Checkpoint.from_directory(tmpdir)
+        report_with_retries(metrics=metrics, checkpoint=checkpoint)
+
+        # Add a barrier to ensure all workers finished reporting here
+        trainer.strategy.barrier()
+
+        if self.local_rank == 0:
+            shutil.rmtree(tmpdir)
+            shutil.rmtree(tmpdir)

lightray/fs.py

@@ -0,0 +1,43 @@
+import os
+from typing import Optional
+
+import pyarrow
+import s3fs
+
+# s3 retry configuration
+retry_config = {"retries": {"total_max_attempts": 10, "mode": "adaptive"}}
+
+
+def setup_filesystem(
+    storage_dir: Optional[str] = None,
+) -> tuple[
+    Optional[pyarrow.fs.PyFileSystem], Optional[pyarrow.fs.PyFileSystem]
+]:
+    internal_fs = external_fs = None
+    if storage_dir and storage_dir.startswith("s3://"):
+        storage_dir = storage_dir.removeprefix("s3://")
+        endpoint_url = os.getenv("AWS_ENDPOINT_URL")
+        # directly use s3 instead of rays pyarrow s3 default due to
+        # this issue https://github.com/ray-project/ray/issues/41137
+        internal_fs = s3fs.S3FileSystem(
+            key=os.getenv("AWS_ACCESS_KEY_ID"),
+            secret=os.getenv("AWS_SECRET_ACCESS_KEY"),
+            endpoint_url=endpoint_url,
+            config_kwargs=retry_config,
+        )
+        internal_fs = pyarrow.fs.PyFileSystem(
+            pyarrow.fs.FSSpecHandler(internal_fs)
+        )
+
+        external_fs = s3fs.S3FileSystem(
+            key=os.getenv("AWS_ACCESS_KEY_ID"),
+            secret=os.getenv("AWS_SECRET_ACCESS_KEY"),
+            endpoint_url=os.getenv(
+                "AWS_EXTERNAL_ENDPOINT_URL", os.getenv("AWS_ENDPOINT_URL")
+            ),
+            config_kwargs=retry_config,
+        )
+        external_fs = pyarrow.fs.PyFileSystem(
+            pyarrow.fs.FSSpecHandler(external_fs)
+        )
+    return internal_fs, external_fs

lightray/tune.py

@@ -0,0 +1,110 @@
+import logging
+import os
+from pathlib import Path
+from typing import TYPE_CHECKING, Literal, Optional, Type, Union
+
+import lightning.pytorch as pl
+import ray
+from ray import tune
+from ray.tune.schedulers import TrialScheduler
+
+from lightray import utils
+from lightray.fs import setup_filesystem
+
+if TYPE_CHECKING:
+    from lightning.pytorch.cli import LightningCLI
+
+
+def run(
+    cli_cls: Type["LightningCLI"],
+    name: str,
+    metric_name: str,
+    objective: Literal["min", "max"],
+    search_space: Union[dict, Path],
+    scheduler: TrialScheduler,
+    storage_dir: Optional[str] = None,
+    address: Optional[str] = None,
+    num_samples: int = 10,
+    workers_per_trial: int = 1,
+    gpus_per_worker: float = 1.0,
+    cpus_per_gpu: float = 1.0,
+    callbacks: Optional[list[pl.callbacks.Callback]] = None,
+    temp_dir: Optional[str] = None,
+    args: Optional[list[str]] = None,
+) -> tune.ResultGrid:
+    # parse the training configuration file, and
+    # any argument overrides
+    # using the user passed LightningCLI class;
+    config = utils.parse_args(cli_cls, args)
+
+    # if specified, connect to a running ray cluster
+    # otherwise, ray will assume one is running locally
+    logging.info("Initializing Ray")
+    if address is not None:
+        logging.info(f"Connecting to Ray cluster at {address}")
+        # ensure the ray adress starts with "ray://"
+        if not address.startswith("ray://"):
+            raise ValueError(
+                f"Address must start with 'ray://', got {address}"
+            )
+    ray.init(address, _temp_dir=temp_dir)
+
+    logging.info("Initializing checkpoint storage filesystems")
+    internal_fs, external_fs = setup_filesystem(str(storage_dir))
+
+    # construct the function that will actually
+    # execute the training loop, and then set it
+    # up for Ray to distribute it over our cluster,
+    # with the desired number of resources allocated
+    # to each running version of the job
+    train_func = utils.configure_deployment(
+        utils.TrainFunc(cli_cls, name, config, callbacks),
+        metric_name=metric_name,
+        workers_per_trial=workers_per_trial,
+        gpus_per_worker=gpus_per_worker,
+        cpus_per_gpu=cpus_per_gpu,
+        objective=objective,
+        storage_dir=storage_dir or None,
+        fs=internal_fs,
+    )
+
+    search_space = utils.get_search_space(search_space)
+
+    # restore from a previous tuning run
+    path = os.path.join(storage_dir, name)
+    if tune.Tuner.can_restore(path, storage_filesystem=external_fs):
+        logging.info(f"Restoring from previous tuning run at {path}")
+        tuner = tune.Tuner.restore(
+            path,
+            train_func,
+            resume_errored=True,
+            storage_filesystem=external_fs,
+        )
+
+    else:
+        tuner = tune.Tuner(
+            train_func,
+            param_space={"train_loop_config": search_space},
+            tune_config=tune.TuneConfig(
+                metric=metric_name,
+                mode=objective,
+                num_samples=num_samples,
+                scheduler=scheduler,
+                reuse_actors=True,
+                trial_name_creator=lambda trial: f"{trial.trial_id}",
+                trial_dirname_creator=lambda trial: f"{trial.trial_id}",
+            ),
+        )
+
+    logging.info("Starting tune job")
+    results = tuner.fit()
+
+    return results
+
+
+def cli():
+    pass
+
+
+if __name__ == "__main__":
+    cli()

lightray/utils.py

@@ -0,0 +1,277 @@
+"""
+Hyperparameter tuning utilities based largely on this tutorial
+https://docs.ray.io/en/latest/tune/examples/tune-pytorch-lightning.html
+"""
+
+
+import importlib
+import math
+import os
+from tempfile import NamedTemporaryFile
+from typing import Any, List, Optional, Type, Union
+
+import lightning.pytorch as pl
+import pyarrow.fs
+import yaml
+from lightning.pytorch.cli import LightningCLI
+from ray import train
+from ray.train import CheckpointConfig, FailureConfig, RunConfig, ScalingConfig
+from ray.train.lightning import (
+    RayDDPStrategy,
+    RayLightningEnvironment,
+    prepare_trainer,
+)
+from ray.train.torch import TorchTrainer
+
+from lightray.callbacks import RayTrainReportCallback
+
+
+def get_host_cli(cli: Type[LightningCLI]):
+    """
+    Return a LightningCLI class that will utilize
+    the `cli` class passed as a parent class
+    for parsing arguments.
+
+    Since this is run on the client,
+    we don't actually want to do anything with the arguments we parse,
+    just record them, so override the couple parent
+    methods responsible for instantiating classes and running
+    subcommands.
+    """
+
+    class HostCLI(cli):
+        def instantiate_classes(self):
+            return
+
+        def _run_subcommand(self):
+            return
+
+    return HostCLI
+
+
+def get_worker_cli(
+    cli: Type[LightningCLI],
+    callbacks: Optional[List[pl.callbacks.Callback]] = None,
+):
+    """
+    Return a LightningCLI class that will actually execute
+    training runs on worker nodes
+    """
+
+    callbacks = callbacks or [RayTrainReportCallback()]
+
+    class WorkerCLI(cli):
+        def instantiate_trainer(self, **kwargs):
+            kwargs = kwargs | dict(
+                enable_progress_bar=False,
+                devices="auto",
+                accelerator="auto",
+                strategy=RayDDPStrategy(),
+                callbacks=callbacks,
+                plugins=[RayLightningEnvironment()],
+            )
+            return super().instantiate_trainer(**kwargs)
+
+    return WorkerCLI
+
+
+def get_search_space(search_space: Union[str, dict]) -> dict[str, callable]:
+    if isinstance(search_space, dict):
+        return search_space
+
+    # determine if the path is a file path or a module path
+    if os.path.isfile(search_space):
+        # load the module from the file
+        module_name = os.path.splitext(os.path.basename(search_space))[0]
+        spec = importlib.util.spec_from_file_location(
+            module_name, search_space
+        )
+        module = importlib.util.module_from_spec(spec)
+        spec.loader.exec_module(module)
+    else:
+        # load the module using importlib
+        module = importlib.import_module(search_space)
+
+    # try to get the 'space' attribute from the module
+    try:
+        space = module.space
+    except AttributeError:
+        raise ValueError(f"Module {module.__name__} has no space dictionary")
+
+    if not isinstance(space, dict):
+        raise TypeError(
+            "Expected search space in module {} to be "
+            "a dictionary, found {}".format(module.__name__, type(space))
+        )
+
+    return space
+
+
+def stop_on_nan(trial_id: str, result: dict) -> bool:
+    return math.isnan(result["train_loss"])
+
+
+class TrainFunc:
+    """
+    Callable wrapper that takes a `LightningCLI` and executes
+    it with the both the `config` passed here at initialization
+    time as well as the arguments supplied by a particular
+    hyperparameter config. Meant for execution on workers during
+    tuning run, which expect a callable that a particular
+    hyperparameter config as its only argument.
+
+    All runs of the function will be given the same
+    Weights & Biases group name of `name` for tracking.
+    The names of individual runs in this group will be
+    randomly chosen by W&B.
+    """
+
+    def __init__(
+        self,
+        cli: Type[LightningCLI],
+        name: str,
+        config: dict,
+        callbacks: Optional[list[pl.callbacks.Callback]] = None,
+    ) -> None:
+        self.cli = cli
+        self.name = name
+        self.config = config
+        self.callbacks = callbacks
+
+    def validate_logger(self, args):
+        try:
+            logger = self.config["trainer"]["logger"]
+        except KeyError:
+            logger = None
+
+        if logger is not None:
+            if "WandbLogger" not in logger:
+                raise ValueError(
+                    "Only W&B logging is supported for distributed tuning"
+                )
+            args.append(f"--trainer.logger.group={self.name}")
+        return args
+
+    def __call__(self, hparams: dict[str, Any]):
+        """
+        Dump the config to a file, then parse it
+        along with the hyperparameter configuration
+        passed here using our CLI.
+        """
+
+        with NamedTemporaryFile(mode="w") as f:
+            # dump the core config,
+            # then add the hyperparameters
+            yaml.dump(self.config, f)
+            args = ["-c", f.name]
+            for key, value in hparams.items():
+                args.append(f"--{key}={value}")
+
+            args = self.validate_logger(args)
+
+            cli_cls = get_worker_cli(self.cli, self.callbacks)
+            cli = cli_cls(
+                run=False, args=args, save_config_kwargs={"overwrite": True}
+            )
+
+        log_dir = cli.trainer.logger.log_dir or cli.trainer.logger.save_dir
+        if not log_dir.startswith("s3://"):
+            ckpt_prefix = ""
+        else:
+            ckpt_prefix = "s3://"
+
+        # restore from checkpoint if available
+        checkpoint = train.get_checkpoint()
+        ckpt_path = None
+        if checkpoint:
+            ckpt_path = os.path.join(
+                ckpt_prefix, checkpoint.path, "checkpoint.ckpt"
+            )
+
+        trainer = prepare_trainer(cli.trainer)
+        trainer.fit(cli.model, cli.datamodule, ckpt_path=ckpt_path)
+
+
+def configure_deployment(
+    train_func: TrainFunc,
+    metric_name: str,
+    workers_per_trial: int,
+    gpus_per_worker: int,
+    cpus_per_gpu: int,
+    objective: str = "max",
+    storage_dir: Optional[str] = None,
+    fs: Optional[pyarrow.fs.FileSystem] = None,
+) -> TorchTrainer:
+    """
+    Set up a training function that can be distributed
+    among the workers in a ray cluster.
+
+    Args:
+        train_func:
+            Function that each worker will execute
+            with a config specifying the hyperparameter
+            configuration for that trial.
+        metric_name:
+            Name of the metric that will be optimized
+            during the hyperparameter search
+        workers_per_trial:
+            Number of training workers to deploy
+        gpus_per_worker:
+            Number of GPUs to train over within each worker
+        cpus_per_gpu:
+            Number of CPUs to attach to each GPU. If gpus_per_worker
+            is 0, this is interpretated as the number of CPUs per worker
+        objective:
+            `"max"` or `"min"`, indicating how the indicated
+            metric ought to be optimized
+        storage_dir:
+            Directory to save ray checkpoints and logs
+            during training.
+        fs: Filesystem to use for storage
+    """
+
+    if gpus_per_worker == 0:
+        cpus_per_worker = cpus_per_gpu
+    else:
+        cpus_per_worker = cpus_per_gpu * gpus_per_worker
+
+    use_gpu = gpus_per_worker > 0
+    scaling_config = ScalingConfig(
+        trainer_resources={"CPU": 0},
+        resources_per_worker={"CPU": cpus_per_worker, "GPU": gpus_per_worker},
+        num_workers=workers_per_trial,
+        use_gpu=use_gpu,
+    )
+
+    run_config = RunConfig(
+        checkpoint_config=CheckpointConfig(
+            num_to_keep=2,
+            checkpoint_score_attribute=metric_name,
+            checkpoint_score_order=objective,
+        ),
+        failure_config=FailureConfig(
+            max_failures=5,
+        ),
+        storage_filesystem=fs,
+        storage_path=storage_dir,
+        name=train_func.name,
+        stop=stop_on_nan,
+    )
+    return TorchTrainer(
+        train_func,
+        scaling_config=scaling_config,
+        run_config=run_config,
+    )
+
+
+def parse_args(
+    cli_cls: Type[LightningCLI], args: Optional[list[str]] = None
+) -> dict:
+    """
+    Use a `LightningCLI` class to parse command line arguments
+    """
+    host_cli = get_host_cli(cli_cls)
+    host_cli = host_cli(run=False, args=args)
+    config = host_cli.parser.dump(host_cli.config, format="yaml")
+    config = yaml.safe_load(config)
+    return config

lightray-0.1.0.dist-info/METADATA

@@ -0,0 +1,96 @@
+Metadata-Version: 2.1
+Name: lightray
+Version: 0.1.0
+Summary: Distribute a LightningCLI hyperparameter search with Ray Tune
+License: MIT
+Author: Ethan Jacob Marx
+Author-email: ethan.marx@ligo.org
+Requires-Python: >=3.9,<4.0
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Requires-Dist: jsonargparse[signatures] (>=4.29,<5.0)
+Requires-Dist: lightning (>=2.4.0,<3.0.0)
+Requires-Dist: pyarrow (>=17.0.0,<18.0.0)
+Requires-Dist: ray[default,tune] (>=2.8.0,<3.0.0)
+Requires-Dist: s3fs (>=2024.9.0,<2025.0.0)
+Description-Content-Type: text/markdown
+
+# lightray
+Easily integrate a `LightningCLI` with `RayTune` hyperparameter optimization
+
+## Getting Started
+Extend your custom `LightningCLI` to a `Trainable` compatible with `RayTune`.
+
+Imagine you have a `LightningCLI` parser for a 
+
+```python
+from lightning.pytorch as pl
+
+class CustomCLI(pl.cli.LightningCLI):
+    def add_arguments_to_parser(self, parser):
+        parser.link_arguments(
+            "data.init_args.parameter", 
+            "model.init_args.parameter", 
+            apply_on="parse"
+        )
+
+class DataModule(pl.LightningDataModule):
+    def __init__(self, hidden_dim: int, learning_rate: float, parameter: int):
+        self.parameter = parameter
+        self.hidden_dim = hidden_dim
+        self.learning_rate = learning_rate
+    
+    def train_dataloader(self):
+        ...
+
+class DataModule(pl.LightningModule):
+    def __init__(self, parameter: int):
+        self.parameter = parameter
+    
+    def training_step(self):
+        ...
+
+```
+
+Launching a hyperparameter tuning job with `RayTune` using this `LightningCLI` is as simple as
+
+```python
+from ray import tune
+from raylightning import run
+
+# define search space
+search_space = {
+    "model.init_args.hidden_dim": tune.choice([32, 64])
+    "model.init_args.learning_rate": tune.loguniform(1e-4, 1e-2)
+}
+
+# define scheduler
+scheduler = ASHAScheduler(max_t=50, grace_period=1, reduction_factor=2)
+
+# pass command line style arguments as if you were
+# running your `LightningCLI` from the command line
+args = ["--config", "/path/to/config.yaml"]
+results = run(
+    args=args
+    cli_cls=simple_cli,
+    name="tune-test",
+    metric="val_loss",
+    objective="min",
+    search_space=search_space,
+    scheduler=scheduler,
+    storage_dir=storage_dir,
+    address=None,
+    num_samples=num_samples,
+    workers_per_trial=1,
+    gpus_per_worker=1.0,
+    cpus_per_gpu=4.0,
+    temp_dir=None,
+)
+```
+
+s3 storage works out of the box. Make sure you have set the `AWS_ENDPOINT_URL`, `AWS_ACCESS_KEY_ID`, and `AWS_SECRET_ACCESS_KEY` environment variables set. Then, simply pass an s3 path (e.g. `s3://{bucket}/{folder}` to the `storage_dir` argument.
+

lightray-0.1.0.dist-info/RECORD

@@ -0,0 +1,8 @@
+lightray/__init__.py,sha256=JgTN3Hzek3t8ejv76E8ffQSuptc3PYtEHjTKinsc9ks,22
+lightray/callbacks.py,sha256=YaTNoz-KP1hS2ZjWs09_0De1SCT-f10-u5NMN_Otivw,2024
+lightray/fs.py,sha256=3WTWtWue8L1rFNusmtJTNrnRQWzFYXJHlf4yjN9Ewgc,1461
+lightray/tune.py,sha256=xmDj6mdKe1xjOY6ya5USWsTJKx9_iSDlux56Z232734,3409
+lightray/utils.py,sha256=qNtWoReHTVSCQJ5hUHJGcci_VmwJnBVyHJ7phHb6pOk,8407
+lightray-0.1.0.dist-info/METADATA,sha256=6onMoD9_uafK11YnUctjFYetjRTF7uAv3Eo9YlPKLpI,2919
+lightray-0.1.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+lightray-0.1.0.dist-info/RECORD,,

lightray-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: poetry-core 1.9.0
+Root-Is-Purelib: true
+Tag: py3-none-any