[BUG] Segfault when batch size is larger than 255 on Atari environments

[imoneoi]

Describe the bug

Segfault when batch size is larger than 255 on Atari environments

MuJoCo environment seems to work well.

To Reproduce

Steps to reproduce the behavior.

import time

import envpool
import numpy as np

batch_size = 256  # set to 255 works

env = envpool.make_gym("Breakout-v5",
                        stack_num=1,

                        num_envs=batch_size * 2,
                        batch_size=batch_size,

                        use_inter_area_resize=False,

                        img_width=88,
                        img_height=88,

                        num_threads=0,
                        thread_affinity_offset=0)
action = np.array(
    [env.action_space.sample() for _ in range(batch_size)]
)

counter = 0

env.async_reset()

last_time = time.time()
while True:
    obs, rew, done, info = env.recv()

    env_id = info["env_id"]
    env.send(action, env_id)

    counter += batch_size
    if counter >= 100000:
        cur_time = time.time()
        print("TPS", counter / (cur_time - last_time))

        counter = 0
        last_time = cur_time

[1]    2959596 segmentation fault (core dumped)  python test_envpool.py

Expected behavior

Can run with large batch size, like 1024, 2048, etc.

System info

Describe the characteristic of your environment:

• Describe how the library was installed (pip, source, ...)
• Python version
• Versions of any other relevant libraries

import envpool, numpy, sys
print(envpool.__version__, numpy.__version__, sys.version, sys.platform)

0.6.1.post1 1.21.2 3.8.12 (default, Oct 12 2021, 13:49:34) 
[GCC 7.5.0] linux

Additional context

Set batch size to 1024 works / segfaults randomly

1024
TPS 49611.30131772514
TPS 57661.12695997062
TPS 52648.235412990536
TPS 52059.6945247295
[1]    2971074 segmentation fault (core dumped)  python test_envpool.py

Reason and Possible fixes

Checklist

• [x] I have checked that there is no similar issue in the repo (required)
• [x] I have read the documentation (required)
• [x] I have provided a minimal working example to reproduce the bug (required)

[Trinkle23897]

Could you please provide your system info? For example, the number of CPU cores and the type of CPU: cat /proc/cpuinfo | grep model | uniq -c.

I ran on DGX-A100 (256 cores) with >255 Atari envs and it works well.

[imoneoi]

AMD Ryzen Threadripper 3960X 24-Core Processor, with 128GB RAM, installed using pip install envpool

Architecture:                    x86_64
CPU op-mode(s):                  32-bit, 64-bit
Byte Order:                      Little Endian
Address sizes:                   43 bits physical, 48 bits virtual
CPU(s):                          48
On-line CPU(s) list:             0-47
Thread(s) per core:              2
Core(s) per socket:              24
Socket(s):                       1
NUMA node(s):                    1
Vendor ID:                       AuthenticAMD
CPU family:                      23
Model:                           49
Model name:                      AMD Ryzen Threadripper 3960X 24-Core Processor
Stepping:                        0
Frequency boost:                 enabled
CPU MHz:                         2200.000
CPU max MHz:                     4568.1641
CPU min MHz:                     2200.0000
BogoMIPS:                        7600.29
Virtualization:                  AMD-V
L1d cache:                       768 KiB
L1i cache:                       768 KiB
L2 cache:                        12 MiB
L3 cache:                        128 MiB
NUMA node0 CPU(s):               0-47
Vulnerability Itlb multihit:     Not affected
Vulnerability L1tf:              Not affected
Vulnerability Mds:               Not affected
Vulnerability Meltdown:          Not affected
Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl and seccomp
Vulnerability Spectre v1:        Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:        Mitigation; Retpolines, IBPB conditional, STIBP conditional, RSB filling
Vulnerability Srbds:             Not affected
Vulnerability Tsx async abort:   Not affected
Flags:                           fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fx
                                 sr sse sse2 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm constant_tsc rep_good nopl
                                  nonstop_tsc cpuid extd_apicid aperfmperf pni pclmulqdq monitor ssse3 fma cx16 sse4_1 
                                 sse4_2 movbe popcnt aes xsave avx f16c rdrand lahf_lm cmp_legacy svm extapic cr8_legac
                                 y abm sse4a misalignsse 3dnowprefetch osvw ibs skinit wdt tce topoext perfctr_core per
                                 fctr_nb bpext perfctr_llc mwaitx cpb cat_l3 cdp_l3 hw_pstate ssbd mba ibpb stibp vmmca
                                 ll fsgsbase bmi1 avx2 smep bmi2 cqm rdt_a rdseed adx smap clflushopt clwb sha_ni xsave
                                 opt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local clzero irp
                                 erf xsaveerptr rdpru wbnoinvd arat npt lbrv svm_lock nrip_save tsc_scale vmcb_clean fl
                                 ushbyasid decodeassists pausefilter pfthreshold avic v_vmsave_vmload vgif v_spec_ctrl 
                                 umip rdpid overflow_recov succor smca sme sev sev_es

Run the code above using batch_size=255

python test_envpool.py

255
TPS 60889.970751004286
TPS 60038.54218780955
TPS 59991.044900352266
TPS 59110.94207866427

Run the code above using batch_size=256

python test_envpool.py

256
[1]    3556131 segmentation fault (core dumped)  python test_envpool.py

[Trinkle23897]

Hi @imoneoi, could you please rerun the test with this wheel?

[imoneoi]

Thanks! It now works with >255 batch_size (tested with 1024 and 2560).

[imoneoi]

BTW, setting num_threads=0 seems to allocate 1024 threads with 1024 environments, resulting in ~10% fps drop due to context switching. I think default num_threads should be min(num_envs, cpu_count() - 1)?

[Trinkle23897]

But you can always set num_threads manually.