fixed various type checking and mismatched type errors

src/algos/MetaL2C.py

@@ -398,7 +398,7 @@ def run_protocol(self) -> None:
         total_rounds = self.config["rounds"]
 
         stats = []
-        avg_alpha = None
+        avg_alpha : Dict[str, torch.Tensor] = {}
         for cur_round in range(start_round, total_rounds):
             self.round = cur_round
             self.log_utils.log_console(f"Starting round {cur_round}")

src/algos/attack_add_noise.py

@@ -19,7 +19,7 @@
 from collections import OrderedDict
 from typing import Dict
 from torch import Tensor
-from utils.types import ConfigType
+from utils.custom_types import ConfigType
 
 
 class AddNoiseAttack:

src/algos/attack_bad_weights.py

@@ -17,7 +17,7 @@
 from collections import OrderedDict
 from typing import Dict
 from torch import Tensor
-from utils.types import ConfigType
+from utils.custom_types import ConfigType
 
 
 class BadWeightsAttack:

src/algos/attack_sign_flip.py

@@ -17,7 +17,7 @@
 from collections import OrderedDict
 from typing import Dict
 from torch import Tensor
-from utils.types import ConfigType
+from utils.custom_types import ConfigType
 
 
 class SignFlipAttack:

src/algos/base_class.py

@@ -36,7 +36,7 @@
     get_dset_balanced_communities,
     get_dset_communities,
 )
-from utils.types import ConfigType
+from utils.custom_types import ConfigType
 from utils.dropout_utils import NodeDropout
 
 import torchvision.transforms as T  # type: ignore

src/algos/topologies/base.py

@@ -4,7 +4,7 @@
 import numpy as np
 import networkx as nx
 
-from utils.types import ConfigType
+from utils.custom_types import ConfigType
 
 class BaseTopology(ABC):
     """

src/algos/topologies/collections.py

@@ -1,5 +1,5 @@
 from algos.topologies.base import BaseTopology
-from utils.types import ConfigType
+from utils.custom_types import ConfigType
 
 from math import ceil
 import networkx as nx

src/configs/algo_config.py

@@ -1,7 +1,7 @@
 from typing import Dict, List
 from .malicious_config import malicious_config_list
 import random
-from utils.types import ConfigType
+from utils.custom_types import ConfigType
 
 
 def get_malicious_types(malicious_config_list: List[ConfigType]) -> Dict[str, str]:

src/configs/algo_config_test.py

@@ -1,4 +1,4 @@
-from utils.types import ConfigType
+from utils.custom_types import ConfigType
 
 fedstatic: ConfigType = {
     # Collaboration setup

src/configs/malicious_config.py

@@ -1,5 +1,5 @@
 # Malicious Configuration
-from utils.types import ConfigType
+from utils.custom_types import ConfigType
 from typing import Dict
 
 # Weight Update Attacks

src/configs/non_iid_clients.py

@@ -101,25 +101,25 @@ def get_domain_support(
     return support
 
 
-CIFAR10_ROT_DMN: List[str] = ["r0", "r90", "r180", "r270"]
+CIFAR10_ROT_DMN: List[str | int] = ["r0", "r90", "r180", "r270"]
 
 
 def get_cifar10_rot_support(
-    num_clients: int, domains: List[str] = CIFAR10_ROT_DMN
+    num_clients: int, domains: List[str | int] = CIFAR10_ROT_DMN
 ) -> Dict[str, str]:
     return get_domain_support(num_clients, "cifar10", domains)
 
 
-DOMAINNET_DMN: List[str] = ["real", "sketch", "clipart"]
+DOMAINNET_DMN: List[str | int] = ["real", "sketch", "clipart"]
 
 
 def get_domainnet_support(
-    num_clients: int, domains: List[str] = DOMAINNET_DMN
+    num_clients: int, domains: List[str | int] = DOMAINNET_DMN
 ) -> Dict[str, str]:
     return get_domain_support(num_clients, "domainnet", domains)
 
 
-DOMAINNET_DMN_FULL: List[str] = [
+DOMAINNET_DMN_FULL: List[str | int] = [
     "real",
     "sketch",
     "clipart",
@@ -130,42 +130,42 @@ def get_domainnet_support(
 
 
 def get_domainnet_support_full(
-    num_clients: int, domains: List[str] = DOMAINNET_DMN_FULL
+    num_clients: int, domains: List[str | int] = DOMAINNET_DMN_FULL
 ) -> Dict[str, str]:
     return get_domain_support(num_clients, "domainnet", domains)
 
 
-DOMAINNET_DMN_V2: List[str] = ["infograph", "quickdraw", "painting"]
+DOMAINNET_DMN_V2: List[str | int] = ["infograph", "quickdraw", "painting"]
 
 
 def get_domainnet_support_v2(
-    num_clients: int, domains: List[str] = DOMAINNET_DMN_V2
+    num_clients: int, domains: List[str | int] = DOMAINNET_DMN_V2
 ) -> Dict[str, str]:
     return get_domain_support(num_clients, "domainnet", domains)
 
 
-IWILDCAM_DMN: List[int] = list(range(1, 5))  # 245 possible
+IWILDCAM_DMN: List[str | int] = list(range(1, 5))  # 245 possible
 
 
 def get_iwildcam_support(
-    num_clients: int, domains: List[int] = IWILDCAM_DMN
+    num_clients: int, domains: List[str | int] = IWILDCAM_DMN
 ) -> Dict[str, str]:
     return get_domain_support(num_clients, "wilds_iwildcam", domains)
 
 
 # 2 classes
 # 3 domains: 0:116'959, 3:132'052, 4:5'3425 in training set
-CAMELYON17_DMN: List[int] = [0, 3, 4]  # + 1, 2 in test set
+CAMELYON17_DMN: List[str | int] = [0, 3, 4]  # + 1, 2 in test set
 
 
 def get_camelyon17_support(
-    num_clients: int, domains: List[int] = CAMELYON17_DMN
+    num_clients: int, domains: List[str | int] = CAMELYON17_DMN
 ) -> Dict[str, str]:
     return get_domain_support(num_clients, "wilds_camelyon17", domains)
 
 
 # Issue every of the 1139 classes has only 1 sample per domain => how to create in domain test set ?
-RXRX1_DMN = [
+RXRX1_DMN: List[str | int] = [
     0,
     1,
     2,
@@ -175,15 +175,15 @@ def get_camelyon17_support(
 ]  # in train set: 0-6, 11-26, 35-41, 46-48, in test set: 7-10, 27-34, 42-45, 49-50
 
 
-def get_rxrx1_support(num_clients, domains=RXRX1_DMN):
+def get_rxrx1_support(num_clients: int, domains: List[str|int] =RXRX1_DMN):
     return get_domain_support(num_clients, "wilds_rxrx1", domains)
 
 
 # 0: 17'809, 1: 34'816, 2: 1'582, 3: 20'973, 4:1'641, 5: 42
-FMOW_DMN = [0, 1, 3]
+FMOW_DMN : List[str|int] = [0, 1, 3]
 
 
-def get_fmow_support(num_clients, domains=FMOW_DMN):
+def get_fmow_support(num_clients: int, domains : List[str|int] =FMOW_DMN):
     return get_domain_support(num_clients, "wilds_fmow", domains)
 
 

src/configs/sys_config.py

@@ -3,7 +3,7 @@
 # is to simulate different real-world scenarios without changing the algorithm configuration.
 from typing import Dict, List, Literal, Optional
 import random
-from utils.types import ConfigType
+from utils.custom_types import ConfigType
 
 # from utils.config_utils import get_sliding_window_support, get_device_ids
 from .algo_config import (
@@ -334,7 +334,7 @@ def get_digit_five_support(num_users: int, domains: List[str] = DIGIT_FIVE):
 
 # for swift or fedavgpush, just modify the algo_configs list
 # for swift, synchronous should preferable be False
-gpu_ids = [2, 3, 5, 6]
+gpu_ids : List[int | Literal['cpu']] = [2, 3, 5, 6]
 
 grpc_system_config: ConfigType = {
     "exp_id": "static",

src/configs/sys_config_test.py

@@ -1,9 +1,9 @@
 from typing import Dict, List, Literal, Optional
 import random
-from utils.types import ConfigType
+from utils.custom_types import ConfigType
 
 from .algo_config_test import (
-    traditional_fl,
+    #traditional_fl,
     fedstatic
 )
 
@@ -99,7 +99,7 @@ def get_algo_configs(
 for i in range(1, num_users + 1):
     dropout_dicts[f"node_{i}"] = dropout_dict
 
-gpu_ids = [2, 3, 5, 6]
+gpu_ids : List[int | Literal['cpu']] = [2, 3, 5, 6]
 
 grpc_system_config: ConfigType = {
     "exp_id": "static",

src/main_grpc.py

@@ -35,7 +35,7 @@
 # Command for opening each process
 command_list: List[str] = ["python", "main.py", "-host", args.host]
 if args.dev == True:
-    command_list: List[str] = ["python", "main.py", "-b", "./configs/algo_config_test.py", "-s", "./configs/sys_config_test.py", "-host", args.host]
+    command_list = ["python", "main.py", "-b", "./configs/algo_config_test.py", "-s", "./configs/sys_config_test.py", "-host", args.host]
 
 # Start process for each user
 for i in range(args.n):