Skip to content

Commit

Permalink
Update serialization code in benchmarks (#1098)
Browse files Browse the repository at this point in the history
  • Loading branch information
galipremsagar authored Dec 11, 2024
1 parent d5e2ae5 commit 9efacc6
Showing 1 changed file with 3 additions and 3 deletions.
6 changes: 3 additions & 3 deletions ucp/benchmarks/cudf_merge.py
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,7 @@

import cudf # noqa: E402
import rmm # noqa: E402
from cudf.core.abc import Serializable # noqa: E402
from rmm.allocators.cupy import rmm_cupy_allocator # noqa: E402


Expand All @@ -48,7 +49,7 @@ def sizeof_cudf_dataframe(df):


async def send_df(ep, df):
header, frames = df.serialize()
header, frames = df.device_serialize()
header["frame_ifaces"] = [f.__cuda_array_interface__ for f in frames]
header = pickle.dumps(header)
header_nbytes = np.array([len(header)], dtype=np.uint64)
Expand All @@ -72,8 +73,7 @@ async def recv_df(ep):
for frame in frames:
await ep.recv(frame)

cudf_typ = pickle.loads(header["type-serialized"])
return cudf_typ.deserialize(header, frames)
return Serializable.device_deserialize(header, frames)


async def barrier(rank, eps):
Expand Down

0 comments on commit 9efacc6

Please sign in to comment.