Ok so answer to my own question, something that didn't work before but works perfectly now is to use the output in the custom data, like so:

def get_cloud_init(param1, param2, param3):
    with open('cloud-init.sh') as f:
        return f.read().format(param1=param1, param2=param2, param3=param3)
def swarm_scaleset(role, count=1, instance_type=None):
    return az.compute.ScaleSet(
        ...
        os_profile={
            ...
            'custom_data': pulumi.Output.all(param1, param2, param3).apply(
                lambda args: get_cloud_init(*args)
            )
            ...
        },
        ...
    )

managers = swarm_scaleset('manager', 3, 'Standard_B1ms')
workers = swarm_scaleset('worker', 10, 'Standard_B1ms')

I'm testing it at the moment and it seems to work ok, I'll test it on a fresh install soon (which was the problem before)