simulation speed with num_replicates versus looping

[grahamgower]

In #1899, @jeromekelleher pointed out just how much faster it is to pass num_replicates versus looping.

import time
import msprime

nreps = 5000
sim_kwargs = dict(
    samples=64,
    random_seed=1,
    record_provenance=False,
    recombination_rate=1e-9,
    sequence_length=1_000_000,
)

t0 = time.perf_counter()

for _ in range(nreps):
    msprime.sim_ancestry(**sim_kwargs)

t1 = time.perf_counter()

for _ in msprime.sim_ancestry(num_replicates=nreps, **sim_kwargs):
    pass

t2 = time.perf_counter()
print("loop", t1 - t0)
print("num_replicates", t2 - t1)

Output:

loop 15.983280896994984
num_replicates 2.054126285001985

Addmittedly, each of these simulations is very small which magnifies the overhead of the looping approach. But such small simulations are actually very common when doing simulation-based inference (e.g. discussion #1834), in which case we can't use num_replicates because the simulation parameters are different in each simulation. So I wonder, is there anything that could be done to reduce overhead here?

[jeromekelleher]

This is a drastic difference, I agree. Any chance you could try running this through the profiler and see where we're spending the time?

[grahamgower]

Yep. Here's the loop version using msprime git HEAD and results sorted by cumulative time.

Thu Nov 11 13:48:34 2021    loop.txt

         22633362 function calls (21538146 primitive calls) in 17.309 seconds

   Ordered by: cumulative time
   List reduced from 348 to 40 due to restriction <40>

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
        1    0.000    0.000   17.343   17.343 {built-in method builtins.exec}
        1    0.000    0.000   17.343   17.343 <string>:1(<module>)
        1    0.255    0.255   17.343   17.343 /home/grg/src/msprime-speed/loop.py:14(loop)
     5000    0.036    0.000   17.087    0.003 /home/grg/src/msprime/msprime/ancestry.py:985(sim_ancestry)
     5000    0.123    0.000   14.183    0.003 /home/grg/src/msprime/msprime/ancestry.py:784(_parse_sim_ancestry)
     5000    0.041    0.000    6.237    0.001 /home/grg/src/msprime/msprime/demography.py:1034(insert_populations)
     5002    0.041    0.000    5.743    0.001 /home/grg/.local/lib/python3.9/site-packages/tskit/metadata.py:594(__init__)
    10002    0.017    0.000    5.299    0.001 /home/grg/.local/lib/python3.9/site-packages/jsonschema/validators.py:291(check_schema)
     5000    3.982    0.001    5.189    0.001 /home/grg/src/msprime/msprime/ancestry.py:1208(__init__)
275052/15002    0.617    0.000    4.830    0.000 /home/grg/.local/lib/python3.9/site-packages/jsonschema/validators.py:296(iter_errors)
45008/15002    0.150    0.000    4.283    0.000 /home/grg/.local/lib/python3.9/site-packages/jsonschema/_validators.py:276(properties)
260050/50008    0.124    0.000    4.164    0.000 /home/grg/.local/lib/python3.9/site-packages/jsonschema/validators.py:343(descend)
90018/40008    0.131    0.000    3.727    0.000 /home/grg/.local/lib/python3.9/site-packages/jsonschema/_validators.py:252(ref)
     5000    0.036    0.000    2.864    0.001 /home/grg/src/msprime/msprime/ancestry.py:645(_wrap_replicates)
    10000    0.066    0.000    2.827    0.000 /home/grg/src/msprime/msprime/ancestry.py:1410(run_replicates)
     5001    0.013    0.000    2.524    0.001 /home/grg/.local/lib/python3.9/site-packages/tskit/metadata.py:133(__init__)
    15002    0.041    0.000    2.073    0.000 /home/grg/.local/lib/python3.9/site-packages/jsonschema/_validators.py:41(additionalProperties)
     5000    0.418    0.000    1.438    0.000 /home/grg/src/msprime/msprime/ancestry.py:1378(run)
    20004    0.070    0.000    1.360    0.000 /home/grg/.local/lib/python3.9/site-packages/jsonschema/_validators.py:319(anyOf)
    80016    0.144    0.000    1.248    0.000 /usr/lib/python3.9/urllib/parse.py:524(urljoin)
    10000    0.019    0.000    1.212    0.000 /home/grg/src/msprime/msprime/ancestry.py:674(_parse_rate_map)
    10000    0.018    0.000    1.170    0.000 /home/grg/src/msprime/msprime/intervals.py:445(uniform)
    10000    0.232    0.000    1.153    0.000 /home/grg/src/msprime/msprime/intervals.py:48(__init__)
    80016    0.089    0.000    1.109    0.000 /home/grg/.local/lib/python3.9/site-packages/jsonschema/validators.py:761(resolve)
    15003    0.022    0.000    0.986    0.000 /home/grg/.local/lib/python3.9/site-packages/jsonschema/validators.py:262(__init__)
    15003    0.023    0.000    0.964    0.000 /home/grg/.local/lib/python3.9/site-packages/jsonschema/validators.py:666(from_schema)
    15003    0.100    0.000    0.933    0.000 /home/grg/.local/lib/python3.9/site-packages/jsonschema/validators.py:636(__init__)
130000/90000    0.151    0.000    0.886    0.000 {built-in method numpy.core._multiarray_umath.implement_array_function}
   120024    0.098    0.000    0.878    0.000 /home/grg/.local/lib/python3.9/site-packages/jsonschema/validators.py:684(push_scope)
     5000    0.111    0.000    0.857    0.000 /home/grg/.local/lib/python3.9/site-packages/tskit/tables.py:2679(fromdict)
   305054    0.104    0.000    0.819    0.000 /home/grg/.local/lib/python3.9/site-packages/jsonschema/validators.py:355(is_type)
   295059    0.469    0.000    0.814    0.000 /usr/lib/python3.9/urllib/parse.py:437(urlsplit)
   160032    0.192    0.000    0.765    0.000 /usr/lib/python3.9/urllib/parse.py:372(urlparse)
     5000    0.040    0.000    0.756    0.000 /home/grg/src/msprime/msprime/ancestry.py:1350(_run_until)
260000/70000    0.280    0.000    0.740    0.000 /usr/lib/python3.9/copy.py:128(deepcopy)
   305054    0.158    0.000    0.715    0.000 /home/grg/.local/lib/python3.9/site-packages/jsonschema/_types.py:66(is_type)
     5000    0.053    0.000    0.700    0.000 /home/grg/.local/lib/python3.9/site-packages/tskit/tables.py:2427(asdict)
     5000    0.697    0.000    0.697    0.000 {function Simulator.run at 0x7f024ec580d0}
   105018    0.146    0.000    0.630    0.000 /home/grg/.local/lib/python3.9/site-packages/jsonschema/_validators.py:269(type)
   105021    0.061    0.000    0.597    0.000 /home/grg/.local/lib/python3.9/site-packages/jsonschema/_utils.py:14(normalize)

And the same thing, sorted by total time.

Thu Nov 11 13:48:34 2021    loop.txt

         22633362 function calls (21538146 primitive calls) in 17.309 seconds

   Ordered by: internal time
   List reduced from 348 to 40 due to restriction <40>

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
     5000    3.982    0.001    5.189    0.001 /home/grg/src/msprime/msprime/ancestry.py:1208(__init__)
     5000    0.697    0.000    0.697    0.000 {function Simulator.run at 0x7f024ec580d0}
275052/15002    0.617    0.000    4.830    0.000 /home/grg/.local/lib/python3.9/site-packages/jsonschema/validators.py:296(iter_errors)
   295059    0.469    0.000    0.814    0.000 /usr/lib/python3.9/urllib/parse.py:437(urlsplit)
     5000    0.418    0.000    1.438    0.000 /home/grg/src/msprime/msprime/ancestry.py:1378(run)
   830166    0.416    0.000    0.509    0.000 /usr/lib/python3.9/urllib/parse.py:114(_coerce_args)
1030074/800074    0.318    0.000    0.377    0.000 {built-in method builtins.getattr}
260000/70000    0.280    0.000    0.740    0.000 /usr/lib/python3.9/copy.py:128(deepcopy)
        1    0.255    0.255   17.343   17.343 /home/grg/src/msprime-speed/loop.py:14(loop)
  2670461    0.236    0.000    0.320    0.000 {built-in method builtins.isinstance}
    10000    0.232    0.000    1.153    0.000 /home/grg/src/msprime/msprime/intervals.py:48(__init__)
  2035285    0.219    0.000    0.219    0.000 {method 'get' of 'dict' objects}
   215043    0.195    0.000    0.380    0.000 /usr/lib/python3.9/urllib/parse.py:505(urlunsplit)
   160032    0.192    0.000    0.765    0.000 /usr/lib/python3.9/urllib/parse.py:372(urlparse)
455000/450000    0.191    0.000    0.210    0.000 /home/grg/.local/lib/python3.9/site-packages/tskit/tables.py:329(__setattr__)
   305054    0.179    0.000    0.240    0.000 /home/grg/.local/lib/python3.9/site-packages/pyrsistent/_pmap.py:53(_get_bucket)
   115000    0.162    0.000    0.506    0.000 /home/grg/.local/lib/python3.9/site-packages/tskit/tables.py:195(__init__)
85000/20000    0.161    0.000    0.546    0.000 /usr/lib/python3.9/dataclasses.py:1078(_asdict_inner)
   305054    0.158    0.000    0.715    0.000 /home/grg/.local/lib/python3.9/site-packages/jsonschema/_types.py:66(is_type)
130000/90000    0.151    0.000    0.886    0.000 {built-in method numpy.core._multiarray_umath.implement_array_function}
45008/15002    0.150    0.000    4.283    0.000 /home/grg/.local/lib/python3.9/site-packages/jsonschema/_validators.py:276(properties)
  1890378    0.146    0.000    0.146    0.000 {method 'replace' of 'str' objects}
   105018    0.146    0.000    0.630    0.000 /home/grg/.local/lib/python3.9/site-packages/jsonschema/_validators.py:269(type)
    80016    0.144    0.000    1.248    0.000 /usr/lib/python3.9/urllib/parse.py:524(urljoin)
     5000    0.141    0.000    0.141    0.000 {method 'finalise_tables' of '_msprime.Simulator' objects}
     5000    0.134    0.000    0.134    0.000 {method 'asdict' of 'LightweightTableCollection' objects}
   305054    0.132    0.000    0.372    0.000 /home/grg/.local/lib/python3.9/site-packages/pyrsistent/_pmap.py:59(_getitem)
    45000    0.132    0.000    0.132    0.000 {method 'reduce' of 'numpy.ufunc' objects}
90018/40008    0.131    0.000    3.727    0.000 /home/grg/.local/lib/python3.9/site-packages/jsonschema/_validators.py:252(ref)
260050/50008    0.124    0.000    4.164    0.000 /home/grg/.local/lib/python3.9/site-packages/jsonschema/validators.py:343(descend)
     5000    0.123    0.000   14.183    0.003 /home/grg/src/msprime/msprime/ancestry.py:784(_parse_sim_ancestry)
     5000    0.111    0.000    0.857    0.000 /home/grg/.local/lib/python3.9/site-packages/tskit/tables.py:2679(fromdict)
   305054    0.104    0.000    0.819    0.000 /home/grg/.local/lib/python3.9/site-packages/jsonschema/validators.py:355(is_type)
     5000    0.104    0.000    0.104    0.000 {method 'reset' of '_msprime.Simulator' objects}
    10000    0.103    0.000    0.274    0.000 /home/grg/.local/lib/python3.9/site-packages/numpy/lib/function_base.py:4567(insert)
   305054    0.101    0.000    0.473    0.000 /home/grg/.local/lib/python3.9/site-packages/pyrsistent/_pmap.py:69(__getitem__)
    15003    0.100    0.000    0.933    0.000 /home/grg/.local/lib/python3.9/site-packages/jsonschema/validators.py:636(__init__)
   120024    0.098    0.000    0.878    0.000 /home/grg/.local/lib/python3.9/site-packages/jsonschema/validators.py:684(push_scope)
    85000    0.095    0.000    0.110    0.000 /usr/lib/python3.9/copy.py:242(_keep_alive)
   105000    0.092    0.000    0.224    0.000 /home/grg/.local/lib/python3.9/site-packages/tskit/tables.py:503(__init__)

[jeromekelleher]

Excellent, thanks @grahamgower. Looks like the majority of the time is spent in _parse_sim_ancestry and in particular in validating metadata in the built table. Weird that we're calling urlib.parse so many times too.

Well, this is good news I guess - the bottlenecks are all in Python, so we should be able to trim quite a bit of time off.

[grahamgower]

Yeah, I wonder if transitioning tskit metadata validation to fastjsonschema might be an option. Thoughts @benjeffery?

[jeromekelleher]

Very interesting! It has conda-forge and pypi packages and has no dependencies, so I'd definitely be up for it.

[jeromekelleher]

Opened tskit issue: tskit-dev/tskit#1907

[jeromekelleher]

#1908 made some good inroads into this. A lot of time is still spend in _parse_sim_ancestry, and we could possibly do a bit better with our memory allocation defaults for the C code. I don't want to optimise to much in favour of the small simulation case, though,and we'd need to keep an eye on what happens to big sims too if we start tweaking memory allocation parameters. (That said, they're probably totally arbitrary at the moment!)