Skip to content

Commit 1ad40e3

Browse files
committed
review: address feedback on RuntimeCacheHandle pickle changes
- ``RuntimeCacheHandle::serialize``: switch the pending-bytes copy from ``std::memcpy`` to ``std::copy`` (with free-function ``std::cbegin`` / ``std::cend``), and use free-function ``std::empty`` for the empty-vector check. Idiomatic C++17 in a hot path that didn't need the C-string-style API. - Drop the explanatory comment block above the pending-bytes branch. - Trim the wrapper ``__getstate__`` docstring: the rationale that named specific picklability hazards (``weakref``, ``threading.Lock``) is implementation noise that doesn't belong in the wrapper's contract. - Drop the cpp-rt skip on ``test_python_handle_pickle_preserves_pending_warm_bytes``. The test instantiates ``_RuntimeCacheHandle`` directly rather than going through ``RuntimeCache._handle`` selection, so the python class is exercisable regardless of which runtime is active. Updated docstring notes the directness. Behavior is unchanged. Addresses review comments on core/runtime/RuntimeSettings.cpp:131, :136 and tests/py/dynamo/runtime/test_000_runtime_cache.py:341.
1 parent c4772af commit 1ad40e3

3 files changed

Lines changed: 10 additions & 16 deletions

File tree

core/runtime/RuntimeSettings.cpp

Lines changed: 4 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
#include "core/runtime/RuntimeSettings.h"
22

3+
#include <algorithm>
34
#include <array>
45
#include <cstring>
56
#include <iterator>
@@ -124,16 +125,12 @@ at::Tensor RuntimeCacheHandle::serialize() const {
124125
#ifdef TRT_MAJOR_RTX
125126
std::lock_guard<std::mutex> lock(state_mu_);
126127
if (!trt_handle_) {
127-
// Pre-materialize: forward any ``pending_warm_bytes_`` so that the
128-
// handle's persistable state survives ``save_to_stream`` and pickle
129-
// even before any engine has triggered ``ensure_materialized``.
130-
// ``save_to_stream`` / ``def_pickle`` round-trip then matches the
131-
// python facade's behavior (it reads ``_pending_warm_bytes`` directly).
132-
if (pending_warm_bytes_.empty()) {
128+
if (std::empty(pending_warm_bytes_)) {
133129
return empty();
134130
}
135131
auto tensor = at::empty({static_cast<int64_t>(pending_warm_bytes_.size())}, opts);
136-
std::memcpy(tensor.data_ptr(), pending_warm_bytes_.data(), pending_warm_bytes_.size());
132+
std::copy(
133+
std::cbegin(pending_warm_bytes_), std::cend(pending_warm_bytes_), static_cast<uint8_t*>(tensor.data_ptr()));
137134
return tensor;
138135
}
139136
auto host_mem = make_trt(trt_handle_->serialize());

py/torch_tensorrt/dynamo/runtime/_TorchTensorRTModule.py

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -547,10 +547,7 @@ def __getstate__(self) -> dict[str, Any]:
547547
cpp-side exclusion (engine bytes never carry these fields).
548548
549549
``_implicit_cache_handle`` is dropped alongside ``_runtime_settings``
550-
because it aliases the same ``RuntimeCache`` instance and would
551-
otherwise drag a ``weakref`` (via the handle's ``atexit`` closure)
552-
and a Python-only ``threading.Lock`` (when the python-runtime path
553-
is active) into pickle -- neither is picklable.
550+
because it aliases the same ``RuntimeCache`` instance.
554551
"""
555552
get_state = getattr(super(), "__getstate__", None)
556553
state = (get_state() if get_state else self.__dict__).copy()

tests/py/dynamo/runtime/test_000_runtime_cache.py

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -334,15 +334,15 @@ def test_user_built_handle_no_autosave_by_default(self):
334334
"User-built handle with autosave_on_del=False should not save on GC",
335335
)
336336

337-
@unittest.skipIf(
338-
ENABLED_FEATURES.torch_tensorrt_runtime,
339-
"exercises the python ``_RuntimeCacheHandle`` directly; cpp-rt "
340-
"path is covered by ``register_jit_hooks.cpp`` ``def_pickle``",
341-
)
342337
def test_python_handle_pickle_preserves_pending_warm_bytes(self):
343338
"""A python ``_RuntimeCacheHandle`` that has bytes loaded but
344339
hasn't materialized them yet must round-trip those bytes through
345340
pickle. Matches the cpp ``def_pickle`` contract (path + bytes).
341+
342+
The test instantiates ``_RuntimeCacheHandle`` directly rather than
343+
going through ``RuntimeCache`` (which would pick the torchbind
344+
backing on cpp rt), so the python class is exercisable regardless
345+
of which runtime is active.
346346
"""
347347
import pickle
348348

0 commit comments

Comments
 (0)