Hello, Recently we created a metaflow extension for our internal use cases (It has a Pytorch dependency) but then when we try to bootstrap the environment, it starts consuming memory like crazy. Does anyone know why this is happening? We suspect that it is because of the PyTorch dependency but maybe there are other reasons. Thanks!

it’s hard for me to answer this without more information. What do you mean “it has a pytorch dependency”. When is it consuming memory?

the extension has code that uses Pytorch, so it is a dependency of the extension

so your extension is doing

import pytorch

somewhere?

yes it is

Just a bit of extra context, we are creating a metaflow extension to share code across our company (following a suggestion in this same Slack channel) The code we are sharing has multiple dependencies (e.g. pytorch, transformers, numpy, pandas). When we run metaflow with our extension the amount of resources consumed skyrockets. By monitoring the resources we noticed that there's a bunch of python processes being executed during the Virtual Environment Bootstrapping step (previous to the execution of the flow). We suspect that it might have to do with the aforementioned dependencies making the Metaflow environment too large during the Virtual Env Bootstrapping process.

Sorry for the delay. If you are loading massive libraries that actually do something when loading, I could see that happening. Metaflow can launch multiple processes and if each process is loading the extension you wrote which INCLUDES large stuff, it may increase resource usage. I would suggest, for large libraries, to delay import until use (ie: the extension, instead of always loading Pytorch, loads it when the function that needs it is called). That may help a bit.

We ended up coming up with the same solution! We took the large libraries and lazy loaded them and that fixed the issue 🙂

cool. Sorry for being late tot he party 😞

Don't worry! Thank you for your help! It was really insightful 🙂