minor runtime fixes + better type annotations

experiments/simbricks/orchestration/instantiation/base.py

@@ -26,6 +26,7 @@ import asyncio
 import enum
 import pathlib
 import shutil
+import typing
 from simbricks.orchestration.utils import base as util_base
 from simbricks.orchestration.system import base as sys_base
 from simbricks.orchestration.system import pcie as sys_pcie
@@ -33,9 +34,11 @@ from simbricks.orchestration.system import mem as sys_mem
 from simbricks.orchestration.system import eth as sys_eth
 from simbricks.orchestration.system.host import base as sys_host
 from simbricks.orchestration.system.host import disk_images
-from simbricks.orchestration.simulation import base as sim_base
 from simbricks.orchestration.runtime_new import command_executor
 
+if typing.TYPE_CHECKING:
+    from simbricks.orchestration.simulation import base as sim_base
+
 
 class SockType(enum.Enum):
     LISTEN = enum.auto()

experiments/simbricks/orchestration/runtime_new/runs/base.py

@@ -24,9 +24,6 @@
 from __future__ import annotations
 
 import itertools
-import pathlib
-import shutil
-import typing as tp
 import abc
 
 from simbricks.orchestration.simulation import output

experiments/simbricks/orchestration/simulation/base.py

@@ -55,10 +55,9 @@ class Simulator(utils_base.IdObj):
         super().__init__()
         self.name: str = name
         self._relative_executable_path: str = relative_executable_path
-        # self.extra_deps: list[Simulator] = []
         self._simulation: sim_base.Simulation = simulation
         self._components: set[sys_conf.Component] = set()
-        self._wait: bool = False
+        self._wait: bool = False # TODO: FIXME
         simulation.add_sim(self)
 
     @staticmethod
@@ -122,12 +121,6 @@ class Simulator(utils_base.IdObj):
         """Full name of the simulator."""
         return ""
 
-    # TODO: move into prepare method
-    # pylint: disable=unused-argument
-    # def prep_cmds(self, inst: inst_base.Instantiation) -> list[str]:
-    # """Commands to prepare execution of this simulator."""
-    # return []
-
     def add(self, comp: sys_conf.Component) -> None:
         if comp in self._components:
             raise Exception("cannot add the same specification twice to a simulator")
@@ -211,6 +204,9 @@ class Simulator(utils_base.IdObj):
     def run_cmd(self, inst: inst_base.Instantiation) -> str:
         """Command to execute this simulator."""
         return ""
+    
+    def checkpoint_commands(self) -> list[str]:
+        return []
 
     # TODO: overwrite in sub-classes to reflect that currently not all adapters support both listening and connecting
     # In future version adapters should support both which would render this method obsolete
@@ -236,6 +232,7 @@ class Simulator(utils_base.IdObj):
     def start_delay(self) -> int:
         return 5
 
+    # TODO: FIXME
     def wait_terminate(self) -> bool:
         return self._wait
 

experiments/simbricks/orchestration/simulation/host.py

@@ -20,20 +20,18 @@
 # TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
 # SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 
-import io
-import tarfile
+from __future__ import annotations
+
 import math
 import asyncio
 import typing as tp
 import simbricks.orchestration.simulation.base as sim_base
-import simbricks.orchestration.system.pcie as system_pcie
 import simbricks.orchestration.system as system
 from simbricks.orchestration.instantiation import base as inst_base
 from simbricks.orchestration.experiment.experiment_environment_new import ExpEnv
 
 if tp.TYPE_CHECKING:
-    from simbricks.orchestration.system.host.base import Host
-
+    from simbricks.orchestration.system import host as sys_host
 
 class HostSim(sim_base.Simulator):
 
@@ -44,22 +42,7 @@ class HostSim(sim_base.Simulator):
     def full_name(self) -> str:
         return "host." + self.name
 
-    def dependencies(self) -> tp.List[sim_base.Simulator]:
-        deps = []
-        for h in self._components:
-            for dev in h.ifs:
-                if not dev.is_connected():
-                    raise Exception("host interface is not connected")
-                else:
-                    if dev.channel.a == dev:
-                        peer_if = dev.channel.b
-                    else:
-                        peer_if = dev.channel.a
-
-                deps.append(self._simulation.find_sim(peer_if.component))
-        return deps
-
-    def add(self, host: "Host"):
+    def add(self, host: sys_host.Host):
         super().add(host)
 
     def config_str(self) -> str:
@@ -139,8 +122,9 @@ class Gem5Sim(HostSim):
 
         return cmd
 
-    def wait_terminate(self) -> bool:
-        return self.wait
+
+    def checkpoint_commands(self) -> list[str]:
+        return ["m5 checkpoint"]
 
 
 class QemuSim(HostSim):
@@ -173,9 +157,9 @@ class QemuSim(HostSim):
             disks = tp.cast(list[system.DiskImage], fsh.disks)
             for disk in disks:
                 prep_cmds.append(
-                    f"{inst.qemu_img_path()} create -f qcow2 -o "
+                    f'{inst.qemu_img_path()} create -f qcow2 -o '
                     f'backing_file="{disk.path(inst=inst, format="qcow2")}" '
-                    f"{inst.hdcopy_path(img=disk, format="qcow2")}"
+                    f'{inst.hdcopy_path(img=disk, format="qcow2")}'
                 )
 
         task = asyncio.create_task(

experiments/simbricks/orchestration/system/base.py

@@ -21,10 +21,13 @@
 # SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 
 from __future__ import annotations
+
 import abc
 import typing as tp
 from simbricks.orchestration.utils import base as util_base
-from simbricks.orchestration.instantiation import base as inst_base
+
+if tp.TYPE_CHECKING:
+    from simbricks.orchestration.instantiation import base as inst_base
 
 
 class System:

experiments/simbricks/orchestration/system/host/app.py

@@ -20,41 +20,43 @@
 # TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
 # SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 
+from __future__ import annotations
+
 import typing as tp
 import abc
 import io
-from simbricks.orchestration.experiment import experiment_environment as expenv
+from simbricks.orchestration.instantiation import base as inst_base
 
 if tp.TYPE_CHECKING:
-    from simbricks.orchestration.system.host import base
+    from simbricks.orchestration.system import host as sys_host
 
 class Application(abc.ABC):
-    def __init__(self, h: 'Host') -> None:
+    def __init__(self, h: sys_host.Host) -> None:
         self.host = h
 
 
 # Note AK: Maybe we can factor most of the duplicate calls with the host out
 # into a separate module.
-class BaseLinuxApplication(abc.ABC): # TODO: FIXME!!!
-    def __init__(self, h: 'LinuxHost') -> None:
+class BaseLinuxApplication(abc.ABC):
+    def __init__(self, h: sys_host.LinuxHost) -> None:
         self.host = h
         self.start_delay: float | None = None
         self.end_delay: float | None = None
         self.wait = True
 
     @abc.abstractmethod
-    def run_cmds(self, env: expenv.ExpEnv) -> list[str]:
+    def run_cmds(self, inst: inst_base.Instantiation) -> list[str]:
         """Commands to run on node."""
         return []
 
-    def cleanup_cmds(self, env: expenv.ExpEnv) -> list[str]:
+    def cleanup_cmds(self, inst: inst_base.Instantiation) -> list[str]:
         """Commands to run to cleanup node."""
         if self.end_delay is None:
             return []
         else:
             return [f'sleep {self.start_delay}']
 
-    def config_files(self, env: expenv.ExpEnv) -> dict[str, tp.IO]:
+    def config_files(self, inst: inst_base.Instantiation) -> dict[str, tp.IO]:
         """
         Additional files to put inside the node, which are mounted under
         `/tmp/guest/`.
@@ -64,11 +66,11 @@ class BaseLinuxApplication(abc.ABC): # TODO: FIXME!!!
         """
         return {}
 
-    def prepare_pre_cp(self, env: expenv.ExpEnv) -> list[str]:
+    def prepare_pre_cp(self, inst: inst_base.Instantiation) -> list[str]:
         """Commands to run to prepare node before checkpointing."""
         return []
 
-    def prepare_post_cp(self, env: expenv.ExpEnv) -> list[str]:
+    def prepare_post_cp(self, inst: inst_base.Instantiation) -> list[str]:
         """Commands to run to prepare node after checkpoint restore."""
         if self.end_delay is None:
             return []
@@ -87,39 +89,39 @@ class BaseLinuxApplication(abc.ABC): # TODO: FIXME!!!
 
 
 class PingClient(BaseLinuxApplication):
-    def __init__(self, h: 'LinuxHost', server_ip: str = '192.168.64.1') -> None:
+    def __init__(self, h: sys_host.LinuxHost, server_ip: str = '192.168.64.1') -> None:
         super().__init__(h)
         self.server_ip = server_ip
 
-    def run_cmds(self, env: expenv.ExpEnv) -> tp.List[str]:
+    def run_cmds(self, inst: inst_base.Instantiation) -> tp.List[str]:
         return [f'ping {self.server_ip} -c 10']
 
 
 class Sleep(BaseLinuxApplication):
-    def __init__(self, h: 'LinuxHost', delay: float = 10) -> None:
+    def __init__(self, h: sys_host.LinuxHost, delay: float = 10) -> None:
         super().__init__(h)
         self.delay = delay
 
-    def run_cmds(self, env: expenv.ExpEnv) -> tp.List[str]:
+    def run_cmds(self, inst: inst_base.Instantiation) -> tp.List[str]:
         return [f'sleep {self.delay}']
 
 
 class NetperfServer(BaseLinuxApplication):
-    def __init__(self, h: 'LinuxHost') -> None:
+    def __init__(self, h: sys_host.LinuxHost) -> None:
         super().__init__(h)
 
-    def run_cmds(self, env: expenv.ExpEnv) -> tp.List[str]:
+    def run_cmds(self, inst: inst_base.Instantiation) -> tp.List[str]:
         return ['netserver', 'sleep infinity']
 
 
 class NetperfClient(BaseLinuxApplication):
-    def __init__(self, h: 'LinuxHost', server_ip: str = '192.168.64.1') -> None:
+    def __init__(self, h: sys_host.LinuxHost, server_ip: str = '192.168.64.1') -> None:
         super().__init__(h)
         self.server_ip = server_ip
         self.duration_tp = 10
         self.duration_lat = 10
 
-    def run_cmds(self, env: expenv.ExpEnv) -> tp.List[str]:
+    def run_cmds(self, inst: inst_base.Instantiation) -> tp.List[str]:
         return [
             'netserver',
             'sleep 0.5',

experiments/simbricks/orchestration/system/host/base.py

@@ -20,6 +20,8 @@
 # TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
 # SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 
+from __future__ import annotations
+
 import typing as tp
 import io
 import asyncio
@@ -30,7 +32,6 @@ from simbricks.orchestration.system import eth as eth
 from simbricks.orchestration.system.host import app
 
 if tp.TYPE_CHECKING:
-    from simbricks.orchestration.system import eth, mem, pcie
     from simbricks.orchestration.system.host import disk_images
 
 
@@ -56,9 +57,9 @@ class FullSystemHost(Host):
         self.memory = 512
         self.cores = 1
         self.cpu_freq = "3GHz"
-        self.disks: list["DiskImage"] = []
+        self.disks: list[disk_images.DiskImage] = []
 
-    def add_disk(self, disk: "DiskImage") -> None:
+    def add_disk(self, disk: disk_images.DiskImage) -> None:
         self.disks.append(disk)
 
     async def prepare(self, inst: instantiation.Instantiation) -> None:
@@ -80,7 +81,7 @@ class BaseLinuxHost(FullSystemHost):
         self,
         inst: instantiation.Instantiation,
         mapper: tp.Callable[
-            ["BaseLinuxApplication", instantiation.Instantiation], list[str]
+            [app.BaseLinuxApplication, instantiation.Instantiation], list[str]
         ],
     ) -> list[str]:
         """
@@ -111,7 +112,7 @@ class BaseLinuxHost(FullSystemHost):
         """
         cfg_files = {}
         for app in self.applications:
-            cfg_files |= self.applications[0].config_files(inst)
+            cfg_files |= app.config_files(inst)
         return cfg_files
 
     def prepare_pre_cp(self, inst: instantiation.Instantiation) -> list[str]:
@@ -122,13 +123,13 @@ class BaseLinuxHost(FullSystemHost):
         """Commands to run to prepare node after checkpoint restore."""
         return self._concat_app_cmds(inst, app.BaseLinuxApplication.prepare_post_cp)
 
-    def _config_str(self, inst: instantiation.Instantiation) -> str:
-        if inst.create_cp:
-            cp_cmd = self.checkpoint_commands()
+    def config_str(self, inst: instantiation.Instantiation) -> str:
+        if inst.create_cp():
+            sim = inst.find_sim_by_spec(spec=self)
+            cp_cmd = sim.checkpoint_commands()
         else:
             cp_cmd = []
 
-        # TODO: FIXME
         es = (
             self.prepare_pre_cp(inst)
             + self.applications[0].prepare_pre_cp(inst)
@@ -205,9 +206,6 @@ class I40ELinuxHost(LinuxHost):
         super().__init__(sys)
         self.drivers.append("i40e")
 
-    def checkpoint_commands(self) -> list[str]:
-        return ["m5 checkpoint"]
-
 
 class CorundumLinuxHost(LinuxHost):
     def __init__(self, sys) -> None:

experiments/simbricks/orchestration/system/host/disk_images.py

@@ -20,6 +20,8 @@
 # TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
 # SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 
+from __future__ import annotations
+
 import abc
 import io
 import os.path
@@ -27,13 +29,13 @@ import tarfile
 import typing as tp
 from simbricks.orchestration.utils import base as utils_base
 from simbricks.orchestration.instantiation import base as inst_base
-from simbricks.orchestration.experiment import experiment_environment as expenv
+
 if tp.TYPE_CHECKING:
-    from simbricks.orchestration.system.host import base
+    from simbricks.orchestration.system import host as sys_host
 
 
 class DiskImage(utils_base.IdObj):
-    def __init__(self, h: 'Host') -> None:
+    def __init__(self, h: sys_host.Host) -> None:
         self.host = None | str
 
     @abc.abstractmethod
@@ -65,7 +67,7 @@ class DiskImage(utils_base.IdObj):
 
 # Disk image where user just provides a path
 class ExternalDiskImage(DiskImage):
-    def __init__(self, h: 'FullSystemHost', path: str) -> None:
+    def __init__(self, h: sys_host.FullSystemHost, path: str) -> None:
         super().__init__(h)
         self.path = path
         self.formats = ["raw", "qcow2"]
@@ -80,7 +82,7 @@ class ExternalDiskImage(DiskImage):
 
 # Disk images shipped with simbricks
 class DistroDiskImage(DiskImage):
-    def __init__(self, h: 'FullSystemHost', name: str) -> None:
+    def __init__(self, h: sys_host.FullSystemHost, name: str) -> None:
         super().__init__(h)
         self.name = name
         self.formats = ["raw", "qcow2"]
@@ -101,7 +103,7 @@ class DistroDiskImage(DiskImage):
 
 # Abstract base class for dynamically generated images
 class DynamicDiskImage(DiskImage):
-    def __init__(self, h: 'FullSystemHost') -> None:
+    def __init__(self, h: sys_host.FullSystemHost) -> None:
         super().__init__(h)
 
     def path(self, inst: inst_base.Instantiation, format: str) -> str:
@@ -113,9 +115,9 @@ class DynamicDiskImage(DiskImage):
 
 # Builds the Tar with the commands to run etc.
 class LinuxConfigDiskImage(DynamicDiskImage):
-    def __init__(self, h: 'LinuxHost') -> None:
+    def __init__(self, h: sys_host.LinuxHost) -> None:
         super().__init__(h)
-        self.host: base.LinuxHost
+        self.host: sys_host.LinuxHost
 
     def available_formats(self) -> list[str]:
         return ["raw"]
@@ -126,7 +128,7 @@ class LinuxConfigDiskImage(DynamicDiskImage):
             # add main run script
             cfg_i = tarfile.TarInfo('guest/run.sh')
             cfg_i.mode = 0o777
-            cfg_f = self.host.strfile(self.host._config_str(inst))
+            cfg_f = self.host.strfile(self.host.config_str(inst))
             cfg_f.seek(0, io.SEEK_END)
             cfg_i.size = cfg_f.tell()
             cfg_f.seek(0, io.SEEK_SET)
@@ -149,7 +151,7 @@ class LinuxConfigDiskImage(DynamicDiskImage):
 # Could of course also have a version that generates the packer config from
 # python
 class PackerDiskImage(DynamicDiskImage):
-    def __init__(self, h: 'FullSystemHost', packer_config_path: str) -> None:
+    def __init__(self, h: sys_host.FullSystemHost, packer_config_path: str) -> None:
         super().__init__(h)
         self.config_path = packer_config_path