task.py 2.98 KB
Newer Older
William Ellis Warriner's avatar
William Ellis Warriner committed
1
2
3
4
5
6
7
8
9
10
11
12
"""
task.py

Written by William Warriner 2021
wwarr@uab.edu

Contains example process code which we want to checkpoint. The general workflow
is to loop through a sample computation which sleeps a random amount of time and
returns a random float. Both values are returned in a dict and printed to
stdout. This process is repeated
"""

William Ellis Warriner's avatar
William Ellis Warriner committed
13
14
15
import argparse
import random
import sys
William Ellis Warriner's avatar
William Ellis Warriner committed
16
17
import time
from typing import Dict
William Ellis Warriner's avatar
William Ellis Warriner committed
18
19

VALUE = "value"
William Ellis Warriner's avatar
William Ellis Warriner committed
20
SLEEP_TIME = "sleep_time"
William Ellis Warriner's avatar
William Ellis Warriner committed
21
22


William Ellis Warriner's avatar
William Ellis Warriner committed
23
24
25
26
27
28
29
30
31
32
def print_step(current_step: int, data: Dict[str, float]) -> None:
    """
    Prints the data dict containing the value and the sleep time, with an
    iteration count. We flush the buffer each time a line is printed to ensure
    real-time updating. We can (and do) use the `-u` argument, i.e. unbuffered
    streams, with the Python interpreter for the same effect.
    """
    print(
        f"Step {current_step+1: >4}: {data[VALUE]: >6.2f} ({data[SLEEP_TIME]: >6.2f}s)"
    )
William Ellis Warriner's avatar
William Ellis Warriner committed
33
34
35
36
    sys.stdout.flush()


def step() -> Dict[str, float]:
William Ellis Warriner's avatar
William Ellis Warriner committed
37
38
39
40
41
42
43
    """
    The kernel computation of interest which is run once per loop step. Computes
    a value using a lognormal distribution. Computes a sleep time using a
    triangular distribution to ensure non-negative values with a central
    tendency, and then sleeps for that amount of time. A dict is returned
    containing both.
    """
William Ellis Warriner's avatar
William Ellis Warriner committed
44
    value = random.lognormvariate(mu=0.0, sigma=1.0)
William Ellis Warriner's avatar
William Ellis Warriner committed
45
46
47
    sleep_time = abs(random.triangular(low=0.0, high=1, mode=0.5))
    time.sleep(sleep_time)
    return {VALUE: value, SLEEP_TIME: sleep_time}
William Ellis Warriner's avatar
William Ellis Warriner committed
48
49
50


def task(step_count: int) -> None:
William Ellis Warriner's avatar
William Ellis Warriner committed
51
52
53
54
55
56
    """
    The primary kernel loop. Runs the kernel `step()` function, then reports the
    values using `print_step()`. The only input is a positive integer which
    limits how many times the loop is executed.
    """
    assert 0 < step_count
William Ellis Warriner's avatar
William Ellis Warriner committed
57
58
59
60
61
62
    for current_step in range(step_count):
        data = step()
        print_step(current_step, data)


def interface() -> None:
William Ellis Warriner's avatar
William Ellis Warriner committed
63
64
65
66
67
68
69
    """
    The external interface. Prepares an argument parser accepting a positive
    integer steps. Steps is clipped to 100 for sanity.
    """
    parser = argparse.ArgumentParser(
        description="Run kernel function and report values. Intended for use as part of DMTCP checkpointing with SLURM tutorial on Cheaha."
    )
William Ellis Warriner's avatar
William Ellis Warriner committed
70
    parser.add_argument(
William Ellis Warriner's avatar
William Ellis Warriner committed
71
72
73
74
75
76
        "steps",
        metavar="N",
        nargs="?",
        default=100,
        type=_check_positive,
        help="positive integer number of steps to run, clipped to 100",
William Ellis Warriner's avatar
William Ellis Warriner committed
77
78
    )
    args = parser.parse_args()
William Ellis Warriner's avatar
William Ellis Warriner committed
79
    steps = args.steps
William Ellis Warriner's avatar
William Ellis Warriner committed
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
    steps = min(steps, 100)
    task(steps)


def _check_positive(value: str) -> int:
    """
    Checks if a string supplied to argparse is a positive integer. Raises
    argparse.ArgumentTypeError if not.
    """
    try:
        ivalue = int(value)
    except ValueError as e:
        raise argparse.ArgumentTypeError(f"{value} must be an integer")
    if ivalue <= 0:
        raise argparse.ArgumentTypeError(f"{value} must be positive")
    return ivalue
William Ellis Warriner's avatar
William Ellis Warriner committed
96
97
98
99


if __name__ == "__main__":
    interface()