TransferQueue/scripts/performance_test.py at dev · TransferQueue/TransferQueue · GitHub

This repository was archived by the owner on Jan 21, 2026. It is now read-only.

History

338 lines (267 loc) · 12.9 KB

Raw

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

157

158

159

160

161

162

163

164

165

166

167

168

169

170

171

172

173

174

175

176

177

178

179

180

181

182

183

184

185

186

187

188

189

190

191

192

193

194

195

196

197

198

199

200

201

202

203

204

205

206

207

208

209

210

211

212

213

214

215

216

217

218

219

220

221

222

223

224

225

226

227

228

229

230

231

232

233

234

235

236

237

238

239

240

241

242

243

244

245

246

247

248

249

250

251

252

253

254

255

256

257

258

259

260

261

262

263

264

265

266

267

268

269

270

271

272

273

274

275

276

277

278

279

280

281

282

283

284

285

286

287

288

289

290

291

292

293

294

295

296

297

298

299

300

301

302

303

304

305

306

307

308

309

310

311

312

313

314

315

316

317

318

319

320

321

322

323

324

325

326

327

328

329

330

331

332

333

334

335

336

337

338

import asyncio

import logging

import math

import random

import sys

import time

from pathlib import Path

import ray

import torch

from omegaconf import OmegaConf

from tensordict import TensorDict

from tensordict.tensorclass import NonTensorData

parent_dir = Path(__file__).resolve().parent.parent.parent

sys.path.append(str(parent_dir))

from transfer_queue import ( # noqa: E402

AsyncTransferQueueClient,

SimpleStorageUnit,

TransferQueueController,

process_zmq_server_info,

)

from transfer_queue.utils.utils import get_placement_group # noqa: E402

logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")

logger = logging.getLogger(__name__)

########################################################################

# Please set up Ray cluster before running this script

########################################################################

HEAD_NODE_IP = "NodeA" # Replace with your head node IP

WORKER_NODE_IP = "NodeB" # Replace with your worker node IP

# This is the Medium setting of the performance test.

# You can modify the parameters according to

# https://www.yuque.com/haomingzi-lfse7/lhp4el/tml8ke0zkgn6roey?singleDoc#

config_str = """

global_batch_size: 1024

seq_length: 8192

field_num: 10

num_global_batch: 1

num_data_storage_units: 8

"""

dict_conf = OmegaConf.create(config_str)

def create_complex_test_case(batch_size=None, seq_length=None, field_num=None):

tensor_field_size_bytes = batch_size * seq_length * 4

tensor_field_size_gb = tensor_field_size_bytes / (1024**3)

num_tensor_fields = (field_num + 1) // 2

num_nontensor_fields = field_num // 2

total_tensor_size_gb = tensor_field_size_gb * num_tensor_fields

total_nontensor_size_gb = (batch_size * 1024 / (1024**3)) * num_nontensor_fields

total_size_gb = total_tensor_size_gb + total_nontensor_size_gb

logger.info(f"Total data size: {total_size_gb:.6f} GB")

fields = {}

for i in range(field_num):

field_name = f"field_{i}"

if i % 2 == 0:

# Tensor

tensor_data = torch.randn(batch_size, seq_length, dtype=torch.float32)

fields[field_name] = tensor_data

else:

# NonTensorData

str_length = 1024

non_tensor_data = [

"".join(random.choices("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789", k=str_length))

for _ in range(batch_size)

]

fields[field_name] = NonTensorData(data=non_tensor_data, batch_size=(batch_size,), device=None)

batch_size_tuple = (batch_size,)

prompt_batch = TensorDict(

fields,

batch_size=batch_size_tuple,

device=None,

)

return prompt_batch, total_size_gb

@ray.remote

class RemoteDataStoreObjStore:

def __init__(self):

pass

def get_data(self, data_handler):

start_get = time.time()

ray.get(data_handler)

end_get = time.time()

get_time = end_get - start_get

return get_time

@ray.remote

class RemoteDataStoreRemote:

def __init__(self):

self.stored_data = None

def put_data(self, data):

self.stored_data = data

def get_data(self):

return self.stored_data

def clear_data(self):

self.stored_data = None

class RayBandwidthTester:

def __init__(self, config, test_mode="obj_store"):

self.config = config

self.test_mode = test_mode

if test_mode == "obj_store":

RemoteDataStore = RemoteDataStoreObjStore

else:

RemoteDataStore = RemoteDataStoreRemote

self.remote_store = RemoteDataStore.options(num_cpus=10, resources={f"node:{WORKER_NODE_IP}": 0.001}).remote()

logger.info(f"Remote data store created on worker node {WORKER_NODE_IP}")

def run_bandwidth_test(self):

start_create_data = time.time()

test_data, total_data_size_gb = create_complex_test_case(

batch_size=self.config.global_batch_size, seq_length=self.config.seq_length, field_num=self.config.field_num

)

end_create_data = time.time()

logger.info(f"Data creation time: {end_create_data - start_create_data:.8f}s")

if self.test_mode == "obj_store":

self._run_obj_store_test(test_data, total_data_size_gb)

else:

self._run_remote_test(test_data, total_data_size_gb)

def _run_obj_store_test(self, test_data, total_data_size_gb):

start_time = time.time()

data_handler = ray.put(test_data)

ray.get(self.remote_store.get_data.remote([data_handler]))

end_time = time.time()

transfer_time = end_time - start_time

throughput = (total_data_size_gb * 8) / transfer_time

logger.info("=" * 60)

logger.info("RAY OBJECT STORE BANDWIDTH TEST SUMMARY")

logger.info("=" * 60)

logger.info(f"Data Size: {(total_data_size_gb):.6f} GB")

logger.info(f"Transfer Time: {transfer_time:.8f}s")

logger.info(f"Throughput: {throughput:.8f} Gb/s")

def _run_remote_test(self, test_data, total_data_size_gb):

logger.info("Starting Ray PUT bandwidth test...")

start_put = time.time()

ray.get(self.remote_store.put_data.remote(test_data))

end_put = time.time()

put_time = end_put - start_put

logger.info(f"PUT Time: {put_time:.8f}s")

time.sleep(2)

logger.info("Starting Ray GET bandwidth test...")

start_get = time.time()

ray.get(self.remote_store.get_data.remote())

end_get = time.time()

get_time = end_get - start_get

logger.info(f"GET Time: {get_time:.8f}s")

ray.get(self.remote_store.clear_data.remote())

put_throughput = (total_data_size_gb * 8) / put_time

get_throughput = (total_data_size_gb * 8) / get_time

logger.info("=" * 60)

logger.info("RAY REMOTE ACTOR BANDWIDTH TEST SUMMARY")

logger.info("=" * 60)

logger.info(f"Data Size: {total_data_size_gb:.6f} GB")

logger.info(f"PUT Time: {put_time:.8f}s")

logger.info(f"GET Time: {get_time:.8f}s")

logger.info(f"PUT Throughput (Head->Worker): {put_throughput:.8f} Gb/s")

logger.info(f"GET Throughput (Worker->Head): {get_throughput:.8f} Gb/s")

logger.info(f"Round-trip Average Throughput: {total_data_size_gb * 16 / (put_time + get_time):.8f} Gb/s")

class TQBandwidthTester:

def __init__(self, config, remote_mode=False):

self.config = config

self.remote_mode = remote_mode

self.data_system_client = self._initialize_data_system()

def _initialize_data_system(self):

total_storage_size = self.config.global_batch_size * self.config.num_global_batch

self.data_system_storage_units = {}

if self.remote_mode:

for storage_unit_rank in range(self.config.num_data_storage_units):

storage_node = SimpleStorageUnit.options(

num_cpus=10,

resources={f"node:{WORKER_NODE_IP}": 0.001},

).remote(storage_unit_size=math.ceil(total_storage_size / self.config.num_data_storage_units))

self.data_system_storage_units[storage_unit_rank] = storage_node

else:

storage_placement_group = get_placement_group(self.config.num_data_storage_units, num_cpus_per_actor=10)

for storage_unit_rank in range(self.config.num_data_storage_units):

storage_node = SimpleStorageUnit.options(

placement_group=storage_placement_group,

placement_group_bundle_index=storage_unit_rank,

).remote(storage_unit_size=math.ceil(total_storage_size / self.config.num_data_storage_units))

self.data_system_storage_units[storage_unit_rank] = storage_node

logger.info(f"TransferQueueStorageSimpleUnit #0 ~ #{storage_unit_rank} has been created.")

self.data_system_controller = TransferQueueController.remote()

logger.info("TransferQueueController has been created.")

self.data_system_controller_info = process_zmq_server_info(self.data_system_controller)

self.data_system_storage_unit_infos = process_zmq_server_info(self.data_system_storage_units)

tq_config = OmegaConf.create({}, flags={"allow_objects": True})

tq_config.controller_info = self.data_system_controller_info

tq_config.storage_unit_infos = self.data_system_storage_unit_infos

self.config = OmegaConf.merge(tq_config, self.config)

self.data_system_client = AsyncTransferQueueClient(

client_id="Trainer", controller_info=self.data_system_controller_info

)

self.data_system_client.initialize_storage_manager(manager_type="AsyncSimpleStorageManager", config=self.config)

return self.data_system_client

def run_bandwidth_test(self):

logger.info("Creating large batch for bandwidth test...")

start_create_data = time.time()

big_input_ids, total_data_size_gb = create_complex_test_case(

batch_size=self.config.global_batch_size, seq_length=self.config.seq_length, field_num=self.config.field_num

)

end_create_data = time.time()

logger.info(f"Data creation time: {end_create_data - start_create_data:.8f}s")

logger.info("Starting PUT operation...")

start_async_put = time.time()

asyncio.run(self.data_system_client.async_put(data=big_input_ids, partition_id="train_0"))

end_async_put = time.time()

put_time = end_async_put - start_async_put

put_throughput_gbps = (total_data_size_gb * 8) / put_time

logger.info(f"async_put cost time: {put_time:.8f}s")

logger.info(f"PUT Throughput: {put_throughput_gbps:.8f} Gb/s")

time.sleep(2)

logger.info("Starting GET_META operation...")

start_async_get_meta = time.time()

prompt_meta = asyncio.run(

self.data_system_client.async_get_meta(

data_fields=list(big_input_ids.keys()),

batch_size=big_input_ids.size(0),

partition_id="train_0",

task_name="generate_sequences",

)

end_async_get_meta = time.time()

logger.info(f"async_get_meta cost time: {end_async_get_meta - start_async_get_meta:.8f}s")

time.sleep(2)

logger.info("Starting GET_DATA operation...")

start_async_get_data = time.time()

asyncio.run(self.data_system_client.async_get_data(prompt_meta))

end_async_get_data = time.time()

get_time = end_async_get_data - start_async_get_data

get_throughput_gbps = (total_data_size_gb * 8) / get_time

logger.info(f"async_get_data cost time: {get_time:.8f}s")

logger.info(f"GET Throughput: {get_throughput_gbps:.8f} Gb/s")

mode_name = "TQ REMOTE" if self.remote_mode else "TQ NORMAL"

logger.info("=" * 60)

logger.info(f"{mode_name} BANDWIDTH TEST SUMMARY")

logger.info("=" * 60)

logger.info(f"Total Data Size: {total_data_size_gb:.6f} GB")

logger.info(f"PUT Time: {put_time:.8f}s")

logger.info(f"GET Time: {get_time:.8f}s")

logger.info(f"PUT Throughput: {put_throughput_gbps:.8f} Gb/s")

logger.info(f"GET Throughput: {get_throughput_gbps:.8f} Gb/s")

logger.info(f"Network Round-trip Throughput: {(total_data_size_gb * 16) / (put_time + get_time):.8f} Gb/s")

def main():

if len(sys.argv) < 2:

print("Usage: python performance_test.py <test_mode>")

print("Available test modes:")

print(" ray-obj-store - Ray Object Store bandwidth test")

print(" ray-remote - Ray Remote Actor bandwidth test")

print(" tq-normal - TQ Normal mode bandwidth test")

print(" tq-remote - TQ Remote mode bandwidth test")

return

test_mode = sys.argv[1]

if test_mode == "ray-obj-store":

logger.info("Starting Ray Object Store bandwidth test")

tester = RayBandwidthTester(config=dict_conf, test_mode="obj_store")

tester.run_bandwidth_test()

logger.info("Ray Object Store bandwidth test completed successfully!")

elif test_mode == "ray-remote":

logger.info("Starting Ray Remote Actor bandwidth test")

tester = RayBandwidthTester(config=dict_conf, test_mode="remote")

tester.run_bandwidth_test()

logger.info("Ray Remote Actor bandwidth test completed successfully!")

elif test_mode in ["tq-normal", "tq-remote"]:

remote_mode = test_mode == "tq-remote"

mode_name = "TQ Remote" if remote_mode else "TQ Normal"

logger.info(f"Starting {mode_name} bandwidth test")

tester = TQBandwidthTester(config=dict_conf, remote_mode=remote_mode)

tester.run_bandwidth_test()

logger.info(f"{mode_name} bandwidth test completed successfully!")

else:

print(f"Unknown test mode: {test_mode}")

print("Available test modes: ray-obj-store, ray-remote, tq-normal, tq-remote")

if __name__ == "__main__":

main()

TransferQueue/scripts/performance_test.py at dev · TransferQueue/TransferQueue · GitHub

Uh oh!

Expand file tree

Breadcrumbs

performance_test.py

Latest commit

History

Breadcrumbs

performance_test.py

File metadata and controls