diff --git a/estela_cli/create/cronjob.py b/estela_cli/create/cronjob.py index 2f71147..fcbcdd5 100644 --- a/estela_cli/create/cronjob.py +++ b/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: diff --git a/estela_cli/create/job.py b/estela_cli/create/job.py index a1ac88b..a46abdd 100644 --- a/estela_cli/create/job.py +++ b/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.") diff --git a/estela_cli/estela_client.py b/estela_cli/estela_client.py index a9b4ac1..149de66 100644 --- a/estela_cli/estela_client.py +++ b/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) diff --git a/estela_cli/list/cronjob.py b/estela_cli/list/cronjob.py index e4bc36f..ff64a2a 100644 --- a/estela_cli/list/cronjob.py +++ b/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")) diff --git a/estela_cli/list/job.py b/estela_cli/list/job.py index b2b1a87..2144638 100644 --- a/estela_cli/list/job.py +++ b/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")) diff --git a/estela_cli/utils.py b/estela_cli/utils.py index 132efa3..52d7e2d 100644 --- a/estela_cli/utils.py +++ b/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: