JonathanColetti commited on
Commit
efc5130
·
verified ·
1 Parent(s): da34457

disable expandable segments: they abort the ZeroGPU worker

Browse files
Files changed (1) hide show
  1. app.py +7 -3
app.py CHANGED
@@ -21,9 +21,13 @@ the GitHub repo locally for free text.
21
 
22
  import os
23
 
24
- # Set before torch is imported anywhere: the decoder allocates and frees a
25
- # pixel-space buffer every block, which is the pattern expandable segments exist for.
26
- os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
 
 
 
 
27
 
28
  import queue
29
  import shutil
 
21
 
22
  import os
23
 
24
+ # Deliberately NOT `expandable_segments:True`. That is the usual fix for allocator
25
+ # trouble under transient spikes, but here it *causes* it: expandable segments grow
26
+ # through the CUDA VMM path, and the first growth inside a ZeroGPU worker aborts with
27
+ # NVML_SUCCESS == r INTERNAL ASSERT FAILED ... CUDACachingAllocator.cpp
28
+ # on an allocation of ~49 MB, while the same allocation succeeds with the default
29
+ # allocator. Set it explicitly so a platform default cannot turn it back on.
30
+ os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:False"
31
 
32
  import queue
33
  import shutil