fosters commited on
Commit
b412294
·
verified ·
1 Parent(s): 3763e7a

fix: use sched_getaffinity for container-aware CPU count

Browse files
Files changed (1) hide show
  1. app.py +4 -1
app.py CHANGED
@@ -32,7 +32,10 @@ from music_detector import _get_ast_runtime, judge_chunk_files_batched
32
  # Constants
33
  # ---------------------------------------------------------------------------
34
 
35
- N_CPUS = os.cpu_count() or 2
 
 
 
36
  torch.set_num_threads(N_CPUS)
37
 
38
  BATCH_SIZE = 32
 
32
  # Constants
33
  # ---------------------------------------------------------------------------
34
 
35
+ try:
36
+ N_CPUS = len(os.sched_getaffinity(0)) # container-aware (Linux cgroup)
37
+ except AttributeError:
38
+ N_CPUS = os.cpu_count() or 2 # macOS / Windows fallback
39
  torch.set_num_threads(N_CPUS)
40
 
41
  BATCH_SIZE = 32