From dd768083510553ff1cb79d5eb564aa0a93efb668 Mon Sep 17 00:00:00 2001
From: Wey Gu <weyl.gu@gmail.com>
Date: Fri, 12 Jan 2024 18:53:30 +0800
Subject: [PATCH] feat: support running inside a container

cloase: #10
---
 README.md                           | 20 ++++++++++++++++++++
 src/nebulagraph_lite/nebulagraph.py | 25 +++++++++++++++++++++++--
 2 files changed, 43 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index 8ee7984..76c8a6b 100644
--- a/README.md
+++ b/README.md
@@ -82,6 +82,26 @@ With udocker, the opinionated subset docker running in user space, we could run
 
 Thus we support running inside docker container, WSL2, Google Colab.
 
+### Can NebulaGraph-Lite run inside a container?
+
+Yes! Say we are in a container that runs Ubuntu, we could run:
+
+```bash
+docker run -it --rm ubuntu:latest bash
+# inside the container
+apt update && apt install python3-pip curl -y
+pip3 install nebulagraph-lite
+python3
+```
+
+In python3:
+
+```python
+from nebulagraph_lite import nebulagraph_let as ng_let
+n = ng_let(in_container=True)
+n.start()
+```
+
 ### Does it support Windows?
 
 Yes, it supports Windows with WSL2 or other Linux VMs with a Hypervisor.
diff --git a/src/nebulagraph_lite/nebulagraph.py b/src/nebulagraph_lite/nebulagraph.py
index 7849cf4..b570de5 100644
--- a/src/nebulagraph_lite/nebulagraph.py
+++ b/src/nebulagraph_lite/nebulagraph.py
@@ -23,6 +23,7 @@ def __init__(
         base_path: str = BASE_PATH,
         debug=False,
         clean_up=False,
+        in_container=False,
     ):
         if clean_up:
             self.clean_up()
@@ -52,6 +53,20 @@ def __init__(
             self._python_bin_path = os.path.dirname(_path[0])
         else:
             self._python_bin_path = os.path.dirname(os.sys.executable)
+            result = subprocess.run(
+                f"{self._python_bin_path}/udocker --help",
+                shell=True,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+            )
+            if result.returncode != 0:
+                _path = subprocess.getoutput("which udocker")
+                if _path:
+                    self._python_bin_path = os.path.dirname(_path)
+                else:
+                    raise Exception(
+                        "udocker not found. Please install or link it manually to your PATH."
+                    )
 
         self._debug = debug
 
@@ -59,6 +74,8 @@ def __init__(
         if self.on_colab:
             self.base_path = COLAB_BASE_PATH
 
+        self.in_container = in_container
+
         self.create_nebulagraph_lite_folders()
 
     def _is_running_on_colab(self):
@@ -137,6 +154,8 @@ def _run_udocker(self, command: str):
         if self.on_colab:
             return self._run_udocker_on_colab(command)
         udocker_command_prefix = os.path.join(self._python_bin_path, "udocker")
+        if self.in_container:
+            udocker_command_prefix = udocker_command_prefix + " --allow-root"
         udocker_command = f"{udocker_command_prefix} {command}"
         result = subprocess.run(
             udocker_command,
@@ -179,6 +198,8 @@ def _run_udocker_background(self, command: str):
             return
 
         udocker_command_prefix = os.path.join(self._python_bin_path, "udocker")
+        if self.in_container:
+            udocker_command_prefix = udocker_command_prefix + " --allow-root"
         udocker_command = f"{udocker_command_prefix} {command} &"
         subprocess.Popen(
             udocker_command,
@@ -199,7 +220,7 @@ def udocker_pull_backgroud(self, image: str):
     def _try_shoot_service(self, service: str):
         try:
             self._run_udocker(
-                f"ps | grep {service} | awk '{{print $1}}' | xargs -I {{}} udocker rm {{}}"
+                f"ps | grep {service} | awk '{{print $1}}' | xargs -I {{}} udocker rm -f {{}}"
             )
             os.system(f"killall nebula-{service} > /dev/null 2>&1")
         except Exception as e:
@@ -347,7 +368,7 @@ def docker_ps(self):
     def stop(self):
         if self.on_colab:
             self._run_udocker(
-                "ps | grep nebula | awk '{print $1}' | xargs -I {} udocker rm {}"
+                "ps | grep nebula | awk '{print $1}' | xargs -I {} udocker rm -f {}"
             )
             self._try_shoot_all_services()
             return