update docs and deps

Author: apoorvkh

docs/conf.py

@@ -1,12 +1,17 @@
 """Configuration file for the Sphinx documentation builder."""
 
+from glob import glob
+import shutil
+
+shutil.copyfile("../README.md", "source/README.md")
+
 project = "torchrunx"
 github_username = "apoorvkh"
 github_repository = "torchrunx"
 html_theme = "furo"
 language = "en"
 
-html_extra_path = ["source/examples/scripts"]
+html_extra_path = list(glob("source/examples/scripts/*.py"))
 
 extensions = [
     "autodoc2",
@@ -27,69 +32,5 @@
     "numpy": ("https://numpy.org/doc/stable", None),
 }
 
-## Copy README.md to docs/source/
-import shutil
-shutil.copyfile("../README.md", "source/README.md")
-## End of "Copy README.md to docs/source/"
-
-## sphinx.ext.linkcode configuration
-# Link code to Github source
-# From: https://github.com/scikit-learn/scikit-learn/blob/main/doc/sphinxext/github_link.py
-
-import inspect
-import os
-import subprocess
-import sys
-from operator import attrgetter
-
-package = project
-
-try:
-    revision = (
-        subprocess.check_output("git rev-parse --short HEAD".split()).strip().decode("utf-8")
-    )
-except (subprocess.CalledProcessError, OSError):
-    print("Failed to execute git to get revision")
-    revision = None
-
-url_fmt = (
-    f"https://github.com/{github_username}/{github_repository}/"
-    "blob/{revision}/src/{package}/{path}#L{lineno}"
-)
-
-def linkcode_resolve(domain, info):
-    if revision is None:
-        return
-    if domain not in ("py", "pyx"):
-        return
-    if not info.get("module") or not info.get("fullname"):
-        return
-
-    class_name = info["fullname"].split(".")[0]
-    module = __import__(info["module"], fromlist=[class_name])
-    obj = attrgetter(info["fullname"])(module)
-
-    # Unwrap the object to get the correct source
-    # file in case that is wrapped by a decorator
-    obj = inspect.unwrap(obj)
-
-    try:
-        fn = inspect.getsourcefile(obj)
-    except Exception:
-        fn = None
-    if not fn:
-        try:
-            fn = inspect.getsourcefile(sys.modules[obj.__module__])
-        except Exception:
-            fn = None
-    if not fn:
-        return
-
-    fn = os.path.relpath(fn, start=os.path.dirname(__import__(package).__file__))
-    try:
-        lineno = inspect.getsourcelines(obj)[1]
-    except Exception:
-        lineno = ""
-    return url_fmt.format(revision=revision, package=package, path=fn, lineno=lineno)
-
-## End of "sphinx.ext.linkcode configuration"
+from docs.linkcode_github import generate_linkcode_resolve_fn
+linkcode_resolve = generate_linkcode_resolve_fn(project, github_username, github_repository)

docs/linkcode_github.py

@@ -0,0 +1,62 @@
+## sphinx.ext.linkcode configuration
+# Link code to Github source
+# From: https://github.com/scikit-learn/scikit-learn/blob/main/doc/sphinxext/github_link.py
+
+import inspect
+import os
+import subprocess
+import sys
+from operator import attrgetter
+
+
+def generate_linkcode_resolve_fn(package, github_username, github_repository):
+
+    try:
+        revision = (
+            subprocess.check_output("git rev-parse --short HEAD".split()).strip().decode("utf-8")
+        )
+    except (subprocess.CalledProcessError, OSError):
+        print("Failed to execute git to get revision")
+        revision = None
+
+    url_fmt = (
+        f"https://github.com/{github_username}/{github_repository}/"
+        "blob/{revision}/src/{package}/{path}#L{lineno}"
+    )
+
+    def linkcode_resolve(domain, info):
+        if revision is None:
+            return
+        if domain not in ("py", "pyx"):
+            return
+        if not info.get("module") or not info.get("fullname"):
+            return
+
+        class_name = info["fullname"].split(".")[0]
+        module = __import__(info["module"], fromlist=[class_name])
+        obj = attrgetter(info["fullname"])(module)
+
+        # Unwrap the object to get the correct source
+        # file in case that is wrapped by a decorator
+        obj = inspect.unwrap(obj)
+
+        try:
+            fn = inspect.getsourcefile(obj)
+        except Exception:
+            fn = None
+        if not fn:
+            try:
+                fn = inspect.getsourcefile(sys.modules[obj.__module__])
+            except Exception:
+                fn = None
+        if not fn:
+            return
+
+        fn = os.path.relpath(fn, start=os.path.dirname(__import__(package).__file__))
+        try:
+            lineno = inspect.getsourcelines(obj)[1]
+        except Exception:
+            lineno = ""
+        return url_fmt.format(revision=revision, package=package, path=fn, lineno=lineno)
+
+    return linkcode_resolve

docs/source/examples/accelerate.md

@@ -1,5 +1,5 @@
 # Accelerate
 
 ```{eval-rst}
-.. literalinclude:: ./scripts/accelerate_example.py
+.. literalinclude:: ./scripts/accelerate_train.py
 ```

docs/source/examples/deepspeed.md

@@ -1,6 +1,6 @@
 # DeepSpeed
 
 ```{eval-rst}
-.. literalinclude:: ./scripts/deepspeed_example.py
+.. literalinclude:: ./scripts/deepspeed_train.py
 .. literalinclude:: ./scripts/deepspeed_config.json
 ```

docs/source/examples/lightning.md

@@ -1,5 +1,5 @@
 # Pytorch Lightning
 
 ```{eval-rst}
-.. literalinclude:: ./scripts/lightning_example.py
+.. literalinclude:: ./scripts/lightning_train.py
 ```

docs/source/examples/scripts/accelerate_example.py renamed to docs/source/examples/scripts/accelerate_train.py

@@ -14,7 +14,7 @@ The following arguments are required: `--model.name`, `--dataset.tokenizer-name`
   <summary><p style="display: inline-block;"><code class="docutils literal notranslate"><span class="pre">python transformers_train.py --help</span></code></p> (expand)</summary>
 
   ```{eval-rst}
-  .. literalinclude:: ./transformers_help.txt
+  .. literalinclude:: ./scripts/transformers_help.txt
   ```
 </details>
 

docs/source/examples/scripts/deepspeed_example.py renamed to docs/source/examples/scripts/deepspeed_train.py

@@ -30,7 +30,7 @@ dependencies = [
   "numpy>=1.20",
 ]
 [dependency-groups]
-dev = ["ruff", "pyright", "pytest", "build", "twine"]
+dev = ["ruff", "pyright[nodejs]", "pytest", "build", "twine"]
 dev-extras = ["submitit", "transformers"]
 docs = ["sphinx==7.4.7", "furo==2024.8.6", "myst-parser==3.0.1", "sphinx-autodoc2==0.5.0", "sphinx-toolbox==3.8.1"]