It’s just been over a week since I blogged last and a lot has been going on, not just in this replication effort. Zoe had a remote, prerecorded Zoom gig (which I’ve been helping out with) and tons of other things. Still haven’t heard about the promotion (though I think the latest they could tell me would be the end of June). I had some exercise energy early in the weak but that’s waned.
Oh, and I’ve been fighting with two things: measuring Python memory consumption and wrestling with Cython. There’s lots of good stuff to write about from those efforts…but it’s another thing and a tiring one. I need a smoother format than WordPress…maybe some Markdowny thingy.
First bit, other team members have made good progress on population density. We can load the appropriate Landscape and a rasterisation derived from ONS local authority based census data. So, getting closer.
Second bit, whoof, it’s super hard to measure the “deep” object memory weight of a Python object. I’ve tried 3 installed things, some manual stuff, and it’s not only not consistent between techniques, even within techniques, things are weird. Really weird.
I quickly found that our strings are almost certainly interned, so no big wins there. I tried the bit twiddling consolidation of booleans and saw memory consumption shrink at 10x the access time and a ton of ugly code. (Would want macros.) I quickly decided that the only productive way forward was Cython. Since a lot of our attribute values are essentially primitive, we could hope that they’d be captured as thin C types and popped in the struct Cython uses to represent extension object instance variables.That also would open the way for sub-word width representation with C speed bit operations. I found this cool feature of Cython where you can leave your Python file untouched and “just” add Cython declarations in a separate file…this seems perfect!
Well, “perfect” heh. Let’s say that I haven’t quite got it all working, esp with inheritance. (It’s a messy story for another day.) But I got it working well enough to see significant benefit in memory consumption. Here’s a guppy “relative” heap dump:
Partition of a set of 9000002 objects. Total size = 980889350 bytes.
Index Count % Size % Cumulative % Kind (class / dict of class)
0 1000000 11 408000000 42 408000000 42 dict of model.AgentsPlain.Person
1 1000000 11 240000000 24 648000000 66 dict (no owner)
2 1000000 11 104000000 11 752000000 77 list
3 3000001 33 84000028 9 836000028 85 int
4 1000000 11 64888890 7 900888918 92 str
5 1000000 11 56000000 6 956888918 98 model.AgentsPlain.Person
6 1000000 11 24000000 2 980888918 100 float
7 1 0 432 0 980889350 100 types.FrameType
Here I just created 1 million Person objects. (Remember for the UK we need 60 million or so.) To a first approximation, the total size is probably about right. So 1GB for a million Persons. Index 5 gives us the “base” overhead of an object, about 56bytes. Index 0 is the dictionary which holds all they instance variables, which weighs in at around 400bytes. Anything that has a count of 1,000,000 is probably part of our Person object. (I created 1 distinct string per object, just to see. This makes the total size a bit of an overstatement but not as much as you might thing as schools, workplaces, and social centres we each one string and there would normally be more. It’s only 64byes for a total of 0.06GB/million or 3.6GB for 60 million. It’s currently a rounding error.)
After a naive Cython conversion, we see:
Partition of a set of 4000002 objects. Total size = 576889350 bytes.
Index Count % Size % Cumulative % Kind (class / dict of class)
0 1000000 25 240000000 42 240000000 42 dict (no owner)
1 1000000 25 168000000 29 408000000 71 Agents.Person
2 1000000 25 104000000 18 512000000 89 list
3 1000000 25 64888890 11 576888890 100 str
4 1 0 432 0 576889322 100 types.FrameType
5 1 0 28 0 576889350 100 int
Only 60% or so the memory needed! That was totally worth the Cython learning curve.
Index 0 can be driven out just by a simple Python level refactor (that makes things a lot nicer) and perhaps the lists can too. If both of those work out (the dict is 100%, the list less so) then we get to around 0.16GB per million agents or 9.6GB for the UK population. This is with almost no changes to the substantive Python code. It looks the same. It runs the same. You can run it with our without the Cython version.
This is still about twice our overall model budget (if we’re going to run on certain sections of Manchester’s Condor setup) and we have some other stuff to represent. Not a lot, though and a lot of them are covered by our too many strs. 4GBs is a tight tight budget. 10GB for a model is low enough to run on my development laptop.
It does mean that the teams efforts at tessellating the map so we can work on parts isn’t so critical. I think it still can help in a variety of ways and for places like the US it will be critical. But assuming I’ve not messed up the measurement, we might be ok.
I do think it’s worth pressing on a bit with layout and alignment. I’m pretty sure all the fields of the Cython struct are 64bit and…we don’t need that for almost anything except population and even there, 32 bits gets us 4 billion. Unless we want to literally run 1 model with the entire world population, we have a lot of wasted space.
Just to hint, here’s the Cython declaration:
def class Person:
cdef public bint is_infected, acquired_infection, dead, is_immune, timetable_may_change, is_at_hospital
cdef public str social_center, workplace, school, location, symptoms
cdef public int id, hours_since_infection, incubation_length, infection_length, timetable_group
cdef public float age, p_infection, p_death
cdef public tuple timetable
cdef public dict hospital
cdef public int household
(That household declaration is a bit of an experiment.)
Here’s the corresponding struct:
struct __pyx_obj_6Agents_Person {
PyObject_HEAD
int is_infected;
int acquired_infection;
int dead;
int is_immune;
int timetable_may_change;
int is_at_hospital;
PyObject *social_center;
PyObject *workplace;
PyObject *school;
PyObject *location;
PyObject *symptoms;
int id;
int hours_since_infection;
int incubation_length;
int infection_length;
int timetable_group;
float age;
float p_infection;
float p_death;
PyObject *timetable;
PyObject *hospital;
int household;
};
So, it’s pretty clear that all the PyObjects except timetable and hospital are dispensable. social_center et al are just ids and could be ints with almost no loss of code clarity. You’d have to add a bit of pretty printing code but that’s it. The names are all generated by interpolating an incremented integer into a name like ‘social_center’.
Most of the ints have very short ranges, even beyond the booleans. We should be able to pack them quite tightly. I’m not sure if alignment is even worth all that much given how much memory we’d drop. I’m no C programmer, even at the novice level. So I’m going to have to research a bit on bit twiddling vs. packed and unaligned struct access.
Back of the envelope suggests that we could comfortably get down to 32 bytes a Person which would be around 2GB for all of the UK.
There is a super modest speed win at the moment at least for tests comfortably in memory. Just compiling Person doesn’t gain all that much and I suspect there’s not huge gains from Cythonising Person alone. But there are places where we iterate over the whole population doing simple stat gathering. That seems like an easy win if we can push the whole thing into C.
Ok, this went on a bit!
I have to get back to the spec extraction. I’m trying to keep the extraction ahead of the team and they are catching up!
But all in all, I’m feeling pretty confident that it will work out ok. I’m going to push, at our Monday meeting, to open the repo, though it won’t be “nice” yet. (I mean, it’s pretty good right now, but the master branch is the ad hoc model we threw together. There’s some but nowhere enough documentation. Etc.) It just bugs me. But maybe opening it early will bug others! I don’t think we can absorb a lot of contributions at the moment (though I don’t see people champing at the bit…I’ve gotten no feedback or help on the extraction.)
I have thoughts there, but that’s another post.
