[DNR][data] Refractor OpRuntimeMetrics

[iamjustinhsu]

Description

Currently, ray data internal metrics were handled by the monolith OpRuntimeMetrics class, which handled task and input/output metrics. While I do think it simplifies abstractions by only thinking about one OpRuntimeMetrics, task metrics is only used for MapOperator and HashShuffleOperator.

Proposal

Split OpRuntimeMetrics into BaseOpMetrics and TaskOpMetrics, where BaseOpMetrics will handle input/outputs, and TaskOpMetrics will inherit from BaseOpMetrics and keep track of task metrics.

Why

1. I think it's pretty hard to understand which metrics are relevant for an operator, and which are not. For example, limit operator launches tasks to split the blocks, but doesn't keep track of task metrics.
2. In the future, I would like to have dedicated metrics on aggregator health. For example, I would be able to extend the TaskOpMetrics(BaseOpMetrics). This will include partition_size, num_partitions_per_aggregator, etc...
3. It's pretty confusing which metrics are exposed publicly (by public, i mean outside the operator, like the StreamingExecutor, ResourceManager), and which metrics are only kept internal for prometheus. By implementing the BaseOpMetrics, you know what gets exposed.
4. Slightly smaller ray-data.log log file because we only export the metrics the operator actually uses.

Why not

1. More abstractions can lead to additional complexity, so need to understand which metrics an operator uses
2. Not as easily extendible to use metrics outside the operator, so would need to implement in the BaseOpMetrics class

Related issues

Additional information

Will run release test, assuming PR is worth merging to make sure I didn't miss anything.

[gemini-code-assist]

Code Review

This pull request is a well-executed refactoring of OpRuntimeMetrics into a more modular and understandable BaseOpMetrics and TaskOpMetrics hierarchy. The changes are applied consistently throughout the codebase, improving clarity and maintainability. I've found one critical issue related to a circular import that would cause a runtime error, which I've commented on.

[cursor]

Bug: Actor metrics only updated on task completion

Actor pool metrics (num_alive_actors, num_pending_actors, num_restarting_actors) are now only updated when on_task_finished is called with actor_info. Previously, these metrics were updated in add_output (in streaming_executor_state.py), which was called more frequently. This causes actor metrics to become stale between task completions, failing to reflect real-time changes in actor pool state like newly started or restarted actors.

 

[github-actions]

This pull request has been automatically marked as stale because it has not had
any activity for 14 days. It will be closed in another 14 days if no further activity occurs.
Thank you for your contributions.

You can always ask for help on our discussion forum or Ray's public slack channel.

If you'd like to keep this open, just leave any comment, and the stale label will be removed.

[github-actions]

This pull request has been automatically marked as stale because it has not had
any activity for 14 days. It will be closed in another 14 days if no further activity occurs.
Thank you for your contributions.

You can always ask for help on our discussion forum or Ray's public slack channel.

If you'd like to keep this open, just leave any comment, and the stale label will be removed.