[RLlib] 2.50.0 regression. Lifetime metrics are not checkpoint compatible anymore
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 43.9k
- Forks
- 8.1k
- PR merge metrics
- PR metrics pending
Description
What happened + What you expected to happen
https://github.com/ray-project/ray/pull/56047 introduced quite a intelligent method to reduce sum metrics. However it relies on the id_ of an incoming stat object. When I load a checkpoint and Stats are regenerated the id changes.
Tracking lifetime metrics like NUM_ENV_STEPS_SAMPLED_LIFETIME are incorrect after a restore as the needed value in _prev_merge_values cannot be accessed (its also not included in the state).
Below is a small repo that simulates the core issue.
Related are of course the general past issues of metrics & restore: #53877 #54174. But, one could at least work around them in <2.50 by updating some MetricsLoggers manually. Now however in 2.50.0 this does not work anymore.
CC: @sven1977 @ArturNiederfahrenhorst
btw. you can find some two standalone tests here that cover quite a few aspects of metric save & restores via Algorithms: https://github.com/Daraan/ray_utilities/blob/main/test/test_standalone.py
Versions / Dependencies
2.50.0+
Reproduction script
Simulate a save and load of metrics during checkpointing.
Correct output would be 114 but is 214.
from ray.rllib.utils.metrics.metrics_logger import *
# Simulate first training
root_logger = MetricsLogger()
sub_logger = MetricsLogger()
sub_logger.log_value("test", 100, reduce="sum")
# Combine
root_logger.aggregate([sub_logger.reduce()])
root_logger.reduce()
# Checkpoint
logger_state = root_logger.get_state()
sub_logger_state = sub_logger.get_state()
# Recreate
restored_root_logger = MetricsLogger()
restored_sub_logger = MetricsLogger()
restored_root_logger.set_state(logger_state)
restored_sub_logger.set_state(sub_logger_state)
# Continue training
sub_logger.log_value("test", 14, reduce="sum")
restored_sub_logger.log_value("test", 14, reduce="sum")
# Reduce
sub_reduced = sub_logger.reduce()
restored_sub_reduced = restored_sub_logger.reduce()
# Aggregate
root_logger.aggregate([sub_reduced])
# Error: Should log only 14, but logs 114
# aggregate on restored logger, does not work neither sub_reduced nor restored_sub_reduced
restored_root_logger.aggregate([restored_sub_reduced])
# Out
log_reduced = root_logger.reduce()
new_log_reduced = restored_root_logger.reduce()
assert sub_reduced == restored_sub_reduced # OK
assert log_reduced == new_log_reduced, f"{log_reduced} != {new_log_reduced}" # FAIL
# 114 (correct) vs 214 (wrong, double counted, last reduce not subtracted)
Issue Severity
High: It blocks me from completing my task.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start with ray/rllib/utils/metrics/metrics_logger.py, the module imported by the reproduction script, and run the supplied save/load simulation. Compare restored and non-restored reductions, expecting 114 rather than 214; use the referenced standalone tests as additional restore coverage. Done means lifetime metrics remain correct after checkpoint restore.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- distributed-systems, machine-learning
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 52/100