From 0b3b3bb2574a4ce07a38842a87d40c6b1f40ac39 Mon Sep 17 00:00:00 2001
From: Evan Shi <14984764+evanyeyeye@users.noreply.github.com>
Date: Sun, 21 Jan 2024 14:57:44 -0500
Subject: [PATCH] Add docker resource limits for jobs (#241)

* Linting changes

* linting

* Use getattr for config
---
 config.template.py         |  4 ++++
 restful_tango/tangoREST.py | 13 ++++++++-----
 vmms/localDocker.py        |  4 ++++
 3 files changed, 16 insertions(+), 5 deletions(-)

diff --git a/config.template.py b/config.template.py
index 091fb3c8..400e51cd 100644
--- a/config.template.py
+++ b/config.template.py
@@ -86,6 +86,10 @@ class Config(object):
     DOCKER_RM_TIMEOUT = 5
     DOCKER_HOST_USER = ""
 
+    # Docker autograding container resource limits
+    DOCKER_CORES_LIMIT = None
+    DOCKER_MEMORY_LIMIT = None  # in MB
+
     # Maximum size for input files in bytes
     MAX_INPUT_FILE_SIZE = 10 * 1024 * 1024 * 1024  # 10GB
 
diff --git a/restful_tango/tangoREST.py b/restful_tango/tangoREST.py
index 31b66b41..62f3ebf5 100644
--- a/restful_tango/tangoREST.py
+++ b/restful_tango/tangoREST.py
@@ -111,16 +111,19 @@ def checkFileExists(self, directory, filename, fileMD5):
                 except IOError:
                     continue
 
-    def createTangoMachine(
-        self, image, vmms=Config.VMMS_NAME, vmObj={"cores": 1, "memory": 512}
-    ):
+    def createTangoMachine(self, image, vmms=Config.VMMS_NAME, vmObj=None):
         """createTangoMachine - Creates a tango machine object from image"""
+        cores = getattr(Config, "DOCKER_CORES_LIMIT", None)
+        memory = getattr(Config, "DOCKER_MEMORY_LIMIT", None)
+        if vmObj and "cores" in vmObj and "memory" in vmObj:
+            cores = vmObj["cores"]
+            memory = vmObj["memory"]
         return TangoMachine(
             name=image,
             vmms=vmms,
             image="%s" % (image),
-            cores=vmObj["cores"],
-            memory=vmObj["memory"],
+            cores=cores,
+            memory=memory,
             disk=None,
             network=None,
         )
diff --git a/vmms/localDocker.py b/vmms/localDocker.py
index c9385f4b..6557237b 100644
--- a/vmms/localDocker.py
+++ b/vmms/localDocker.py
@@ -157,6 +157,10 @@ def runJob(self, vm, runTimeout, maxOutputFileSize):
             )
         args = ["docker", "run", "--name", instanceName, "-v"]
         args = args + ["%s:%s" % (volumePath, "/home/mount")]
+        if vm.cores:
+            args = args + [f"--cpus={vm.cores}"]
+        if vm.memory:
+            args = args + ["-m", f"{vm.memory}m"]
         args = args + [vm.image]
         args = args + ["sh", "-c"]