fix: collect system metrics from all agents
This change corrects usage of rank, cross_rank and local_rank in the profiler so that it correctly knows it should collect system metrics when the local rank is 0.
@@ -39,8 +39,8 @@ class TrialController(metaclass=abc.ABCMeta):
self.prof = profiler.ProfilerAgent.from_env(
env,
- context.distributed.cross_rank,
- context.distributed.rank,
+ global_rank=context.distributed.rank,
+ local_rank=context.distributed.local_rank,
)
distributed_backend = _DistributedBackend()
