diff --git a/build.sh b/build.sh
index c1e1498..720184b 100755
--- a/build.sh
+++ b/build.sh
@@ -14,6 +14,4 @@
     ossDiskSetup \
     deployPolicyEngine \
     mdtStorageSku \
-    mdtNumDisks \
-    ostStorageSku \
-    ostNumDisks
+    ostStorageSku
diff --git a/scripts/setup_lustre.sh b/scripts/setup_lustre.sh
index 2ce74ea..bfe5a40 100755
--- a/scripts/setup_lustre.sh
+++ b/scripts/setup_lustre.sh
@@ -14,17 +14,13 @@ storage_container="$5"
 oss_disk_setup="$6"
 deploy_policy_engine="$7"
 mdt_storage_sku="$8"
-num_mdts="$9"
-ost_storage_sku="${10}"
-num_osts="${11}"
+ost_storage_sku="$9"
 
 
 if [ "$HOSTNAME" = "$mds" ]; then
 	storage_sku="$mdt_storage_sku"
-	num_disks="$num_mdts"
 else
 	storage_sku="$ost_storage_sku"
-	num_disks="$num_osts"
 fi
 
 echo mds="$1"
@@ -35,11 +31,8 @@ echo storage_container="$5"
 echo oss_disk_setup="$6"
 echo deploy_policy_engine="$7"
 echo mdt_storage_sku="$8"
-echo num_msts="$9"
-echo ost_storage_sku="${10}"
-echo num_osts="${11}"
+echo ost_storage_sku="$9"
 echo storage_sku="$storage_sku"
-echo num_disks="$num_disks"
 
 
 rbh="${mds}rbh"
@@ -74,29 +67,19 @@ else
 			devices='/dev/nvme*n1'
 			n_devices=$(echo $devices | wc -w)
 			echo "Using $n_devices NVME devices"
-		elif [ -e /dev/sdb ]; then
-			devices='/dev/sdb'
+		elif [ -e /dev/disk/azure/resource ]; then
+			devices='/dev/disk/azure/resource'
 			n_devices=1
-			echo "Using ephemeral disk on /dev/sdb"
+			echo "Using ephemeral disk on /dev/disk/azure/resource"
 		else
 			echo "ERROR: cannot find devices for storage"
 			exit 1
 		fi
 	else
-		if [ -e /dev/nvme0n1 ]; then
-			devices='/dev/nvme*n1'
-			n_devices=$(echo $devices | wc -w)
-			echo "Using $n_devices NVME devices"
-		elif [ -e /dev/sdb ]; then
-			devices='/dev/sd[b-m]'
+		if [ -d /dev/disk/azure/scsi1 ]; then
+			devices='/dev/disk/azure/scsi1/*'
 			n_devices=$(echo $devices | wc -w)
 
-			if [ "$n_devices" -gt "$num_disks" ]; then
-				echo "VM has ephemeral disk, ignoring /dev/sdb"
-				devices='/dev/sd[c-m]'
-				n_devices=$(echo $devices | wc -w)
-			fi
-
 			echo "Using $n_devices managed disks"
 		else
 			echo "ERROR: cannot find devices for storage"