Documentation updates for 0.3.0 (#138)

* Adding torch clean up function between examples * Adding v0.3.0 switcher * Changelog * Adding header
NVIDIA · Sep 20, 2024 · ccd2311 · ccd2311
1 parent 730460a
commit ccd2311
Show file tree

Hide file tree

Showing 6 changed files with 42 additions and 9 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -53,6 +53,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Dependencies
 
 - Restrict torch_harmonics version to >=0.5.0, <0.7.1
+- Removed specific ONNX version requirement, newer ORT-gpu versions appear to operate
+  fine with CUDA 12
 
 ## [0.2.0] - 2024-07-23
 

diff --git a/docs/_static/switcher.json b/docs/_static/switcher.json
@@ -5,6 +5,11 @@
         "url": "https://nvidia.github.io/earth2studio/",
         "preferred": true
     },
+    {
+        "name": "0.3.0",
+        "version": "0.3.0",
+        "url": "https://nvidia.github.io/earth2studio/v/0.3.0/"
+    },
     {
         "name": "0.2.0",
         "version": "0.2.0",

diff --git a/docs/conf.py b/docs/conf.py
@@ -42,7 +42,8 @@
 release = version("earth2studio")
 
 sys.path.insert(0, root.parent.as_posix())
-sys.path.insert(0, modulus.as_posix())
+# Add current folder to use sphinxext.py
+sys.path.insert(0, os.path.dirname(__file__))
 
 # -- Project information -----------------------------------------------------
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#project-information
@@ -69,7 +70,7 @@
 source_suffix = [".rst", ".md"]
 myst_enable_extensions = ["colon_fence"]
 templates_path = ["_templates"]
-exclude_patterns = ["_build", "Thumbs.db", ".DS_Store"]
+exclude_patterns = ["_build", "sphinxext.py", "Thumbs.db", ".DS_Store"]
 autodoc_typehints = "description"
 
 # -- Options for HTML output -------------------------------------------------
@@ -81,7 +82,6 @@
     "css/custom.css",
 ]
 html_theme_options = {
-    "announcement": "Earth2Studio is now OSS!",
     "logo": {
         "text": "Earth2Studio",
         "image_light": "_static/NVIDIA-Logo-V-ForScreen-ForLightBG.png",
@@ -131,5 +131,8 @@
     "run_stale_examples": run_stale_examples,
     "backreferences_dir": "modules/backreferences",
     "doc_module": ("earth2studio"),
+    "reset_modules": ("matplotlib", "sphinxext.reset_torch"),
+    "show_memory": True,
     "exclude_implicit_doc": {r"load_model", r"load_default_package"},
+    "log_level": {"backreference_missing": "debug"},
 }
diff --git a/docs/sphinxext.py b/docs/sphinxext.py
@@ -0,0 +1,25 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+def reset_torch(gallery_conf, fname):
+    """Function to clean up torch between sphinx examples"""
+    import gc
+
+    import torch
+
+    torch.cuda.empty_cache()
+    gc.collect()
diff --git a/docs/userguide/about/install.md b/docs/userguide/about/install.md
@@ -150,7 +150,6 @@ Note that for ONNX models to work we will need a [specific install](https://onnx
 ```bash
 docker run -i -t nvcr.io/nvidia/pytorch:24.01-py3
 
->>> pip install onnx onnxruntime-gpu==1.17.1 --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-12/pypi/simple/
 >>> pip install "makani[all] @ git+https://github.com/NVIDIA/[email protected]"
 
 >>> pip install earth2studio[all]
@@ -168,7 +167,6 @@ conda activate earth2studio
 
 conda install pytorch pytorch-cuda=12.1 -c pytorch -c nvidia
 conda install eccodes python-eccodes -c conda-forge
-pip install onnx onnxruntime-gpu==1.17.1 --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-12/pypi/simple/
 pip install "makani[all] @ git+https://github.com/NVIDIA/[email protected]"
 
 pip install earth2studio[all]

diff --git a/pyproject.toml b/pyproject.toml
@@ -58,13 +58,13 @@ Changelog = "https://github.com/NVIDIA/earth2studio/blob/main/CHANGELOG.md"
 
 [project.optional-dependencies]
 fengwu = [
-    "onnxruntime-gpu==1.17.1",
+    "onnxruntime-gpu",
 ]
 fuxi = [
-    "onnxruntime-gpu==1.17.1",
+    "onnxruntime-gpu",
 ]
 pangu = [
-    "onnxruntime-gpu==1.17.1",
+    "onnxruntime-gpu",
 ]
 sfno = [
     # makani needs to be installed manually
@@ -81,7 +81,7 @@ corrdiff = [
     "einops>=0.7.0",
 ]
 all = [
-    "onnxruntime-gpu>=1.14.0",
+    "onnxruntime-gpu",
     "jsbeautifier",
     "pynvml",
     "numba",