Skip to content

Commit 4f20636

Browse files
piotr-topnotchpiotr-topnotch
committed
Added the AIStor plugin support to nixlbench
1 parent 7288fd8 commit 4f20636

File tree

3 files changed

+46
-10
lines changed

3 files changed

+46
-10
lines changed

‎benchmark/nixlbench/src/utils/utils.cpp‎

Lines changed: 29 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,7 @@ DEFINE_string(runtime_type, XFERBENCH_RT_ETCD, "Runtime type to use for communic
4343
DEFINE_string(worker_type, XFERBENCH_WORKER_NIXL, "Type of worker [nixl, nvshmem]");
4444
DEFINE_string(backend,
4545
XFERBENCH_BACKEND_UCX,
46-
"Name of communication backend [UCX, UCX_MO, GDS, POSIX, GPUNETIO, OBJ] \
46+
"Name of communication backend [UCX, UCX_MO, GDS, POSIX, GPUNETIO, OBJ, AIStor] \
4747
(only used with nixl worker)");
4848
DEFINE_string(initiator_seg_type, XFERBENCH_SEG_TYPE_DRAM, "Type of memory segment for initiator \
4949
[DRAM, VRAM]");
@@ -77,7 +77,7 @@ DEFINE_int32(num_target_dev, 1, "Number of device in target process");
7777
DEFINE_bool(enable_pt, false, "Enable Progress Thread (only used with nixl worker)");
7878
DEFINE_bool(enable_vmm, false, "Enable VMM memory allocation when DRAM is requested");
7979

80-
// Storage backend(GDS, POSIX, HF3FS, OBJ) options
80+
// Storage backend(GDS, POSIX, HF3FS, OBJ, AIStor) options
8181
DEFINE_string (filepath, "", "File path for storage operations");
8282
DEFINE_int32 (num_files, 1, "Number of files used by benchmark");
8383
DEFINE_bool (storage_enable_direct, false, "Enable direct I/O for storage operations");
@@ -102,7 +102,7 @@ DEFINE_string (posix_api_type,
102102
DEFINE_string(gpunetio_device_list, "0", "Comma-separated GPU CUDA device id to use for \
103103
communication (only used with nixl worker)");
104104

105-
// OBJ options - only used when backend is OBJ
105+
// OBJ options - only used when backend is OBJ or AIStor
106106
DEFINE_string(obj_access_key, "", "Access key for S3 backend");
107107
DEFINE_string(obj_secret_key, "", "Secret key for S3 backend");
108108
DEFINE_string(obj_session_token, "", "Session token for S3 backend");
@@ -236,6 +236,16 @@ xferBenchConfig::loadFromFlags() {
236236
return -1;
237237
}
238238
}
239+
240+
241+
// Load AIStor-specific configurations if backend is AIStor
242+
if (backend == XFERBENCH_BACKEND_AISTOR) {
243+
obj_access_key = FLAGS_obj_access_key;
244+
obj_secret_key = FLAGS_obj_secret_key;
245+
obj_bucket_name = FLAGS_obj_bucket_name;
246+
obj_region = FLAGS_obj_region;
247+
obj_endpoint_override = FLAGS_obj_endpoint_override;
248+
}
239249
}
240250

241251
initiator_seg_type = FLAGS_initiator_seg_type;
@@ -352,7 +362,7 @@ void xferBenchConfig::printConfig() {
352362
}
353363
printOption ("Worker type (--worker_type=[nixl,nvshmem])", worker_type);
354364
if (worker_type == XFERBENCH_WORKER_NIXL) {
355-
printOption("Backend (--backend=[UCX,UCX_MO,GDS,POSIX,OBJ])", backend);
365+
printOption("Backend (--backend=[UCX,UCX_MO,GDS,POSIX,OBJ,AIStor])", backend);
356366
printOption ("Enable pt (--enable_pt=[0,1])", std::to_string (enable_pt));
357367
printOption ("Device list (--device_list=dev1,dev2,...)", device_list);
358368
printOption ("Enable VMM (--enable_vmm=[0,1])", std::to_string (enable_vmm));
@@ -385,6 +395,17 @@ void xferBenchConfig::printConfig() {
385395
obj_req_checksum);
386396
}
387397

398+
399+
// Print AIStor options if backend is AIStor
400+
if (backend == XFERBENCH_BACKEND_AISTOR) {
401+
printOption("AIStor access key (--obj_access_key=key)", obj_access_key);
402+
printOption("AIStor secret key (--obj_secret_key=key)", obj_secret_key);
403+
printOption("AIStor bucket name (--obj_bucket_name=nixlbench-bucket)", obj_bucket_name);
404+
printOption("AIStor region (--obj_region=region)", obj_region);
405+
printOption("AIStor endpoint override (--obj_endpoint_override=endpoint)",
406+
obj_endpoint_override);
407+
}
408+
388409
if (xferBenchConfig::isStorageBackend()) {
389410
printOption ("filepath (--filepath=path)", filepath);
390411
printOption ("Number of files (--num_files=N)", std::to_string (num_files));
@@ -452,8 +473,10 @@ xferBenchConfig::isStorageBackend() {
452473
return (XFERBENCH_BACKEND_GDS == xferBenchConfig::backend ||
453474
XFERBENCH_BACKEND_HF3FS == xferBenchConfig::backend ||
454475
XFERBENCH_BACKEND_POSIX == xferBenchConfig::backend ||
455-
XFERBENCH_BACKEND_OBJ == xferBenchConfig::backend);
476+
XFERBENCH_BACKEND_OBJ == xferBenchConfig::backend ||
477+
XFERBENCH_BACKEND_AISTOR == xferBenchConfig::backend);
456478
}
479+
457480
/**********
458481
* xferBench Utils
459482
**********/
@@ -516,7 +539,7 @@ void xferBenchUtils::checkConsistency(std::vector<std::vector<xferBenchIOV>> &io
516539
} else if (xferBenchConfig::op_type == XFERBENCH_OP_WRITE) {
517540
addr = calloc(1, len);
518541
is_allocated = true;
519-
if (xferBenchConfig::backend == XFERBENCH_BACKEND_OBJ) {
542+
if ((xferBenchConfig::backend == XFERBENCH_BACKEND_OBJ) || (xferBenchConfig::backend == XFERBENCH_BACKEND_AISTOR)) {
520543
if (!getObjS3(iov.metaInfo)) {
521544
std::cerr << "Failed to get S3 object: " << iov.metaInfo << std::endl;
522545
exit(EXIT_FAILURE);

‎benchmark/nixlbench/src/utils/utils.h‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -72,6 +72,7 @@
7272
#define XFERBENCH_BACKEND_MOONCAKE "Mooncake"
7373
#define XFERBENCH_BACKEND_HF3FS "HF3FS"
7474
#define XFERBENCH_BACKEND_OBJ "OBJ"
75+
#define XFERBENCH_BACKEND_AISTOR "AIStor"
7576

7677
// POSIX API types
7778
#define XFERBENCH_POSIX_API_AIO "AIO"

‎benchmark/nixlbench/src/worker/nixl/nixl_worker.cpp‎

Lines changed: 16 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -177,6 +177,18 @@ xferBenchNixlWorker::xferBenchNixlWorker(int *argc, char ***argv, std::vector<st
177177
}
178178

179179
std::cout << "OBJ backend" << std::endl;
180+
} else if (0 == xferBenchConfig::backend.compare(XFERBENCH_BACKEND_AISTOR)) {
181+
// Using default param values for AIStor backend
182+
backend_params["access_key"] = xferBenchConfig::obj_access_key;
183+
backend_params["secret_key"] = xferBenchConfig::obj_secret_key;
184+
backend_params["bucket"] = xferBenchConfig::obj_bucket_name;
185+
backend_params["region"] = xferBenchConfig::obj_region;
186+
187+
if (xferBenchConfig::obj_endpoint_override != "") {
188+
backend_params["url"] = xferBenchConfig::obj_endpoint_override;
189+
}
190+
191+
std::cout << "AIStor backend" << std::endl;
180192
} else {
181193
std::cerr << "Unsupported backend: " << xferBenchConfig::backend << std::endl;
182194
exit(EXIT_FAILURE);
@@ -538,7 +550,7 @@ xferBenchNixlWorker::allocateMemory(int num_lists) {
538550

539551
opt_args.backends.push_back(backend_engine);
540552

541-
if (xferBenchConfig::backend == XFERBENCH_BACKEND_OBJ) {
553+
if ((xferBenchConfig::backend == XFERBENCH_BACKEND_OBJ) || (xferBenchConfig::backend == XFERBENCH_BACKEND_AISTOR)) {
542554
struct timeval tv;
543555
gettimeofday(&tv, nullptr);
544556
uint64_t timestamp = tv.tv_sec * 1000000ULL + tv.tv_usec;
@@ -656,7 +668,7 @@ xferBenchNixlWorker::deallocateMemory(std::vector<std::vector<xferBenchIOV>> &io
656668
CHECK_NIXL_ERROR(agent->deregisterMem(desc_list, &opt_args), "deregisterMem failed");
657669
}
658670

659-
if (xferBenchConfig::backend == XFERBENCH_BACKEND_OBJ) {
671+
if ((xferBenchConfig::backend == XFERBENCH_BACKEND_OBJ) || (xferBenchConfig::backend == XFERBENCH_BACKEND_AISTOR)) {
660672
for (auto &iov_list : remote_iovs) {
661673
for (auto &iov : iov_list) {
662674
cleanupBasicDescObj(iov);
@@ -740,7 +752,7 @@ xferBenchNixlWorker::exchangeIOV(const std::vector<std::vector<xferBenchIOV>> &l
740752
std::vector<xferBenchIOV> remote_iov_list;
741753
for (auto &iov : iov_list) {
742754
std::optional<xferBenchIOV> basic_desc;
743-
if (XFERBENCH_BACKEND_OBJ == xferBenchConfig::backend) {
755+
if ((XFERBENCH_BACKEND_OBJ == xferBenchConfig::backend) || (XFERBENCH_BACKEND_AISTOR == xferBenchConfig::backend)) {
744756
basic_desc = initBasicDescObj(iov.len, iov.devId, iov.metaInfo);
745757
} else {
746758
basic_desc = initBasicDescFile(iov.len, remote_fds[0], iov.devId);
@@ -823,7 +835,7 @@ execTransfer(nixlAgent *agent,
823835
nixl_xfer_dlist_t local_desc(GET_SEG_TYPE(true));
824836
nixl_xfer_dlist_t remote_desc(GET_SEG_TYPE(false));
825837

826-
if (XFERBENCH_BACKEND_OBJ == xferBenchConfig::backend) {
838+
if ((XFERBENCH_BACKEND_OBJ == xferBenchConfig::backend) || (XFERBENCH_BACKEND_AISTOR == xferBenchConfig::backend)) {
827839
remote_desc = nixl_xfer_dlist_t(OBJ_SEG);
828840
} else if (xferBenchConfig::isStorageBackend()) {
829841
remote_desc = nixl_xfer_dlist_t(FILE_SEG);

0 commit comments

Comments
 (0)