[air/tune] Internal resource management 1 - Ray AIR resource manager implementation

[krfricke]

Why are these changes needed?

Prerequisite to #30016

This PR adds a new Ray AIR resource manager to replace the PlacementGroupManager of Ray Tune. Details can be found in #30016.

Specifically, this PR

• Adds the main resource manager abstractions
• Renames (and moves) PlacementGroupFactory to ResourceRequest
• Adds implementations and tests for a placement group based manager and a budget based manager

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

[matthewdeng]

Thanks for taking the time to explain to me all the additional context around this, I feel like I learned a lot 😄

In this review, I focused primarily on making sure the APIs/documentation for the existing Placement Group flow are really clear and representative of the resource management lifecycle.

Consequently, I still think the FixedResourceManager flow doesn't adhere to the APIs very well, but I'm okay with not overindexing on this in this PR since it is not being exposed yet.

[krfricke]

Thanks for the review @matthewdeng. I've updated the docstrings throughout. ptal

[gjoliver]

can I give a suggestion, things like head_bundle_is_empty should not live on request class.
ResourceRequest class will be purer if it doesn't care about how applications plan to use the resources.
instead, we can make it a parameter to this annotate_remote_objects() function:

def annotate_remote_objects(
    self, objects: List[RemoteRayObject], empty_head_bundle=False,
) -> List[Union[RemoteRayObject]]:
    annotated = []

    def _annotate_obj(obj, bundle, bundle_idx):
        num_cpus = bundle.pop("CPU", 0)
        num_gpus = bundle.pop("GPU", 0)
        memory = bundle.pop("memory", 0.0)

        annotated.append(
            obj.options(
                scheduling_strategy=PlacementGroupSchedulingStrategy(
                    placement_group=self.placement_group,
                    # Max ensures that empty head bundles are correctly placed
                    placement_group_bundle_index=bundle_idx,
                    placement_group_capture_child_tasks=True,
                ),
                num_cpus=num_cpus,
                num_gpus=num_gpus,
                memory=memory,
                resources=bundle,
            )
        )

    first_obj = 0
    if empty_head_bundle:
        _annotate_obj(objects[0], {}, 0)
        first_obj = 1

    for i, (obj, bundle) in enumerate(
        zip(objects[first_obj:], self.resource_request.bundles)
    ):
        _annotate_obj(obj, bundle, i)

[krfricke]

Hm, _head_bundle_is_empty is automatically derived from the bundles in the __init__ of the resource request, i.e. it's no something the user passes. When a user passes bundles=[{}, {"CPU": 1}] the property is set to True. So this is just an internal representation, not an API.

I would not want to introduce this to the annotation API as then the caller will have to analyze the bundles to decide what to pass to the API - this logic should remain in the acquired resource class, which has access to all this information.

I'd be open to leaving the empty bundle in the self._bundles, but this would just be a refactor of the internal representation without any major effect. So for now I think it's best to keep it. (Note that this is also just copied over from PlacementGroupFactory, there is no new logic here).

[gjoliver]

ah, sorry, I looked at ResourceRequest more closely and understand what you are trying to do now.
Is there a background story about why we give head node, and head node only, such special treatment?
I would imagine an otherwise more generic framework for remote actors that don't require resources, regardless of where they are in the list.

[krfricke]

There is a background story!

This is mostly about distributed training. Let's say we have a Ray Train trainer that starts 2 distributed workers a 4 CPUs, so 8 CPUs in total. Ever since we've been using Ray Tune as the execution backend, this requires us to start a remote "head" worker (comparable to the RLlib learner) that starts these 2 distributed workers. But because the head worker is also a remote actor, it requires resources.

So now we're in a situation where we actually would need 9 CPUs (1 for the trainable actor, 2*4 for the distributed workers). That's undesirable, as it means we need bigger clusters (especially when tuning). Also (unlike the rllib trainer), the trainable worker is not doing much at all, it just does a bit of coordination between the distributed workers.

So the solution is to pass num_cpus=0 to this actor.

So now we need to somehow specify resources (a PlacementGroupFactory) that schedules the first actor with 0 CPUs. So what we want is [{"CPU": 0}, {"CPU": 1}, {"CPU": 1}].

At first, we wanted to make this work without any special casing. Placement groups however don't accept empty bundles. The first workaround was to pass [{"CPU": 0.01}, {"CPU": 0.99}, {"CPU": 1}] or so which is obviously not great from a user perspective. The second workaround is what we're currently doing: We don't schedule the first "bundle" but just schedule the first actor onto the first bundle with empty resources.

I agree this is not ideal - e.g. what happens if we pass another empty bundle in one of the workers? At the moment this will just raise an error.

The current solution works well from a user perspective, though. So I'd say let's keep it for now until we feel the need to come up with something better.

[matthewdeng]

Thanks for the many iterations on this! A few more naming/documentation clarifications but otherwise looks good to me.

[gjoliver]

thanks for your patience, 1 minor suggestion, your call :)

[gjoliver]

ah, sorry, I looked at ResourceRequest more closely and understand what you are trying to do now.
Is there a background story about why we give head node, and head node only, such special treatment?
I would imagine an otherwise more generic framework for remote actors that don't require resources, regardless of where they are in the list.