feat: use name_resolve for worker discovery and fix perf_tracer in the single-controller mode

[garrett4wade]

Description

1. Use name_resolve for RPC worker discovery in the distributed setting.

2. Fix the usage of perf_tracer in the single-controller mode:

   • Each process has its own role and rank. Traces are saved in sub folders: {root}/{role}/traces_{rank}.jsonl.
   • The trace export tool will merge data from different roles and ranks.
   • Add new engine methods for tracer configuration and saving for distributed control.

This is a prologue PR for the slurm scheduler.

Type of Change

• Bug fix (non-breaking change that fixes an issue)
• New feature (non-breaking change that adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not
  work as expected)
• Documentation update
• Code refactoring (no functional changes)
• Performance improvement
• Test coverage improvement

Checklist

• I have read the Contributing Guide
• I have run formatting tools (pre-commit or manual)
• I have run relevant unit tests and they pass
• I have added tests for new functionality
• I have updated documentation if needed
• My branch is up to date with main
• This PR introduces breaking changes (if yes, fill out details below)
• If this PR changes documentation, I have built and previewed it locally with
  jb build docs
• No critical issues raised by AI reviewers (/gemini review)

[gemini-code-assist]

Summary of Changes

Hello @garrett4wade, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request introduces crucial improvements to the system's distributed computing infrastructure. It refines the mechanism for RPC worker discovery by adopting name_resolve, ensuring more robust communication between components. Concurrently, the performance tracing utility has been overhauled to better support distributed operations, allowing for detailed, role-aware performance analysis. These changes are essential prerequisites for future scalability and integration with external job schedulers like Slurm.

Highlights

• RPC Worker Discovery: The system now utilizes name_resolve for RPC worker discovery in distributed environments, enhancing how workers locate each other.
• Performance Tracer Enhancements: The perf_tracer has been significantly updated for single-controller mode, including saving traces in role-specific and rank-specific subfolders, and introducing new engine methods (save_perf_tracer, config_perf_tracer) for distributed control. The trace export tool will now merge data from different roles and ranks.
• Scheduler Integration Foundation: This pull request serves as a foundational step for integrating with the Slurm scheduler, preparing the system for more advanced distributed job management.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[garrett4wade]

What the new trace looks like:

[gemini-code-assist]

Code Review

This pull request introduces significant enhancements for worker discovery using name_resolve and refactors the performance tracer (perf_tracer) to support roles in single-controller mode. The changes are comprehensive, touching API definitions, engine implementations, controllers, schedulers, and utility tools. The introduction of role-based tracing is a great improvement for debugging and performance analysis in distributed settings. I've identified a few issues: a potential for log file collisions, a bug where None could be passed to a function expecting a configuration object, and an inconsistency in how asynchronous tasks are executed. Overall, this is a solid contribution that improves the system's observability and distributed capabilities.

[garrett4wade]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces two main features: using name_resolve for worker discovery and enhancing perf_tracer for distributed environments. The changes for perf_tracer are well-implemented, adding new API methods and correctly propagating role and rank information to produce structured, per-role traces. The move to name_resolve for worker discovery is a good step towards more flexible scheduler integration.

My review focuses on a critical issue with log file naming that could lead to corrupted logs, and a couple of medium-severity suggestions regarding code duplication and the use of asyncio.run() in the controller classes, which could impact maintainability and reusability. Overall, the changes are well-aligned with the PR's goals, but the identified issues should be addressed.

[Copilot]

Pull request overview

This PR introduces worker discovery via name_resolve in distributed settings and fixes perf_tracer usage in single-controller mode by adding role-based trace organization.

Key Changes

• Added role parameter throughout the perf_tracer system to organize traces by role in subdirectories
• Implemented RPC worker discovery using name_resolve with new worker_discovery() name function
• Updated trace conversion tools to recursively search subdirectories and handle role information in visualizations

Reviewed changes

Copilot reviewed 18 out of 18 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
areal/utils/perf_tracer.py	Added role parameter to PerfTracer, SessionTracer, and related functions; traces now saved in role-specific subdirectories
areal/utils/names.py	Added worker_discovery() function for name-based RPC worker discovery
areal/tools/plot_session_trace.py	Updated session trace visualization to display role information in labels and hover text
areal/tools/perf_trace_converter.py	Enhanced trace converter to handle role metadata, use rglob for recursive file discovery, and properly sort by role
areal/tests/test_perf_tracer.py	Added comprehensive tests for role parameter in PerfTracer and SessionTracer
areal/tests/test_local_scheduler.py	Updated tests to use experiment_name and trial_name instead of BaseExperimentConfig
areal/scheduler/rpc/rpc_server.py	Integrated name_resolve for worker registration; added perf_tracer wrapping for RPC method calls; changed to auto-assign ports
areal/scheduler/local.py	Added name_resolve configuration and worker identity args to RPC server command; changed log file naming to role-based
areal/experimental/trainer/sft.py	Refactored perf_tracer configuration to separate method with role assignment
areal/experimental/trainer/rl.py	Extended perf_tracer configuration to all engine components with role-specific setup
areal/engine/vllm_remote.py	Implemented save_perf_tracer and config_perf_tracer methods
areal/engine/sglang_remote.py	Implemented save_perf_tracer and config_perf_tracer methods
areal/engine/megatron_engine.py	Implemented save_perf_tracer and config_perf_tracer methods
areal/engine/fsdp_engine.py	Implemented save_perf_tracer and config_perf_tracer methods
areal/controller/train_controller.py	Added config_perf_tracer and save_perf_tracer methods for distributed control
areal/controller/rollout_controller.py	Added config_perf_tracer and save_perf_tracer methods for rollout workers
areal/api/engine_api.py	Added save_perf_tracer and config_perf_tracer abstract methods to TrainEngine and InferenceEngine
examples/single-controller/gsm8k_grpo.yaml	Consolidated scheduling_spec references; added NCCL environment variables; added perf_tracer configuration

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.