@@ -196,7 +196,7 @@ def bind_cores_for_best_perf():
196
196
197
197
if not is_numa_available ():
198
198
logger .error ("'numa' module not found" )
199
- raise ImportError ("'numa' module not found, install with 'pip install numa '" )
199
+ raise ImportError ("'numa' module not found, install with 'pip install py-libnuma '" )
200
200
import numa
201
201
202
202
local_size = get_int_from_env (
@@ -205,7 +205,7 @@ def bind_cores_for_best_perf():
205
205
rank_id = get_int_from_env (
206
206
["LOCAL_RANK" , "MPI_LOCALRANKID" , "OMPI_COMM_WORLD_LOCAL_RANK" , "MV2_COMM_WORLD_LOCAL_RANK" ], 0
207
207
)
208
- nodes = numa .get_max_node () + 1
208
+ nodes = numa .info . get_max_node () + 1
209
209
rank_per_node = math .ceil (local_size / nodes )
210
210
num_cpus_per_nodes = int (psutil .cpu_count (logical = False ) / nodes )
211
211
node_id = int (rank_id / rank_per_node )
@@ -216,17 +216,17 @@ def bind_cores_for_best_perf():
216
216
else :
217
217
num_cpus_per_rank = int (os .getenv ("OMP_NUM_THREADS" ))
218
218
logger .info (f"OMP_NUM_THREADS already set to { num_cpus_per_rank } " )
219
- if len (numa .get_membind ()) == nodes :
219
+ if len (numa .memory . get_membind_nodes ()) == nodes :
220
220
# if numa memory binding is not set, set it to the node where the rank is running
221
- numa .set_membind ([ node_id ] )
221
+ numa .memory . set_membind_nodes (( node_id ) )
222
222
223
223
torch .set_num_threads (num_cpus_per_rank )
224
224
225
- if len (numa .get_affinity (0 )) == psutil .cpu_count (logical = True ):
225
+ if len (numa .schedule . get_affinitive_cpus (0 )) == psutil .cpu_count (logical = True ):
226
226
# if numa affinity is unset (default value is set to all logical cores) set it to the physical cores assigned to the rank
227
227
cpu_start = num_cpus_per_rank * rank_offset_per_node
228
- numa .set_affinity (
228
+ numa .schedule . run_on_cpus (
229
229
0 ,
230
- list (numa .node_to_cpus (node_id )) [cpu_start : cpu_start + num_cpus_per_rank ],
230
+ * (numa .info . node_to_cpus (node_id )[cpu_start : cpu_start + num_cpus_per_rank ]) ,
231
231
)
232
- logger .info (f"affinity={ numa .get_affinity (0 )} , membind = { numa .get_membind ()} " )
232
+ logger .info (f"affinity={ numa .schedule . get_affinitive_cpus (0 )} , membind = { numa .memory . get_membind_nodes ()} " )
0 commit comments