Merge pull request #86 from apoorvkh/support-type-checking

Support type checking

Author: apoorvkh

README.md

@@ -28,6 +28,7 @@ Requires: Linux (+ SSH & shared filesystem if using multiple machines)
 Dummy distributed training function:
 
 ```python
+from __future__ import annotations
 import os
 import torch
 import torch.nn as nn
@@ -59,26 +60,24 @@ Launching training with `torchrunx`:
 ```python
 import torchrunx
 
-results = torchrunx.launch(
-    func = train,
-    kwargs = dict(
-        model = nn.Linear(10, 10),
-        num_steps = 10
-    ),
-    #
+results = torchrunx.Launcher(
     hostnames = ["localhost", "second_machine"],
     workers_per_host = 2
+).run(
+    train,
+    model = nn.Linear(10, 10),
+    num_steps = 10
 )
 
 trained_model: nn.Module = results.rank(0)
 torch.save(trained_model.state_dict(), "output/model.pth")
 ```
 
 **See examples where we fine-tune LLMs (e.g. GPT-2 on WikiText) using:**
-  - [Accelerate](https://torchrun.xyz/examples/accelerate.html)
-  - [HF Transformers](https://torchrun.xyz/examples/transformers.html)
+  - [Transformers](https://torchrun.xyz/examples/transformers.html)
   - [DeepSpeed](https://torchrun.xyz/examples/deepspeed.html)
   - [PyTorch Lightning](https://torchrun.xyz/examples/lightning.html)
+  - [Accelerate](https://torchrun.xyz/examples/accelerate.html)
 
 **Refer to our [API](https://torchrun.xyz/api.html) and [Advanced Usage Guide](https://torchrun.xyz/advanced.html) for many more capabilities!**
 
@@ -118,4 +117,4 @@ torch.save(trained_model.state_dict(), "output/model.pth")
 > - Automatic detection of SLURM environments.
 > - Start multi-node training from Python notebooks!
 
-**On our [roadmap](https://github.com/apoorvkh/torchrunx/issues?q=is%3Aopen+is%3Aissue+label%3Aenhancement): higher-order parallelism, support for debuggers, fuller typing, and more!**
+**On our [roadmap](https://github.com/apoorvkh/torchrunx/issues?q=is%3Aopen+is%3Aissue+label%3Aenhancement): higher-order parallelism, support for debuggers, and more!**

docs/conf.py

@@ -24,11 +24,12 @@
     "sphinx_toolbox.github",
 ]
 
+maximum_signature_line_length = 90
 autodoc_member_order = "bysource"
-autodoc_typehints = "description"
-autodoc_typehints_description_target = "documented"
 
-intersphinx_mapping = {'python': ('https://docs.python.org/3', None)}
+intersphinx_mapping = {
+    'python': ('https://docs.python.org/3.9', 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,29 +1,6 @@
 # API
 
 ```{eval-rst}
-.. autofunction:: torchrunx.launch(func, args, kwargs, ...)
-```
-
-We provide the {obj}`torchrunx.Launcher` class as an alias to {obj}`torchrunx.launch`.
-
-```{eval-rst}
-.. autoclass:: torchrunx.Launcher
-  :members:
-```
-
-## Results
-
-```{eval-rst}
-.. autoclass:: torchrunx.LaunchResult
+.. automodule:: torchrunx
   :members:
 ```
-
-## Exceptions
-
-```{eval-rst}
-.. autoexception:: torchrunx.AgentFailedError
-```
-
-```{eval-rst}
-.. autoexception:: torchrunx.WorkerFailedError
-```

docs/source/features/cli.md

@@ -1,16 +1,16 @@
 # CLI Integration
 
-We can use {mod}`torchrunx.Launcher` to populate arguments from the CLI (e.g. with [tyro](https://brentyi.github.io/tyro/)):
+We can automatically populate {mod}`torchrunx.Launcher` arguments using most CLI tools (those that generate interfaces from Data Classes, e.g. [tyro](https://brentyi.github.io/tyro/)):
 
 ```python
-import torchrunx as trx
+import torchrunx
 import tyro
 
 def distributed_function():
-    pass
+    ...
 
 if __name__ == "__main__":
-    launcher = tyro.cli(trx.Launcher)
+    launcher = tyro.cli(torchrunx.Launcher)
     launcher.run(distributed_function)
 ```
 

pyproject.toml

@@ -11,7 +11,7 @@ authors = [
 ]
 description = "Automatically initialize distributed PyTorch environments"
 readme = "README.md"
-license = {file = "LICENSE"}
+license = { file = "LICENSE" }
 urls = { Repository = "https://github.com/apoorvkh/torchrunx.git", Documentation = "https://torchrun.xyz" }
 requires-python = ">=3.9"
 dependencies = [
@@ -21,12 +21,17 @@ dependencies = [
   # torch.distributed depends on numpy
   # torch<=2.2 needs numpy<2
   "numpy>=1.20",
+  "typing-extensions>=4.9.0",
 ]
 [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-toolbox==3.8.2"]
-
+docs = [
+  "sphinx==7.4.7",
+  "furo==2024.8.6",
+  "myst-parser==3.0.1",
+  "sphinx-toolbox==3.8.2",
+]
 
 [tool.ruff]
 include = ["pyproject.toml", "src/**/*.py", "tests/**/*.py"]
@@ -36,6 +41,8 @@ src = ["src", "tests"]
 [tool.ruff.lint]
 select = ["ALL"]
 ignore = [
+  "TC003",   # no type checking blocks for stdlib
+  "D104",    # package docstrings
   "ANN401",  # self / cls / Any annotations
   "BLE001",  # blind exceptions
   "TD",      # todo syntax

src/torchrunx/__init__.py

@@ -1,12 +1,10 @@
-"""API for our torchrunx library."""
-
-from .launcher import Launcher, LaunchResult, launch
+from .launcher import DEFAULT_ENV_VARS_FOR_COPY, Launcher, LaunchResult
 from .utils.errors import AgentFailedError, WorkerFailedError
 
-__all__ = [
-    "AgentFailedError",
-    "LaunchResult",
+__all__ = [  # noqa: RUF022
+    "DEFAULT_ENV_VARS_FOR_COPY",
     "Launcher",
+    "LaunchResult",
+    "AgentFailedError",
     "WorkerFailedError",
-    "launch",
 ]

src/torchrunx/integrations/__init__.py

@@ -1 +0,0 @@
-"""Utilities for integrations with other libraries."""