From 267d1086bc59b7f6458a3d7aadf0f7e8ebab900e Mon Sep 17 00:00:00 2001
From: kshitijrajsharma <skshitizraj@gmail.com>
Date: Thu, 10 Aug 2023 15:32:38 +0545
Subject: [PATCH] Added validation for epochs and batchsize

---
 backend/aiproject/settings.py |  5 +++++
 backend/core/views.py         | 12 ++++++++++++
 2 files changed, 17 insertions(+)

diff --git a/backend/aiproject/settings.py b/backend/aiproject/settings.py
index 22de17ef..c4824a0c 100644
--- a/backend/aiproject/settings.py
+++ b/backend/aiproject/settings.py
@@ -53,6 +53,11 @@
 OSM_SECRET_KEY = env("OSM_SECRET_KEY")
 
 
+# Limiter
+EPOCHS_LIMIT = env("EPOCHS_LIMIT", default=30)
+BATCH_SIZE_LIMIT = env("BATCH_SIZE_LIMIT", default=8)
+
+
 # Application definition
 
 INSTALLED_APPS = [
diff --git a/backend/core/views.py b/backend/core/views.py
index ddbf8357..500df8d9 100644
--- a/backend/core/views.py
+++ b/backend/core/views.py
@@ -118,6 +118,18 @@ def create(self, validated_data):
                 "Another training is already running or submitted for this model."
             )
 
+        epochs = validated_data["epochs"]
+        batch_size = validated_data["batch_size"]
+
+        if epochs > settings.EPOCHS_LIMIT:
+            raise ValidationError(
+                f"Epochs can't be greater than {settings.EPOCHS_LIMIT} on this server"
+            )
+        if batch_size > settings.BATCH_SIZE_LIMIT:
+            raise ValidationError(
+                f"Batch size can't be greater than {settings.BATCH_SIZE_LIMIT} on this server"
+            )
+
         user = self.context["request"].user
         validated_data["created_by"] = user
         # create the model instance