add lightning hook

.github/workflows/test.yaml

@@ -7,7 +7,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ["3.7", "3.8", "3.9", "3.10", "3.11"]
+        python-version: ["3.8", "3.9", "3.10", "3.11"]
 
     steps:
       - uses: actions/checkout@v3
@@ -18,14 +18,9 @@ jobs:
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
-      - name: Install dependencies != 3.11
-        if: ${{ matrix.python-version < '3.11' }}
-        run: |
-          pip install .[testing,ray]
       - name: Install dependencies >= 3.11
-        if: ${{ matrix.python-version >= '3.11' }}
         run: |
-          pip install .[testing]
+          pip install .[testing,ray,lightning]
       - name: Test with pytest
         run: |
           pytest

.pre-commit-config.yaml

@@ -50,7 +50,7 @@ repos:
     rev: v3.3.2
     hooks:
       - id: pyupgrade
-        args: ["--py37-plus"]
+        args: ["--py38-plus"]
   - repo: https://github.com/asottile/setup-cfg-fmt
     rev: "v2.2.0"
     hooks:

CHANGELOG.md

@@ -3,6 +3,16 @@
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## 1.1.0 (04.07.2023)
+
+### Added
+
+- Added support for pytorch lightning
+
+### Changed
+
+- Removed python 3.7 support
+
 ## 1.0.4 (18.01.2023)
 
 ### Fixed

README.md

@@ -104,6 +104,27 @@ for batch_idx, (data, target) in enumerate(train_loader):
         trigger_sync()  # <-- New!
 ```
 
+#### With pytorch lightning
+
+```python
+from wandb_osh.lightning_hooks import TriggerWandbSyncLightningCallback  # <-- New!
+from pytorch_lightning.loggers import WandbLogger
+from pytorch_lightning import Trainer
+
+logger = WandbLogger(
+    project="project",
+    group="group",
+    offline=True,
+)
+
+model = MyLightningModule()
+trainer = Trainer(
+    logger=logger,
+    callbacks=[TriggerWandbSyncLightningCallback()]  # <-- New!
+)
+trainer.fit(model, train_dataloader, val_dataloader)
+```
+
 #### With ray tune
 
 > **Note**

setup.cfg

@@ -24,7 +24,7 @@ project_urls =
     Bug Tracker =   https://github.com/klieret/wandb-offline-sync-hook/issues
     Documentation = https://wandb_osh.readthedocs.io/
     Source Code =   https://github.com/klieret/wandb-offline-sync-hook
-python_requires = >=3.7
+python_requires = >=3.8
 
 [options]
 packages = find:
@@ -46,6 +46,8 @@ console_scripts =
     wandb-osh = wandb_osh.cli:main
 
 [options.extras_require]
+lightning =
+    pytorch-lightning
 ray =
     ray[tune]
 testing =

src/wandb_osh/lightning_hooks.py

@@ -0,0 +1,33 @@
+from __future__ import annotations
+
+from os import PathLike
+
+from pytorch_lightning.callbacks import Callback
+
+from wandb_osh.hooks import TriggerWandbSyncHook, _comm_default_dir
+
+
+class TriggerWandbSyncLightningCallback(Callback):
+    def __init__(self, communication_dir: PathLike = _comm_default_dir):
+        """Hook to be used when interfacing wandb with pytorch lightning.
+
+        Args:
+            communication_dir: Directory used for communication with wandb-osh.
+
+        Usage
+
+        .. code-block:: python
+
+            from wandb_osh.lightning_hooks import TriggerWandbSyncLightningCallback
+
+            trainer = Trainer(callbacks=[TriggerWandbSyncLightningCallback()])
+
+        """
+        super().__init__()
+        self._hook = TriggerWandbSyncHook(communication_dir=communication_dir)
+
+    def on_validation_epoch_end(
+        self,
+        *args,
+    ):
+        self._hook()

tests/conftest.py

@@ -0,0 +1,7 @@
+from __future__ import annotations
+
+import logging
+
+from wandb_osh.util.log import logger
+
+logger.setLevel(logging.DEBUG)

tests/test_lightning_hooks.py

@@ -0,0 +1,12 @@
+from __future__ import annotations
+
+import wandb
+
+from wandb_osh.lightning_hooks import TriggerWandbSyncLightningCallback
+
+
+def test_manual_trigger(tmp_path):
+    wandb.init(project="test", mode="offline", dir=tmp_path)
+    lh = TriggerWandbSyncLightningCallback(tmp_path)
+    lh.on_validation_epoch_end(None, None)  # type: ignore
+    assert len([f for f in tmp_path.iterdir() if f.suffix == ".command"]) == 1