:wave: My pipeline is processing 200k documents wi...
# ask-metaflow
b
👋 My pipeline is processing 200k documents with LLM. I'm thinking to run the Metaflow flow code in a single GPU-supported machine. steps: 1) fetch document data; 2) do information extraction with LLM; 3) save results to a database Consider I have to do processing in batches of 1000 documents: fetch 1000 documents, do information extraction with LLM all 1000 at once (batched LLM calls), save to DB, and then repeat with the next batch of 1000 documents. Q: How to implement such batching? 1) Do I have a flow that processes a single batch and I just call that flow for every batch? Feels like I'm missing something; 2) I could have a batching in
start
and then use
foreach
for every batch. However, how do I then ensure that the information extraction step works for only 1 batch at a time? Do I use internal semaphore?
1
e
I believe one solution may be in using
foreach
together with safeguards like
--max-workers
(see https://docs.metaflow.org/scaling/remote-tasks/introduction#safeguard-flags). This would provide you some flexibility in case you intend to scale-out in the future. Although you'd have to evaluate the potential overhead of running each batch one at a time in a separate worker, as compared with sequentially processing the batches one at a time from the same worker.
👍 1
b
@elegant-lamp-40213 thank you, I've tried
--max-workers=1
but it seems by default Metaflow wants to end one step completely before proceeding with the next. Do you think there is any way to enforce task execution like this: processing 0 -> saving 0 -> processing 1 -> saving 1 -> ... I do understand worst case I could merge the two steps into one
f
is there a reason that the execution ordering across each of the parallel steps (each branch of the foreach fanout) needs to be done in a certain way? If you think about the DAG structure, that would effectively be adding more dependencies into how those task executions can happen within each of those parallel branches of the DAG -- metaflow will make sure that the processing/saving happens in that order (e.g. within the scope of that batch). In terms of how those steps are actually scheduled/executed, I believe that may vary depending on the backend you're using the only time the scheduler will "end one step completely before proceeding with the next" is when it's required by the DAG, e.g. the
join
step requires all the previous parallel steps to have finished. In real-world use cases on Step Functions or Argo, those orchestrators will schedule the
save
after the
process
to be executed. If you think of a basic first-in-first-out job queue, you can see how this can result in your screenshot. There's of course a ton of ways those job queues/orchestrators can be configured, and it also depends on the compute requirements of the steps (e.g. some steps are CPU only, others might require GPUs) and if each of those tasks can be executed immediately or if scaling up new nodes is required before placement, etc etc as for how to do the batching itself, that really depends on the specifics of your problem – e.g. how large the documents are will likely inform how many you process at once, if a single "failed" batch should cause the whole "set of batches" to fail (if a step fails, the whole flow will fail), and so on
b
@fresh-laptop-72652 the
join
step is added because Metaflow flow validator complains if there is no join step before
end
. I'm not allowed to remove this, even though I don't need it. The reason I want execution like this: processing 0 -> saving 0 -> processing 1 -> saving 1 -> ... is because of memory constraints. Processing each batch of documents will produce a result object and it's just a risk for out of memory if I'm processing >200k documents with a pipeline. So I want to finish one batch and begin other batch. In my case, I only have 1 GPU machine (which will execute other short steps too) and it can only process 1000 documents at once (1 batch). I was starting my test like this:
python preprocessing/tests/test_metaflow_ratelimit.py run --max-workers=1
analogous use case would be a step that has API rate limit. For example, if API doesn't support concurrency, must make only one call at a time. So metaflow step needs to be configured to not start same step concurrently. Is there support for such a thing in metaflow?
e
Couldn't you rely on creating a data artifact at each step? This way you have some backup and can still chain all processing and then saving. And it would not impact memory (at least RAM or VRAM) as artifacts are serialized on disk. Otherwise, as you mentioned, if you want processing and saving to be tightly coupled, you probably want the whole thing to be a single step.
b
@elegant-lamp-40213 okay but do consider another use case: API rate limit. Consider such pipeline with 1 million items: 1. fetch from API (no concurrency allowed, 1 request per second) 2. do processing (concurrency allowed) 3. save to db (concurrency allowed) it'll take so much longer if first we have to complete step #1 a million times and only then begin the next steps.
e
It sounds like it would be beneficial to rethink what is for you an atomic workflow. The definition of a workflow is that it will execute one step at a time, you cannot change that without altering the contract with users. What you're looking for, it sounds like, is orchestration of distinct runs of a workflow, one per batch.
You may be able to further optimize runtime by erasing the latency of rate limiting steps. E.g. if you leverage distinct flows for processing and saving, and each run going through processing would trigger a saving flow when it's done, so saving0 would run in parallel with processing1. This may work, although only if you do not need to monitor the status of "saving" from the flow "processing". And also, this may take a queuing system to distribute batches to processing, only one at a time.
b
@elegant-lamp-40213 so you think I should have like a sub-flow inside a flow? The sub-flow processes a particular batch, the outter flow splits the data into batches and initiates the sub-flow?
e
Does this sound like this architecture could work for you? Please note that I am currently biased towards workflow of workflows because a project of mine. Likely other architectures could work, although this one sounds nice to me 😁
f
The reason I want execution like this:
processing 0 -> saving 0 -> processing 1 -> saving 1 -> ...
is because of memory constraints. Processing each batch of documents will produce a result object and it's just a risk for out of memory if I'm processing >200k documents with a pipeline. So I want to finish one batch and begin other batch.
In my case, I only have 1 GPU machine (which will execute other short steps too) and it can only process 1000 documents at once (1 batch).
you can control the CPU/RAM/GPU resources allocated at the step level, so that's easy to line up with your desired batching strategy you can configure your backend compute environments so nodes with GPUs only process steps that request GPUs (e.g. create separate Batch job queues if deploying Step Functions, or similar taints/tolerations/etc in the k8s/argo stack) -- that's what I meant by there being some scheduler considerations that might vary depending how you've configured your deployment you can absolutely have a single GPU machine that will execute the
process
tasks, and then CPU machines that will handle the
save
tasks as they become available. The subsequent
save
tasks (within each of the parallel branches) will not wait until all
process
tasks have finished. The reason you're seeing the execution ordering in the screenshot is not due to the DAG/parallel fanout, but is due to how the toy example is constructed whether or not you want to have a single flow, with a single step that iterates through the batches to rate limit, or a single flow with a parallel foreach fanout to process in parallel, or separate flow executions per batch really just depends on what you want and the specifics of your problem. the main considerations that come to my mind for that are that changing the number of parallel steps or flow executions doesn't change the specifics of your compute environment – if you only have a single GPU node, you aren't addressing bottlenecks through those changes in the DAG structure. Similarly if the API you're hitting has tight limits, then it may not make sense to try to increase throughput with parallel requests on your end – really just depends how/why those limits are in place to begin with. The larger question I think is around what you are considering a "unit of work" in the case of writes to the DB and when addressing failures -- e.g. if a run fails and you need to retry it, does it need to retry only a single batch or all the batches. Answering that question can help inform how you structure the flow and what makes sense to persist as each step's artifacts
b
@fresh-laptop-72652 if a single batch fails, restart should be done only on that batch because the batches are independent in my use case. Here is the test code (https://pastebin.com/FAKw6jFr), if you'd suggest how to make the execution order like I suggested. I'm experimenting with a single GPU and CPU machine but in future might expect to have a few GPU machines, single CPU machine.
f
I'm still not quite following what ordering constraints you have -- but if the goal is to ensure that within each
foreach
split the process/save happen in the correct order, that is guaranteed to happen by the structure of the DAG at the "workflow" level, the steps will be executed as they become available (and depending on some job orchestrator specifics) as an example of having some of the
save
happen in specific branches of the parallel fanout, while
process
are still going on in other branches of the fanout, give it a go yourself! check out the outerbounds sandbox to take it for a spin https://outerbounds.com/sandbox/ I copied your code into a
test.py
then I deployed the flow to a proper orchestrator (Argo in this case) via
Copy code
python test.py argo-workflows create
and then triggered an execution via
Copy code
python test.py argo-workflows trigger
and in the UI you can see the run, along with the ordering of each of those steps