refactor: Improve efficiency of .scan (and .load) if limit is set

[tomaarsen]

Hello!

Pull Request overview

• Refactor Datasets.scan to never request more samples than required.
• Implements a batch_size variable which is now set to self.DEFAULT_SCAN_SIZE. This parameter can easily be exposed in the future for batch_size support, as discussed previously in Slack etc.

Details

The core of the changes is that the URL is now built on-the-fly at every request, rather than just once. For each request, the limit for that request is computed using min(limit, batch_size), and limit is decremented, allowing it to represent a "remaining samples to fetch". When limit reaches 0, i.e. 0 more samples to fetch, we return from scan.

I've also added a check to ensure that the limit passed to scan can't negative.

Tests

For tests, I've created a test_scan_efficient_limiting function which verifies the new and improved behaviour. It contains two monkeypatches:

1. DEFAULT_SCAN_SIZE is set to 10. Because our dataset in this test only has 100 samples, we want to ensure that we can't just sample the entire dataset with one request.
2. The http_client.post method is monkeypatched to allow us to track the calls to the server.

We test the following scenarios:

• limit=23 -> 3 requests: for 10, 10 and 3 samples.
• limit=20 -> 2 requests: for 10, 10 samples.
• limit=6 -> 1 request: for 6 samples.

There's also a test to cover the new ValueError if limit < 0.

Effects

Consider the following script:

import argilla as rg
import time

def test_scan_records(fields):
    client = rg.active_client()

    data = client.datasets.scan(
        name="banking77-topics-setfit",
        projection=fields,
        limit=1,  # <- Note, just one sample
    )
    start_t = time.time()
    list(data)
    print(f"load time for 1 sample with fields={fields}: {time.time() - start_t:.4f}s")

test_scan_records(set())
test_scan_records({"text"})
test_scan_records({"tokens"})

On this PR, this outputs:

load time for 1 sample with fields=set(): 0.0774s
load time for 1 sample with fields={'text'}: 0.0646s
load time for 1 sample with fields={'tokens'}: 0.0669s

On the develop branch, this outputs:

load time for 1 sample with fields=set(): 0.1355s
load time for 1 sample with fields={'text'}: 0.1347s
load time for 1 sample with fields={'tokens'}: 0.1173s

This can be repeated for rg.load(..., limit=1), as that relies on .scan under the hood.

Note that this is the most extreme example of performance gains. The performance increases in almost all scenarios if limit is set, but the effects are generally not this big. Going from fetching 250 samples 8 times to fetching 250 samples 7 times and 173 once doesn't have as big of an impact.

---

Type of change

• Refactor (change restructuring the codebase without changing functionality)

How Has This Been Tested

See above.

Checklist

• I have merged the original branch into my forked branch
• I added relevant documentation
• follows the style guidelines of this project
• I did a self-review of my code
• I made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works

---

• Tom Aarsen

[tomaarsen]

I've taken it a little further than we originally discussed, i.e. I've also improved the performance if limit > DEFAULT_SCAN_SIZE. However, no new parameters (e.g. batch_size) are introduced yet, as I imagine there should be a discussion on how to introduce that throughout Argilla as a whole.

[codecov]

Codecov Report

Patch coverage: 98.20% and project coverage change: +0.14 🎉

Comparison is base (5a8bb28) 91.86% compared to head (e7c356f) 92.00%.

Additional details and impacted files

@@             Coverage Diff             @@
##           develop    #2434      +/-   ##
===========================================
+ Coverage    91.86%   92.00%   +0.14%     
===========================================
  Files          161      162       +1     
  Lines         7869     7896      +27     
===========================================
+ Hits          7229     7265      +36     
+ Misses         640      631       -9     

Flag	Coverage Δ
pytest	92.00% <98.20%> (+0.14%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

Impacted Files	Coverage Δ
src/argilla/__init__.py	86.66% <ø> (ø)
src/argilla/server/daos/backend/search/model.py	100.00% <ø> (ø)
src/argilla/server/daos/datasets.py	95.45% <ø> (-0.06%)	⬇️
...ver/services/tasks/text_classification/__init__.py	100.00% <ø> (ø)
...gilla/labeling/text_classification/label_errors.py	90.12% <83.33%> (-0.13%)	⬇️
src/argilla/client/datasets.py	74.72% <93.33%> (+0.09%)	⬆️
src/argilla/utils/dependency.py	96.15% <96.15%> (ø)
src/argilla/client/__init__.py	100.00% <100.00%> (ø)
src/argilla/client/apis/datasets.py	97.24% <100.00%> (+1.05%)	⬆️
src/argilla/client/client.py	87.13% <100.00%> (-0.42%)	⬇️
... and 18 more

Help us with your feedback. Take ten seconds to tell us how you rate us. Have a feature suggestion? Share it here.

☔ View full report at Codecov.
📢 Do you have feedback about the report comment? Let us know in this issue.

[frascuchon]

Good work, @tomaarsen !!!

[davidberenstein1957]

Hi, @frascuchon @tomaarsen, currently rg.log uses chunk_size. I am in favour of deprecating this variable and moving over to the more intuitive batch_size when reviewing this discussion. If you don't agree, we should also use chunk_size within this PR.

[tomaarsen]

I'm in favour of batch_size over chunk_size. I think (in a future PR?) we should expose batch_size in both functions.

[dvsrepo]

I'm in favor of batch_size too

[davidberenstein1957]

Yes, for me it can be the PR where/when/if we actually start using the batch_size within rg.load

[tomaarsen]

I think we're best off keeping that separate from this PR for now, but we should work on that soon. Do you want to raise an issue, @davidberenstein1957? (i.e. one issue with two parts: rename chunk_size to batch_size for rg.log and introduce batch_size for rg.load)

[frascuchon]

Yes, I prefer to keep PRs as simple as possible. We can tackle the batch_size in 2 separate PRs:

• Deprecate the current chunk_size in favour of batch_size for rg.load
• Add a batch_size argument to rg.load to load data (adding also this new argument to datasets.scan method

UPDATE.

It's the same thing that wrote @tomaarsen one hour ago 😄

[tomaarsen]

Great minds think alike ;)