-
Notifications
You must be signed in to change notification settings - Fork 0
/
submit-job.py
110 lines (94 loc) · 2.89 KB
/
submit-job.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
import os
from datetime import datetime
import boto3
import click
params = {}
region = os.getenv("AWS_REGION", os.getenv("AWS_DEFAULT_REGION"))
profile = os.getenv("AWS_PROFILE")
if region:
params["region_name"] = region
if profile:
boto3.setup_default_session(profile_name=profile, **params)
client = boto3.client("batch", **params)
def _parse_env(env: list):
environment = []
for env_var in env:
if "=" not in env_var:
raise ValueError(
"Environment variable is not properly formed. Use --env 'KEY=value'"
)
key, value = env_var.split("=", maxsplit=1)
environment.append({"name": key, "value": value})
return environment
def _submit_job(job_queue: str, job_definition: str, job_details: dict, environ: list):
experiment_conf = job_details["path"]
cur_time = datetime.now().strftime("%Y-%m-%d_%H%M%S")
experiment_last_path = experiment_conf.split("/")[-2]
name = f"{cur_time}-{experiment_last_path}"
environment = [
{
"name": "CHAOS_CONFIG_FILE",
"value": experiment_conf,
},
]
environment.extend(environ)
client.submit_job(
jobName=name,
jobQueue=job_queue,
jobDefinition=job_definition,
containerOverrides={
"environment": environment,
"resourceRequirements": [
{"value": str(job_details.get("vcpu", 1)), "type": "VCPU"},
{"value": str(job_details.get("memory", 2048)), "type": "MEMORY"},
],
},
timeout={"attemptDurationSeconds": 3600},
)
@click.command()
@click.option(
"--vcpu",
type=click.FLOAT,
default=1.0,
help="The number of vcpus for the chaos job",
)
@click.option(
"--memory",
type=click.INT,
default=2048,
help="The memory (in MB) for the chaos job",
)
@click.option(
"--env",
"-e",
multiple=True,
help="Override environment variable for job execution. Format: --env 'KEY=value'",
)
@click.option(
"--queue",
required=True,
help="The name of the AWS batch job queue",
)
@click.option(
"--job-definition",
required=True,
help="The name of the AWS batch job definition",
)
@click.argument("experiment-configurations", nargs=-1)
def cli(vcpu, memory, env, queue, job_definition, experiment_configurations):
# pylint: disable=too-many-arguments
for experiment_config_file in list(experiment_configurations):
print(f"Submitting experiment with config file: {experiment_config_file}")
job_details = {
"path": experiment_config_file,
"vcpu": vcpu,
"memory": memory,
}
_submit_job(
job_queue=queue,
job_definition=job_definition,
job_details=job_details,
environ=_parse_env(env),
)
if __name__ == "__main__":
cli() # pylint: disable=no-value-for-parameter