updated socket handling to retrieve sockets at later stage + pcap argument for...

updated socket handling to retrieve sockets at later stage + pcap argument for some network simulators

experiments/simbricks/orchestration/instantiation/base.py

@@ -46,12 +46,14 @@ class InstantiationEnvironment:
         workdir: str = pathlib.Path(),
         cpdir: str = pathlib.Path(),
         shm_base: str = pathlib.Path(),
+        output_base: str = pathlib.Path(),
     ):
         # TODO: add more parameters that wont change during instantiation
         self._repodir: str = pathlib.Path(repo_path).absolute()
         self._workdir: str = pathlib.Path(workdir).absolute()
         self._cpdir: str = pathlib.Path(cpdir).absolute()
-        self._shm_base = pathlib.Path(workdir).joinpath(shm_base).absolute()
+        self._shm_base: str = pathlib.Path(workdir).joinpath(shm_base).absolute()
+        self._output_base: str = pathlib.Path(workdir).joinpath(output_base).absolute()
 
 
 class Instantiation:
@@ -61,7 +63,11 @@ class Instantiation:
     ):
         self._simulation = simulation
         self._env: InstantiationEnvironment = env
-        self._socket_tracker: dict[simulation.channel.Channel, Socket] = {}
+        # we track sockets per channel and interface:
+        # - per channel tracking is for ensuring that listening as well as connecting simulator use the same socket path
+        # - per interface mapping is helpful for simulators to access their particular socket objects for cleanup and there alike
+        self._socket_per_channel: dict[simulation.channel.Channel, Socket] = {}
+        self._socket_per_interface: dict[system.base.Interface, Socket] = {}
 
     @staticmethod
     def is_absolute_exists(path: str) -> bool:
@@ -109,27 +115,37 @@ class Instantiation:
     def _get_socket_by_channel(
         self, channel: simualtion.channel.Channel
     ) -> Socket | None:
-        if not channel in self._socket_tracker:
+        if not channel in self._socket_per_channel:
             return None
-        return self._socket_tracker[channel]
+        return self._socket_per_channel[channel]
 
     def _updated_tracker_mapping(
         self, interface: system.base.Interface, socket: Socket
     ) -> None:
+        # update channel mapping
         channel = self._get_chan_by_interface(interface=interface)
-        if channel in self._socket_tracker:
-            raise Exception(
-                "cannot update socket tracker mapping, channel is already mapped"
-            )
-        self._socket_tracker[channel] = socket
+        if channel not in self._socket_per_channel:
+            self._socket_per_channel[channel] = socket
 
-    def _get_socket_by_interface(
+        # update interface mapping
+        if interface in self._socket_per_interface:
+            raise Exception("an interface cannot be associated with two sockets")
+        self._socket_per_interface[interface] = socket
+
+    def _get_opposing_socket_by_interface(
         self, interface: system.base.Interface
     ) -> Socket | None:
         channel = self._get_chan_by_interface(interface=interface)
         socket = self._get_socket_by_channel(channel=channel)
         return socket
 
+    def _get_socket_by_interface(
+        self, interface: system.base.Interface
+    ) -> Socket | None:
+        if interface not in self._socket_per_interface:
+            return None
+        return self._socket_per_interface[interface]
+
     def _interface_to_sock_path(self, interface: system.base.Interface) -> str:
         basepath = pathlib.Path(self._env._workdir)
 
@@ -160,12 +176,16 @@ class Instantiation:
         return new_socket
 
     def get_socket(self, interface: system.base.Interface) -> Socket:
-        # The other side already created a socket, we just create the opposing
-        # side (i.e. connect or listening depending on the already created type)
-        # and return
+        # check if already a socket is associated with this interface
         socket = self._get_socket_by_interface(interface=interface)
+        if socket is not None:
+            return socket
+
+        # Check if other side already created a socket, and create an opposing one
+        socket = self._get_opposing_socket_by_interface(interface=interface)
         if socket is not None:
             new_socket = self._create_opposing_socket(socket=socket)
+            self._updated_tracker_mapping(interface=interface, socket=new_socket)
             return new_socket
 
         # neither connecting nor listening side already created a socket, thus we
@@ -180,13 +200,21 @@ class Instantiation:
 
     # TODO: add more methods constructing paths as required by methods in simulators or image handling classes
 
-    def join_repo_base(self, relative_path: str) -> str:
-        path = pathlib.Path(self._env._repodir)
+    def _join_paths(self, base: str = "", relative_path: str = "") -> str:
+        path = pathlib.Path(base)
         path.joinpath(relative_path)
         if not path.exists():
-            raise Exception(f"couldn't join {self._env._repodir} and {relative_path}")
+            raise Exception(f"couldn't join {base} and {relative_path}")
         return path.absolute()
 
+    def join_repo_base(self, relative_path: str) -> str:
+        return self._join_paths(base=self._env._repodir, relative_path=relative_path)
+
+    def join_output_base(self, relative_path: str) -> str:
+        return self._join_paths(
+            base=self._env._output_base, relative_path=relative_path
+        )
+
     def hd_path(self, hd_name_or_path: str) -> str:
         if Instantiation.is_absolute_exists(hd_name_or_path):
             return hd_name_or_path

experiments/simbricks/orchestration/simulation/base.py

@@ -94,7 +94,6 @@ class Simulator(abc.ABC):
         """Commands to prepare execution of this simulator."""
         return []
 
-    # TODO: call this in subclasses
     def _add_component(self, comp: sys_base.Component) -> None:
         if comp in self._components:
             raise Exception("cannot add the same specification twice to a simulator")
@@ -115,29 +114,29 @@ class Simulator(abc.ABC):
             if inter.component in self._components:
                 assert interface is None
                 interface = inter
-
         if interface is None:
             raise Exception(
                 "unable to find channel interface for simulators specification"
             )
-
         return interface
 
-    # TODO: change method to take interface
+    def _get_sys_chan(self, interface: sys_conf.Interface) -> sys_conf.Channel:
+        if not interface.is_connected():
+            raise Exception("interface does not need a channel as it is not connected")
+        return interface.channel
+
     def _get_socket_and_chan(
-        self, inst: inst_base.Instantiation, chan: sys_conf.Channel
+        self, inst: inst_base.Instantiation, interface: sys_conf.Interface
     ) -> tuple[sys_conf.Channel, inst_base.Socket] | tuple[None, None]:
-        # check if this channel is simulator internal, i.e. doesn't need a shared memory queue
+        # get the channel associated with this interface
+        chan = self._get_sys_chan(interface=interface)
+        # check if interfaces channel is simulator internal, i.e. doesnt need an instanciation
         if not self._chan_needs_instance(chan):
             return None, None
-
         # create channel simualtion object
         channel = self.experiment.retrieve_or_create_channel(chan)
-
         # create the socket to listen on or connect to
-        my_interface = self._get_my_interface(chan)
-        socket = inst.get_socket(my_interface)
-
+        socket = inst.get_socket(interface=interface)
         return (channel, socket)
 
     def _get_channels_and_sockets(
@@ -148,12 +147,11 @@ class Simulator(abc.ABC):
         sockets = []
 
         for comp_spec in self._components:
+            for interface in comp_spec.interfaces():
 
-            # TODO: use interfaces() method instead of channels
-            for chan in comp_spec.channels():
-
-                channel, socket = self._get_socket_and_chan(inst=inst, chan=chan)
-
+                channel, socket = self._get_socket_and_chan(
+                    inst=inst, interface=interface
+                )
                 if channel is None or socket is None:
                     continue
 
@@ -172,10 +170,17 @@ class Simulator(abc.ABC):
         """Other simulators to execute before this one."""
         return []
 
-    # Sockets to be cleaned up
+    # Sockets to be cleaned up: always the CONNECTING sockets
     # pylint: disable=unused-argument
-    def sockets_cleanup(self, env: exp_env.ExpEnv) -> list[str]:
-        return []
+    def sockets_cleanup(self, inst: inst_base.Instantiation) -> list[inst_base.Socket]:
+        sockets = []
+        for comp_spec in self._components:
+            for interface in comp_spec.interfaces():
+                socket = inst.get_socket(interface=interface)
+                if socket._type == inst_base.SockType.CONNECT:
+                    sockets.append(socket)
+
+        return sockets
 
     # sockets to wait for indicating the simulator is ready
     # pylint: disable=unused-argument

experiments/simbricks/orchestration/simulation/net/base.py

@@ -49,10 +49,6 @@ class NetSim(base.Simulator):
                 deps.append(n.net[0].sim)
         return deps
 
-    # TODO
-    def sockets_cleanup(self, env: exp_env.ExpEnv) -> list[str]:
-        pass
-
     # TODO
     def sockets_wait(self, env: exp_env.ExpEnv) -> list[str]:
         pass
@@ -73,10 +69,13 @@ class WireNet(NetSim):
 
     def __init__(self, simulation: base.Simulation) -> None:
         super().__init__(
-            simulation=simulation, relative_executable_path="/sims/net/wire/net_wire"
+            simulation=simulation,
+            relative_executable_path="/sims/net/wire/net_wire",
+            relative_pcap_file_path=None,
         )
         # TODO: probably we want to store these in a common base class...
         self._wire_comp: eth.EthWire | None = None
+        self._relative_pcap_file_path: str | None = relative_pcap_file_path
 
     def add_wire(self, wire: eth.EthWire):
         assert self._wire_comp is None
@@ -105,9 +104,11 @@ class WireNet(NetSim):
         cmd = inst.join_repo_base(self._relative_executable_path)
         cmd += f"{sockets[0]} {sockets[1]} {run_sync} {sync_period} {eth_latency}"
 
-        # TODO
-        if len(env.pcap_file) > 0:
-            cmd += " " + env.pcap_file
+        if self._relative_pcap_file_path is not None:
+            pcap_file = inst.join_output_base(
+                relative_path=self._relative_pcap_file_path
+            )
+            cmd += " " + pcap_file
         return cmd
 
 
@@ -117,12 +118,14 @@ class SwitchNet(NetSim):
         self,
         simulation: base.Simulation,
         relative_executable_path="/sims/net/switch/net_switch",
+        relative_pcap_file_path=None,
     ) -> None:
         super().__init__(
             simulation=simulation, relative_executable_path=relative_executable_path
         )
         # TODO: probably we want to store these in a common base class...
         self._switch_spec: eth.EthSwitch | None = None
+        self._relative_pcap_file_path: str | None = relative_pcap_file_path
 
     def add_switch(self, switch_spec: eth.EthSwitch):
         assert self._switch_spec is None
@@ -155,9 +158,11 @@ class SwitchNet(NetSim):
         if not run_sync:
             cmd += " -u"
 
-        # TODO: pcap_file --> no env!!!
-        if len(env.pcap_file) > 0:
-            cmd += " -p " + env.pcap_file
+        if self._relative_pcap_file_path is not None:
+            pcap_file = inst.join_output_base(
+                relative_path=self._relative_pcap_file_path
+            )
+            cmd += " " + pcap_file
 
         listen, connect = base.Simulator.split_sockets_by_type(sockets)
 
@@ -169,23 +174,16 @@ class SwitchNet(NetSim):
 
         return cmd
 
-    # TODO
-    def sockets_cleanup(self, env: exp_env.ExpEnv) -> list[str]:
-        # cleanup here will just have listening eth sockets, switch also creates
-        # shm regions for each with a "-shm" suffix
-        cleanup = []
-        for s in super().sockets_cleanup(env):
-            cleanup.append(s)
-            cleanup.append(s + "-shm")
-        return cleanup
-
 
 class MemSwitchNet(SwitchNet):
 
-    def __init__(self, simulation: base.Simulation) -> None:
+    def __init__(
+        self, simulation: base.Simulation, relative_pcap_file_path=None
+    ) -> None:
         super().__init__(
             simulation=simulation,
             relative_executable_path="/sims/mem/memswitch/memswitch",
+            relative_pcap_file_path=relative_pcap_file_path,
         )
         """AS_ID,VADDR_START,VADDR_END,MEMNODE_MAC,PHYS_START."""
         self.mem_map = []