Update metrics_utils for future global metrics aggregation in controller.

[Stack-Attack]

Why are these changes needed?

These changes modify the autoscaler metrics collection and aggregation functions in preparation for global aggregation in the controller.

Related issue number

Partial for #46497

Required for #41135 #51905

Checks

• I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
• I've run scripts/format.sh to lint the changes in this PR.
• I've included any doc changes needed for https://docs.ray.io/en/master/.
  • I've added any new APIs to the API Reference. For example, if I added a
    method in Tune, I've added it in doc/source/tune/api/ under the
    corresponding .rst file.
• I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
• Testing Strategy
  • Unit tests
  • Release tests
  • This PR is not tested :(

[gemini-code-assist]

Code Review

This pull request refactors the metrics aggregation logic for autoscaling by introducing new generic aggregation methods in InMemoryMetricsStore and a consolidate_metrics_stores function. The changes are a good step towards improving the separation of concerns. However, I've identified a few critical issues where None return values from aggregation functions are not handled, which will lead to TypeError exceptions. Additionally, there's a bug in consolidate_metrics_stores that could cause data loss for a specific metric, and a minor API design concern. Please review the detailed comments for suggestions on how to address these points.

[Stack-Attack]

@zcin Partial PR for #51905 is here.

[abrarsheikh]

can you explain which serve component would call this function?

[Stack-Attack]

@abrarsheikh This would be called in the 'autoscaler_state_manager' to consolidate metrics from multiple handles. AFAIK there is no guarantee that replica-handle metrics are one-to-one (based on current logic), so I wanted to handle the case of the same replica reporting through different handles with a different cadence.

[zcin]

can we add a docstring explaining inputs and return value? Since there are 2 axes, handles and replicas, it would be helpful if the docstring could explain with that context in mind.

[Stack-Attack]

Added here, as well as the aggregate functions as the tuple return is new.

You mentioned in the other PR to use a single aggregate function which takes an enum. Let me know if you want that here. I'd need to import the enum to test_metrics_utils and, replica.py and router.py, which I'm not sure is preferable.

[abrarsheikh]

finding it hard to follow this logic. can you deconstruct the expression to better readability

[zcin]

each timeseries here is the number of requests (over e.g. past 30 seconds) that handle X send to replica Y right? are we choosing only the data from one out of all handles here, instead of aggregating across handles?

[Stack-Attack]

We're going to have a lot of handle_metrics or replica_metrics ending up in the controller. We will store all of these as raw InMemoryMetricsStore objects. Each key stores one timeseries. Assuming that:

1. Different metric_store's might contain the same key.
2. All key's map to one unique replica.
   This function allows us to feed all the metrics_stores in, and only get the latest metrics for each key. So if a key is found, we overwrite it's data if the new timeseries newer. So this doesn't aggregate, it just consolidates the data and returns a single object representing the global state of all metrics_stores.

[Stack-Attack]

I'll break out that large conditional for readability, and add more documentation to make this clear!

[zcin]

@Stack-Attack I'm not convinced that's what we want. If we have 2 handles, A and B, sending requests to a single replica X, then handle A can have time series (3, 4, 5), and handle B can have time series (7, 8, 9). This means that handle A most recently recorded 5 requests that were sent to replica X, and handle B recorded 9 requests that were sent to replica X. This means that 14 requests in total are executing on replica X.

However this code seems to only be taking either 5 or 9, instead of 14.

[Stack-Attack]

Alternatively, we could use the replica-level metric collection for running requests, and only collect the queued requests from handles? If custom metrics will be recorded at the replica anyway, this might be better? @abrarsheikh

[Stack-Attack]

If I recall correctly, one reason metrics collection was moved to handles was to strictly guarantee that replicas do not exceed their max_running limit. Handles currently report the number of requests they send to each replica, and this is the metric source for running_requests.

To support custom metrics in the future, however, we are going to need to send metrics from the replica-level right? So it might make sense to unify the metrics collection back at the replica, and send only queue length from the handles.

I think this is the only major blocker. Since either solution above has pro's and con's, I think I should leave it to maintainers to decide on the design.

[abrarsheikh]

I think @Stack-Attack, your proposal to send queued requests from the handle and the rest of the metrics from replicas makes sense to me.

[abrarsheikh]

we'll have to apply a windowed sum across all replica metrics?

that is correct

[Stack-Attack]

Changed the consolidate function to a merge function using a time window. In practice this window will be metrics_interval_s, which will be sent from the controller.

[zcin]

can we add a docstring explaining inputs and return value? Since there are 2 axes, handles and replicas, it would be helpful if the docstring could explain with that context in mind.

[zcin]

each timeseries here is the number of requests (over e.g. past 30 seconds) that handle X send to replica Y right? are we choosing only the data from one out of all handles here, instead of aggregating across handles?

[Stack-Attack]

Error in rebasing on master. Fixing now.

[zcin]

CI is failing because the example you added in the docstring for _aggregate_reduce is tested and failing: https://buildkite.com/ray-project/microcheck/builds/23995#0198e5ec-e7b9-49a2-b299-2226bf66710e/189-869

[zcin]

OK, let's table this for when you raise the next PR!

[zcin]

why is there data at timestamp=4? I don't see data added for m1 at timestamp=4

[Stack-Attack]

The windows can shift based on ordering. The only way to avoid this would be to compute the global start_timestamp before any merges.

merge(s1,s2, window=2):
  t=1 => window==[0, 2) => sum(1) => TimeStampedValue(1,1)
  t=3 => window==[2,4) => sum(2) => TimeStampedValue(3,2)
  result = (1,1),(3,2)
  
merge(result, s4, window=2):
  t=0 => window==[-1,1) => sum(100) => TimeStampedValue(0,100)
  t=2 => window==[1,3) => sum(1) => TimeStampedValue(2,1)
  t=4 => window==[3,5) => sum(2) => TimeStampedValue(4,2)
  result = (0,100),(2,1),(4,2)

Note, this was still passing incorrectly due to:

@dataclass(order=True)
class TimeStampedValue:
    timestamp: float
    value: float = field(compare=False)

I added a comparison to the tests to force check the value as well. I think we should remove the compare=False, but it's not required.

[zcin]

@Stack-Attack Hmm it seems pretty strange that the windows can shift. can this be fixed if we instead key windows by start of the window instead of using ts_center = start + w * window_s + (window_s / 2.0)? aka directly use start + w * window_s?

[zcin]

@Stack-Attack actually I'm not sure if that's the exact issue. We just need the windows to be consistent - aka if on the previous merge the windows were (0,2), (2,4) etc, it shouldn't shift to (1,3),(3,5) etc at the next merge. How about we make sure that the windows are always multiples of window_s, aka it's always (0,window_s), (window_s,2*window_s)?

[Stack-Attack]

@zcin Good call! Modified to snap start to window_s grid. I maintained the window_s / 2.0 as it should still be helpful when window_s matches the polling rate of the metrics.

[abrarsheikh]

@Stack-Attack could you please address @zcin 's comment? I think we are very close to closing this.

[abrarsheikh]

i made 3 changes

1. added a validation for window_s
2. apply the bucketing logic even when a ensure that if a single timeseries is passed
3. added more unit tests

[zcin]

LGTM with one more comment! I think we're good to merge after this.