[Train][Observability] Track Train Run Info with TrainStateActor

[woshiyyya]

Why are these changes needed?

This PR added StateActor to collect Train Run metadata for Train Dashboard.

The main components include:

• StatsActor:
  • A singleton, detached actor initiated on the head node.
  • It receives metrics and metadata from all train runs.
  • The Train dashboard periodically retrieves information from this actor.
  • Currently it stores only static metadata, but we can also extend it to support real-time metrics in the future.
• TrainRunStatsManager:
  • A management class instantiated on the controller layer(BackendExecutor here) of each train run.
  • It aggregates information from each worker and collects data about the trainer.
• Schemas (TrainRunInfo, TrainWorkerInfo, TrainDatasetInfo):
  • Pydantic schemas defined for the training run.
  • These schemas are immutable, which is constructed and reported only once before training starts.

How to launch the StateActor?

We decide to launch StateActor in the driver instead of the trainable (controller). Below are the reasons behind it:

Resource Limitation
Ray Tune turns placement_group_capture_child_tasks=True, which limit the restricts the total resource a trainable can use.

ray/python/ray/air/execution/resources/placement_group.py

Line 35 in 3606da8

placement_group_capture_child_tasks=True,

In Ray Train, workers and controllers already used up all the resources specified in ScalingConfig, thus no additional actors can be launched in TrainTrainable. Therefore, we cannot launch TrainStateActor in controller, since it requires "node:__head__": 0.001.

Force Cleanup
Ray Tune will force clean up all the actors and subactors launched inside the trainable. Therefore, if we launched the detached actor in the trainable, it will be purged after trainer.fit finished.

Related issue number

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 :(

[woshiyyya]

Trying to make sure all the pydantic related code is encapsulated in TrainRunStatsManager. So that OSS users who are not using ray[default] will not get an error.

[woshiyyya]

Create a unique ID that differentiate each run.

• Cannot use trial id because trainer.restore will reuse trial id.
• Cannot use job id because there could be multiple train runs in one job.
• Cannot use trial id + job id because one can restore a run multiple times in one job.

[justinvyu]

Thanks, looking good!

[matthewdeng]

I do think it feels a bit weird to pass the WorkerGroup here, but not sure if there is another cleaner way to organize it.

[woshiyyya]

The consideration here is: when we do elastic/fault-tolerant training, we can avoid using an old workergroup if we pass it through the function arguments.

[justinvyu]

I think it is actually better to pass it in register_train_run as before so we don't keep unnecessary state.

[woshiyyya]

What about having a api to update the worker group? We can re-register the worker group in the TrainRunStateManager when we it's updated (e.g. for elastic training)

[raulchen]

• let's name the above 2 as "dataset_name" and "dataset_uuid", we don't want to expose the concept of plan.
• also, I think we can prefix train-level dataset name to the data-level dataset name. So it's easier to identify them on the data dashboard.

dataset_key = "train"
dataset._set_name(dataset_key + "_" + dataset._name)

[woshiyyya]

OK. I'll update it with dataset_name and dataset_uuid.

For setting the prefix, probably we should do it when we initialize the trainer. I'll post another PR to do this.

[raulchen]

instead of forcing it on the head node, it'd be better to force it on the current node. Because we may launch a job driver on a worker node.

        scheduling_strategy = NodeAffinitySchedulingStrategy(
            ray.get_runtime_context().get_node_id(),
            soft=False,
        )

[woshiyyya]

Oh, actually the StateActor is tracking all Train Runs, we don't want the dashboard crash because of a worker node crash.

For example, we have a 2 train runs

• run A on node 1
• run B on node 2.

If we launch StateActor on node 1, and node 1 died, we will not be able to track run B.

[justinvyu]

Good work! I just left some nits.

[justinvyu]

Nit: we can just create a WorkerGroup directly instead of using BackendExecutor.

[woshiyyya]

Changed to WorkerGroup.

I was originally trying to also test the session initialization in this test. But since we already covered that in the end-2-end test, let's go with WorkerGroup instead.

[justinvyu]

def get_state_actor() -> Optional[ActorHandle[TrainStateActor]]:

[woshiyyya]

Nice catch. But seems that ActorHandle is not subscriptable since it's not a generic class. I've changed to def get_state_actor() -> Optional[ActorHandle]:

[justinvyu]

Nit: We don't need this one after removing BackendExecutor, but we should use monkeypatch.setenv(...) fixture so that the env var is only set within the test.

Otherwise it will spill over to the next one. We should also change the next test's env var setting to use monkeypatch.

[woshiyyya]

Ah got it! Good to see that after removing BackendExecutor, we don't need to setenv anymore.