Add stubs for upcoming documentation (#19)

* Setup a readme and a documentation page

Signed-off-by: Fabrice Normandin <[email protected]>

* Add more pages to the docs

Signed-off-by: Fabrice Normandin <[email protected]>

* Add a plugin to fix md formatting issues in mkdocs

Signed-off-by: Fabrice Normandin <[email protected]>

* Playing around with mkdocs to generate docs

Signed-off-by: Fabrice Normandin <[email protected]>

* Ignore `SUMMARY.md` files in mdformat hook

Signed-off-by: Fabrice Normandin <[email protected]>

* Fix structure a bit, remove unused files

Signed-off-by: Fabrice Normandin <[email protected]>

* Remove need for preexisting $SLURM_TMPDIR env var

Signed-off-by: Fabrice Normandin <[email protected]>

* Fix pre-commit issue

Signed-off-by: Fabrice Normandin <[email protected]>

* Small doc fixes

Signed-off-by: Fabrice Normandin <[email protected]>

* Add a workflow to deploy the docs

Signed-off-by: Fabrice Normandin <[email protected]>

* Add material theme to the docs

Signed-off-by: Fabrice Normandin <[email protected]>

* Tweak indent of install.md

Signed-off-by: Fabrice Normandin <[email protected]>

* Add a (wip) in the title of the docs

Signed-off-by: Fabrice Normandin <[email protected]>

* Add a link to the intro section

Signed-off-by: Fabrice Normandin <[email protected]>

---------

Signed-off-by: Fabrice Normandin <[email protected]>

.devcontainer/devcontainer.json

@@ -74,7 +74,7 @@
 		// FIXME: This assumes that either the NETWORK_DIR environment variable is set on the host, or
 		// that the /network directory exists.
 		"source=${localEnv:NETWORK_DIR:/network},target=/network,type=bind,readonly",
-		// Mount a /tmp on the host machine to /tmp/slurm_tmpdir in the container.
+		// Mount $SLURM_TMPDIR on the host machine to /tmp/slurm_tmpdir in the container.
 		// note: there's also a SLURM_TMPDIR env variable set to /tmp/slurm_tmpdir in the container.
 		// NOTE: this assumes that either $SLURM_TMPDIR is set on the host machine (e.g. a compute node)
 		// or that `/tmp/slurm_tmpdir` exists on the host machine.
@@ -89,7 +89,7 @@
 	// doesn't fail.
 	"initializeCommand": {
 		"create pdm install cache": "mkdir -p ${SCRATCH?need the SCRATCH environment variable to be set.}/.cache/pdm",      // todo: put this on $SCRATCH on the host (e.g. compute node)
-		"create fake SLURM_TMPDIR": "mkdir -p ${SLURM_TMPDIR?need the SLURM_TMPDIR environment variable to be set.}"  // this is fine on compute nodes
+		"create fake SLURM_TMPDIR": "mkdir -p ${SLURM_TMPDIR:-/tmp/slurm_tmpdir}"  // this is fine on compute nodes
 	},
 	// NOTE: Getting some permission issues with the .cache dir if mounting .cache/pdm to
 	// .cache/pdm in the container. Therefore, here I'm making a symlink from ~/.cache/pdm to

.github/workflows/docs.yml

@@ -0,0 +1,19 @@
+name: Publish docs via GitHub Pages
+on:
+  push:
+    branches:
+      - master
+
+jobs:
+  build:
+    name: Deploy docs
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout main
+        uses: actions/checkout@v4
+
+      - name: Deploy docs
+        uses: mhausenblas/[email protected]
+        # Or use mhausenblas/mkdocs-deploy-gh-pages@nomaterial to build without the mkdocs-material theme
+        env:
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}

.pre-commit-config.yaml

@@ -72,16 +72,20 @@ repos:
 
   # md formatting
   - repo: https://github.com/executablebooks/mdformat
-    rev: 0.7.16
+    rev: 0.7.17
     hooks:
       - id: mdformat
+        exclude: 'SUMMARY.md'
         args: ["--number"]
         additional_dependencies:
           - mdformat-gfm
           - mdformat-tables
           - mdformat_frontmatter
-          # - mdformat-toc
-          # - mdformat-black
+          - mdformat-toc
+          - mdformat-config
+          - mdformat-black
+          # see https://github.com/KyleKing/mdformat-mkdocs
+          - mdformat-mkdocs[recommended]>=2.1.0
         require_serial: true
 
 

README.md

@@ -1,4 +1,8 @@
-# research_template
+# Research Project Template
 
 ![Build](https://github.com/mila-iqia/ResearchTemplate/workflows/build.yml/badge.svg)
 [![codecov](https://codecov.io/gh/mila-iqia/ResearchTemplate/graph/badge.svg?token=I2DYLK8NTD)](https://codecov.io/gh/mila-iqia/ResearchTemplate)
+
+Please note: This is a **Work-in-Progress**. The goal is to make a first release by the end of summer 2024.
+
+For now, feel free to take a look at the [documentation page](https://mila-iqia.github.io/ResearchTemplate/) if you want more information about this project.

docs/.gitignore

@@ -0,0 +1,2 @@
+reference.md
+reference/*

docs/SUMMARY.md

@@ -0,0 +1,13 @@
+* [Home](index.md)
+  * Overview
+    * overview/*.md
+  * Getting Started
+    * getting_started/*.md
+  * Reference
+    * reference/*
+  * Examples
+    * examples/*
+* [Tests](tests.md)
+* [Related projects](related.md)
+* [Getting Help](help.md)
+* [Contributing](contributing.md)

docs/contributing.md

@@ -0,0 +1,5 @@
+# Contributing
+
+TODOs:
+
+- [ ] Describe how to contribute to the project.

docs/examples/examples.md

@@ -0,0 +1,22 @@
+# Examples
+
+TODOs:
+
+- [ ] Show examples (that are also to be tested with doctest or similar) of how to add a new algo.
+- [ ] Show examples of how to add a new datamodule.
+- [ ] Add a link to the RL example once [#13](https://github.com/mila-iqia/ResearchTemplate/issues/13) is done.
+- [ ] Add a link to the NLP example once [#14](https://github.com/mila-iqia/ResearchTemplate/issues/14) is done.
+- [ ] Add an example of how to use Jax for the dataset/dataloading:
+    - Either through an RL example, or with `tfds` in [#18](https://github.com/mila-iqia/ResearchTemplate/issues/18)
+
+## Simple run
+
+```bash
+python project/main.py algorithm=example_algo datamodule=mnist network=fcnet
+```
+
+## Running a Hyper-Parameter sweep on a SLURM cluster
+
+```bash
+python project/main.py experiment=cluster_sweep_example
+```

docs/examples/jax.md

@@ -0,0 +1,18 @@
+# Using Jax
+
+You can use Jax for your dataloading, your network, or the learning algorithm, all while still benefiting from the nice stuff that comes from using PyTorch-Lightning.
+
+How does this work?
+Well, we use [torch-jax-interop](https://www.github.com/lebrice/torch_jax_interop), another package developed here at Mila, which allows easy interop between torch and jax code. See the readme on that repo for more details.
+
+## Example Algorithm that uses Jax
+
+You can use Jax for your training step, but not the entire training loop (since that is handled by Lightning).
+There are a few good reasons why you should let Lightning handle the training loop, most notably the fact that it handles all the logging, checkpointing, and other stuff that you'd lose if you swapped out the entire training framework for something based on Jax.
+
+In this [example Jax algorithm](https://www.github.com/mila-iqia/ResearchTemplate/tree/master/project/algorithms/jax_algo.py),
+a Neural network written in Jax (using [flax](https://flax.readthedocs.io/en/latest/)) is wrapped using the `torch_jax_interop.JaxFunction`, so that its parameters are learnable. The parameters are saved on the LightningModule as nn.Parameters (which use the same underlying memory as the jax arrays). In this example, the loss function is written in PyTorch, while the network forward and backward passes are written in Jax.
+
+## Example datamodule that uses Jax
+
+(todo)

docs/generate_reference_docs.py

@@ -0,0 +1,128 @@
+#!/usr/bin/env python
+# based on https://github.com/mkdocstrings/mkdocstrings/blob/5802b1ef5ad9bf6077974f777bd55f32ce2bc219/docs/gen_doc_stubs.py#L25
+
+
+import textwrap
+from pathlib import Path
+
+import mkdocs_gen_files
+import mkdocs_gen_files.nav
+
+from project.utils.env_vars import REPO_ROOTDIR
+
+nav = mkdocs_gen_files.nav.Nav()
+
+
+package = "project"
+module = "project/main.py"
+submodules = ["project.datamodules", "project.utils", "project.networks", "project.algorithms"]
+
+
+def _get_import_path(module_path: Path) -> str:
+    """Returns the path to use to import a given (internal) module."""
+    return ".".join(module_path.relative_to(REPO_ROOTDIR).with_suffix("").parts)
+
+
+def add_doc_for_module(module_path: Path) -> None:
+    a = "reference" / (module_path.relative_to(REPO_ROOTDIR).with_suffix(".md"))
+    module_import_path = _get_import_path(module_path)
+
+    with mkdocs_gen_files.open(a, "w") as f:
+        print(
+            textwrap.dedent(f"""\
+            ::: {module_import_path}
+            """),
+            file=f,
+        )
+    docs_dir = REPO_ROOTDIR / "docs"
+    module_path_relative_to_docs_dir = module_path.relative_to(docs_dir, walk_up=True)
+    mkdocs_gen_files.set_edit_path(a, module_path_relative_to_docs_dir)
+
+
+def get_modules(package: Path) -> list[Path]:
+    return [
+        p
+        for p in package.glob("*.py")
+        if not p.name.endswith("_test.py") and not p.name == "__init__.py"
+    ]
+
+
+def get_subpackages(package: Path) -> list[Path]:
+    return [
+        p
+        for p in package.iterdir()
+        if p.is_dir() and not p.name.startswith("__") and (p / "__init__.py").exists()
+    ]
+
+
+project_nav = mkdocs_gen_files.nav.Nav()
+with mkdocs_gen_files.open("reference/project/main.md", "w") as f:
+    print(
+        textwrap.dedent("""\
+        ::: project.main
+        """),
+        file=f,
+    )
+nav["project", "main"] = "project/main.md"
+mkdocs_gen_files.set_edit_path("reference/project/main.md", "../project/main.py")
+
+with mkdocs_gen_files.open("reference/project/experiment.md", "w") as f:
+    print(
+        textwrap.dedent("""\
+        ::: project.experiment
+        """),
+        file=f,
+    )
+nav["project", "experiment"] = "reference/project/experiment.md"
+mkdocs_gen_files.set_edit_path("reference/project/experiment.md", "../project/experiment.py")
+
+project_utils_nav = mkdocs_gen_files.nav.Nav()
+with mkdocs_gen_files.open("reference/project/utils/types.md", "w") as f:
+    print(
+        textwrap.dedent("""\
+        ::: project.utils.types
+            options:
+                show_source: true
+        """),
+        file=f,
+    )
+nav["project", "utils", "types"] = "reference/project/utils/types.md"
+mkdocs_gen_files.set_edit_path("reference/project/utils/types.md", "../project/utils/types.py")
+
+
+with mkdocs_gen_files.open("reference.md", "w") as nav_file:
+    # assert False, "\n".join(nav.build_literate_nav())
+    nav_file.writelines(nav.build_literate_nav())
+
+# with mkdocs_gen_files.open("reference/SUMMARY.md", "w") as project_nav_file:
+#     project_nav_file.writelines(project_nav.build_literate_nav())
+
+
+# project_root = REPO_ROOTDIR / "project"
+# for python_module_path in sorted(
+#     f
+#     # for f in project_root.glob("*.py")
+#     for f in [(project_root / "project")]
+#     if not f.name.endswith("_test.py") and not f.name == "__init__.py"
+# ):
+#     doc_path = python_module_path.relative_to(REPO_ROOTDIR).with_suffix(".md")
+
+#     full_doc_path = Path("reference") / doc_path
+
+#     nav[full_doc_path.with_suffix("").parts] = str(full_doc_path)
+
+#     with mkdocs_gen_files.open(full_doc_path, "w") as f:
+#         module_import_path = ".".join(
+#             python_module_path.relative_to(REPO_ROOTDIR).with_suffix("").parts
+#         )
+#         print(f"::: {module_import_path}", file=f)
+
+#     mkdocs_gen_files.set_edit_path(
+#         full_doc_path, python_module_path.relative_to(REPO_ROOTDIR / "docs", walk_up=True)
+#     )
+
+# nav["mkdocs_autorefs", "references"] = "autorefs/references.md"
+# nav["mkdocs_autorefs", "plugin"] = "autorefs/plugin.md"
+
+# with mkdocs_gen_files.open("reference.md", "w") as nav_file:
+#     nav_file.writelines(nav.build_literate_nav())

docs/getting_started/install.md

@@ -0,0 +1,73 @@
+# Installation instructions
+
+There are two ways to install this project
+
+1. Using Conda (recommended for newcomers)
+2. Using a development container (recommended if you are able to install Docker on your machine)
+
+## Using Conda and pip
+
+### Prerequisites
+
+You need to have [Conda](https://docs.conda.io/en/latest/) installed on your machine.
+
+### Installation
+
+1. Clone the repository and navigate to the root directory:
+
+    ```bash
+    git clone https://www.github.com/mila-iqia/ResearchTemplate
+    cd ResearchTemplate
+    ```
+
+2. Create a conda environment
+
+    ```bash
+    conda create -n research_template python=3.12
+    conda activate research_template
+    ```
+
+    Notes:
+
+    - If you don't Conda installed, you can download it from [here](https://docs.conda.io/en/latest/miniconda.html).
+    - If you'd rather use a virtual environment instead of Conda, you can totally do so, as long as you have a version of Python >= 3.12.
+
+    <!-- TODO: - If you're on the `mila` cluster, you can run this setup script: (...) -->
+
+3. Install the package using pip:
+
+    ```bash
+    pip install -e .
+    ```
+
+    Optionally, you can also install the package using [PDM](https://pdm-project.org/en/latest/). This makes it easier to add or change the dependencies later on:
+
+    ```bash
+    pip install pdm
+    pdm install
+    ```
+
+## Using a development container
+
+This repo provides a [Devcontainer](https://code.visualstudio.com/docs/remote/containers) configuration for [Visual Studio Code](https://code.visualstudio.com/) to use a Docker container as a pre-configured development environment. This avoids struggles setting up a development environment and makes them reproducible and consistent.  and make yourself familiar with the [container tutorials](https://code.visualstudio.com/docs/remote/containers-tutorial) if you want to use them. In order to use GPUs, you can enable them within the `.devcontainer/devcontainer.json` file.
+
+1. Setup Docker on your local machine
+
+    On an Linux machine where you have root access, you can install Docker using the following commands:
+
+    ```bash
+    curl -fsSL https://get.docker.com -o get-docker.sh
+    sudo sh get-docker.sh
+    ```
+
+    On Windows or Mac, follow [these installation instructions](https://code.visualstudio.com/docs/remote/containers#_installation)
+
+2. (optional) Install the [nvidia-container-toolkit](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html) to use your local machine's GPU(s).
+
+3. Install the [Dev Containers extension](vscode:extension/ms-vscode-remote.remote-containers) for Visual Studio Code.
+
+4. When opening repository in Visual Studio Code, you should be prompted to reopen the repository in a container:
+
+    ![VsCode popup image](https://github.com/mila-iqia/ResearchTemplate/assets/13387299/37d00ce7-1214-44b2-b1d6-411ee286999f)
+
+    Alternatively, you can open the command palette (Ctrl+Shift+P) and select `Dev Containers: Rebuild and Reopen in Container`.

docs/help.md

@@ -0,0 +1,5 @@
+# Help and Support
+
+## FAQ
+
+## How to get help

docs/index.md

@@ -0,0 +1,51 @@
+# Research Project Template
+
+[![Build](https://github.com/mila-iqia/ResearchTemplate/actions/workflows/build.yml/badge.svg?branch=master)](https://github.com/mila-iqia/ResearchTemplate/actions/workflows/build.yml)
+[![codecov](https://codecov.io/gh/mila-iqia/ResearchTemplate/graph/badge.svg?token=I2DYLK8NTD)](https://codecov.io/gh/mila-iqia/ResearchTemplate)
+[![hydra](https://img.shields.io/badge/Config-Hydra_1.3-89b8cd)](https://hydra.cc/)
+[![license](https://img.shields.io/badge/License-MIT-green.svg?labelColor=gray)](https://github.com/mila-iqia/ResearchTemplate#license)
+
+Please note: This is a Work-in-Progress. The goal is to make a first release by the end of summer 2024.
+
+This is a research project template. It is meant to be a starting point for ML researchers at [Mila](https://mila.quebec/en).
+
+For more context, see [this  introduction to the project.](overview/intro.md).
+
+## Overview
+
+This project makes use of the following libraries:
+
+- [Hydra](https://hydra.cc/) is used to configure the project. It allows you to define configuration files and override them from the command line.
+- [PyTorch Lightning](https://lightning.ai/docs/pytorch/stable/) is used to as the training framework. It provides a high-level interface to organize ML research code.
+    - 🔥 Please note: You can also use [Jax](https://jax.readthedocs.io/en/latest/) with this repo, as is shown in the [Jax example](examples/examples.md#using-jax) 🔥
+- [Weights & Biases](https://wandb.ai) is used to log metrics and visualize results.
+- [pytest](https://docs.pytest.org/en/stable/) is used for testing.
+
+## Usage
+
+To see all available options:
+
+```bash
+python project/main.py --help
+```
+
+For a detailed list of examples, see the [examples page](examples/examples.md).
+
+<!-- * `mkdocs new [dir-name]` - Create a new project.
+* `mkdocs serve` - Start the live-reloading docs server.
+* `mkdocs build` - Build the documentation site.
+* `mkdocs -h` - Print help message and exit. -->
+
+## Project layout
+
+```
+pyproject.toml   # Project metadata and dependencies
+project/
+    main.py      # main entry-point
+    algorithms/  # learning algorithms
+    datamodules/ # datasets, processing and loading
+    networks/    # Neural networks used by algorithms
+    configs/     # configuration files
+docs/            # documentation
+conftest.py      # Test fixtures and utilities
+```