[Bugfix][Disco] Handle NDArray larger than OS buffer for pipe #16992
Conversation
Prior to this commit, using `disco.Session` methods to transfer `NDArray` instances to workers could raise an exception if the `NDArray` is larger than the buffer allocated by the OS for the controller/worker pipe. In these case, the first call to the `Read` method of `tvm::support::Pipe` would successfully return, but only with the initial bytes of the `NDArray`. Receiving the full `NDArray` requires repeatedly calling the POSIX `read` function. This commit updates the `Read` and `Write` methods of `tvm::support::Pipe` to repeatedly call the underlying read/write methods, until the full `NDArray` has been transferred. This commit does not add any unit tests, as the existing unit test `tests/python/disco/test_ccl.py::test_attention[nccl-ProcessSession]` requires this change to pass.
|
Thanks @Lunderberg ! Just to keep Socket and pipe behavior consistent with each other, low-level read/write can remain partial while readall and writeall contains convenient method for the intended behavior |
|
I agree. This is a stop-gap measure, as the longer-term fix will require updating the |
This resolves a conflict between two recent changes. In apache#16989, reads of size zero are used to identify hangups in `ProcessSession`. In apache#16992, reads of size zero are treated as an error to avoid infinite loops while waiting for data to be ready. For a long-term resolution, the `dmlc::Stream` interface will need to be updated, so that the `Write` method returns the number of bytes written, just as the `Read` method currently does. This will allow the calling scope to verify the number of bytes received.
|
@tqchen The changes for the long-term fix are implemented in dmlc/dmlc-core#686, with TVM compatibility changes implemented in #16998. Once those changes land, we will be able to remove the stop-gap implementation, and provide |
|
Ah I see, the Stream interface was actually created to be in analogy of FileStream interface(which do not have partial writes). After seeing your fix, i realized that this PR actually did the right thing given pipe inheritated from the stream interface |
|
The socket interface was mainly meant to accomodate for the non-blocking case where full read is not always possible |
Thank you, and that history makes a lot of sense. I'm thinking it would still be a good change (though low priority) to expose partial writes in
Good point. Looking into it, it looks like the |
This resolves a conflict between two recent changes. In #16989, reads of size zero are used to identify hangups in `ProcessSession`. In #16992, reads of size zero are treated as an error to avoid infinite loops while waiting for data to be ready. For a long-term resolution, the `dmlc::Stream` interface will need to be updated, so that the `Write` method returns the number of bytes written, just as the `Read` method currently does. This will allow the calling scope to verify the number of bytes received.
Prior to this commit, using
disco.Sessionmethods to transferNDArrayinstances to workers could raise an exception if theNDArrayis larger than the buffer allocated by the OS for the controller/worker pipe. In these case, the first call to theReadmethod oftvm::support::Pipewould successfully return, but only with the initial bytes of theNDArray. Receiving the fullNDArrayrequires repeatedly calling the POSIXreadfunction.This commit updates the
ReadandWritemethods oftvm::support::Pipeto repeatedly call the underlying read/write methods, until the fullNDArrayhas been transferred.This commit does not add any unit tests, as the existing unit test
tests/python/disco/test_ccl.py::test_attention[nccl-ProcessSession]requires this change to pass.