Code Style

Style Guides

• Python: Google Python Style Guide
• Shell: Google Shell Style Guide

This repository is Python-first.

uv

Use uv run to execute scripts. Do not activate a virtual environment and call python directly.

Do:

uv run examples/run_grpo.py

Don't:

source .venv/bin/activate
python examples/run_grpo.py

Exception: Dockerfile.ngc_pytorch is exempt from this rule.

Python Standard

Code must conform to Python 3.13.13+.

Indentation

Indent with 4 spaces. Do not use tabs.

Naming

• Avoid shadowing variables declared in an outer scope.
• Initialize all externally visible members of a class in the constructor.

Comments

• For interfaces used outside a file, prefer docstrings over comments.
• Comments are for code within a function or file-local interfaces.
• Commented-out code must have a comment explaining why it is commented out. Otherwise remove it before merging.

Docstrings

Use Google style docstrings (parseable by Sphinx).

Enforce Keyword Arguments for Ambiguous Parameters

When a function has multiple parameters of the same type that could easily be swapped by mistake, use a bare * to force keyword-only arguments starting from where the ambiguity begins. This prevents callers from accidentally transposing arguments.

Don't:

def loss_fn(input: Tensor, cp_group: ProcessGroup, tp_group: ProcessGroup, cp_rank: int, tp_rank: int):
    ...

# Caller can silently swap cp_group/tp_group or cp_rank/tp_rank
loss_fn(x, tp_group, cp_group, tp_rank, cp_rank)  # wrong order, no error

Do:

def loss_fn(input: Tensor, *, cp_group: ProcessGroup, tp_group: ProcessGroup, cp_rank: int, tp_rank: int):
    ...

# Caller must name every argument — swaps are impossible
loss_fn(x, cp_group=cp_group, tp_group=tp_group, cp_rank=cp_rank, tp_rank=tp_rank)

Avoid Reflection

Do not use reflection when functionality can be achieved without it.

Don't:

def make_complex(*args):
    x, y = args
    return dict(**locals())

Do:

def make_complex(x, y):
    return {'x': x, 'y': y}

Type Annotations

Annotate new functions and methods — both parameters and return type. When you add a parameter to an existing signature, type it, and match the type already used at the call site (don't leave a new arg untyped while the caller already declares it, e.g. def __init__(self, teacher_worker_groups=None) when the caller passes teacher_worker_groups: Optional[dict[str, Any]]).

When you add a new module under a type-checked area, add it to pyrefly.toml project-includes. pyrefly checks an explicit allow-list of files, so a new file that isn't listed silently escapes type-checking and its annotations are never verified.

Imports

Put imports at module top. Defer an import into a function body ONLY to break a circular import or to avoid loading a heavy/optional dependency in a path that shouldn't need it — and when you do, add a one-line comment saying which. An in-function import with no such reason should move to the top. In particular, deferring stdlib (concurrent.futures, collections, …) or a module that is already imported at module top buys nothing — hoist it.