adding support for assigning gpu ids based on host (#141)

src/algos/base_class.py

@@ -169,9 +169,16 @@ def setup_logging(self, config: Dict[str, ConfigType]) -> None:
     def setup_cuda(self, config: Dict[str, ConfigType]) -> None:
         """add docstring here"""
         # Need a mapping from rank to device id
-        device_ids_map = config["device_ids"]
-        node_name = f"node_{self.node_id}"
-        self.device_ids = device_ids_map[node_name]
+        if (config.get("assign_based_on_host", False)):
+            device_ids_map = config["device_ids"]
+            node_name = f"node_{self.node_id}"
+            self.device_ids = device_ids_map[node_name]
+        else:
+            hostname_to_device_ids = config["hostname_to_device_ids"]
+            hostname = os.uname().nodename
+            # choose a random one of the available devices
+            ind = self.node_id % len(hostname_to_device_ids[hostname])
+            self.device_ids = [hostname_to_device_ids[hostname][ind]]
         gpu_id = self.device_ids[0]
 
         if isinstance(gpu_id, int) and torch.cuda.is_available():

src/configs/sys_config.py

@@ -354,6 +354,14 @@ def get_digit_five_support(num_users: int, domains: List[str] = DIGIT_FIVE):
     "exp_keys": [],
     "dropout_dicts": dropout_dicts,
     "test_samples_per_user": 200,
+    "log_memory": True,
+    "assign_based_on_host": True,
+    "hostname_to_device_ids": {
+        "matlaber1": [2, 3, 4, 5, 6, 7],
+        "matlaber12": [0, 1, 2, 3],
+        "matlaber3": [0, 1, 2, 3],
+        "matlaber4": [0, 2, 3, 4, 5, 6, 7],
+    }
 }