gh-101178: refactor base64.b85encode to be memory friendly

[romuald]

Current description

Rewrote the base64._85encode method logic in C, by plugging in the binascii module (already taking care of the bae64 methods)

By using C and a single buffer, the memory use is reduced to a minimum, addressing the initial issue.

It also greatly improves performance as a bonus:

main

SMALL (11 bytes): 1575427 iterations (1.27 µs per call, 115.41 ns per byte) MEDIUM (200 bytes): 204909 iterations (9.76 µs per call, 48.80 ns per byte) BIG (5000 bytes): 8623 iterations (231.94 µs per call, 46.39 ns per byte) VERYBIG (500000 bytes): 81 iterations (24.69 ms per call, 49.38 ns per byte) 

branch

SMALL (11 bytes): 11230718 iterations (178.08 ns per call, 16.19 ns per byte) MEDIUM (200 bytes): 6004721 iterations (333.07 ns per call, 1.67 ns per byte) BIG (5000 bytes): 458005 iterations (4.37 µs per call, 873.35 ps per byte) VERYBIG (500000 bytes): 4772 iterations (419.11 µs per call, 838.22 ps per byte) 

Script used to test: https://gist.github.com/romuald/7aeba5f40693bb351da4abe62ad7321d

Previous description (python refactor)

not up to date with current PR

Refactor code to make use of generators instead of allocating 2 potentially huge lists for large datasets

Memory gain only measured using macOS and a 5Mb input.

Using main:

Before encoding Physical footprint: 16.3M Physical footprint (peak): 21.3M After encoding Physical footprint: 45.0M Physical footprint (peak): 244.1M 

With refactor:

Before encoding Physical footprint: 14.6M Physical footprint (peak): 19.6M After encoding Physical footprint: 28.5M Physical footprint (peak): 34.4M 

The execution time is more than doubled, which may not be acceptable. However the memory used is reduced by more than 90%
edit: changed the algorithm to be more efficient, the performance decrease now seems to be negligible

I also have no idea how (and if) I should test this

Here is the script I've used to measure the execution time, the memdebug can probably be adapted to read /proc/{pid} on Linux
edit: updated to work on Linux too

importosimportsysimportrandomimporthashlibimportplatformimportsubprocessfromtimeimporttimefrombase64importb85encodedefmemdebug(): ifplatform.system=="Darwin": ifnotos.environ.get("MallocStackLogging"): returnres=subprocess.check_output(["malloc_history", str(os.getpid()), "-highWaterMark", "-allBySize"]) forlineinres.splitlines(): ifline.startswith(b"Physical"): print(line.decode()) elifplatform.system() =="Linux": withopen(f"/proc/{os.getpid()}/status") asreader: forlineinreader: ifline.startswith("VmPeak:"): print(line, end="") defmain(): # use a stable inputrnd=random.Random() rnd.seed(42) data=rnd.randbytes(5_000_000) memdebug() start=time() importpdbtry: res=b85encode(data) exceptException: # pdb.post_mortem()raiseend=time() memdebug() print("Data length:", len(data)) print("Output length:", len(res)) print(f"Decode time: {end-start:.3f}s") h=hashlib.md5(res).hexdigest() print("Hashed result", h) asserth=="ad97e45ba085865e70f7aa05c9a31388"if__name__=='__main__': main()

• Issue: base64.b85encode uses significant amount of RAM #101178

[romuald]

I've changed the algorithm to use a dedicated generator, using unpack("!512I") unstead of 128 unpack("!I") seems to be far more efficient. The issue is now that this may not be as readable due to the additional generator

[arhadthedev]

cc @sobolevn as a commenter in the issue, @pitrou an an author of the original _85encode.

[romuald]

Tested on a Linux VM the new code is actually faster with the 5Mb dataset 🤔

Memory gain is as follow:

branch: VmPeak: 31532 kB -> 45344 kB
main: VmPeak: 33608 kB -> 253752 kB

[romuald]

@pitrou any chance you could spare a little time to review this?

[pitrou]

Could you post timeit numbers for both small and large inputs? (please be sure to compile in non-debug mode)

[pitrou]

This can be simplified to:

Suggested change

	forcinunpack512(b[offset:offset+512]):
	yieldc
	yieldfromunpack512(b[offset:offset+512])

[romuald]

Changed (I'm not yet used to yield from ^^)

[pitrou]

I'm not sure md5 is always available. WDYT @gpshead ?

[romuald]

Good catch, forgot about this
According to https://docs.python.org/3/library/hashlib.html#hashlib.algorithms_guaranteed md5 may not be present, I'll switch to a sha1

[romuald]

@pitrou here is a timeit script / results on a Macbook M1 (I don't have access to a Linux version right now)

importtimeitimportrandomfrombase64importb85encodeREPEAT=5COUNT=10_000SMALL_INPUT: bytes=b"hello world"MEDIUM_INPUT: bytes# 200 bytesBIG_INPUT: bytes# 5000 bytesSMALL_COUNT=500_000MEDIUM_COUNT=100_000BIG_COUNT=20_000definit(): globalMEDIUM_INPUT, BIG_INPUTrnd=random.Random() rnd.seed(42) MEDIUM_INPUT=rnd.randbytes(200) BIG_INPUT=rnd.randbytes(5000) defmain(): init() fornamein"SMALL", "MEDIUM", "BIG": timer=timeit.Timer(f"b85encode({name}_INPUT)", globals=globals()) count=globals()[f"{name}_COUNT"] values=timer.repeat(REPEAT, count) values=", ".join("%.3fs"%xforxinvalues) print(f"Timeit {name} ({count} iterations: {values}") if__name__=='__main__': main()

Results:

main branch Timeit SMALL (500000 iterations: 0.617s, 0.607s, 0.605s, 0.605s, 0.604s Timeit MEDIUM (100000 iterations: 1.000s, 0.999s, 0.999s, 0.999s, 0.999s Timeit BIG (20000 iterations: 4.789s, 4.788s, 4.794s, 4.800s, 4.782s gh-101178-b58encode-memuse branch Timeit SMALL (500000 iterations: 1.193s, 1.186s, 1.184s, 1.190s, 1.174s Timeit MEDIUM (100000 iterations: 1.748s, 1.701s, 1.701s, 1.701s, 1.700s Timeit BIG (20000 iterations: 5.705s, 5.675s, 5.673s, 5.668s, 5.672s 

[pitrou]

The performance decrease is a bit unfortunate. Instead of defining a separate _85buffer_iter_words generator, perhaps we can look for a hybrid approach, something like (just a sketch):

def_85encode(b, chars, chars2, pad=False, foldnuls=False, foldspaces=False): # Helper function for a85encode and b85encodeifnotisinstance(b, bytes_types): # TODO can this be `memoryview(b).cast('B')` instead?b=memoryview(b).tobytes() defencode_words(words): # Encode a sequence of 32-bit words, excluding paddingchunks= [b'z'iffoldnulsandnotwordelseb'y'iffoldspacesandword==0x20202020else (chars2[word//614125] +chars2[word//85%7225] +chars[word%85]) forwordinwords] returnb''.join(chunks) n1=len(b) //512# number of 512 bytes unpackn2= (len(b) -n1*512) //4# number of 4 bytes unpackpadding= (-len(b)) %4unpack512=struct.Struct("!128I").unpackunpack4=struct.Struct("!I").unpackoffset=0blocks= [] for_inrange(n1): blocks.append(encode_words(unpack512(b[offset:offset+512]))) offset+=512for_inrange(n2): blocks.append(encode_words(unpack4(b[offset:offset+4]))) offset+=4ifpadding: # TODO deal with last bytes and padding...returnb''.join(blocks)

[romuald]

Performance regression also on Linux amd64:

Timeit SMALL (500000 iterations: 0.874s, 0.880s, 0.913s, 0.868s, 0.869s Timeit MEDIUM (100000 iterations: 1.208s, 1.201s, 1.211s, 1.211s, 1.212s Timeit BIG (20000 iterations: 5.753s, 5.737s, 5.736s, 5.773s, 5.954s Timeit SMALL (500000 iterations: 1.581s, 1.600s, 1.596s, 1.583s, 1.542s Timeit MEDIUM (100000 iterations: 2.200s, 2.243s, 2.293s, 2.303s, 2.340s Timeit BIG (20000 iterations: 7.443s, 7.478s, 7.678s, 7.502s, 7.457s 

I find it a bit strange because my initial test a few months back the execution was slower on macOS but slightly faster on Linux. I attribued that to the fact that the memory allocation was clostly enough to be a factor

[romuald]

@pitrou sorry for the huge lag

I gave up after spending a lot of time trying to find a way to have be both CPU and memory friendly for small and large dataset

Until last week when I realized that I could simply rewrite the function in C (which I have done), to have both performance and memory improvements

My question is, should I create a new PR or push -f on this one?

[sobolevn]

You surely can do both :)
But I prefer to re-use PRs.

[romuald]

@pitrou / @sobolevn I rewrote this PR from scratch to use the a C implementation instead of a python one

Note that I do not consider myself a seasoned C developer so the implementation may be lacking.

I've tried to maximize compatibility with previous implementation even if the _85encode method is private, we could drop chars2 for example, and possibly change the type of XXchars to bytes

I've also added a test to check for a regression I found while testing the new implementation with random data

[picnixz]

It would feel weird not to have binascii.a2b_base85 so I would suggest keeping it private for now. Ideally, base64 should have its own C accelerator module but maybe it's an overkill.

[romuald]

Renamed to private

That was what delayed me initially because I had no idea how to add a module dedicated to base64, I found out that it did use the binascii one only last week

[picnixz]

By the way, I didn't look at the implementation in detail yet, but if you want to compare your implementation with a popular one, you can have a look at https://github.com/git/git/blob/master/base85.c#L40.

[romuald]

By the way, I didn't look at the implementation in detail yet, but if you want to compare your implementation with a popular one, you can have a look at https://github.com/git/git/blob/master/base85.c#L40.

Thanks, I didn't know where too look when I started (the base algorithm originally came from a LLM :/)

Inspiring from git's code, this could be used instead:

 size_t i = 0 ; size_t underflow = 0; // was overflow, but underflow may be a better name since `i` will not go over bin_len while (i < bin_len){// translate each 4 byte chunk to 32bit integer uint32_t value = 0; for (int cnt = 24; cnt >= 0; cnt -= 8){value |= bin_data[i] << cnt; if (++i == bin_len){// Number of bytes under the 4 bytes rounded value underflow = cnt / 8; break} } // ... } 

There is a 20% performance gain for large inputs (starting a 5kb) since the potential padding does not need to be computed for each chunk.

Shall I use this method instead?

[picnixz]

There is a 20% performance gain for large inputs [...] Shall I use this method instead?

Yes, since base85 can be used for encoding large inputs it's worth it I think.

[picnixz]

I would also credit the git implementation for this one as it's heavily based on it.

[romuald]

Credit added. I don't know if I should phrase it differently?

I've also added some other comments to try to explain the logic more