Skip to content

Commit aa3c018

Browse files
committed
add env var to enalbe use of original mem info
1 parent 1cd2c29 commit aa3c018

File tree

1 file changed

+9
-12
lines changed

1 file changed

+9
-12
lines changed

src/hook.c

Lines changed: 9 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -40,8 +40,8 @@
4040
#include "client.h"
4141
#include "utlist.h"
4242

43-
#define ENV_NVSHARE_ENABLE_SINGLE_OVERSUB "NVSHARE_ENABLE_SINGLE_OVERSUB"
44-
#define ENV_NVSHARE_HIDE_MEMINFO_RESERVE_MIB "NVSHARE_HIDE_MEMINFO_RESERVE_MIB"
43+
#define ENV_NVSHARE_ENABLE_SINGLE_OVERSUB "NVSHARE_ENABLE_SINGLE_OVERSUB"
44+
#define ENV_NVSHARE_USE_ORIGINAL_MEMINFO "NVSHARE_USE_ORIGINAL_MEMINFO"
4545

4646
#define MEMINFO_RESERVE_MIB 1536 /* MiB */
4747
#define KERN_SYNC_DURATION_BIG 10 /* seconds */
@@ -85,7 +85,7 @@ int pending_kernel_window = 1;
8585
pthread_mutex_t kcount_mutex;
8686

8787
int enable_single_oversub = 0;
88-
int hide_meminfo_reserve_mib = 0;
88+
int use_original_meminfo = 0;
8989
int nvml_ok = 1;
9090

9191
/* Representation of a CUDA memory allocation */
@@ -109,7 +109,6 @@ static void bootstrap_cuda(void)
109109
void *cuda_handle;
110110
void *nvml_handle;
111111

112-
113112
true_or_exit(pthread_mutex_init(&kcount_mutex, NULL) == 0);
114113

115114
nvml_handle = dlopen("libnvidia-ml.so.1", RTLD_LAZY);
@@ -332,10 +331,10 @@ static void initialize_libnvshare(void)
332331
log_warn("Enabling GPU memory oversubscription for this"
333332
" application");
334333
}
335-
value = getenv(ENV_NVSHARE_HIDE_MEMINFO_RESERVE_MIB);
334+
value = getenv(ENV_NVSHARE_USE_ORIGINAL_MEMINFO);
336335
if (value != NULL) {
337-
hide_meminfo_reserve_mib = 1;
338-
log_warn("Hiding %d MiB from CUDA memory info", MEMINFO_RESERVE_MIB);
336+
use_original_meminfo = 1;
337+
log_warn("Using original mem info without removing %d MiB", MEMINFO_RESERVE_MIB);
339338
}
340339

341340
bootstrap_cuda();
@@ -786,11 +785,9 @@ CUresult cuMemGetInfo(size_t *free, size_t *total)
786785
* To avoid internal thrashing, we empirically choose a sane value for
787786
* MEMINFO_RESERVE_MIB.
788787
*/
789-
reserve_mib = (MEMINFO_RESERVE_MIB) MiB;
790-
*free = *total - (size_t) reserve_mib;
791-
792-
if (hide_meminfo_reserve_mib == 1) {
793-
*total -= reserve_mib;
788+
if (use_original_meminfo == 0) {
789+
reserve_mib = (MEMINFO_RESERVE_MIB) MiB;
790+
*free = *total - (size_t) reserve_mib;
794791
}
795792

796793
log_debug("nvshare's cuMemGetInfo returning free=%.2f MiB,"

0 commit comments

Comments
 (0)