Just-in-time deserialization #353
Conversation
Codecov Report
@@ Coverage Diff @@
## branch-0.17 #353 +/- ##
==============================================
Coverage ? 58.36%
==============================================
Files ? 19
Lines ? 1561
Branches ? 0
==============================================
Hits ? 911
Misses ? 650
Partials ? 0 Continue to review full report at Codecov.
|
|
@madsbk so it seems that the reason it fails to |
Make sense, thanks @pentschev for investigating this! |
madsbk
force-pushed
the
jit_deserialization
branch
from
edf503e to
40058cd
Compare
|
Well Pandas is required by |
There was a problem hiding this comment.
Mads, this is some great work you put here, it's great to see the incredible speedup you've achieved!
Functionality-wise, I don't have much to comment on, some parts are much beyond my understanding, and it seems like there was enough testing by you and others with TPCx-BB. Therefore, my comments/requests revolve around style and documentations only.
| return sys.getsizeof(self._obj_pxy_deserialize()) | ||
|
|
||
| def __len__(self): | ||
| return len(self._obj_pxy_deserialize()) |
There was a problem hiding this comment.
Will this deserialize the entire object just to check its length? If so, wouldn't it make sense to store an attribute with the length during serialization and just return it to avoid this?
There was a problem hiding this comment.
Yes currently len(x) will deserialize x. What you suggest, we are already doing with x.name, which I added because worker will access x.name before execution tasks.
My plan is to address this issue in a following PR when we have some more experience using JIT deserialization. I suspect we will need to handle a range of attributes.
There was a problem hiding this comment.
Makes sense, thanks for the continued work on it.
| return complex(self._obj_pxy_deserialize()) | ||
|
|
||
| def __index__(self): | ||
| return operator.index(self._obj_pxy_deserialize()) |
There was a problem hiding this comment.
Are all these dunder method implementations here because they are going to be needed in practice, or are they here just to match a full implementation?
There was a problem hiding this comment.
Some are needed for ProxyObject to pass-through cudf and tpcx-bb workflows but many of them are just here for completeness.
There was a problem hiding this comment.
I see, given they're already there I won't suggest we remove them, but for future reference I think we can try to keep code shorter instead of solving all possible cases that are probably unnecessary, this helps with future maintainability.
There was a problem hiding this comment.
I agree but since ProxyObject is exposed to the end user, I think it is reasonable to support most common operations. Say a user writes a Dask tasks that uses Numpy arrays, the user should be able to use most, if not all, Numpy operations.
| "subclass": subclass, | ||
| "serializers": serializers, | ||
| } | ||
| self._obj_pxy_lock = threading.RLock() |
There was a problem hiding this comment.
I'm curious, why isn't a simple threading.Lock sufficient here?
There was a problem hiding this comment.
It is to handle all the methods that acquires the lock before calling _obj_pxy_deserialize(), which also acquires the lock.
dask_cuda/proxy_object.py
Outdated
| List of frames that makes up the serialized object | ||
| """ | ||
| with self._obj_pxy_lock: | ||
| assert serializers is not None |
There was a problem hiding this comment.
Should we replace runtime assertions to raising more appropriate exceptions? For example here I think it would make more sense to raise a ValueError.
There was a problem hiding this comment.
I am not sure, the check is there to assert the internal logic of the class. The _obj_pxy_* methods are not supposed to be called by the user.
There was a problem hiding this comment.
I understand that, but nevertheless I think of assert as a debug statement, so I'm not sure it makes much sense in runtime code. I know Python doesn't distinguish between "debug" and "release" builds, like we can with C, but regardless I think it's clearer to have errors be specific, one could also argue that assertions would substitute virtually any exceptions we could raise, but the exceptions let us be clearer about our intent.
In the interest of avoiding nitpicking, I will leave the final decision up to you, but I think there's real value in being more specific instead of using assertions, whenever we want to raise an exception if something goes wrong.
There was a problem hiding this comment.
I am not sure I agree completely. I like to use assert in internal logic in order to both catch bugs while developing and document the expected state of an object.
However, in this case I agree with you. Since it can be triggered by user input, it is preferable to raise ValueError.
|
We'll give this a test in a fresh environment on TPCxBB workloads on constrained memory systems as well |
|
@pentschev thanks for the review! |
|
Also, I think we can leave this open until EOD or Monday so people who are interested can still review or test it before merging, but otherwise it looks good to me! Thanks again @madsbk for the work you put into this, it's really great! |
Co-authored-by: Peter Andreas Entschev <[email protected]>
It is all in :)
Yes, no hurry. I also like to wait on a verdict from @beckernick :P |
|
I am running this currently :) |
|
In a fresh environment today, I ran tpcxbb q02 several times with the following configuration (the default):
By default, q02 takes ~300 seconds with the above. With With With DASK_JIT_UNSPILL=False, the peak memory after each time q02 finishes is the following: Peak memory spiked on GPU 0 due to the client process, but went back down as expected after completion. After additional runs, the memory profile looks identical. With DASK_JIT_UNSPILL=True, the peak memory after each time q02 finishes looks different, due to the OOM. When it succeeded, we had to allocate more memory so the active memory does not go back down to 29352 MB on the GPUs that needed more memory.
Going down to a device memory limit of 10GB doesn't stop the OOM, but can also sometimes succeed. When it succeeds, we still see the huge speed boost. Finally, when it succeeds, we don't ever need to allocate more memory for the pool. It seems like if the workload needs to allocate additional memory due to JIT_UNSPILL=True, we will end up failing. This is a GPU memory screenshot from the middle of a q02 run that failed, which was kicked off after one that succeeded:
As a result, I think this PR poses limited risk. When enabled, JIT unspilling provides significant performance gains but consistently increase memory pressure (and usually OOMs on these 8 GPU tests). If the memory increases could be avoided, I think it would be worth enabling this by default. |
|
Thanks for all of your work on this Mads! And everyone for the reviews! 😄 |
|
Thanks @beckernick, I don't know exactly why we see the memory spikes but I will make a follow up PR to address it when I find the time :) |
This PR implements Just-in-time deserialization of device memory by wrapping
DeviceHostFileitems in a proxy class:ProxyObject. The deserializaion of the items are then delayed until the items are accessed.Fixes #342
Ref. #57
Use
In order to enable JIT deserializaion, use the new
jit_unspillargument when creatingLocalCUDACluster, set--enable-jit-unspillwhen starting a CUDAWorker, or set the environment variableDASK_JIT_UNSPILL=TrueTODO
ProxyObject(should come up with a better name)ProxyObjectin order to allow the communication spilled dataProxyObjectin order to avoid re-spilling of data in the case whereProxyObjectwraps un-spilled device data.ProxyObject