I found a good solution that works both locally and in the cloud. Copy pasting the code in case it helps someone.

This is the higher level function in charge of partitioning the data and sending the data and the function fn to each node. 

def decrypt_data(df: SparkDataFrame) -> SparkDataFrame:
    """Decrypts the data by partitioning and sending to each Spark node a pandas
    DataFrame. Note that the GroupBy does nothing but it's necessary with the
    current Spark DataFrame API.

    Parameters
    ----------
    df : SparkDataFrame
        The initial Spark DataFrame

    Returns
    -------
    SparkDataFrame
        The decrypted Spark DataFrame
    """
    # import here to avoid error in databricks - NameError: name 'spark' is not defined
    from databricks.sdk.runtime import dbutils

    KEY_ARN = dbutils.secrets.get(
        scope=f"mlc_{Cts.ENV}", key=f"KEY_ARN_{Cts.ENV.upper()}"
    )
    fn = decrypt_on_partition(key_arn=KEY_ARN)
    return df.groupBy().applyInPandas(fn, schema=get_schema())

 See that decrypt_on_partition returns a function with they key embedded into it:

def decrypt_on_partition(key_arn: str) -> Callable[[pd.DataFrame], pd.DataFrame]:
    """This function is just a wrapper around the actual function used in each Spark
    partition and it's used to embedd the ARN KEY into each Spark node in a clean way.

    Parameters
    ----------
    key_arn : str
        The key

    Returns
    -------
    Callable[[pd.DataFrame], pd.DataFrame]
        The function that runs on each partition
    """

    def _decrypt_on_partition(partition_df: pd.DataFrame) -> pd.DataFrame:
        """The decrypt function called on each DataFrame partition. Note that it's only
        in this function that makes sense to instantiate the decryption client and
        decryption cache, so that each Spark node has its own client and cache (it
        can't be share at the cluster level).

        Parameters
        ----------
        partition_df : pd.DataFrame
            The partitioned data as a pandas DataFrame

        Returns
        -------
        pd.DataFrame
            The resulting pandas DataFrame
        """
        key_prov = aws_encryption_sdk.StrictAwsKmsMasterKeyProvider(key_ids=[key_arn])
        cache = aws_encryption_sdk.LocalCryptoMaterialsCache(LOCAL_CACHE)
        caching_cmm = aws_encryption_sdk.CachingCryptoMaterialsManager(
            master_key_provider=key_prov,
            cache=cache,
            max_age=MAX_AGE_IN_SECONDS,
            max_messages_encrypted=MAX_ENTRY_MESSAGES,
        )
        client = EncryptionSDKClient(
            commitment_policy=aws_encryption_sdk.CommitmentPolicy.REQUIRE_ENCRYPT_REQUIRE_DECRYPT
        )

        partition_df = (
            partition_df.assign(
                data=partition_df["data"].apply(
                    lambda x: _decrypt_string(x, client, caching_cmm)
                )
            )
        )

        return partition_df

    return _decrypt_on_partition