summaryrefslogtreecommitdiff
path: root/buildscripts/evergreen_gen_powercycle_tasks.py
blob: 9d949c7068a53d5a5606ee57746e7e120264397f (plain)
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
111
112
113
114
115
116
117
118
119
#!/usr/bin/env python3
"""Generate multiple powercycle tasks to run in evergreen."""
from collections import namedtuple
from typing import Any, List, Tuple, Set

import click
from shrub.v2 import BuildVariant, FunctionCall, ShrubProject, Task, TaskDependency, ExistingTask
from shrub.v2.command import BuiltInCommand

from buildscripts.util.fileops import write_file
from buildscripts.util.read_config import read_config_file
from buildscripts.util.taskname import name_generated_task

Config = namedtuple("config", [
    "current_task_name",
    "task_names",
    "num_tasks",
    "timeout_params",
    "remote_credentials_vars",
    "set_up_ec2_instance_vars",
    "run_powercycle_vars",
    "build_variant",
    "distro",
])


def make_config(expansions_file: Any) -> Config:
    """Group expansions into config."""
    expansions = read_config_file(expansions_file)
    current_task_name = expansions.get("task_name", "powercycle")
    task_names = expansions.get("task_names", "powercycle_smoke_skip_compile")
    # Avoid duplicated task names
    task_names = {task_name for task_name in task_names.split(" ")}
    num_tasks = int(expansions.get("num_tasks", 10))
    timeout_params = {
        "exec_timeout_secs": int(expansions.get("exec_timeout_secs", 7200)),
        "timeout_secs": int(expansions.get("timeout_secs", 1800)),
    }
    remote_credentials_vars = {
        "private_key_file": "src/powercycle.pem",
        "private_key_remote": "${__project_aws_ssh_key_value}",
    }
    set_up_ec2_instance_vars = {
        "set_up_retry_count": int(expansions.get("set_up_retry_count", 2)),
    }
    run_powercycle_vars = {
        "run_powercycle_args": expansions.get("run_powercycle_args"),
    }
    build_variant = expansions.get("build_variant")
    distro = expansions.get("distro_id")

    return Config(current_task_name, task_names, num_tasks, timeout_params, remote_credentials_vars,
                  set_up_ec2_instance_vars, run_powercycle_vars, build_variant, distro)


def get_setup_commands() -> Tuple[List[FunctionCall], Set[TaskDependency]]:
    """Return setup commands."""
    return [
        FunctionCall("do setup"),
    ], {TaskDependency("archive_dist_test_debug")}


def get_skip_compile_setup_commands() -> Tuple[List[FunctionCall], set]:
    """Return skip compile setup commands."""
    return [
        BuiltInCommand("manifest.load", {}),
        FunctionCall("git get project"),
        FunctionCall("f_expansions_write"),
        FunctionCall("kill processes"),
        FunctionCall("cleanup environment"),
        FunctionCall("set up venv"),
        FunctionCall("upload pip requirements"),
        FunctionCall("configure evergreen api credentials"),
        FunctionCall("get compiled binaries"),
    ], set()


@click.command()
@click.argument("expansions_file", type=str, default="expansions.yml")
@click.argument("output_file", type=str, default="powercycle_tasks.json")
def main(expansions_file: str = "expansions.yml",
         output_file: str = "powercycle_tasks.json") -> None:
    """Generate multiple powercycle tasks to run in evergreen."""

    config = make_config(expansions_file)
    build_variant = BuildVariant(config.build_variant)

    sub_tasks = set()
    for task_name in config.task_names:
        if "skip_compile" in task_name:
            commands, task_dependency = get_skip_compile_setup_commands()
        else:
            commands, task_dependency = get_setup_commands()

        commands.extend([
            FunctionCall("set up remote credentials", config.remote_credentials_vars),
            BuiltInCommand("timeout.update", config.timeout_params),
            FunctionCall("set up EC2 instance", config.set_up_ec2_instance_vars),
            FunctionCall("run powercycle test", config.run_powercycle_vars),
        ])

        sub_tasks.update({
            Task(
                name_generated_task(task_name, index, config.num_tasks, config.build_variant),
                commands, task_dependency)
            for index in range(config.num_tasks)
        })

    build_variant.display_task(
        config.current_task_name.replace("_gen", ""), sub_tasks, distros=[config.distro],
        execution_existing_tasks={ExistingTask(config.current_task_name)})
    shrub_project = ShrubProject.empty()
    shrub_project.add_build_variant(build_variant)

    write_file(output_file, shrub_project.json())


if __name__ == '__main__':
    main()