Set SLURM_PARTITION env var for all job types

Set the environment variable SLURM_PARTITION to the partition in
which a job is running. Set for salloc, sbatch and srun.