Is there a way to limit node-wise parallelization of certain memory-intensive UDF computations in PySpark?

BACKGROUND

PySpark allows you to apply arbitrary pandas functions to groups of a Spark DataFrame, as follows:

import pandas as pd
from pyspark.sql.functions import pandas_udf, PandasUDFType

df = spark.createDataFrame(
    [(1, 1.0), (1, 2.0), (2, 3.0), (2, 5.0), (2, 10.0)], ("id", "v"))

@pandas_udf(df.schema, PandasUDFType.GROUPED_MAP)
def subtract_mean(pdf):
    v = pdf.v
    return pdf.assign(v=v - v.mean())

df.groupby("id").apply(subtract_mean).show()

Behind the scenes, Spark distributes data from df to worker nodes, partitioned according to 'id', applies the user-defined pandas function to each group in parallel, and then takes the union of the results of these.

PROBLEM The pandas UDF I want to run is very memory intensive. When Spark tries to run many instances of the UDF on the same node in parallel (say 64 for example), the node runs out of memory, and the UDF throws a memory error.

QUESTION Is there a way to tell Spark: "Parallelize this Grouped UDF across nodes, but only run at most k instances of the UDF in parallel on each node."? (e.g., k=1 means each node only executes one instance of the UDF at a time).

PREVIOUS ATTEMPTS AT SOLUTION Break up df into smaller chunks, and process them in series using a for-loop. Why is this bad? Because we would like Spark to take advantage of parallelization across nodes!