Modify sim-baobab connector to allow for parallelization

When initiating a new sim environment, we create also a baobab instance with unique namespace for environment. This allows to properly run multiple sim in parallel.
Azure · Aug 14, 2023 · 04c7c5d · 04c7c5d
1 parent f2952b5
commit 04c7c5d
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 13 deletions.
diff --git a/examples/getting-started-anylogic-sim/src/sim.py b/examples/getting-started-anylogic-sim/src/sim.py
@@ -1,4 +1,5 @@
 import json
+import os
 import subprocess
 import time
 from pathlib import Path
@@ -9,12 +10,13 @@
 from platotk.logger import log
 from platotk.serialize import GymEncoder, check_and_transform
 
-BASE_URL = "http://localhost:8000"
-
 
 class SimWrapper(Env):
     def __init__(self, env_config):
-        self.base_url = BASE_URL
+        self.base_host = "localhost"
+        self.env_id = int(f"{env_config.worker_index}{env_config.vector_index}")
+        self.base_port = 8000 + self.env_id
+        self.base_url = f"http://{self.base_host}:{self.base_port}"
 
         self.action_space = spaces.Dict(
             {
@@ -29,24 +31,47 @@ def __init__(self, env_config):
             "arrivalRate": 0.5,
             "sizeBufferQueues": 45,
         }
-        self.start_anylogic_sim()
+        self.start_sim_framework()
+
+    @staticmethod
+    def find_unique_port(worker, vector):
+        return 8000 + int(f"{worker}{vector}")
 
-    def start_anylogic_sim(self):
+    def start_sim_framework(self):
+        """Start Baobab API and external sim."""
+
+        # Find the sim executable
         scripts = [script for script in Path(__file__).parent.rglob("*_linux.sh")]
         if len(scripts) > 1:
             raise RuntimeError(f"Too many Anylogic sims found: {scripts}")
         elif len(scripts) < 1:
             raise RuntimeError("No Anylogic sim found.")
-
         sim_exec = scripts.pop()
+
+        os.environ["BAOBAB_NAMESPACE"] = str(self.env_id)
+
+        # Launch Baobab
+        subprocess.Popen(
+            [
+                "gunicorn",
+                "--worker-class",
+                "uvicorn.workers.UvicornWorker",
+                "--bind",
+                f"{self.base_host}:{self.base_port}",
+                "platotk.baobab:app",
+            ]
+        )
+        time.sleep(2)
+
+        # Launch the sim that will connect to Baobab
         penv = {
-            "SIM_API_HOST": "http://localhost:8000",
+            "SIM_API_HOST": self.base_url,
             "SIM_CONTEXT": "{}",
             "SIM_WORKSPACE": "dummy",
             "SIM_ACCESS_KEY": "dummy",
         }
         subprocess.Popen([sim_exec], env=penv)
-        time.sleep(10)
+        time.sleep(5)
 
     def reset(self, *, seed=None, options=None):
         log.debug("Reset send.")

diff --git a/examples/getting-started-anylogic-sim/src/start.sh b/examples/getting-started-anylogic-sim/src/start.sh
@@ -2,9 +2,4 @@
 set -xe
 memcached -p 11211 -u memcache -l 127.0.0.1 -d
 
-gunicorn \
-	--worker-class uvicorn.workers.UvicornWorker \
-	--bind '127.0.0.1:8000' \
-	platotk.baobab:app &
-
 python -u main.py --test-local