Debugging a slow PyTorch GPU Inference Pipeline on Beam/Google Cloud Dataflow

We are attempting to use Google Cloud Dataflow to build a simple GPU-based classification pipeline that looks like this: Pub/Sub request comes in with link to a file on GCS → Read data from GCS → Chop up and batch data → Run inference in PyTorch.

Background

We deploy our pipeline on Dataflow with a custom Docker image adapted from the pytorch-minimal sample.

We ingest Pub/Sub messages and download data audio files from GCS using pathy, then chop audio into chunks for classification.

We've adapted Beam's relatively new RunInference function. Currently, there is no GPU support for RunInference on Dataflow (see open issue https://issues.apache.org/jira/browse/BEAM-13986). Upon building the Beam pipeline locally before deploy to Dataflow, the model initialization step doesn't recognize a CUDA environment and defaults to a CPU device for inference. This configuration gets propagated to the Dataflow execution environment that is properly GPU-enabled. So, we force a GPU device if requested without a CUDA device check. Other than that, the code is the same as the general RunInference code: A BatchElements operation followed by a ParDo that calls the model.

Problem

Everything is working-ish, but GPU inference is very slow – much slower than what we can time the same GPU-instance on processing batches on Google Cloud Compute Engine.

We're looking for advice on how to debug and speed up the pipeline. We suspect that the issue might have to do with threading as well as how Beam/Dataflow manages load across the pipeline stages. We kept running into CUDA OOM issues with multiple threads trying to access the GPU in the ParDo function. We launch our jobs with --num_workers=1 --experiment="use_runner_v2" --experiment="no_use_multiple_sdk_containers" to avoid multi-processing altogether. We saw that this 2021 beam summit talk on using Dataflow for local ML batch inference recommended going even further and just using a single worker thread --number_of_worker_harness_threads=1. However, we ideally don't want to do this: it's pretty common practice in ML pipelines like these to have multiple threads doing the I/O work of downloading data from the bucket and preparing batches so that the GPU never sits idle. Unfortunately, it seems that there is no way to tell beam to use a certain max number of threads per stage (?), so the best solution we could come up with is to protect the GPU with a Semaphore like so:

class _RunInferenceDoFn(beam.DoFn, Generic[ExampleT, PredictionT]):
    ...

    def _get_semaphore(self):
        def get_semaphore():
            logging.info('intializing semaphore...')
            return Semaphore(1)

        return self._shared_semaphore.acquire(get_semaphore)

    def setup(self):
        ...
        self._model = self._load_model()
        self._semaphore = self._get_semaphore()

    def process(self, batch, inference_args):
        ...
        logging.info('trying to acquire semaphore...')
        self._semaphore.acquire()
        logging.info('semaphore acquired')
        start_time = _to_microseconds(self._clock.time_ns())
        result_generator = self._model_handler.run_inference(
            batch, self._model, inference_args)
        end_time = _to_microseconds(self._clock.time_ns())
        self._semaphore.release()
        ...

We make three odd observations in that setup:

1. Beam always uses the minimum possible batch size we allow; if we specify a batch size of min 8 max 32, it'll always choose a batch size of at most 8, sometimes lower.
2. The inference timed here is still much much slower when allow multiple threads (--number_of_worker_harness_threads=10) than when we single-thread (--number_of_worker_harness_threads=1). 2.7s per batch vs. 0.4s per batch, both of which are a bit slower than running on compute engine directly.
3. In the multi-threaded setup, we keep seeing occasional CUDA OOM errors despite using a conservative batch size.

Would appreciate any and all debugging guidance for how to make this work! Right now, the whole pipeline is so slow that we have resorted to just running things in batches on Compute Engine again :/ – but there must be a way to make this work on Dataflow, right?

For reference:

• Single-threaded job: catalin-debug-classifier-test-1660143139 (Job ID: 2022-08-10_07_53_06-5898402459767488826)
• Multi-threaded job: catalin-debug-classifier-10threads-32batch-1660156741 (Job ID: 2022-08-10_11_39_50-2452382118954657386)