switch to sphinx.ext.autodoc

Author: apoorvkh

.gitignore

@@ -1,4 +1,5 @@
 docs/source/README.md
+docs/source/contributing.md
 torchrunx_logs/
 .pixi/
 .ruff_cache/

docs/conf.py

@@ -1,11 +1,12 @@
 """Configuration file for the Sphinx documentation builder."""
-
 from glob import glob
 import shutil
 
 shutil.copyfile("../README.md", "source/README.md")
+shutil.copyfile("../CONTRIBUTING.md", "source/contributing.md")
 
 project = "torchrunx"
+copyright = 'Apoorv Khandelwal and Peter Curtin'
 github_username = "apoorvkh"
 github_repository = "torchrunx"
 html_theme = "furo"
@@ -14,7 +15,7 @@
 html_extra_path = list(glob("source/examples/scripts/*.py"))
 
 extensions = [
-    "autodoc2",
+    "sphinx.ext.autodoc",
     "myst_parser",  # support markdown
     "sphinx.ext.intersphinx",  # link to external docs
     "sphinx.ext.napoleon",  # for google style docstrings
@@ -23,14 +24,11 @@
     "sphinx_toolbox.github",
 ]
 
-maximum_signature_line_length = 100
-autodoc2_render_plugin = "myst"
-intersphinx_mapping = {
-    "python": ("https://docs.python.org/3.9", None),
-    "fabric": ("https://docs.fabfile.org/en/stable", None),
-    'torch': ('https://pytorch.org/docs/stable', None),
-    "numpy": ("https://numpy.org/doc/stable", None),
-}
+autodoc_member_order = "bysource"
+autodoc_typehints = "description"
+autodoc_typehints_description_target = "documented"
+
+intersphinx_mapping = {'python': ('https://docs.python.org/3', None)}
 
 from docs.linkcode_github import generate_linkcode_resolve_fn
 linkcode_resolve = generate_linkcode_resolve_fn(project, github_username, github_repository)

docs/source/api.md

@@ -1,10 +1,10 @@
 # API
 
 ```{eval-rst}
-.. autofunction:: torchrunx.launch
+.. autofunction:: torchrunx.launch(func, args, kwargs, ...)
 ```
 
-We provide the {mod}`torchrunx.Launcher` class as an alias to {mod}`torchrunx.launch`.
+We provide the {obj}`torchrunx.Launcher` class as an alias to {obj}`torchrunx.launch`.
 
 ```{eval-rst}
 .. autoclass:: torchrunx.Launcher
@@ -21,9 +21,9 @@ We provide the {mod}`torchrunx.Launcher` class as an alias to {mod}`torchrunx.la
 ## Exceptions
 
 ```{eval-rst}
-.. autoclass:: torchrunx.AgentFailedError
+.. autoexception:: torchrunx.AgentFailedError
 ```
 
 ```{eval-rst}
-.. autoclass:: torchrunx.WorkerFailedError
+.. autoexception:: torchrunx.WorkerFailedError
 ```

docs/source/contributing.md

@@ -23,13 +23,13 @@ We forward all logs (i.e. from {mod}`logging` and {mod}`sys.stdout`/{mod}`sys.st
 We provide some utilities to help:
 
 ```{eval-rst}
-.. autofunction:: torchrunx.file_handler
+.. autofunction:: torchrunx.utils.file_handler
 ```
 
 ```{eval-rst}
-.. autofunction:: torchrunx.stream_handler
+.. autofunction:: torchrunx.utils.stream_handler
 ```
 
 ```{eval-rst}
-.. autofunction:: torchrunx.add_filter_to_handler
+.. autofunction:: torchrunx.utils.add_filter_to_handler
 ```

docs/source/features/customization.md

@@ -25,7 +25,7 @@ dependencies = [
 [dependency-groups]
 dev = ["ruff==0.9.5", "pyright[nodejs]==1.1.393", "pytest==8.3.4"]
 test-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"]
+docs = ["sphinx==7.4.7", "furo==2024.8.6", "myst-parser==3.0.1", "sphinx-toolbox==3.8.2"]
 
 
 [tool.ruff]

pyproject.toml

@@ -2,15 +2,11 @@
 
 from .launcher import Launcher, LaunchResult, launch
 from .utils.errors import AgentFailedError, WorkerFailedError
-from .utils.logging_utilities import add_filter_to_handler, file_handler, stream_handler
 
 __all__ = [
     "AgentFailedError",
     "LaunchResult",
     "Launcher",
     "WorkerFailedError",
-    "add_filter_to_handler",
-    "file_handler",
     "launch",
-    "stream_handler",
 ]

src/torchrunx/__init__.py

@@ -19,7 +19,7 @@
     LauncherAgentGroup,
     get_open_port,
 )
-from .utils.logging_utilities import log_records_to_socket, redirect_stdio_to_logger
+from .utils.logging import log_records_to_socket, redirect_stdio_to_logger
 from .worker import WorkerArgs, worker_entrypoint
 
 

src/torchrunx/agent.py

@@ -34,7 +34,7 @@
     ExceptionFromWorker,
     WorkerFailedError,
 )
-from .utils.logging_utilities import LoggingServerArgs, start_logging_server
+from .utils.logging import LoggingServerArgs, start_logging_server
 
 
 def launch(
@@ -65,25 +65,33 @@ def launch(
     """Distribute and parallelize a function onto specified nodes and workers.
 
     Arguments:
-        func: Function to launch on each node and replicate for each worker.
-        args: Positional arguments for ``func``.
-        kwargs: Keyword arguments for ``func``.
+        func: Function to replicate on each node/worker.
+        args: Positional arguments for ``func``. Default: :py:obj:`None`.
+        kwargs: Keyword arguments for ``func``. Default: :py:obj:`None`.
         hostnames: Nodes on which to launch the function.
-            Default: infer from localhost or SLURM.
+            Default: ``"auto"`` (infer from localhost or SLURM).
         workers_per_host: Number of processes to run (e.g. # of GPUs) per node.
-            Default: use number of GPUs on each host.
+            Default: ``"auto"`` (number of GPUs per host).
         ssh_config_file: Path to an SSH configuration file for connecting to nodes.
-            Default: ``~/.ssh/config`` or ``/etc/ssh/ssh_config``.
+            Default: ``"~/.ssh/config"`` or ``"/etc/ssh/ssh_config"``.
         backend: `Backend <https://pytorch.org/docs/stable/distributed.html#torch.distributed.Backend>`_
-            for worker process group. Set `None` to disable. Default: NCCL (GPU) or GLOO (CPU).
+            for worker process group. Set `None` to disable.
+            Default: ``"auto"`` (NCCL if GPU or GLOO if CPU).
         timeout: Worker process group timeout (seconds).
+            Default: ``600``.
         default_env_vars: Environment variables to copy from the launcher process to workers.
             Supports bash pattern matching syntax.
+            Default: ``("PATH", "LD_LIBRARY", "LIBRARY_PATH", "PYTHON*", "CUDA*", "TORCH*",
+            "PYTORCH*", "NCCL*")``.
         extra_env_vars: Additional user-specified environment variables to copy.
+            Default: ``()``.
         env_file: Path to a file (e.g., ``.env``) with additional environment variables to copy.
+            Default: :py:obj:`None`.
         propagate_exceptions: Raise exceptions from worker processes in the launcher.
-            If false, raises :obj:`WorkerFailedError` instead.
+            If false, raises :exc:`WorkerFailedError` instead.
+            Default: :py:obj:`True`.
         handler_factory: Function to customize processing of agent and worker logs with handlers.
+            Default: ``"auto"`` (see `custom logging <https://torchrun.xyz/features/customization.html#logging>`_).
 
     Raises:
         RuntimeError: If there are configuration issues.
@@ -115,7 +123,7 @@ def launch(
 
 @dataclass
 class Launcher:
-    """Useful for sequential invocations or for specifying arguments via CLI."""
+    """Alias class for :func:`launch`. Refer to that function for documentation."""
 
     hostnames: list[str] | Literal["auto", "slurm"] = "auto"
     workers_per_host: int | list[int] | Literal["auto"] = "auto"
@@ -153,7 +161,7 @@ def run(  # noqa: C901, PLR0912
         args: tuple | None = None,
         kwargs: dict[str, Any] | None = None,
     ) -> LaunchResult:
-        """Run a function using the :mod:`torchrunx.Launcher` configuration."""
+        """Launch a function using class configuration."""
         if not dist.is_available():
             msg = "The torch.distributed package is not available."
             raise RuntimeError(msg)

src/torchrunx/launcher.py

@@ -1 +1,5 @@
 """Utility classes and functions."""
+
+from .logging import add_filter_to_handler, file_handler, stream_handler
+
+__all__ = ["add_filter_to_handler", "file_handler", "stream_handler"]