fix: H100+H200 CoreWeave fix (#308)

* add shm path mount

* add fix to h200 cw

* change path

* change path back

Author: cquil11

runners/launch_h100-cw.sh

@@ -7,14 +7,17 @@ SQUASH_FILE="/mnt/vast/squash/$(echo "$IMAGE" | sed 's/[\/:@#]/_/g').sqsh"
 salloc --partition=$PARTITION --gres=gpu:$TP --exclusive --time=180 --no-shell
 JOB_ID=$(squeue -u $USER -h -o %A | head -n1)
 
+SAGEMAKER_SHM_PATH=$(mktemp -d /mnt/vast/shm-XXXXXX)
+
 set -x
 srun --jobid=$JOB_ID bash -c "enroot import -o $SQUASH_FILE docker://$IMAGE"
 srun --jobid=$JOB_ID \
 --container-image=$SQUASH_FILE \
---container-mounts=$GITHUB_WORKSPACE:/workspace/,$HF_HUB_CACHE_MOUNT:$HF_HUB_CACHE \
+--container-mounts=$GITHUB_WORKSPACE:/workspace/,$HF_HUB_CACHE_MOUNT:$HF_HUB_CACHE,$SAGEMAKER_SHM_PATH:/dev/shm/sagemaker_sessions \
 --container-mount-home \
 --container-workdir=/workspace/ \
 --no-container-entrypoint --export=ALL,PORT=8888 \
 bash benchmarks/${EXP_NAME%%_*}_${PRECISION}_h100_slurm.sh
 
+rmdir $SAGEMAKER_SHM_PATH
 scancel $JOB_ID

runners/launch_h200-cw.sh

@@ -9,6 +9,8 @@ FRAMEWORK_SUFFIX=$([[ "$FRAMEWORK" == "trt" ]] && printf '_trt' || printf '')
 PARTITION="h200"
 SQUASH_FILE="/mnt/vast/squash/$(echo "$IMAGE" | sed 's/[\/:@#]/_/g').sqsh"
 
+SAGEMAKER_SHM_PATH=$(mktemp -d /mnt/vast/shm-XXXXXX)
+
 salloc --partition=$PARTITION --gres=gpu:$TP --exclusive --time=180 --no-shell
 JOB_ID=$(squeue -u $USER -h -o %A | head -n1)
 
@@ -25,10 +27,11 @@ fi
 # This seems to have been introduced in vLLM 0.11.2, but the issue is specific to CoreWeave runners.
 srun --jobid=$JOB_ID \
 --container-image=$CONTAINER_IMAGE \
---container-mounts=$GITHUB_WORKSPACE:/workspace/,$HF_HUB_CACHE_MOUNT:$HF_HUB_CACHE \
+--container-mounts=$GITHUB_WORKSPACE:/workspace/,$HF_HUB_CACHE_MOUNT:$HF_HUB_CACHE,$SAGEMAKER_SHM_PATH:/dev/shm/sagemaker_sessions \
 --container-mount-home \
 --container-workdir=/workspace/ \
 --no-container-entrypoint --export=ALL \
-bash -c "bash benchmarks/${MODEL_CODE}_${PRECISION}_h200${FRAMEWORK_SUFFIX}_slurm.sh; rm -rf /dev/shm/sagemaker_sessions"
+bash benchmarks/${MODEL_CODE}_${PRECISION}_h200${FRAMEWORK_SUFFIX}_slurm.sh
 
+rmdir $SAGEMAKER_SHM_PATH
 scancel $JOB_ID