chore: update contribution guidelines

Author: mrapplexz

CONTRIBUTING.md

@@ -65,6 +65,22 @@ We provide a `Makefile` to automate common development tasks.
 
 We enforce strict quality standards to keep the codebase maintainable.
 
+### Design Principles
+
+They are not enforced by tooling, but PRs that violate them may be asked to change.
+
+* **Composition over inheritance**: components are small single-responsibility classes wired together (see `d9d/loop/component/`). Avoid "God classes" that accumulate unrelated responsibilities, avoid speculative base classes that exist only to hoard "common" code.
+* **Define contracts structurally.** Use a `typing.Protocol` for a *trait* - a secondary capability bolted onto a type that already has its own base class (e.g. `ModuleLateInit` on an `nn.Module`), where you want duck-typed conformance without forcing inheritance. Use an `abc.ABC` when the interface *is* the object's primary identity and the hierarchy is the "main" type (e.g. `PipelineSchedule`).
+* **No reflection where it can be avoided.** Avoid `getattr` / `hasattr` / `inspect` and string-name dispatch. Prefer an explicit `match`-`case`, a proper interface, or a factory. Reflection is acceptable *only* when introspection is intrinsic to the feature itself - i.e. declarative registration APIs that cannot work without it, such as a `@subscribe`/`@register` decorator wiring handlers by signature.
+* **Inject dependencies; don't reach for them.** Components receive their collaborators as constructor arguments and store them as private fields. Don't pull them from globals/singletons or construct them internally - wiring happens at the edges (`d9d/loop/run/`).
+* **Reuse before reinventing.** If PyTorch or the stdlib already solves it, use it, rather than hand-rolling an equivalent.
+* **No needless indirection.** Don't add a wrapper that only forwards to another function/object without adding meaning. Inline it instead.
+* **Validate eagerly, fail fast.** Validate constructor args up front; raise if a method is called outside its required lifecycle scope rather than silently misbehaving.
+* **Decide behavior from explicit inputs, not inferred state.** Drive branching with an explicit parameter, not by sniffing the shape/dtype/contents of the data. Inferred checks silently encode invariants the caller and the next reader won't know are there - make them part of the signature instead.
+* **Validate at the boundary; trust within it**. Data crossing an untrusted boundary (user config, deserialized state) is validated once at the edge into a model that guarantees its own invariants — that's the validation layer, and we use `pydantic` for it. Pass trusted internal data as plain `dataclasses` and assume it is already valid. Don't re-validate trusted internal data, and don't pass unvalidated raw input deeper than the edge.
+* **Separate configuration from behavior.** Config objects describe; classes behave. Don't merge them into one dataclass that needs `__post_init__` magic.
+* **Polymorphism for configurable objects via discriminated unions.** When a configurable object has selectable behavior, model the choices as a Pydantic discriminated union and resolve them in a `build_*()` factory with an exhaustive `match (case _: raise)`.
+
 ### Linting & Formatting
 We use [Ruff](https://docs.astral.sh/ruff/) for both linting and formatting. 
 Configuration is strict (see `pyproject.toml` for enabled rules).
@@ -84,6 +100,15 @@ We have two tiers of tests:
 
 **Requirement:** All PRs must pass `make test`. If you add a feature, you must add corresponding tests.
 
+### Docstrings
+
+We follow the [Google Python style](https://google.github.io/styleguide/pyguide.html#38-comments-and-docstrings) for docstrings.
+
+*   **Style:** Use Google-style docstrings (`Args:`, `Returns:`, `Raises:`, etc.).
+*   **No type annotations in docstrings:** Types are already declared in the signature and checked by `ty`. Do not repeat them in the docstring.
+*   **Document `__init__`:** Write a docstring even for `__init__`, but keep it short and to the point, e.g. `"""Constructs the ``Trainer`` object."""`.
+*   **Public API coverage:** Always write docstrings for everything considered public API.
+
 ## Documentation
 
 Documentation is built with **Zensical**.

d9d/core/autograd/grad_context.py

@@ -3,8 +3,7 @@
 
 
 class GradDirection(StrEnum):
-    """
-    Enum representing the specific gradient edges to compute.
+    """Enum representing the specific gradient edges to compute.
 
     This is used to manually control gradient flow in custom autograd functions
     during split backward passes.
@@ -19,8 +18,7 @@ class GradDirection(StrEnum):
 
 
 class GlobalGradContext:
-    """
-    Global state manager for controlling gradient computation in custom autograd functions.
+    """Global state manager for controlling gradient computation in custom autograd functions.
 
     This context addresses a limitation in PyTorch where custom `torch.autograd.Function`
     implementations set `ctx.needs_input_grad` to True for all edges requiring grad,
@@ -38,13 +36,11 @@ class GlobalGradContext:
 
     def __init__(self):
         """Constructs a GlobalGradContext object with all directions enabled by default."""
-
         # both directions by default
         self._enabled_directions: set[GradDirection] = {GradDirection.inputs, GradDirection.weight}
 
     def check_direction(self, direction: GradDirection | None) -> bool:
-        """
-        Checks if the gradient calculation for the given direction is currently enabled.
+        """Checks if the gradient calculation for the given direction is currently enabled.
 
         Args:
             direction: The direction to check (inputs or weights). If None,
@@ -53,16 +49,14 @@ def check_direction(self, direction: GradDirection | None) -> bool:
         Returns:
             True if the direction is enabled or None is passed, False otherwise.
         """
-
         if direction is None:
             return True
 
         return direction in self._enabled_directions
 
     @contextmanager
     def with_directions(self, *directions: GradDirection):
-        """
-        Context manager that sets the enabled gradient directions.
+        """Context manager that sets the enabled gradient directions.
 
         This overrides the current state for the duration of the context
         and restores the previous state afterwards.

d9d/core/dist_context/__init__.py

@@ -1,6 +1,4 @@
-"""
-This package configures the distributed environment and device meshes.
-"""
+"""This package configures the distributed environment and device meshes."""
 
 from .configured import DistributedContext
 from .device_mesh_domains import BATCH_DOMAIN, DENSE_DOMAIN, EXPERT_DOMAIN, FLAT_DOMAIN, REGULAR_DOMAIN

d9d/core/dist_context/configured.py

@@ -32,8 +32,7 @@ def _build_mesh_domains(params: "DeviceMeshParameters") -> dict[str, DeviceMesh]
 
 
 class DistributedContext:
-    """
-    Acts as the single source of truth for the distributed execution environment.
+    """Acts as the single source of truth for the distributed execution environment.
 
     It acts as the central repository for the distributed configuration, managing the creation
     and synchronization of PyTorch DeviceMeshes for different domains (Regular domain, Expert Parallel domain, ...).
@@ -78,12 +77,10 @@ def __init__(self, params: "DeviceMeshParameters", log_level: int):
     @property
     def logger(self) -> logging.Logger:
         """Returns the logger instance configured for distributed logging."""
-
         return self._logger
 
     def mesh_for(self, domain: str) -> DeviceMesh:
-        """
-        Returns the device mesh view associated with a specific logical domain.
+        """Returns the device mesh view associated with a specific logical domain.
 
         Available Domains and Dimensions:
             *   `regular` (`REGULAR_DOMAIN`): The most granular mesh for fully decomposed parallelism.
@@ -106,38 +103,32 @@ def mesh_for(self, domain: str) -> DeviceMesh:
         Raises:
             ValueError: If the specified domain does not exist.
         """
-
         if domain not in self._meshes:
             raise ValueError(f"Domain {domain} does not exist")
         return self._meshes[domain]
 
     @property
     def is_main_process(self) -> bool:
         """Checks if the current process is the global rank 0."""
-
         return self._global_rank == 0
 
     @property
     def is_local_main_process(self) -> bool:
         """Checks if the current process is the rank 0 on the specific node."""
-
         return self._local_rank == 0
 
     def wait_world(self):
         """Blocks process execution until all ranks reach this point."""
-
         if self._params.is_distributed:
             torch.distributed.barrier(device_ids=[torch.cuda.current_device()])
         torch.cuda.synchronize()
 
     def set_timeout(self, timeout_seconds: float):
-        """
-        Updates the NCCL/process group timeout for all underlying meshes.
+        """Updates the NCCL/process group timeout for all underlying meshes.
 
         Args:
             timeout_seconds: New timeout duration in seconds.
         """
-
         if not self._params.is_distributed:  # does nothing for local setups
             return
 
@@ -154,8 +145,7 @@ def set_timeout(self, timeout_seconds: float):
 
     @contextmanager
     def local_main_process_first(self):
-        """
-        Context manager that executes the block on the local main process first.
+        """Context manager that executes the block on the local main process first.
 
         Other local ranks wait at the entrance. The local main process waits at the
         exit to synchronize before continuing.
@@ -170,13 +160,11 @@ def local_main_process_first(self):
 
     @contextmanager
     def main_process_first(self):
-        """
-        Context manager that executes the block on the global main process first.
+        """Context manager that executes the block on the global main process first.
 
         All other ranks wait at the entrance. The global main process waits at the
         exit to synchronize before continuing.
         """
-
         if not self.is_main_process:
             self.wait_world()
 
@@ -188,35 +176,29 @@ def main_process_first(self):
     @property
     def current_device(self) -> torch.device:
         """Returns the CUDA device associated with this rank."""
-
         return self._current_device
 
     @property
     def mesh_params(self) -> "DeviceMeshParameters":
         """Returns the parameters used to initialize this context."""
-
         return self._params
 
     @property
     def master_addr(self) -> str:
         """Returns the IP address or domain name of the master node."""
-
         return self._master_addr
 
     @property
     def node_rank(self) -> int:
         """Returns the index of the node this process is running on."""
-
         return self._node_rank
 
     @property
     def local_rank(self) -> int:
         """Returns the rank of the current process within its node."""
-
         return self._local_rank
 
     @property
     def num_nodes(self) -> int:
         """Returns the total number of nodes in the cluster."""
-
         return self._num_nodes

d9d/core/dist_context/device_mesh_domains.py

@@ -8,8 +8,7 @@
 
 
 class DeviceMeshDomain(abc.ABC):
-    """
-    Abstract base class for a Device Mesh provider.
+    """Abstract base class for a Device Mesh provider.
 
     A Domain defines a specific strategy for organizing available GPUs into a
     multidimensional grid (Mesh) to support specific parallelism techniques.
@@ -19,21 +18,18 @@ class DeviceMeshDomain(abc.ABC):
     @abc.abstractmethod
     def name(self) -> str:
         """Returns the unique identifier for this mesh domain."""
-
         ...
 
     @abc.abstractmethod
     def build_mesh(self, params: "DeviceMeshParameters") -> DeviceMesh:
-        """
-        Constructs the device mesh configuration.
+        """Constructs the device mesh configuration.
 
         Args:
             params: Global configuration parameters for the distributed environment.
 
         Returns:
             The initialized PyTorch DeviceMesh for this specific domain.
         """
-
         ...
 
 

d9d/core/dist_context/log.py

@@ -3,8 +3,7 @@
 
 
 def build_dist_logger(qualifier: str, level: int) -> logging.Logger:
-    """
-    Configures and returns a logger instance for d9d.
+    """Configures and returns a logger instance for d9d.
 
     The logger is configured to write to stdout with a formatter that includes
     the provided rank qualifier, allowing for easier debugging in distributed logs.
@@ -16,7 +15,6 @@ def build_dist_logger(qualifier: str, level: int) -> logging.Logger:
     Returns:
         A configured logging.Logger instance.
     """
-
     dist_logger = logging.getLogger("d9d")
     dist_logger.setLevel(level)
     dist_logger.handlers.clear()

d9d/core/dist_context/params.py

@@ -7,8 +7,7 @@
 
 
 class DeviceMeshParameters(BaseModel):
-    """
-    Configuration parameters for initializing Distributed Device Meshes.
+    """Configuration parameters for initializing Distributed Device Meshes.
 
     Attributes:
         pipeline_parallel: Degree of pipeline parallelism (PP).
@@ -37,19 +36,16 @@ class DeviceMeshParameters(BaseModel):
     @property
     def has_pipeline_parallel(self) -> bool:
         """Checks if pipeline parallelism is enabled (degree > 1)."""
-
         return self.pipeline_parallel > 1
 
     @property
     def has_data_parallel_replicate(self) -> bool:
         """Checks if data parallel replication is enabled (degree > 1)."""
-
         return self.data_parallel_replicate > 1
 
     @property
     def has_data_parallel_shard(self) -> bool:
         """Checks if data parallel sharding is enabled (degree > 1)."""
-
         return self.data_parallel_shard > 1
 
     @property
@@ -72,7 +68,6 @@ def has_expert_parallel(self) -> bool:
     @property
     def is_distributed(self) -> bool:
         """Checks if any form of parallelism is enabled."""
-
         return (
             self.has_pipeline_parallel
             or self.has_data_parallel_replicate
@@ -85,7 +80,6 @@ def is_distributed(self) -> bool:
 
     @model_validator(mode="after")
     def _check_ep_divisibility(self) -> Self:
-        """Validates that DP/CP/TP dimensions can support the requested EP/ETP degrees."""
         dp_cp_tp_degree = (
             self.data_parallel_shard
             * self.data_parallel_replicate
@@ -103,11 +97,9 @@ def _check_ep_divisibility(self) -> Self:
         return self
 
     def build(self, log_level: int = logging.INFO) -> "DistributedContext":
-        """
-        Initializes the DistributedContext using these parameters.
+        """Initializes the DistributedContext using these parameters.
 
         Returns:
             A new DistributedContext instance containing the initialized device meshes.
         """
-
         return DistributedContext(self, log_level)

d9d/core/dist_ops/__init__.py

@@ -1,6 +1,4 @@
-"""
-This module provides high-level wrappers around `torch.distributed` collective operations.
-"""
+"""This module provides high-level wrappers around `torch.distributed` collective operations."""
 
 from .object import all_gather_object, gather_object
 from .tensor import all_gather, all_gather_variadic_shape, gather, gather_variadic_shape

d9d/core/dist_ops/object.py

@@ -6,8 +6,7 @@
 
 
 def gather_object(obj: T, group: dist.ProcessGroup, group_dst: int) -> list[T] | None:
-    """
-    Gathers picklable objects from the whole process group to a specific destination rank.
+    """Gathers picklable objects from the whole process group to a specific destination rank.
 
     This acts as a wrapper around torch.distributed.gather_object that automatically
     initializes the output buffer list on the destination rank.
@@ -20,7 +19,6 @@ def gather_object(obj: T, group: dist.ProcessGroup, group_dst: int) -> list[T] |
     Returns:
         A list of objects from all ranks on the destination rank; None on other ranks.
     """
-
     if group.rank() == group_dst:
         # We initialize with None, but we cast to list[T] because we know
         # dist.gather_object will populate these slots with actual objects.
@@ -32,8 +30,7 @@ def gather_object(obj: T, group: dist.ProcessGroup, group_dst: int) -> list[T] |
 
 
 def all_gather_object(obj: T, group: dist.ProcessGroup) -> list[T]:
-    """
-    Gathers picklable objects from the whole process group to all ranks.
+    """Gathers picklable objects from the whole process group to all ranks.
 
     This acts as a wrapper around torch.distributed.all_gather_object that automatically
     initializes the output buffer list on all ranks.