Bench/plot sample access times across data and across formats

[knighton]

Description of changes:

Benchmark and plot sample access times across different kinds of data (text, tokens) and different shard formats (MDS, JSONL, CSV).

On my system, this distribution reliably looks like one big spike (as expected) followed by seemingly a second much smaller hump, continuing to fade out very rapidly. I attribute this second rise to some kind of small linear-ish slowdown somewhere affecting the main spike plus the log-scaled bucketing and plotting: the buckets get wider in linear terms as you ascend to the right on the plot.

We probably want to break this into bench_text.py, bench_tokens.py, plot.py?

Issue #, if available:

Merge Checklist:

Put an x without space in the boxes that apply. If you are unsure about any checklist, please don't hesitate to ask. We are here to help! This is simply a reminder of what we are going to look for before merging your pull request.

General

• I have read the contributor guidelines
• This is a documentation change or typo fix. If so, skip the rest of this checklist.
• I certify that the changes I am introducing will be backward compatible, and I have discussed concerns about this, if any, with the MosaicML team.
• I have updated any necessary documentation, including README and API docs (if appropriate).

Tests

• I ran pre-commit on my change. (check out the pre-commit section of prerequisites)
• I have added tests that prove my fix is effective or that my feature works (if appropriate).
• I ran the tests locally to make sure it pass. (check out testing)
• I have added unit and/or integration tests as appropriate to ensure backward compatibility of the changes.

[knighton]

[knighton]

Bench: text
  Generate: 14.735 s                                                                                                    
  Write
    MDS: 1.434 s
    JSONL: 2.855 s
    CSV: 0.856 s
  Read
    Round 0/5
      MDS
        Sequential: 9.839 s
        Random: 9.947 s
      JSONL
        Sequential: 10.011 s
        Random: 11.368 s
      CSV
        Sequential: 9.247 s
        Random: 9.539 s
    Round 1/5
      MDS
        Sequential: 9.339 s
        Random: 10.460 s
      JSONL
        Sequential: 10.343 s
        Random: 10.587 s
      CSV
        Sequential: 9.153 s
        Random: 9.999 s
    Round 2/5
      MDS
        Sequential: 9.655 s
        Random: 10.446 s
      JSONL
        Sequential: 9.942 s
        Random: 11.336 s
      CSV
        Sequential: 9.692 s
        Random: 9.813 s
    Round 3/5
      MDS
        Sequential: 11.019 s
        Random: 10.896 s
      JSONL
        Sequential: 9.718 s
        Random: 10.111 s
      CSV
        Sequential: 9.510 s
        Random: 9.473 s
    Round 4/5
      MDS
        Sequential: 9.912 s
        Random: 10.570 s
      JSONL
        Sequential: 9.752 s
        Random: 10.547 s
      CSV
        Sequential: 9.358 s
        Random: 9.317 s
  Plot
    All sequential
      Min: 29.432 μs
      Max: 7385.989 μs
    All random
      Min: 29.278 μs
      Max: 7991.268 μs
    Bounds
      Min: 28.380 μs
      Max: 660.988 μs
      Min log10: 4.453
      Max log10: 5.820
  Stats
    MDS
      Sequential
        Min: 31.584 μs
        Mode: 38.482 μs
        Median: 39.999 μs
        Mean: 49.152 μs
        Max: 3569.917 μs
      Random
        Min: 31.374 μs
        Mode: 36.591 μs
        Median: 42.543 μs
        Mean: 51.387 μs
        Max: 6281.184 μs
    JSONL
      Sequential
        Min: 31.488 μs
        Mode: 37.525 μs
        Median: 40.048 μs
        Mean: 49.175 μs
        Max: 5863.511 μs
      Random
        Min: 31.848 μs
        Mode: 38.482 μs
        Median: 43.149 μs
        Mean: 52.957 μs
        Max: 7991.268 μs
    CSV
      Sequential
        Min: 29.432 μs
        Mode: 33.716 μs
        Median: 37.972 μs
        Mean: 46.408 μs
        Max: 7385.989 μs
      Random
        Min: 29.278 μs
        Mode: 34.794 μs
        Median: 38.652 μs
        Mean: 47.251 μs
        Max: 7349.424 μs
Bench: tokens
  Generate: 6.218 s                                                                                                     
  Write
    MDS: 7.915 s
  Read
    Round 0/5
      MDS
        Sequential: 13.407 s
        Random: 14.415 s
    Round 1/5
      MDS
        Sequential: 13.090 s
        Random: 14.467 s
    Round 2/5
      MDS
        Sequential: 13.215 s
        Random: 14.483 s
    Round 3/5
      MDS
        Sequential: 12.965 s
        Random: 14.117 s
    Round 4/5
      MDS
        Sequential: 13.742 s
        Random: 13.843 s
  Plot
    All sequential
      Min: 43.872 μs
      Max: 7572.538 μs
    All random
      Min: 44.692 μs
      Max: 5071.478 μs
    Bounds
      Min: 42.807 μs
      Max: 511.750 μs
      Min log10: 4.632
      Max log10: 5.709
  Stats
    MDS
      Sequential
        Min: 43.872 μs
        Mode: 50.792 μs
        Median: 55.832 μs
        Mean: 65.762 μs
        Max: 7572.538 μs
      Random
        Min: 44.692 μs
        Mode: 51.810 μs
        Median: 60.259 μs
        Mean: 70.285 μs
        Max: 5071.478 μs

real	8m1.801s
user	6m13.920s
sys	1m45.858s

[karan6181]

Overall looks great.

• Can you add a README under scripts/samples/ directory explaining how to run this script and how to interpret the output ? This would help internal and external user to run the script in future !!
• As you have mentioned in the description, it would be good idea to split this into bench_text.py, bench_tokens.py, plot.py files. Not a blocker for this PR, but definitely an improvement.