Skip to content

Commit 5620e61

Browse files
pytorchmergebotpytorchmergebot
committed
Revert "[BE] Modify PyObjectSlot the assume only a single interpreter is in use (pytorch#158407)"
This reverts commit 255c054. Reverted pytorch#158407 on behalf of https://github.com/ZainRizvi due to Reverting as per offline discussion to fix internal breaks. @PaliC will reland this as a codev diff. Instructions here: https://fburl.com/fixing-ghfirst-reverts ([comment](pytorch#158288 (comment)))
1 parent ee84ba4 commit 5620e61

File tree

3 files changed

+92
-7
lines changed

3 files changed

+92
-7
lines changed

‎c10/core/impl/PyObjectSlot.cpp‎

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,19 @@ PyInterpreter& PyObjectSlot::load_pyobj_interpreter() const {
4444
if (interpreter) {
4545
return *interpreter;
4646
}
47-
TORCH_CHECK(false, "cannot access PyObject for Tensor - no interpreter set");
47+
TORCH_CHECK(
48+
false,
49+
"cannot access PyObject for Tensor on interpreter ",
50+
(*pyobj_interpreter_.load())->name());
51+
}
52+
53+
bool PyObjectSlot::check_interpreter(PyInterpreter* interpreter) {
54+
return interpreter == pyobj_interpreter();
55+
}
56+
57+
bool PyObjectSlot::has_pyobj_nonhermetic() {
58+
return check_pyobj(pyobj_interpreter(), /*ignore_hermetic_tls=*/true)
59+
.has_value();
4860
}
4961

5062
bool PyObjectSlot::owns_pyobj() {

‎c10/core/impl/PyObjectSlot.h‎

Lines changed: 68 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,48 @@ struct C10_API PyObjectSlot {
2828
PyInterpreter* self_interpreter,
2929
PyObject* pyobj,
3030
PyInterpreterStatus status) {
31-
pyobj_interpreter_.store(self_interpreter, std::memory_order_relaxed);
31+
impl::PyInterpreter* expected = nullptr;
32+
switch (status) {
33+
case impl::PyInterpreterStatus::DEFINITELY_UNINITIALIZED:
34+
// caller guarantees there is no multithreaded access; if there is
35+
// no data race OK to do a relaxed store
36+
pyobj_interpreter_.store(self_interpreter, std::memory_order_relaxed);
37+
break;
38+
case impl::PyInterpreterStatus::TAGGED_BY_US:
39+
// no tagging is necessary, the tag is already correct
40+
break;
41+
case impl::PyInterpreterStatus::MAYBE_UNINITIALIZED:
42+
// attempt to claim this TensorImpl with the specified interpreter
43+
// tag
44+
if (pyobj_interpreter_.compare_exchange_strong(
45+
expected, self_interpreter, std::memory_order_acq_rel)) {
46+
break;
47+
}
48+
// test if, actually, it was already tagged by us! this situation can't
49+
// be caused by a race, but it could be caused by a situation
50+
// where someone conservatively tagged the tensor as MAYBE_UNINITIALIZED
51+
// (because they didn't pre-check the tag) when actually it was
52+
// owned by the interpreter
53+
if (expected == self_interpreter) {
54+
break;
55+
}
56+
// fallthrough, we lost the race. We are guaranteed not to lose the
57+
// race with ourself, as calls to init_pyobj with the same interpreter
58+
// ID must be sequentialized by the GIL
59+
[[fallthrough]];
60+
case impl::PyInterpreterStatus::TAGGED_BY_OTHER:
61+
TORCH_CHECK(
62+
false,
63+
"cannot allocate PyObject for Tensor on interpreter ",
64+
self_interpreter,
65+
" that has already been used by another torch deploy interpreter ",
66+
pyobj_interpreter_.load());
67+
}
68+
69+
// we are the ONLY thread that can have gotten to this point. It is not
70+
// possible to conflict with another zero interpreter as access is protected
71+
// by GIL
72+
// NB: owns_pyobj tag is initially false
3273
pyobj_ = pyobj;
3374
}
3475

@@ -56,16 +97,30 @@ struct C10_API PyObjectSlot {
5697
std::optional<PyObject*> check_pyobj(
5798
PyInterpreter* self_interpreter,
5899
bool ignore_hermetic_tls = false) const {
100+
// Note [Memory ordering on Python interpreter tag]
59101
impl::PyInterpreter* interpreter =
60102
pyobj_interpreter_.load(std::memory_order_acquire);
61103
if (interpreter == nullptr) {
104+
// NB: This never returns DEFINITELY_UNINITIALIZED because there is
105+
// always the possibility that another thread races to initialize
106+
// after we query here. The only time when we can conclude a tensor
107+
// is definitely uninitialized is when we have just allocated it and
108+
// it cannot have escaped to other threads yet
62109
return std::nullopt;
63-
}
64-
65-
if (!ignore_hermetic_tls && c10::impl::HermeticPyObjectTLS::get_state()) {
66-
return std::nullopt;
110+
} else if (interpreter == self_interpreter) {
111+
// NB: pyobj_ could still be null!
112+
if (!ignore_hermetic_tls && c10::impl::HermeticPyObjectTLS::get_state()) {
113+
return std::nullopt;
114+
} else {
115+
return _unchecked_untagged_pyobj();
116+
}
67117
} else {
68-
return _unchecked_untagged_pyobj();
118+
TORCH_CHECK(
119+
false,
120+
"cannot access PyObject for Tensor on interpreter ",
121+
(*self_interpreter)->name(),
122+
" that has already been used by another torch deploy interpreter ",
123+
(*pyobj_interpreter_.load())->name());
69124
}
70125
}
71126

@@ -75,6 +130,13 @@ struct C10_API PyObjectSlot {
75130

76131
PyInterpreter& load_pyobj_interpreter() const;
77132

133+
// Check if the PyObjectSlot's interpreter is the same as the specified
134+
// interpreter
135+
bool check_interpreter(PyInterpreter* interpreter);
136+
137+
// Check if the PyObjectSlot is holding a PyObject, owned or non-owned
138+
bool has_pyobj_nonhermetic();
139+
78140
bool owns_pyobj();
79141

80142
void set_owns_pyobj(bool b);

‎torch/csrc/Storage.cpp‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -98,6 +98,17 @@ PyObject* THPStorage_Wrap(c10::Storage storage) {
9898
}
9999
c10::impl::PyObjectSlot* pyobj_slot = storage_impl->pyobj_slot();
100100

101+
// If the StorageImpl has a PyObject that is managed by a different
102+
// interpreter than the current one, create a new StorageImpl that points to
103+
// the same data and then create the Python storage from that.
104+
// NOTE: This is only supposed to happen in MultiPy // codespell:ignore
105+
if (pyobj_slot->has_pyobj_nonhermetic() &&
106+
!pyobj_slot->check_interpreter(getPyInterpreter())) {
107+
return THPStorage_NewWithStorage(
108+
THPStorageClass,
109+
c10::newStorageImplFromRefcountedDataPtr(storage),
110+
c10::impl::PyInterpreterStatus::DEFINITELY_UNINITIALIZED);
111+
}
101112
std::optional<PyObject*> maybe_pyobj = pyobj_slot->check_pyobj(
102113
getPyInterpreter(), /*ignore_hermetic_tls=*/false);
103114
c10::impl::PyInterpreterStatus status =

0 commit comments

Comments
 (0)