[data] implement per-op resource reservation #43026

raulchen · 2024-02-07T00:15:03Z

Why are these changes needed?

Implement per-op resource reservation. This PR doesn't enable it by default. As it depends on #42930 and we also need to integrate it with backpressure policies.

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

Signed-off-by: Hao Chen <chenh1024@gmail.com>

bveeramani

Overall LGTM

bveeramani · 2024-02-12T19:15:57Z

python/ray/data/_internal/execution/interfaces/execution_options.py

+    @classmethod
+    def zero(cls) -> "ExecutionResources":
+        """Returns an ExecutionResources object with zero resources."""
+        if cls._ZERO is None:
+            cls._ZERO = ExecutionResources(0.0, 0.0, 0)
+        return cls._ZERO


What's our motivation for caching this?

Suggested change

@classmethod

def zero(cls) -> "ExecutionResources":

"""Returns an ExecutionResources object with zero resources."""

if cls._ZERO is None:

cls._ZERO = ExecutionResources(0.0, 0.0, 0)

return cls._ZERO

@classmethod

def zero(cls) -> "ExecutionResources":

"""Returns an ExecutionResources object with zero resources."""

return ExecutionResources(0.0, 0.0, 0)

One concern I have is that ExecutionResources are mutable.

>>> ExecutionResources.zero().cpu = 1 >>> ExecutionResources.zero() ExecutionResources(cpu=1, gpu=0.0, object_store_memory=0)

This might become a gotcha.

good catch. we should make it immutable in the future.

bveeramani · 2024-02-12T19:23:27Z

python/ray/data/_internal/execution/interfaces/execution_options.py

+    def max(self, other: "ExecutionResources") -> "ExecutionResources":
+        """Returns the maximum for each resource type."""
+        return ExecutionResources(
+            cpu=max(self.cpu or 0.0, other.cpu or 0.0),


Nit: if both fields are None, do we expected the max to be zero? I'm wondering it it'd be more intuitive if the field would still be None

For example:

>>> result = ExecutionResources().max(ExecutionResources()) >>> result.cpu 0.0 # You might expect this to still be `None`?

this is super confusing and annoying. but I think defaulting to 0 makes sense here.
The reason is that this ExecutionResources is used for 2 different purposes simultaneously:

representing "resource usage": values should default to 0 in this case.

representing "resource limits": values should default to inf in this case.

the min and max API should only be used for (1). I think maybe we should think more on how to sanitize this API and introduce (potentially breaking) changes before 2.10

I will use a follow-up PR to make these values non-nullable.

bveeramani · 2024-02-12T19:24:40Z

python/ray/data/_internal/execution/interfaces/execution_options.py

+    def min(self, other: "ExecutionResources") -> "ExecutionResources":
+        """Returns the minimum for each resource type."""
+        return ExecutionResources(
+            cpu=min(self.cpu or float("inf"), other.cpu or float("inf")),


Similar to max(), I'm wondering what the behavior should be if both self and other are None

>>> ExecutionResources().min(ExecutionResources()) ExecutionResources(cpu=inf, gpu=inf, object_store_memory=inf)

bveeramani · 2024-02-12T19:26:31Z

python/ray/data/context.py

@@ -239,6 +239,10 @@ def __init__(
        # The max number of blocks that can be buffered at the streaming generator of
        # each `DataOpTask`.
        self._max_num_blocks_in_streaming_gen_buffer = None
+        # Whether to enable ReservationOpResourceLimiter.
+        self.op_resource_reservation_enabled = False


Are we planning on adding an environment variable for this? More generally, when don't we want to add an environment variable for a setting? Seems like most do right now.

We should add environment varilable and add into parameter as other configs.

will add an env var for this.
but regarding parameter of __init__, I don't think it's useful anymore. because there are too many parameters already. and now we already don't construct a DataContext with __init__ params.

but regarding parameter of init, I don't think it's useful anymore. because there are too many parameters already. and now we already don't construct a DataContext with init params.

We can clean it up later. @raulchen - is that you cannot add to __init__ because smoething is failing?

oh no, nothing was failing. I just think we should stop adding new params, and clean up the existing ones later.

bveeramani · 2024-02-12T19:28:34Z

python/ray/data/_internal/execution/resource_manager.py

+        return self._op_resource_limiter.get_op_limits(op)
+
+
+class OpResourceLimiter(metaclass=ABCMeta):


Suggested change

class OpResourceLimiter(metaclass=ABCMeta):

class OpResourceLimiter(abc.ABC):

Using ABC as a base class has essentially the same effect as specifying metaclass=abc.ABCMeta, but is simpler to type and easier to read.

bveeramani · 2024-02-12T19:37:22Z

python/ray/data/_internal/execution/resource_manager.py

+
+        return


Suggested change

return

bveeramani · 2024-02-12T19:39:42Z

python/ray/data/_internal/execution/resource_manager.py

@@ -108,6 +121,8 @@ def get_global_limits(self) -> ExecutionResources:
            gpu=gpu,
            object_store_memory=object_store_memory,
        )
+        if self._op_resource_limiter is not None:
+            self._op_resource_limiter.on_global_limits_updated(self._global_limits)


Felt confused at first about why we're calling on_global_limits_updated in get_global_limits. get_and_update_global_limits might be a clearer name, although maybe that's too verbose

I think the name of get_global_limits makes sense for the users of this function.
Basically this function is responsible for returning the current global limits. The internal cache is an implementation detail that users don't need to know.
And regarding on_global_limits_updated, from OpResourceLimiter's perspective, ResourceManager is responsible for calling this callback. When to call it is internal to ResourceManager.

actually there is an issue, on_global_limits_updated may not get called before update_resources.
I've changed the code to call on_global_limits_updated within the OpResourceLimiter itself.

bveeramani · 2024-02-12T19:41:11Z

python/ray/data/_internal/execution/resource_manager.py

+                op.incremental_resource_usage()
+            )
+
+    def update_usages(self):


Is update_usages always called whenever the global limits change?

it doesn't need to. It only needs to be updated when ReousrceManager.update_usages is called.

c21 · 2024-02-09T17:08:33Z

python/ray/data/_internal/execution/interfaces/execution_options.py

+        """Returns an ExecutionResources object with zero resources."""
+        if cls._ZERO is None:
+            cls._ZERO = ExecutionResources(0.0, 0.0, 0)
+        return cls._ZERO


shall we always return a new object return ExecutionResources(0.0, 0.0, 0)? Worried about if the caller changes the value accidentally, and affects other callers.

c21 · 2024-02-09T17:10:21Z

python/ray/data/_internal/execution/interfaces/execution_options.py

+        """
+        total = ExecutionResources()
+        if self.cpu is not None or other.cpu is not None:
+            total.cpu = (self.cpu or 0.0) - (other.cpu or 0.0)


should we add a check when value is negative? what does it mean when the value is negative?

good point. I'll add a post-create check.

oops, I misunderstood your question. If the value is negative, it just means that self is less than other. I think it's okay here. For the concrete use cases, I'll do something like a.subtract(b).max(ExecutionResources.zero()) to make it non-negative.

ok, let's make sure do the .max() part. The negative execution resource is quite confusing.

c21 · 2024-02-12T21:03:50Z

python/ray/data/_internal/execution/resource_manager.py

+        self._eligible_ops = [
+            op for op in self._resource_manager._topology if isinstance(op, MapOperator)
+        ]


Let's add a TODO to handle sub-plan of map operators? i.e. map().map().repartition().map().map()...

discussed offline, we'll only enable this feature when the dataset only has map/limit/streaming_split. Because memory accounting is not correct for other ops yet.

c21 · 2024-02-12T21:12:53Z

python/ray/data/tests/test_resource_manager.py

+
+
+class TestReservationOpResourceLimiter:
+    def test_basic(self, restore_data_context):


can we also add a unit test when the DAG has non-map operator as well?

Signed-off-by: Hao Chen <chenh1024@gmail.com>

bveeramani

Nice! LGTM!

bveeramani · 2024-02-13T21:32:58Z

python/ray/data/_internal/execution/resource_manager.py

+        if op in self._op_limits:
+            return self._op_limits[op]
+        else:
+            return ExecutionResources.inf()


Nit: Optionally,

Suggested change

if op in self._op_limits:

return self._op_limits[op]

else:

return ExecutionResources.inf()

return self._op_limits.get(op, ExecutionResources.inf())

Not a big deal, but the current way will be slightly more efficient, as inf() now still returns a new object each time.

python/ray/data/tests/test_resource_manager.py

bveeramani · 2024-02-13T21:52:41Z

python/ray/data/tests/test_resource_manager.py

+            2.5, float("inf"), 100
+        )
+
+        # Test global_limits exceeded.


Nit: global_limits isn't exceeded in this test, right? 500 + 125 + 150 = 775 < 800? This is testing what happens when the shared memory is full?

bveeramani · 2024-02-13T21:54:38Z

python/ray/data/tests/test_resource_manager.py

@@ -138,3 +145,145 @@ def test_calculating_usage(self):
        assert resource_manager.get_downstream_object_store_memory(o1) == 1700
        assert resource_manager.get_downstream_object_store_memory(o2) == 1700
        assert resource_manager.get_downstream_object_store_memory(o3) == 1000
+
+
+class TestReservationOpResourceLimiter:


Nit: Not sure if I already missed this, but we could also add a test that we always reserve at minimum incremental_resource_usage.

yes, I missed that.

c21

LGTM

c21 · 2024-02-13T22:00:58Z

python/ray/data/_internal/execution/interfaces/execution_options.py

+        """
+        total = ExecutionResources()
+        if self.cpu is not None or other.cpu is not None:
+            total.cpu = (self.cpu or 0.0) - (other.cpu or 0.0)


ok, let's make sure do the .max() part. The negative execution resource is quite confusing.

Signed-off-by: Hao Chen <chenh1024@gmail.com>

raulchen added 3 commits February 5, 2024 19:19

implement ReservationOpResourceLimiter

35851d6

Signed-off-by: Hao Chen <chenh1024@gmail.com>

fix

fbf40dd

Signed-off-by: Hao Chen <chenh1024@gmail.com>

comment

129ead5

Signed-off-by: Hao Chen <chenh1024@gmail.com>

raulchen requested review from ericl, scv119, c21, amogkam, scottjlee, bveeramani, stephanie-wang and omatthew98 as code owners February 7, 2024 00:15

raulchen added 4 commits February 6, 2024 16:20

comment

d41c99a

Signed-off-by: Hao Chen <chenh1024@gmail.com>

add test

9dd7203

Signed-off-by: Hao Chen <chenh1024@gmail.com>

fix

5240a67

Signed-off-by: Hao Chen <chenh1024@gmail.com>

revert unrelated

7e8381b

bveeramani reviewed Feb 12, 2024

View reviewed changes

c21 reviewed Feb 12, 2024

View reviewed changes

bveeramani mentioned this pull request Feb 12, 2024

[Data] Account internal inqueue to previous operator's memory usage #42930

Merged

8 tasks

raulchen added 10 commits February 12, 2024 16:54

fix zero() and inf()

4ea3ae1

Signed-off-by: Hao Chen <chenh1024@gmail.com>

env vars

7292bbc

Signed-off-by: Hao Chen <chenh1024@gmail.com>

small fixes

8f63c1f

Signed-off-by: Hao Chen <chenh1024@gmail.com>

only enable for ops with accurate accouting

6b7e007

Signed-off-by: Hao Chen <chenh1024@gmail.com>

fix

a906669

Signed-off-by: Hao Chen <chenh1024@gmail.com>

fix on_global_limits_updated

cf84dcb

Signed-off-by: Hao Chen <chenh1024@gmail.com>

fix no map

d5df089

Signed-off-by: Hao Chen <chenh1024@gmail.com>

fix test

43e146f

Signed-off-by: Hao Chen <chenh1024@gmail.com>

test_no_eligible_ops

d3653cb

Signed-off-by: Hao Chen <chenh1024@gmail.com>

add pytest entry

a51f4e0

Signed-off-by: Hao Chen <chenh1024@gmail.com>

bveeramani self-assigned this Feb 13, 2024

raulchen added 2 commits February 13, 2024 12:03

simplify type checking

c170fbb

Signed-off-by: Hao Chen <chenh1024@gmail.com>

lint

e6c38b2

Signed-off-by: Hao Chen <chenh1024@gmail.com>

bveeramani approved these changes Feb 13, 2024

View reviewed changes

c21 approved these changes Feb 13, 2024

View reviewed changes

raulchen added 3 commits February 13, 2024 15:59

Merge branch 'master' into per-op-resource-limit2

65f14b4

test incremental

725aa8f

Signed-off-by: Hao Chen <chenh1024@gmail.com>

lint

4a49b7f

Signed-off-by: Hao Chen <chenh1024@gmail.com>

raulchen merged commit ccf79b2 into ray-project:master Feb 14, 2024
8 of 9 checks passed

raulchen deleted the per-op-resource-limit2 branch February 14, 2024 02:02

kevin85421 pushed a commit to kevin85421/ray that referenced this pull request Feb 17, 2024

[data] implement per-op resource reservation (ray-project#43026)

638f340

Signed-off-by: Hao Chen <chenh1024@gmail.com>

c21 mentioned this pull request Feb 26, 2024

[data] Enable per-op resource reservation #43171

Merged

8 tasks

		return self._op_resource_limiter.get_op_limits(op)


		class OpResourceLimiter(metaclass=ABCMeta):

	class OpResourceLimiter(metaclass=ABCMeta):
	class OpResourceLimiter(abc.ABC):



		class TestReservationOpResourceLimiter:
		def test_basic(self, restore_data_context):

[data] implement per-op resource reservation #43026

[data] implement per-op resource reservation #43026

Conversation

raulchen commented Feb 7, 2024 • edited

Why are these changes needed?

Related issue number

Checks

bveeramani left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

raulchen Feb 13, 2024 • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

bveeramani left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

c21 left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

raulchen commented Feb 7, 2024 •

edited

raulchen Feb 13, 2024 •

edited