Slurm can do this with a `topology.conf` file, and this can be auto-generated for an IB switch. See https://slurm.schedmd.com/topology.html.