From 7966f52faf870fee134033cca38e0eadd4144a98 Mon Sep 17 00:00:00 2001
From: you-n-g <you-n-g@users.noreply.github.com>
Date: Sat, 21 Sep 2024 21:31:56 +0800
Subject: [PATCH] update lint (#290)

---
 rdagent/utils/env.py   |  2 ++
 test/utils/test_env.py | 22 +++++++++++++++++++---
 2 files changed, 21 insertions(+), 3 deletions(-)

diff --git a/rdagent/utils/env.py b/rdagent/utils/env.py
index 5e2524fd..85cca343 100644
--- a/rdagent/utils/env.py
+++ b/rdagent/utils/env.py
@@ -132,6 +132,7 @@ class DockerConf(BaseSettings):
     network: str | None = "bridge"  # the network mode for the docker
     shm_size: str | None = None
     enable_gpu: bool = True  # because we will automatically disable GPU if not available. So we enable it by default.
+    mem_limit: str | None = "48g"  # Add memory limit attribute
 
 
 class QlibDockerConf(DockerConf):
@@ -300,6 +301,7 @@ class DockerEnv(Env[DockerConf]):
                 # auto_remove=True, # remove too fast might cause the logs not to be get
                 network=self.conf.network,
                 shm_size=self.conf.shm_size,
+                mem_limit=self.conf.mem_limit,  # Set memory limit
                 **self._gpu_kwargs(client),
             )
             logs = container.logs(stream=True)
diff --git a/test/utils/test_env.py b/test/utils/test_env.py
index 3572a5ee..54d005e5 100644
--- a/test/utils/test_env.py
+++ b/test/utils/test_env.py
@@ -6,7 +6,7 @@ from pathlib import Path
 sys.path.append(str(Path(__file__).resolve().parent.parent))
 import shutil
 
-from rdagent.utils.env import LocalConf, LocalEnv, QTDockerEnv
+from rdagent.utils.env import LocalConf, LocalEnv, QlibDockerConf, QTDockerEnv
 
 DIRNAME = Path(__file__).absolute().resolve().parent
 
@@ -37,8 +37,7 @@ class EnvUtils(unittest.TestCase):
         self.assertTrue(mlrun_p.exists(), f"Expected output file {mlrun_p} not found")
 
     def test_docker(self):
-        """
-        We will mount `env_tpl` into the docker image.
+        """We will mount `env_tpl` into the docker image.
         And run the docker image with `qrun conf.yaml`
         """
         qtde = QTDockerEnv()
@@ -54,6 +53,23 @@ class EnvUtils(unittest.TestCase):
         result = qtde.run(local_path=str(DIRNAME / "env_tpl"), entry="python read_exp_res.py")
         print(result)
 
+    def test_docker_mem(self):
+        cmd = 'python -c \'print("start"); import numpy as np;  size_mb = 500; size = size_mb * 1024 * 1024 // 8; array = np.random.randn(size).astype(np.float64); print("success")\''
+
+        qtde = QTDockerEnv(QlibDockerConf(mem_limit="10m"))
+        qtde.prepare()
+        result = qtde.run(local_path=str(DIRNAME / "env_tpl"), entry=cmd)
+        self.assertTrue(not result.strip().endswith("success"))
+
+        qtde = QTDockerEnv(QlibDockerConf(mem_limit="1g"))
+        qtde.prepare()
+        result = qtde.run(local_path=str(DIRNAME / "env_tpl"), entry=cmd)
+        self.assertTrue(result.strip().endswith("success"))
+
+        # The above command equals to the follow commands with dockr cli.sh
+        # docker run  --memory=10m  -it --rm local_qlib:latest python -c 'import numpy as np; print(123);  size_mb = 1; size = size_mb * 1024 * 1024 // 8; array = np.random.randn(size).astype(np.float64); array[0], array[-1] = 1.0, 1.0; print(321)'
+        # docker run  --memory=10g  -it --rm local_qlib:latest python -c 'import numpy as np; print(123);  size_mb = 1; size = size_mb * 1024 * 1024 // 8; array = np.random.randn(size).astype(np.float64); array[0], array[-1] = 1.0, 1.0; print(321)'
+
 
 if __name__ == "__main__":
     unittest.main()