updates how we pick partitions

now works on both prince and rusty

Snakefile

@@ -550,9 +550,30 @@ def get_windows(wildcards):
                                                   t_width=t_width, window_type=window_type))
         return windows
 
-# def get_gpu_type(wildcards):
+def get_partition(wildcards, cluster):
     # if our V1 scaling value is small enough, we need a V100 and must specify
-    # it. otherwise, we can use any GPU, because they'll all have enough memory
+    # it. otherwise, we can use any GPU, because they'll all have enough
+    # memory. The partition name depends on the cluster (prince or rusty), so
+    # we have two different params, one for each, and the cluster config grabs
+    # the right one
+    if cluster not in ['prince', 'rusty']:
+        raise Exception(f"Don't know how to handle cluster {cluster}")
+    if int(wildcards.gpu) == 0:
+        if cluster == 'rusty':
+            return 'ccn'
+        elif cluster == 'prince':
+            return None
+    else:
+        scaling = float(wildcards.scaling)
+        if cluster == 'rusty':
+            return 'gpu'
+        elif cluster == 'prince':
+            part = 'simoncelli_gpu,v100_sxm2_4,v100_pci_2'
+            if scaling >= .18:
+                part += ',p40_4'
+            if scaling >= .4:
+                part += ',p100_4'
+            return part
 
 rule create_metamers:
     input:
@@ -582,7 +603,8 @@ rule create_metamers:
     params:
         cache_dir = lambda wildcards: op.join(config['DATA_DIR'], 'windows_cache'),
         time = lambda wildcards: {'V1': '12:00:00', 'RGC': '5-00:00:00'}[wildcards.model_name.split('_')[0]],
-        partition = lambda wildcards: {0: 'ccn'}.get(int(wildcards.gpu), 'gpu')
+        rusty_partition = lamda wilcards: get_partition(wildcards, 'rusty'),
+        prince_partition = lamda wilcards: get_partition(wildcards, 'prince'),
     run:
         import foveated_metamers as met
         import contextlib
@@ -667,7 +689,8 @@ rule continue_metamers:
     params:
         cache_dir = lambda wildcards: op.join(config['DATA_DIR'], 'windows_cache'),
         time = lambda wildcards: {'V1': '12:00:00', 'RGC': '5-00:00:00'}[wildcards.model_name.split('_')[0]],
-        partition = lambda wildcards: {0: 'ccn'}.get(int(wildcards.gpu), 'gpu')
+        rusty_partition = lamda wilcards: get_partition(wildcards, 'rusty'),
+        prince_partition = lamda wilcards: get_partition(wildcards, 'prince'),
     run:
         import foveated_metamers as met
         import contextlib

prince.json

@@ -14,12 +14,14 @@
     {
 	    "gres": "gpu:{resources.gpu}",
 	    "mem": "{resources.mem}GB",
+        "partition": "{params.prince_partition}",
         "time": "{params.time}"
     },
     "continue_metamers":
     {
         "gres": "gpu:{resources.gpu}",
         "mem": "{resources.mem}GB",
+        "partition": "{params.prince_partition}",
         "time": "{params.time}"
     },
     "cache_windows":

rusty.json

@@ -15,14 +15,14 @@
     {
         "gres": "gpu:v100-32gb:{resources.gpu}",
         "mem": "{resources.mem}GB",
-        "partition": "{params.partition}",
+        "partition": "{params.rusty_partition}",
         "time": "{params.time}"
     },
     "continue_metamers":
     {
         "gres": "gpu:v100-32gb:{resources.gpu}",
         "mem": "{resources.mem}GB",
-        "partition": "{params.partition}",
+        "partition": "{params.rusty_partition}",
         "time": "{params.time}"
     },
     "cache_windows":