setup basic otel logging (meta-pytorch#281)

Summary:
Pull Request resolved: meta-pytorch#281

- setup 3 basic structured logs
  - quorums: every time a rank changes quorum id
  - commits: every time a rank commits a step
  - errors: every time a rank calls abort on process group
- allow `otel.py` to initialize loggers in multiple namespaces for each structured loggers
- pass `replica_id` to process group so that it can be logged
- flag to enable or disable otel logging

Reviewed By: d4l3k

Differential Revision: D84571482

fbshipit-source-id: 18dcfef40d8e9cfc1a11bb90d167dd3bf0271dd2

Author: tushar00jain

torchft/__init__.py

@@ -8,6 +8,7 @@
 from torchft.ddp import DistributedDataParallel
 from torchft.manager import Manager
 from torchft.optim import OptimizerWrapper as Optimizer
+from torchft.otel import setup_logger
 from torchft.process_group import (
     ProcessGroupBabyNCCL,
     ProcessGroupBabyXCCL,
@@ -16,6 +17,10 @@
     ProcessGroupXCCL,
 )
 
+setup_logger("torchft_quorums")
+setup_logger("torchft_commits")
+setup_logger("torchft_errors")
+
 __all__ = (
     "DistributedDataParallel",
     "DistributedSampler",

torchft/checkpointing/pg_transport_bench.py

@@ -47,7 +47,7 @@ def run(rank: int) -> None:
 
         with _timeit("init_pg"):
             pg = ProcessGroupBabyNCCL(timeout=timeout)
-            pg.configure(store_addr=store_addr, rank=rank, world_size=2)
+            pg.configure(store_addr=store_addr, replica_id="0", rank=rank, world_size=2)
 
             t = torch.zeros(10, device=device, dtype=torch.float32)
             pg.allreduce([t], dist.ReduceOp.SUM).wait(timeout=timeout)

torchft/checkpointing/pg_transport_test.py

@@ -27,6 +27,7 @@ def init(rank: int, world_size: int) -> CheckpointTransport[dict[str, object]]:
             pg = ProcessGroupGloo()
             pg.configure(
                 store_addr=f"localhost:{store.port}/prefix",
+                replica_id="0",
                 rank=rank,
                 world_size=world_size,
             )
@@ -52,6 +53,7 @@ def init(rank: int, world_size: int) -> CheckpointTransport[dict[str, object]]:
             pg = ProcessGroupBabyNCCL(timeout=timeout)
             pg.configure(
                 store_addr=f"localhost:{store.port}/prefix",
+                replica_id="0",
                 rank=rank,
                 world_size=world_size,
             )
@@ -78,6 +80,7 @@ def init(rank: int, world_size: int) -> CheckpointTransport[dict[str, object]]:
             pg = ProcessGroupBabyNCCL(timeout=timeout)
             pg.configure(
                 store_addr=f"localhost:{store.port}/prefix",
+                replica_id="0",
                 rank=rank,
                 world_size=world_size,
             )

torchft/device_mesh.py

@@ -62,19 +62,19 @@ def __init__(
             raise ValueError(
                 "ManagedDeviceMesh doesn't support both mesh and parent are None."
             )
-        self.mesh = mesh
-        self.mesh_dim_names = mesh_dim_names
+        self._mesh = mesh
+        self._mesh_dim_names = mesh_dim_names
         self.replicate_pg = replicate_pg
         self.replicate_dim = replicate_dim
         self.replicate_dim_name: str = mesh_dim_names[replicate_dim]
         self.parent = parent
         self.flatten_meshes: Dict[str, DeviceMesh] = {}
-        self.device_type: str
+        self._device_type: str
         if mesh is not None:
-            self.device_type = mesh.device_type
+            self._device_type = mesh.device_type
         else:
             assert parent is not None
-            self.device_type = parent.device_type
+            self._device_type = parent.device_type
         self._flatten_mesh_list: tuple[DeviceMesh, ...] = tuple()
         self._thread_id: Optional[int] = None
         self._hash: Optional[int] = None
@@ -102,20 +102,20 @@ def __getitem__(self, mesh_dim_names: Union[str, tuple[str, ...]]) -> DeviceMesh
             elif mesh_dim_names in self.flatten_meshes:
                 res_submesh = self.flatten_meshes[mesh_dim_names]
             else:
-                assert self.mesh is not None
-                res_submesh = self.mesh[mesh_dim_names]
+                assert self._mesh is not None
+                res_submesh = self._mesh[mesh_dim_names]
         else:
             assert isinstance(mesh_dim_names, tuple)
             if self.replicate_dim_name not in mesh_dim_names:
-                assert self.mesh is not None
-                res_submesh = self.mesh[mesh_dim_names]
+                assert self._mesh is not None
+                res_submesh = self._mesh[mesh_dim_names]
             else:
                 mesh_dim_names_wo_replicate = tuple(
                     n for n in mesh_dim_names if n != self.replicate_dim_name
                 )
-                assert self.mesh is not None
+                assert self._mesh is not None
                 res_submesh = ManagedDeviceMesh(
-                    self.mesh[mesh_dim_names_wo_replicate],
+                    self._mesh[mesh_dim_names_wo_replicate],
                     mesh_dim_names,
                     self.replicate_pg,
                     mesh_dim_names.index(self.replicate_dim_name),
@@ -125,7 +125,7 @@ def __getitem__(self, mesh_dim_names: Union[str, tuple[str, ...]]) -> DeviceMesh
         # TODO: find a better way to do this that doesn't depend on device mesh
         # internals
         root = _mesh_resources.get_root_mesh(self)
-        _mesh_resources.child_to_root_mapping[res_submesh] = root
+        res_submesh._root_mesh = root
 
         return res_submesh
 
@@ -134,7 +134,7 @@ def _real_mesh_dim(self, mesh_dim: int) -> int:
 
     def get_group(self, mesh_dim: Optional[Union[int, str]] = None) -> BaseProcessGroup:
         if isinstance(mesh_dim, str):
-            dim = self.mesh_dim_names.index(mesh_dim)
+            dim = self._mesh_dim_names.index(mesh_dim)
         else:
             dim = 0 if mesh_dim is None else int(mesh_dim)
 
@@ -143,8 +143,8 @@ def get_group(self, mesh_dim: Optional[Union[int, str]] = None) -> BaseProcessGr
         elif dim == self.replicate_dim:
             return self.replicate_pg
         else:
-            assert self.mesh is not None
-            return self.mesh.get_group(self._real_mesh_dim(dim))
+            assert self._mesh is not None
+            return self._mesh.get_group(self._real_mesh_dim(dim))
 
     def _flatten(
         self,
@@ -168,64 +168,64 @@ def size(self, mesh_dim: Optional[int] = None) -> int:
         # This is possible during the initialization stage of training.
         replicate_pg_size = 1 if replicate_pg_size == 0 else replicate_pg_size
         if mesh_dim is None:
-            if self.mesh is None:
+            if self._mesh is None:
                 return replicate_pg_size
             else:
-                assert self.mesh is not None
-                return self.mesh.size() * replicate_pg_size
+                assert self._mesh is not None
+                return self._mesh.size() * replicate_pg_size
         elif mesh_dim == self.replicate_dim:
             return replicate_pg_size
         else:
-            assert self.mesh is not None
-            return self.mesh.size(self._real_mesh_dim(mesh_dim))
+            assert self._mesh is not None
+            return self._mesh.size(self._real_mesh_dim(mesh_dim))
 
     @property
     def ndim(self) -> int:
-        assert self.mesh is not None
-        return self.mesh.ndim + 1
+        assert self._mesh is not None
+        return self._mesh.ndim + 1
 
     @property
     def shape(self) -> tuple[int, ...]:
-        assert self.mesh is not None
-        ret: list[int] = list(self.mesh.shape)
+        assert self._mesh is not None
+        ret: list[int] = list(self._mesh.shape)
         ret.insert(self.replicate_dim, self.replicate_pg.size())
         return tuple(ret)
 
     def get_rank(self) -> int:
-        assert self.mesh is not None
-        return self.mesh.get_rank()
+        assert self._mesh is not None
+        return self._mesh.get_rank()
 
     def get_local_rank(self, mesh_dim: Optional[Union[int, str]] = None) -> int:
         if isinstance(mesh_dim, str):
-            dim = self.mesh_dim_names.index(mesh_dim)
+            dim = self._mesh_dim_names.index(mesh_dim)
         else:
             dim = 0 if mesh_dim is None else int(mesh_dim)
 
         if mesh_dim is None:
-            if self.mesh is None:
+            if self._mesh is None:
                 return get_rank(self.replicate_pg)
 
             assert self.replicate_dim == 0, "replicate_dim must be the first one"
-            assert self.mesh is not None
-            other_dim_size = self.mesh.size()
-            assert self.mesh is not None
-            other_dim_rank = self.mesh.get_local_rank()
+            assert self._mesh is not None
+            other_dim_size = self._mesh.size()
+            assert self._mesh is not None
+            other_dim_rank = self._mesh.get_local_rank()
             replicate_pg_rank = get_rank(self.replicate_pg)
             return other_dim_size * replicate_pg_rank + other_dim_rank
         elif dim == self.replicate_dim:
             return get_rank(self.replicate_pg)
         else:
-            assert self.mesh is not None
-            return self.mesh.get_local_rank(self._real_mesh_dim(dim))
+            assert self._mesh is not None
+            return self._mesh.get_local_rank(self._real_mesh_dim(dim))
 
     def get_coordinate(self) -> Optional[list[int]]:
         """
         Return the relative indices of this rank relative to all
         dimensions of the mesh. If this rank is not part of the mesh, return None.
         """
-        assert self.mesh is not None
+        assert self._mesh is not None
         coordinate = (
-            self.mesh._coordinate_on_dim if self.mesh._coordinate_on_dim else None
+            self._mesh._coordinate_on_dim if self._mesh._coordinate_on_dim else None
         )
         if not coordinate:
             return coordinate
@@ -239,20 +239,20 @@ def get_all_groups(self) -> list[BaseProcessGroup]:
         raise NotImplementedError
 
     def __repr__(self) -> str:
-        return f"ManagedDeviceMesh(mesh={self.mesh})"
+        return f"ManagedDeviceMesh(mesh={self._mesh})"
 
     def __hash__(self) -> int:
         # lazily compute hash
         if not self._hash:
             self._hash = hash(
                 (
-                    self.mesh,
-                    self.mesh_dim_names,
+                    self._mesh,
+                    self._mesh_dim_names,
                     self.replicate_pg,
                     self.replicate_dim,
                     self.replicate_dim_name,
                     self.parent,
-                    self.device_type,
+                    self._device_type,
                 )
             )
         return self._hash

torchft/manager.py

@@ -216,9 +216,15 @@ def __init__(
                 before raising an exception. If None, will retry indefinitely.
             quorum_retries: the number of times to retry the quorum before crashing
         """
+        self.quorum_logger: logging.Logger = logging.getLogger("torchft_quorums")
+        self.commits_logger: logging.Logger = logging.getLogger("torchft_commits")
+        self.errors_logger: logging.Logger = logging.getLogger("torchft_errors")
+
         self._load_state_dict_fns: Dict[str, Callable[[object], None]] = {}
         self._user_state_dicts: Dict[str, Callable[[], object]] = {}
 
+        self._replica_id = replica_id
+
         # Protects state dict
         self._state_dict_lock = RWLock(timeout=timeout.total_seconds())
 
@@ -642,6 +648,16 @@ def _async_quorum(
                 self._participating_replica_rank = None
 
         if quorum_id != self._quorum_id:
+            self.quorum_logger.info(
+                "",
+                extra={
+                    "job_id": os.environ.get("JOB_ID", "unknown"),
+                    "replica_id": self._replica_id,
+                    "rank": self._group_rank,
+                    "quorum_id": quorum_id,
+                    "step": max_step,
+                },
+            )
             store_prefixed_addr = (
                 f"{store_address}/torchft/{quorum_id}/{self._group_rank}"
             )
@@ -653,7 +669,10 @@ def _async_quorum(
                     if torch.accelerator.is_available():
                         torch.accelerator.synchronize()
                     self._pg.configure(
-                        store_prefixed_addr, replica_rank, replica_world_size
+                        store_prefixed_addr,
+                        self._replica_id if self._replica_id is not None else "0",
+                        replica_rank,
+                        replica_world_size,
                     )
                 self._quorum_id = quorum_id
             except Exception as e:
@@ -817,6 +836,18 @@ def should_commit(self, timeout: Optional[timedelta] = None) -> bool:
             f"should_commit={should_commit} enough_replicas={enough_replicas}, errored={self._errored}"
         )
 
+        self.commits_logger.info(
+            "",
+            extra={
+                "job_id": os.environ.get("JOB_ID", "unknown"),
+                "replica_id": self._replica_id,
+                "rank": self._group_rank,
+                "quorum_id": self._quorum_id,
+                "step": self._step,
+                "commit_result": should_commit,
+            },
+        )
+
         self._checkpoint_transport.disallow_checkpoint()
 
         # decide whether we're in a healthy state to increase the step count