refactor: flatten sub-module imports to use parent package re-exports

[NJX-njx]

Description

Add __init__.py with fully lazy re-exports (__getattr__ + __all__) for areal/api, areal/engine, areal/reward, areal/dataset, and areal/workflow, then rewrite all external imports across ~100 files to use the shorter parent-package form.

For example: from areal.api.engine_api import TrainEngine → from areal.api import TrainEngine.

Related Issue

N/A

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

Key Changes

• Fully lazy __init__.py re-exports: all symbols use __getattr__ — no eager top-level imports, making areal.api completely side-effect free at import time
• areal/api (+23 lazy symbols), areal/engine (+11), areal/reward (+3), areal/workflow (+3), areal/dataset (__all__ only)
• Flatten imports in areal/trainer/, areal/utils/, areal/models/, areal/engine/, areal/workflow/, areal/infra/, areal/experimental/, tests/, examples/
• Preserve cli_args deep imports — dozens of config classes including SchedulingSpec; explicit path is clearer
• Preserve intra-package imports — avoids circular dependency issues
• Reward submodules use sibling-relative imports (from . import get_math_verify_worker) instead of absolute package imports, eliminating self-import coupling
• Preserve non-exported symbols (DeviceRuntimeInfo, HttpRequest, HttpGenerationResult, etc.) that remain as deep imports

What is NOT changed

Category	Reason
from areal.api.cli_args import ... (including SchedulingSpec)	Dozens of config classes — explicit sub-module path is more readable
from areal.engine.core.*, fsdp_utils.*, megatron_utils.*	Internal engine utilities, not part of public API
from areal.api.io_struct import DeviceRuntimeInfo, HttpRequest, ...	Not added to __all__ — internal symbols
Intra-package imports (e.g., areal/api/engine_api.py importing from areal/api/io_struct)	Prevents circular imports

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)

Breaking Change Details (if applicable):

N/A — all old import paths still work. This is a purely additive refactoring.

Additional Context

All re-exports use PEP 562 lazy __getattr__ to ensure zero side effects at module load time. Symbols are only resolved when first accessed, eliminating circular import risks.

106 files changed, 425 insertions(+), 265 deletions(-)

[chatgpt-codex-connector]

You have reached your Codex usage limits for code reviews. You can see your limits in the Codex usage dashboard.

[gemini-code-assist]

Summary of Changes

Hello, 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 significantly enhances the usability of the areal library by flattening its import structure for the api, workflow, and engine modules. This change allows developers to import key classes and functions directly from the module root (e.g., from areal.workflow import RLVRWorkflow) rather than requiring deep, verbose paths. The implementation utilizes a lazy import pattern to avoid loading heavy dependencies unnecessarily, ensuring performance is not impacted, and maintains full backward compatibility for existing codebases.

Highlights

• Flattened Import Structure: The import structure for areal/api, areal/workflow, and areal/engine modules has been flattened, allowing for shorter and more intuitive import paths (e.g., from areal.workflow import RLVRWorkflow).
• Lazy Import Pattern: A __getattr__ lazy import pattern was implemented to avoid eagerly loading heavy dependencies (like torch) when modules are imported, improving startup performance.
• Clean API Surface: Only the most commonly used public classes are re-exported in the __init__.py files to maintain a clean and focused API surface.
• Backward Compatibility: The changes are fully backward compatible, meaning existing deep import paths will continue to function without issues.
• Comprehensive Testing: New comprehensive tests (tests/test_flatten_imports.py) were added to validate the new import paths, __all__ definitions, and the lazy loading mechanism.

🧠 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.

Changelog

• areal/api/init.py
  • Added a new __init__.py file to the areal/api module.
  • Configured __init__.py to re-export commonly used API types and abstractions.
  • Implemented a __getattr__ lazy import mechanism for these re-exported types.
• areal/engine/init.py
  • Modified the existing __init__.py file in areal/engine.
  • Configured __init__.py to re-export engine implementations like FSDPEngine and MegatronEngine.
  • Integrated a __getattr__ lazy import pattern for the engine classes.
• areal/workflow/init.py
  • Added a new __init__.py file to the areal/workflow module.
  • Configured __init__.py to re-export workflow classes such as RLVRWorkflow and MultiTurnWorkflow.
  • Implemented a __getattr__ lazy import mechanism for these workflow classes.
• tests/test_flatten_imports.py
  • Added a new test file to verify the functionality of the flattened import structure.
  • Included tests for areal.api, areal.workflow, and areal.engine to ensure correct re-exports and lazy loading.
  • Verified that AttributeError is raised for invalid attribute access and that __all__ is correctly defined.
  • Added a specific test to confirm that importing the workflow module does not eagerly load torch.

Activity

• No human activity has been recorded on this pull request yet.

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. ↩

[gemini-code-assist]

Code Review

This pull request introduces a flattened import structure for the api, engine, and workflow modules, which is a great usability improvement. The use of __getattr__ for lazy loading is well-implemented to avoid unnecessary imports of heavy dependencies. My review includes suggestions to improve maintainability by reducing code duplication, enhance developer experience with better support for static analysis tools, and make the new tests more robust.

[Copilot]

Pull request overview

This PR implements the "Flatten the import structure of areal modules" item from the 2026 Q1 roadmap (#907). It adds __init__.py files that allow commonly used classes to be imported via short paths (e.g., from areal.api import AllocationMode) instead of requiring deep module paths (e.g., from areal.api.alloc_mode import AllocationMode). The implementation uses PEP 562 module-level __getattr__ for lazy imports to avoid eagerly pulling in heavy dependencies like torch.

Changes:

• New areal/api/__init__.py with lazy re-exports for engine, scheduler, allocation, and IO types (with AsyncRewardWrapper and RolloutWorkflow eagerly imported since they don't pull in torch)
• New areal/workflow/__init__.py with lazy re-exports for RLVRWorkflow, MultiTurnWorkflow, and VisionRLVRWorkflow
• Modified areal/engine/__init__.py (previously empty) with lazy re-exports for all four engine classes
• New tests/test_flatten_imports.py with tests for all new import paths

Reviewed changes

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

File	Description
areal/api/__init__.py	Adds short-path imports for public API types using lazy __getattr__ pattern
areal/workflow/__init__.py	Adds short-path imports for all workflow classes using lazy __getattr__ pattern
areal/engine/__init__.py	Adds short-path imports for all engine classes (previously empty file)
tests/test_flatten_imports.py	Tests the new import paths across all three modules

---

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

[rchardx]

LGTM