Merge pull request #96 from seung-lab/localssd

Use local scratch ssd as main storage for nfs server

No as flexible as pd-ssd and the data are wiped out when vm is stopped or resets, on the other hand has better performance and slightly cheaper per GB.

cloud/google/cloud-deployment.yaml

@@ -126,13 +126,19 @@ resources:
 # NFS server for meshing and skeletonization tasks. The intermediate files are
 # saved to the nfs server, and the "forge" step will read the input from it.
 # This can reduce the read access/cost siginificantly when the segmentation
-# we want to mesh/skeletonize is large. Make sure the nfsVolumeSizeGB is large
-# enough to accommodate the data
+# we want to mesh/skeletonize is large. Set pdSSDSizeGB to use pd-ssd as the
+# main storage, or numLocalScratchSSD to set the number of local scratch disks
+# for temporary storage. (The data are wipe out once the nfs server is stopped).
+# pdSSDSizeGB and numLocalScratchSSD are mutually exclusive, you can use only
+# choose one of them. You must make sure the storage is large enough to
+# accommodate the data, for example meshing FlyWire segments at 32nmx32nmx40nm
+# resolution generates ~ 1 TB of fragments
           # nfsServer:
           #   zone: us-east1-c
           #   subnetwork: us-east1
-          #   machineType: e2-highmem-16
-          #   nfsVolumeSizeGB: 1000
+          #   machineType: n2d-highmem-16
+          #   pdSSDSizeGB: 1000
+          #   numLocalScratchSSD: 4 # each scratch disk is 375 GB
           easysegWorker:
             zone: us-east1-c
             subnetwork: us-east1

cloud/google/common.py

@@ -118,6 +118,23 @@ def GenerateDisk(diskSizeGb, diskType=None):
     return disk
 
 
+def GenerateScratchDisk(numDisks, diskType):
+    disks = []
+    for i in range(numDisks):
+        disks.append({
+                        "autoDelete": True,
+                        "deviceName": f"local-ssd-{i}",
+                        "initializeParams": {
+                          "diskSizeGb": "375",
+                          "diskType": diskType,
+                        },
+                        "interface": "NVME",
+                        "mode": "READ_WRITE",
+                        "type": "SCRATCH"
+                    })
+
+    return disks
+
 def GenerateNetworkInterface(context, subnetwork, ipAddr=None):
     network_interface = {
         'network': f'$(ref.{context.env["deployment"]}-network.selfLink)',

cloud/google/nfs_server.py

@@ -1,8 +1,9 @@
-from common import ZonalComputeUrl, GenerateDisk, GenerateBootDisk, GenerateNetworkInterface, INSTALL_DOCKER_CMD, GenerateAirflowVar
+from common import ZonalComputeUrl, GenerateDisk, GenerateBootDisk, GenerateScratchDisk, GenerateNetworkInterface, INSTALL_DOCKER_CMD, GenerateAirflowVar
 from workers import GenerateDockerCommand, GenerateCeleryWorkerCommand, GenerateEnvironVar
 
 
 def GenerateNFSServerStartupScript(context, hostname_manager):
+    nfs_server_param = context.properties["nfsServer"]
     env_variables = GenerateAirflowVar(context, hostname_manager)
 
     docker_env = [f'-e {k}' for k in env_variables]
@@ -39,7 +40,16 @@ def GenerateNFSServerStartupScript(context, hostname_manager):
     }
 }
 '''
-
+    scratch_nvme_config = '''
+DRIVES=($(lsblk | grep -oE 'nvme[a-z0-9A-Z]*' | cut -d' ' -f1 | awk '{ print "/dev/"$1 }'))
+if [ ${#DRIVES[@]} -ne 0 ]; then
+mdadm --create /dev/md0 --level=0 --force --raid-devices=${#DRIVES[@]} ${DRIVES[@]}
+mkfs.ext4 -F /dev/md0
+mount /dev/md0 /share
+chmod 777 /share
+mkdir -p /share/mariadb
+fi
+'''
     startup_script = f'''
 #!/bin/bash
 set -e
@@ -50,15 +60,17 @@ def GenerateNFSServerStartupScript(context, hostname_manager):
 
 DEBIAN_FRONTEND=noninteractive apt-get -y -o Dpkg::Options::="--force-confdef" -o Dpkg::Options::="--force-confold" dist-upgrade
 {INSTALL_DOCKER_CMD}
+{scratch_nvme_config}
 
 {GenerateEnvironVar(context, env_variables)}
 
 if [ ! -f "/etc/bootstrap_done" ]; then
-
+if ! mountpoint -q "/share"; then
 mkfs.ext4 -F /dev/sdb
 mount /dev/sdb /share
 chmod 777 /share
 mkdir -p /share/mariadb
+fi
 apt-get install nfs-kernel-server nginx -y
 echo "/share 172.31.0.0/16(insecure,rw,async,no_subtree_check)" >> /etc/exports
 echo "ALL: 172.31.0.0/16" >> /etc/hosts.allow
@@ -79,8 +91,10 @@ def GenerateNFSServerStartupScript(context, hostname_manager):
 
 sysctl -w net.netfilter.nf_conntrack_max=$(awk '/MemAvailable/ {{print int($2/16)}}' /proc/meminfo)
 echo $(awk '/MemAvailable/ {{print int($2/64)}}' /proc/meminfo) > /sys/module/nf_conntrack/parameters/hashsize
+if ! mountpoint -q "/share"; then
 mount /dev/sdb /share
 chmod 777 /share
+fi
 systemctl restart nfs-kernel-server.service
 export INNODB_POOL_SIZE_GB=$(awk '/MemAvailable/ {{print int($2/1024/1024/2)}}' /proc/meminfo)
 export INNODB_LOG_SIZE_GB=$(awk '/MemAvailable/ {{print int($2/1024/1024/8)}}' /proc/meminfo)
@@ -95,23 +109,34 @@ def GenerateNFSServerStartupScript(context, hostname_manager):
 def GenerateNFSServer(context, hostname_manager, hostname_nfs_server):
     nfs_server_param = context.properties["nfsServer"]
 
+    if ('pdSSDSizeGB' in nfs_server_param) == ('numLocalScratchSSD' in nfs_server_param):
+        raise ValueError("You must set one and only one of 'pdSSDSizeGB' and 'numLocalScratchSSD'")
+
     startup_script = GenerateNFSServerStartupScript(context, hostname_manager)
 
-    diskType = ZonalComputeUrl(
-        context.env['project'],
-        nfs_server_param['zone'],
-        'diskTypes', 'pd-ssd')
+    disks = [GenerateBootDisk(diskSizeGb=10),]
+
+    if 'pdSSDSizeGB' in nfs_server_param:
+        diskType = ZonalComputeUrl(
+            context.env['project'],
+            nfs_server_param['zone'],
+            'diskTypes', 'pd-ssd')
+        disks.append(GenerateDisk(diskSizeGb=nfs_server_param['pdSSDSizeGB'], diskType=diskType))
+    elif 'numLocalScratchSSD' in nfs_server_param:
+        diskType = ZonalComputeUrl(
+            context.env['project'],
+            nfs_server_param['zone'],
+            'diskTypes', 'local-ssd')
+        disks += GenerateScratchDisk(numDisks=nfs_server_param['numLocalScratchSSD'], diskType=diskType)
+
 
     instance_resource = {
         'zone': nfs_server_param['zone'],
         'machineType': ZonalComputeUrl(
                       context.env['project'], nfs_server_param['zone'],
                       'machineTypes', nfs_server_param['machineType']
         ),
-        'disks': [
-                  GenerateBootDisk(diskSizeGb=10),
-                  GenerateDisk(diskSizeGb=nfs_server_param['nfsVolumeSizeGB'], diskType=diskType),
-        ],
+        'disks': disks,
         'labels': {
             'vmrole': 'nfs-server',
             'location': nfs_server_param['zone'],

dags/google_api_helper.py

@@ -360,6 +360,7 @@ def start_instance(instance_name, zone):
 def stop_instance(instance_name, zone):
     service = discovery.build('compute', 'v1')
     request = service.instances().stop(
+        discardLocalSsd=True,
         project=get_project_id(),
         zone=zone,
         instance=instance_name