fix more pytype issues

experiments/pyexps/dist_memcache.py

@@ -112,7 +112,8 @@ for host_type in host_types:
                     ip_prefix=16)
             for s in servers:
                 e.assign_sim_host(s, h_i)
-                e.assign_sim_host(s.pcidevs[0].multinic, h_i)
+                s_multisubnic = next(pcidev for pcidev in s.pcidevs if isinstance(pcidev, sim.MultiSubNIC))
+                e.assign_sim_host(s_multisubnic.multinic, h_i)
 
             clients = create_multinic_hosts(e, m, 'client_%d' % (i,),
                     switch, host_class, node.I40eLinuxNode,
@@ -121,7 +122,8 @@ for host_type in host_types:
             for c in clients:
                 c.wait = True
                 e.assign_sim_host(c, h_i)
-                e.assign_sim_host(c.pcidevs[0].multinic, h_i)
+                c_multisubnic = next(pcidev for pcidev in c.pcidevs if isinstance(pcidev, sim.MultiSubNIC))
+                e.assign_sim_host(c_multisubnic.multinic, h_i)
 
             racks.append((servers, clients))
 

experiments/run.py

@@ -30,9 +30,12 @@ import pickle
 import fnmatch
 import typing as tp
 
+from simbricks.runtime.common import *
+from simbricks.runtime.distributed import *
+from simbricks.runtime.local import *
+from simbricks.runtime.slurm import *
 import simbricks.exectools as exectools
 import simbricks.experiments as exp
-import simbricks.runtime as runtime
 
 def mkdir_if_not_exists(path):
     if not os.path.exists(path):
@@ -139,19 +142,19 @@ def warn_multi_exec():
 # initialize runtime
 if args.runtime == 'parallel':
     warn_multi_exec()
-    rt = runtime.LocalParallelRuntime(cores=args.cores, mem=args.mem,
+    rt = LocalParallelRuntime(cores=args.cores, mem=args.mem,
             verbose=args.verbose, exec=executors[0])
 elif args.runtime == 'slurm':
-    rt = runtime.SlurmRuntime(args.slurmdir, args, verbose=args.verbose)
+    rt = SlurmRuntime(args.slurmdir, args, verbose=args.verbose)
 elif args.runtime == 'dist':
-    rt = runtime.DistributedSimpleRuntime(executors, verbose=args.verbose)
+    rt = DistributedSimpleRuntime(executors, verbose=args.verbose)
 else:
     warn_multi_exec()
-    rt = runtime.LocalSimpleRuntime(verbose=args.verbose, exec=executors[0])
+    rt = LocalSimpleRuntime(verbose=args.verbose, exec=executors[0])
 
 
 def add_exp(
-    e: exp.Experiment, run: int, prereq: tp.Optional[runtime.Run],
+    e: exp.Experiment, run: int, prereq: tp.Optional[Run],
     create_cp: bool, restore_cp: bool, no_simbricks: bool
 ):
     outpath = '%s/%s-%d.json' % (args.outdir, e.name, run)
@@ -174,7 +177,7 @@ def add_exp(
     if args.shmdir is not None:
         env.shm_base = os.path.abspath(shmdir)
 
-    run = runtime.Run(e, run, env, outpath, prereq)
+    run = Run(e, run, env, outpath, prereq)
     rt.add_run(run)
     return run
 
@@ -197,7 +200,7 @@ if not args.pickled:
 
     for e in experiments:
         if args.auto_dist and not isinstance(e, exp.DistributedExperiment):
-            e = runtime.auto_dist(e, executors, args.proxy_type)
+            e = auto_dist(e, executors, args.proxy_type)
         # apply filter if any specified
         if (args.filter) and (len(args.filter) > 0):
             match = False

experiments/simbricks/experiments.py

@@ -30,7 +30,7 @@ from simbricks.experiment.experiment_environment import ExpEnv
 from simbricks.experiment.experiment_output import ExpOutput
 from simbricks.exectools import Executor, SimpleComponent
 from simbricks.proxy import NetProxyConnecter, NetProxyListener, SimProxy
-from simbricks.simulators import HostSim, NICSim, NetSim, PCIDevSim, Simulator
+from simbricks.simulators import HostSim, I40eMultiNIC, NICSim, NetSim, PCIDevSim, Simulator
 import simbricks.utils.graphlib as graphlib
 
 class Experiment(object):
@@ -50,13 +50,13 @@ class Experiment(object):
     simulator is restored in the accurate mode using this checkpoint."""
     no_simbricks = False
     """`true` - No simbricks adapters are used in the simulators."""
+    hosts: tp.List[HostSim] = []
+    pcidevs: tp.List[PCIDevSim] = []
+    networks: tp.List[NetSim] = []
+    metadata = {}
 
     def __init__(self, name: str):
         self.name = name
-        self.hosts: tp.List[HostSim] = []
-        self.pcidevs: tp.List[PCIDevSim] = []
-        self.networks: tp.List[NetSim] = []
-        self.metadata = {}
 
     def add_host(self, sim: HostSim):
         for h in self.hosts:
@@ -64,7 +64,7 @@ class Experiment(object):
                 raise Exception('Duplicate host name')
         self.hosts.append(sim)
 
-    def add_nic(self, sim: NICSim):
+    def add_nic(self, sim: tp.Union[NICSim, I40eMultiNIC]):
         self.add_pcidev(sim)
 
     def add_pcidev(self, sim: PCIDevSim):
@@ -124,7 +124,7 @@ class DistributedExperiment(Experiment):
         return itertools.chain(super().all_simulators(),
                 self.proxies_listen, self.proxies_connect)
 
-    def assign_sim_host(self, sim, host):
+    def assign_sim_host(self, sim: Simulator, host: int):
         """ Assign host ID (< self.num_hosts) for a simulator. """
         assert(host >= 0 and host < self.num_hosts)
         self.host_mapping[sim] = host

experiments/simbricks/runtime/distributed.py

@@ -110,14 +110,14 @@ def auto_dist(
     for h in e.hosts:
         de.add_host(h)
         de.assign_sim_host(h, k)
-        for nic in h.nics:
+        for nic in h.nics:  # TODO h.nics does not exist in class HostSim
             de.assign_sim_host(nic, k)
 
             if k != 0:
                 cp.add_nic(nic)
         k = (k + 1) % 2
 
-    for nic in e.nics:
+    for nic in e.nics:  # TODO: e.nics does not exist in class Experiment
         de.add_nic(nic)
 
     return de
\ No newline at end of file

experiments/simbricks/runtime/local.py

@@ -22,7 +22,7 @@
 
 import asyncio
 import pathlib
-import typing as tp
+from typing import Optional
 
 from simbricks.runtime.common import *
 import simbricks.experiments as exp
@@ -71,7 +71,7 @@ class LocalParallelRuntime(Runtime):
     def __init__(
         self,
         cores: int,
-        mem: tp.Optional[int] = None,
+        mem: Optional[int] = None,
         verbose=False,
         exec: exectools.Executor = exectools.LocalExecutor()
     ):

experiments/simbricks/simulator_utils.py

@@ -122,10 +122,10 @@ def create_dctcp_hosts(
     num: int,
     name_prefix: str,
     net: NetSim,
-    nic_class: NICSim,
-    host_class: HostSim,
-    nc_class: NodeConfig,
-    app_class: AppConfig,
+    nic_class: tp.Type[NICSim],
+    host_class: tp.Type[HostSim],
+    nc_class: tp.Type[NodeConfig],
+    app_class: tp.Type[AppConfig],
     cpu_freq: str,
     mtu: int,
     ip_start: int = 1

experiments/simbricks/simulators.py

@@ -103,7 +103,7 @@ class NICSim(PCIDevSim):
     def __init__(self):
         super().__init__()
 
-    def set_network(self, net):
+    def set_network(self, net: NetSim):
         self.network = net
         net.nics.append(self)
 
@@ -419,7 +419,7 @@ class MultiSubNIC(NICSim):
     def start_delay(self):
         return 0
 
-class I40eMultiNIC(Simulator):  # TODO Fix typing, e.g., by making this a sublcass of NICSim
+class I40eMultiNIC(Simulator):
     name = ''
 
     def __init__(self):