Merge pull request #4 from f-dangel/mkdocs

[DOC] Set up `mkdocs` to for generating documentation

README.md

@@ -1,4 +1,4 @@
-# Summary
+# wandb_preempt
 
 This repository contains a tutorial on how to combine `wandb` sweeps with
 Slurm's pre-emption, i.e. how to automatically re-queue and resume runs from a
@@ -8,58 +8,10 @@ This is work in progress:
 
 - TODO Debug failure scenarios
 
-- TODO Freeze interface
-
 - TODO Write a self-contained walk-through.
 
 ## Installation
 
 ```bash
 pip install git+https://github.com/f-dangel/wandb_preempt.git@main
 ```
-
-# Developer guide
-
-This guide describes principles and workflows for developers.
-
-## Setup
-
-We recommend programming in a fresh virtual environment. You can set up the
-`conda` environment and activate it
-
-```bash
-make conda-env
-conda activate wandb_preempt
-```
-
-If you don't use `conda`, set up your preferred environment and run
-
-```bash
-pip install -e ."[lint,test]"
-```
-to install the package in editable mode, along with all required development dependencies
-(the quotes are for OS compatibility, see
-[here](https://github.com/mu-editor/mu/issues/852#issuecomment-498759372)).
-
-## Continuous integration
-
-To standardize code style and enforce high quality, checks are carried out with
-Github actions when you push. You can also run them locally, as they are managed
-via `make`:
-
-- Run tests with `make test`
-
-- Run all linters with `make lint`, or separately with:
-
-    - Run auto-formatting and import sorting with `make black` and `make isort`
-
-    - Run linting with `make flake8`
-
-    - Run docstring checks with `make pydocstyle-check` and `make darglint-check`
-
-## Documentation
-
-We use the [Google docstring
-convention](https://sphinxcontrib-napoleon.readthedocs.io/en/latest/example_google.html)
-and `mkdocs` which allows using markdown syntax in a docstring to achieve
-formatting.

docs/api.md

@@ -0,0 +1,8 @@
+::: wandb_preempt.get_resume_value
+
+::: wandb_preempt.Checkpointer
+    options:
+        members:
+            - __init__
+            - load_latest_checkpoint
+            - step

docs/develop.md

@@ -0,0 +1,45 @@
+# Developer guide
+
+This guide describes principles and workflows for developers.
+
+## Setup
+
+We recommend programming in a fresh virtual environment. You can set up the
+`conda` environment and activate it
+
+```bash
+make conda-env
+conda activate wandb_preempt
+```
+
+If you don't use `conda`, set up your preferred environment and run
+
+```bash
+pip install -e ."[lint,test]"
+```
+to install the package in editable mode, along with all required development dependencies
+(the quotes are for OS compatibility, see
+[here](https://github.com/mu-editor/mu/issues/852#issuecomment-498759372)).
+
+## Continuous integration
+
+To standardize code style and enforce high quality, checks are carried out with
+Github actions when you push. You can also run them locally, as they are managed
+via `make`:
+
+- Run tests with `make test`
+
+- Run all linters with `make lint`, or separately with:
+
+    - Run auto-formatting and import sorting with `make black` and `make isort`
+
+    - Run linting with `make flake8`
+
+    - Run docstring checks with `make pydocstyle-check` and `make darglint-check`
+
+## Documentation
+
+We use the [Google docstring
+convention](https://sphinxcontrib-napoleon.readthedocs.io/en/latest/example_google.html)
+and `mkdocs` which allows using markdown syntax in a docstring to achieve
+formatting. To build the docs, run `mkdocs serve` in the repository root.

docs/index.md

@@ -0,0 +1 @@
+../README.md

mkdocs.yml

@@ -0,0 +1,45 @@
+site_name: wandb_preempt
+site_url: https://https://readthedocs.org/projects/wandb_preempt/
+repo_url: https://github.com/f-dangel/wandb_preempt/
+repo_name: f-dangel/wandb_preempt
+site_author: Felix Dangel, Scott Lowe
+watch:
+  - wandb_preempt
+nav:
+    - Getting Started: index.md
+    - API Documentation: api.md
+    - Developer Notes: develop.md
+theme:
+    name: material
+    features:
+        - content.code.copy
+copyright: Copyright © 2024 Felix Dangel, Scott Lowe
+markdown_extensions:
+    - pymdownx.arithmatex: # LaTeX math
+          generic: true
+    - pymdownx.highlight: # code highlighting
+          anchor_linenums: true
+          line_spans: __span
+          pygments_lang_class: true
+    - pymdownx.inlinehilite # code highlighting
+    - pymdownx.snippets # code highlighting
+    - pymdownx.superfences # code highlighting
+    - footnotes
+plugins:
+    - mkdocstrings:
+        default_handler: python
+        handlers:
+            python:
+                options:
+                      show_root_heading: true
+                      show_source: true
+                      show_bases: true
+                      show_signature_annotations: true
+                      separate_signature: true
+                      docstring_section_style: list
+                      merge_init_into_class: true
+    - search
+extra_javascript:
+    - javascripts/mathjax.js # LaTeX math
+    - https://polyfill.io/v3/polyfill.min.js?features=es6 # LaTeX math
+    - https://cdn.jsdelivr.net/npm/mathjax@3/es5/tex-mml-chtml.js # LaTeX math

pyproject.toml

@@ -70,6 +70,14 @@ example = [
     "torchvision",
 ]
 
+# Dependencies needed to generate the documentation (comma/line-separated), uses
+# pinned versions to avoid future breakdowns
+doc = [
+    "mkdocs==1.4.3",
+    "mkdocs-material==9.1.17",
+    "mkdocstrings[python]==0.22.0",
+]
+
 [tool.setuptools_scm]
 
 [tool.isort]

wandb_preempt/checkpointer.py

@@ -13,6 +13,8 @@
 import wandb
 from torch import cuda, device, get_rng_state, load, save, set_rng_state
 from torch.cuda.amp import GradScaler
+from torch.nn import Module
+from torch.optim import Optimizer
 from torch.optim.lr_scheduler import LRScheduler
 from wandb import Api
 
@@ -71,8 +73,8 @@ class Checkpointer:
     def __init__(
         self,
         run_id: str,
-        model,
-        optimizer,
+        model: Module,
+        optimizer: Optimizer,
         lr_scheduler: Optional[LRScheduler] = None,
         scaler: Optional[GradScaler] = None,
         metadata: Optional[Dict] = None,