add gpu & deep-learning support

examples/D00-tensorflow/seed-compose.py

@@ -0,0 +1,53 @@
+from seedemu.core import Emulator, Binding, Filter, Action
+from seedemu.layers import Base
+from seedemu.compiler import Docker
+
+# Initialize the emulator and layers
+emu = Emulator()
+base = Base()
+
+# Create an autonomous system
+as100 = base.createAutonomousSystem(100)
+
+# Create a network
+as100.createNetwork('net0')
+
+# Create a host in the autonomous system and connect it to the network
+host = as100.createHost('custom-host').joinNetwork('net0')
+# use the first GPU only
+host.setGPUDevices([0])
+
+# download miniconda
+host.addBuildCommand("mkdir -p ~/miniconda3 && curl https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh --output ~/miniconda3/miniconda.sh")
+# # install miniconda
+host.addBuildCommand("bash ~/miniconda3/miniconda.sh -b -u -p ~/miniconda3 && rm -rf ~/miniconda3/miniconda.sh")
+# host.addBuildCommand("curl -Ls https://micro.mamba.pm/api/micromamba/linux-64/latest | tar -xvj bin/micromamba")
+host.addBuildCommand("~/miniconda3/bin/conda install -n base conda-libmamba-solver && ~/miniconda3/bin/conda config --set solver libmamba")
+
+host.addBuildCommand("~/miniconda3/bin/conda init bash && ~/miniconda3/bin/conda init zsh")
+# for chinese users
+host.addBuildCommand("~/miniconda3/bin/conda config --add channels https://mirrors.ustc.edu.cn/anaconda/pkgs/free/ && ~/miniconda3/bin/conda config --add channels https://mirrors.ustc.edu.cn/anaconda/pkgs/main/ && ~/miniconda3/bin/conda config --set show_channel_urls yes")
+host.addBuildCommand("~/miniconda3/bin/conda config --add channels https://mirrors.ustc.edu.cn/anaconda/cloud/conda-forge/")
+# install deps
+host.addBuildCommand("CONDA_OVERRIDE_CUDA=11.8 ~/miniconda3/bin/conda install tensorflow-gpu cudnn cudatoolkit=11.8")
+# test is_gpu_available
+host.appendStartCommand("~/miniconda3/bin/conda run python3 -c \"import tensorflow as tf; print('is_gpu_available': tf.test.is_gpu_available())\"")
+
+# Bind the base layer to the emulator
+emu.addLayer(base)
+
+# Render the emulation
+emu.render()
+
+# Compile the emulation
+emu.compile(Docker(), './output', override=True)
+
+print("""
+Before running docker-compose, it is required to add GPU support to your Docker by:
+    Following Nvidia Container Toolkit Official Guide: 
+        https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/user-guide.html
+    In some linux system, you can driectly install the toolkit via community package manager (e.g. AUR in arch) 
+        https://wiki.archlinux.org/title/docker#With_NVIDIA_Container_Toolkit_.28recommended.29
+
+When the container running, you are able to run your tensorflow code with GPU
+""")

seedemu/compiler/Docker.py

@@ -131,13 +131,42 @@
             - net.ipv4.ip_forward=1
             - net.ipv4.conf.default.rp_filter=0
             - net.ipv4.conf.all.rp_filter=0
+{gpuDevices}
         privileged: true
         networks:
 {networks}{ports}{volumes}
         labels:
 {labelList}
 """
 
+# DockerCompilerFileTemplates['compose_service_resources_limits'] = """\
+#         deploy:
+#             resources:
+#                 limits:
+#                     cpus: {cpuLimit}
+#                     memory: {memLimit}
+#                     pids: {pidLimit}
+# """ # TODO: try writing with yaml lib?
+
+# DockerCompilerFileTemplates['compose_service_resources_reservations'] = """\
+#         deploy:
+#             resources:
+#                 reservations:
+#                     cpus: {cpuLimit}
+#                     memory: {memLimit}
+#                     pids: {pidLimit}
+# """ # TODO: try writing with yaml lib?
+
+DockerCompilerFileTemplates['compose_service_resources_gpus'] = """\
+        deploy:
+            resources:
+                reservations:
+                    devices:
+                      - driver: nvidia
+                        capabilities: [gpu]
+                        device_ids: [{deviceIDList}]
+"""
+
 DockerCompilerFileTemplates['compose_label_meta'] = """\
             org.seedsecuritylabs.seedemu.meta.{key}: "{value}"
 """
@@ -853,6 +882,14 @@ def _compileNode(self, node: Node) -> str:
                 netId = real_netname,
                 address = address
             )
+
+        _gpuDevices = node.getGPUDevices()
+        gpuDevices = ""
+        if len(_gpuDevices) > 0:
+            gpuDevices = DockerCompilerFileTemplates['compose_service_resources_gpus'].format(
+                deviceIDList = ', '.join(f'"{device}"' for device in _gpuDevices)
+            )
+
 
         _ports = node.getPorts()
         ports = ''
@@ -967,7 +1004,8 @@ def _compileNode(self, node: Node) -> str:
             # privileged = 'true' if node.isPrivileged() else 'false',
             ports = ports,
             labelList = self._getNodeMeta(node),
-            volumes = volumes
+            volumes = volumes,
+            gpuDevices = gpuDevices
         )
 
     def _compileNet(self, net: Network) -> str:

seedemu/core/Node.py

@@ -229,6 +229,8 @@ class Node(Printable, Registrable, Configurable, Vertex):
 
     __name_servers: List[str]
 
+    __gpus: List[str]
+
     def __init__(self, name: str, role: NodeRole, asn: int, scope: str = None):
         """!
         @brief Node constructor.
@@ -268,6 +270,8 @@ def __init__(self, name: str, role: NodeRole, asn: int, scope: str = None):
 
         self.__name_servers = []
 
+        self.__gpus = []
+
     def configure(self, emulator: Emulator):
         """!
         @brief configure the node. This is called when rendering.
@@ -779,6 +783,26 @@ def getInterfaces(self) -> List[Interface]:
         """
         return self.__interfaces
 
+    def setGPUDevices(self, devices: List[str]):
+        """!
+        @brief Set the GPU device for the node.
+
+        @param devices the new list of GPU ID.
+
+        @returns self, for chaining API calls.
+        """
+        self.__gpus = devices
+
+        return self
+
+    def getGPUDevices(self) -> List[str]:
+        """!
+        @brief Get the GPU device for the node.
+
+        @returns list of GPU ID.
+        """
+        return self.__gpus
+
     def addSharedFolder(self, nodePath: str, hostPath: str) -> Node:
         """!
         @@brief Add a new shared folder between the node and host.