BITMAKER-1901: Allow setting RAM limits to SpiderJobs

estela_cli/create/cronjob.py

@@ -3,12 +3,12 @@
 from estela_cli.login import login
 from estela_cli.utils import (
     get_estela_settings,
-    validate_key_value_format,
     set_tag_format,
+    validate_key_value_format,
+    validate_limit,
     validate_positive,
 )
 
-
 SHORT_HELP = "Create a new cronjob"
 
 
@@ -32,6 +32,13 @@
     callback=validate_key_value_format,
     help="Set spider cronjob environment variable NAME=VALUE (may be repeated)",
 )
+@click.option(
+    "--memory",
+    "-m",
+    type=click.STRING,
+    callback=validate_limit,
+    help="Set spider job memory limit (e.g. 256Mi, 1Gi).",
+)
 @click.option(
     "--tag",
     "-t",
@@ -47,7 +54,7 @@
     callback=validate_positive,
     help="Set spider cronjob data expiry days",
 )
-def estela_command(sid, pid, schedule, arg, env, tag, day):
+def estela_command(sid, pid, schedule, arg, env, memory, tag, day):
     """Create a new cronjob
 
     \b
@@ -67,7 +74,7 @@ def estela_command(sid, pid, schedule, arg, env, tag, day):
             )
     try:
         response = estela_client.create_spider_cronjob(
-            pid, sid, schedule, arg, env, tag, day
+            pid, sid, schedule, arg, env, memory, tag, day
         )
         click.echo("cronjob/{} created.".format(response["name"]))
     except Exception as ex:

estela_cli/create/job.py

@@ -3,12 +3,12 @@
 from estela_cli.login import login
 from estela_cli.utils import (
     get_estela_settings,
-    validate_key_value_format,
     set_tag_format,
+    validate_key_value_format,
+    validate_limit,
     validate_positive,
 )
 
-
 SHORT_HELP = "Create a new job"
 
 
@@ -31,6 +31,13 @@
     callback=validate_key_value_format,
     help="Set spider job environment variable NAME=VALUE (may be repeated)",
 )
+@click.option(
+    "--memory",
+    "-m",
+    type=click.STRING,
+    callback=validate_limit,
+    help="Set spider job memory limit (e.g. 256Mi, 1Gi).",
+)
 @click.option(
     "--tag",
     "-t",
@@ -46,7 +53,7 @@
     callback=validate_positive,
     help="Set spider job data expiry days",
 )
-def estela_command(sid, pid, arg, env, tag, day):
+def estela_command(sid, pid, arg, env, memory, tag, day):
     """Create a new job
 
     \b
@@ -64,7 +71,7 @@ def estela_command(sid, pid, arg, env, tag, day):
                 "No active project in the current directory. Please specify the PID."
             )
     try:
-        response = estela_client.create_spider_job(pid, sid, arg, env, tag, day)
+        response = estela_client.create_spider_job(pid, sid, arg, env, memory, tag, day)
         click.echo("job/{} created.".format(response["name"]))
     except Exception as ex:
         raise click.ClickException("Cannot create the job for given SID and PID.")

estela_cli/estela_client.py

@@ -181,7 +181,9 @@ def get_spider_job_data(self, pid, sid, jid, last_chunk=None):
         self.check_status(response, 200)
         return response.json()
 
-    def create_spider_job(self, pid, sid, args=[], env_vars=[], tags=[], day=None):
+    def create_spider_job(
+        self, pid, sid, args=[], env_vars=[], memory=None, tags=[], day=None
+    ):
         endpoint = "projects/{}/spiders/{}/jobs".format(pid, sid)
         data = {
             "args": args,
@@ -191,6 +193,8 @@ def create_spider_job(self, pid, sid, args=[], env_vars=[], tags=[], day=None):
         }
         if day:
             data["data_expiry_days"] = f"{date.today() + timedelta(days=day)}"
+        if memory:
+            data["limits"] = {"memory": memory}
 
         response = self.post(endpoint, data=data)
         self.check_status(response, 201)
@@ -203,7 +207,7 @@ def stop_spider_job(self, pid, sid, jid):
         return response.json()
 
     def create_spider_cronjob(
-        self, pid, sid, schedule="", args=[], env_vars=[], tags=[], day=None
+        self, pid, sid, schedule="", args=[], env_vars=[], memory=None, tags=[], day=None
     ):
         endpoint = "projects/{}/spiders/{}/cronjobs".format(pid, sid)
         data = {
@@ -215,6 +219,8 @@ def create_spider_cronjob(
         }
         if day:
             data["data_expiry_days"] = f"0/{day}"
+        if memory:
+            data["limits"] = {"memory": memory}
 
         response = self.post(endpoint, data=data)
         self.check_status(response, 201)

estela_cli/list/cronjob.py

@@ -62,9 +62,10 @@ def estela_command(sid, pid, tag):
             format_tags(cronjob["ctags"]),
             format_key_value_pairs(cronjob["cargs"]),
             format_key_value_pairs(cronjob["cenv_vars"]),
+            cronjob["limits"].get("memory", 0),
         ]
         for cronjob in cronjobs
     ]
 
-    headers = ["CJID", "STATUS", "SCHEDULE", "TAGS", "ARGS", "ENV VARS"]
+    headers = ["CJID", "STATUS", "SCHEDULE", "TAGS", "ARGS", "ENV VARS", "MEMORY_LIMIT"]
     click.echo(tabulate(cronjobs, headers, numalign="left", tablefmt="plain"))

estela_cli/list/job.py

@@ -61,10 +61,11 @@ def estela_command(sid, pid, tag):
             format_tags(job["tags"]),
             format_key_value_pairs(job["args"]),
             format_key_value_pairs(job["env_vars"]),
+            job["limits"].get("memory", 0),
             format_time(job["created"]),
         ]
         for job in jobs
     ]
 
-    headers = ["JID", "STATUS", "TAGS", "ARGS", "ENV VARS", "CREATED"]
+    headers = ["JID", "STATUS", "TAGS", "ARGS", "ENV VARS", "MEMORY_LIMIT", "CREATED"]
     click.echo(tabulate(jobs, headers, numalign="left", tablefmt="plain"))

estela_cli/utils.py

@@ -1,4 +1,5 @@
 import os
+import re
 import yaml
 import csv
 import json
@@ -148,6 +149,16 @@ def validate_positive(ctx, param, value):
     return value
 
 
+def validate_limit(ctx, param, value):
+    if value is not None:
+        match = re.match(r"^(\d+)(Mi|Gi)$", value)
+        if not match:
+            raise click.BadParameter(
+                "must be in a valid limit format (e.g. 256Mi, 1Gi)."
+            )
+    return value
+
+
 def set_tag_format(ctx, param, value):
     tags = []
     for tag_name in value: