Skip to content

Assert on empty PJRT buffers #9062

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: master
Choose a base branch
from

Conversation

rpsilva-aws
Copy link
Collaborator

@rpsilva-aws rpsilva-aws commented Apr 29, 2025

Simple CR to avoid a segmentation fault when there are placeholder tensors involved, as we attempt to de-reference the device from the buffer. It fixes the seg fault for #9049.

ERROR:root:Caught an exception when exiting the process. Exception: 
Traceback (most recent call last):
  File "/ansible/pytorch/xla/torch_xla/__init__.py", line 209, in _prepare_to_exit
    _XLAC._prepare_to_exit()
RuntimeError: torch_xla/csrc/runtime/pjrt_computation_client.cc:741 : Check failed: pjrt_data->buffer != nullptr 
*** Begin stack trace ***
        tsl::CurrentStackTrace[abi:cxx11]()
        torch_xla::runtime::PjRtComputationClient::ExecuteComputation(torch_xla::runtime::ComputationClient::Computation const&, absl::lts_20230802::Span<std::shared_ptr<torch_xla::runtime::ComputationClient::Data> const>, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const&, torch_xla::runtime::ComputationClient::ExecuteComputationOptions const&)




        torch::lazy::MultiWait::Complete(std::function<void ()> const&)
        std::function<void ()>::operator()() const
        tsl::thread::EigenEnvironment::ExecuteTask(tsl::thread::EigenEnvironment::Task const&)
        Eigen::ThreadPoolTempl<tsl::thread::EigenEnvironment>::WorkerLoop(int)
        Eigen::ThreadPoolTempl<tsl::thread::EigenEnvironment>::ThreadPoolTempl(int, bool, tsl::thread::EigenEnvironment)::{lambda()#1}::operator()() const
        void std::__invoke_impl<void, Eigen::ThreadPoolTempl<tsl::thread::EigenEnvironment>::ThreadPoolTempl(int, bool, tsl::thread::EigenEnvironment)::{lambda()#1}&>(std::__invoke_other, Eigen::ThreadPoolTempl<tsl::thread::EigenEnvironment>::ThreadPoolTempl(int, bool, tsl::thread::EigenEnvironment)::{lambda()#1}&)
        std::enable_if<is_invocable_r_v<void, Eigen::ThreadPoolTempl<tsl::thread::EigenEnvironment>::ThreadPoolTempl(int, bool, tsl::thread::EigenEnvironment)::{lambda()#1}&>, std::enable_if>::type std::__invoke_r<void, Eigen::ThreadPoolTempl<tsl::thread::EigenEnvironment>::ThreadPoolTempl(int, bool, tsl::thread::EigenEnvironment)::{lambda()#1}&>(void&&, (Eigen::ThreadPoolTempl<tsl::thread::EigenEnvironment>::ThreadPoolTempl(int, bool, tsl::thread::EigenEnvironment)::{lambda()#1}&)...)
        std::_Function_handler<void (), Eigen::ThreadPoolTempl<tsl::thread::EigenEnvironment>::ThreadPoolTempl(int, bool, tsl::thread::EigenEnvironment)::{lambda()#1}>::_M_invoke(std::_Any_data const&)
        std::function<void ()>::operator()() const
        tsl::thread::EigenEnvironment::CreateThread(std::function<void ()>)::{lambda()#1}::operator()() const
        void std::__invoke_impl<void, tsl::thread::EigenEnvironment::CreateThread(std::function<void ()>)::{lambda()#1}&>(std::__invoke_other, tsl::thread::EigenEnvironment::CreateThread(std::function<void ()>)::{lambda()#1}&)
        std::__invoke_result<tsl::thread::EigenEnvironment::CreateThread(std::function<void ()>)::{lambda()#1}&>::type std::__invoke<tsl::thread::EigenEnvironment::CreateThread(std::function<void ()>)::{lambda()#1}&>(std::__invoke_result&&, (tsl::thread::EigenEnvironment::CreateThread(std::function<void ()>)::{lambda()#1}&)...)
        std::invoke_result<tsl::thread::EigenEnvironment::CreateThread(std::function<void ()>)::{lambda()#1}&>::type std::invoke<tsl::thread::EigenEnvironment::CreateThread(std::function<void ()>)::{lambda()#1}&>(std::invoke_result&&, (tsl::thread::EigenEnvironment::CreateThread(std::function<void ()>)::{lambda()#1}&)...)
        void absl::lts_20230802::internal_any_invocable::InvokeR<void, tsl::thread::EigenEnvironment::CreateThread(std::function<void ()>)::{lambda()#1}&, , void>(tsl::thread::EigenEnvironment::CreateThread(std::function<void ()>)::{lambda()#1}&)
        void absl::lts_20230802::internal_any_invocable::RemoteInvoker<false, void, tsl::thread::EigenEnvironment::CreateThread(std::function<void ()>)::{lambda()#1}&>(absl::lts_20230802::internal_any_invocable::TypeErasedState*)
        absl::lts_20230802::internal_any_invocable::Impl<void ()>::operator()()

@rpsilva-aws rpsilva-aws marked this pull request as ready for review April 29, 2025 23:58
@rpsilva-aws rpsilva-aws requested a review from ysiraichi April 29, 2025 23:59
@bhavya01
Copy link
Collaborator

bhavya01 commented May 2, 2025

Is it possible to add a test for this?

Looking at #9049 quickly, I didn't follow the path how calling _get_xla_tensors_hlo leads to ExecuteComputation call on the PjRtClient

@bhavya01 bhavya01 self-requested a review May 2, 2025 17:17
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants