← Back to blog

Scaling an ML algorithm with Spark as distributed compute

I had a machine learning algorithm that was already solved correctly. The logic lived in a Python module and produced good results. The open question was not whether it worked. It was how to run it over tens of millions of items without it taking days.

That distinction matters, and it is where the interesting engineering decision lived. The same correct algorithm can be trivial or nearly impossible to operate depending on how you choose to run it at scale.

The obvious path: expose it as an API

The default instinct was to wrap the algorithm in an API so a backend could call it to run. It is a clean interface, and for a handful of requests it works fine.

At tens of millions of items it falls apart. A request-and-response service turns a batch problem into an avalanche of calls, and around those calls you end up hand-building everything: autoscaling, load balancing, backpressure, timeouts, retries, idempotency, and orchestration to keep the whole thing moving. The service becomes the bottleneck, and you spend your time operating plumbing rather than running the work. At that volume, an API was the wrong shape for the problem.

The second idea: a Kafka event stream

The next proposal was to turn each item into an event and process the stream with workers that ran the algorithm. That is a real improvement over a synchronous API. It decouples producers from consumers and helps with throughput.

But it still leaves you owning and operating the consumers, the scaling, the retry and dead-letter handling, and the coordination. It is a streaming solution to what is fundamentally a large batch problem, and it carries the operational weight to match.

What I proposed: run the algorithm as a distributed Spark job

I work across software, data, and ML engineering, and looking at it from all three angles at once changed the framing for me. The algorithm did not need to become a service. It needed a distribution layer, and Spark already is one.

The plan was straightforward:

  • Package the Python algorithm as an SDK so it could be imported and called cleanly, rather than wrapped behind a network boundary.
  • Ship a virtual environment to the Spark executors so the exact Python interpreter and dependencies ran inside the cluster, identical to how they ran locally.
  • Run the algorithm inside the Spark executors, applying it across partitions of the input so tens of millions of items were processed in parallel across the whole cluster.

With that, “run this algorithm over tens of millions of items” became a single distributed batch job instead of a fleet of service calls.

Why Spark was the right call

The reason this worked so well is that everything the API approach would have forced me to build by hand, Spark provides out of the box:

  • Scaling is a matter of partitions and executors. Need more throughput, add more of the cluster.
  • Fault tolerance and retries are built in. If a task fails, Spark reruns it, so I did not write my own retry and recovery logic.
  • Parallelism and data locality come for free, so the work runs where the data is and spreads across every core available.

The result was the headline: the job went from something that would have taken days as a call-per-item service to something that finished in hours as a distributed batch run. Same algorithm, a completely different runtime, and far less operational machinery to own.

The takeaway

The lesson I keep from this one is to match the tool to the shape of the problem, not to the first interface that comes to mind. This was an embarrassingly parallel batch problem wearing the disguise of an API request. Once I saw it that way, Spark was the obvious home, and reframing it from “expose a service” to “distribute a job” is what turned days into hours. Sitting across software, data, and ML engineering is what let me see that the algorithm was fine and only the runtime needed to change.