File: wq_alloc_test.py

package info (click to toggle)
cctools 1%3A7.14.5-1
  • links: PTS, VCS
  • area: main
  • in suites: trixie
  • size: 36,956 kB
  • sloc: ansic: 114,614; python: 29,532; cpp: 20,313; sh: 13,675; perl: 4,056; xml: 3,688; makefile: 1,436
file content (130 lines) | stat: -rwxr-xr-x 4,700 bytes parent folder | download | duplicates (2)
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
120
121
122
123
124
125
126
127
128
129
130
#! /usr/bin/env python

# work queue python binding tests
# tests for missing/recursive inputs/outputs.

import sys
import math
import ndcctools.work_queue as wq

def check_task(category, category_mode, max, min, expected):
    q.specify_category_max_resources(category, max)
    q.specify_category_min_resources(category, min)
    q.specify_category_mode(category, category_mode)

    t = wq.Task('/bin/echo hello')
    t.specify_category(category)
    q.submit(t)

    t = q.wait(30)
    if not t:
        print("Task did not complete in the allotted time.")
        sys.exit(1)
    else:
        print("{}:".format(category))
        print("expected: {}".format(expected))
        print("allocated: {}".format({
            'cores': t.resources_allocated.cores,
            'memory': t.resources_allocated.memory,
            'disk': t.resources_allocated.disk,
            'gpus': t.resources_allocated.gpus}))

        assert(t.resources_allocated.cores == expected['cores'])
        assert(t.resources_allocated.memory == expected['memory'])
        assert(t.resources_allocated.disk == expected['disk'])
        assert(t.resources_allocated.gpus == expected['gpus'])


port_file = None
try:
    port_file = sys.argv[1]
    (worker_cores, worker_memory, worker_disk, worker_gpus) = (int(sys.argv[i]) for i in range(2,6))
except IndexError:
    sys.stderr.write("Usage: {} PORTFILE WORKER_CORES WORKER_MEMORY WORKER_DISK\n".format(sys.argv[0]))
    raise

q = wq.WorkQueue(0)
with open(port_file, 'w') as f:
    print('Writing port {port} to file {file}'.format(port=q.port, file=port_file))
    f.write(str(q.port))

worker = wq.Factory("local", manager_host_port="localhost:{}".format(q.port))
worker.max_workers = 1
worker.min_workers = 1

worker.cores=worker_cores
worker.memory=worker_memory
worker.disk=worker_disk
worker.gpus=worker_gpus

worker.debug="all"

with worker:
    q.tune("force-proportional-resources", 0)

    r = {'cores': 1, 'memory': 2, 'disk': 3, 'gpus': 4}
    check_task('all_specified', wq.WORK_QUEUE_ALLOCATION_MODE_FIXED, max = r, min = {}, expected = r)

    check_task('all_specified_no_gpu',
            wq.WORK_QUEUE_ALLOCATION_MODE_FIXED,
            max = {'cores': 1, 'memory': 2, 'disk': 3},
            min = {},
            expected = {'cores': 1, 'memory': 2, 'disk': 3, 'gpus': 0})

    check_task('all_specified_no_cores',
            wq.WORK_QUEUE_ALLOCATION_MODE_FIXED,
            max = {'gpus': 4, 'memory': 2, 'disk': 3},
            min = {},
            expected = {'cores': 0, 'memory': 2, 'disk': 3, 'gpus': 4})

    check_task('all_zero',
            wq.WORK_QUEUE_ALLOCATION_MODE_FIXED,
            max = {'cores': 0, 'memory': 0, 'disk': 0, 'gpus': 0},
            min = {},
            expected = {'cores': worker_cores, 'memory': worker_memory, 'disk': worker_disk, 'gpus': 0})

    q.tune("force-proportional-resources", 1)
    check_task('only_memory',
            wq.WORK_QUEUE_ALLOCATION_MODE_FIXED,
            max = {'memory': worker_memory/2},
            min = {},
            expected = {'cores': worker_cores/2, 'memory': worker_memory/2, 'disk': worker_disk/2, 'gpus': 0})

    check_task('only_memory_w_minimum',
            wq.WORK_QUEUE_ALLOCATION_MODE_FIXED,
            max = {'memory': worker_memory/2},
            min = {'cores': 3, 'gpus': 2},
            expected = {'cores': 3, 'memory': worker_memory/2, 'disk': worker_disk/2, 'gpus': 2})

    check_task('only_cores',
            wq.WORK_QUEUE_ALLOCATION_MODE_FIXED,
            max = {'cores': worker_cores},
            min = {},
            expected = {'cores': worker_cores, 'memory': worker_memory, 'disk': worker_disk, 'gpus': 0})

    check_task('auto_whole_worker',
            wq.WORK_QUEUE_ALLOCATION_MODE_MIN_WASTE,
            max = {},
            min = {},
            expected = {'cores': worker_cores, 'memory': worker_memory, 'disk': worker_disk, 'gpus': 0})

    p = 1/worker_cores
    r = {'cores': 1}
    e = {'cores': 1, 'memory': math.floor(worker_memory * p), 'disk': math.floor(worker_disk * p), 'gpus': 0}
    check_task('only_cores_proportional', wq.WORK_QUEUE_ALLOCATION_MODE_FIXED, max = r, min = {}, expected = e)

    p = 2/worker_cores
    e = {'cores': 2, 'memory': math.floor(worker_memory * p), 'disk': math.floor(worker_disk * p), 'gpus': 0}
    check_task('exhaustive_bucketing',
            wq.WORK_QUEUE_ALLOCATION_MODE_EXHAUSTIVE_BUCKETING,
            max = {},
            min = {},
            expected = e)

    check_task('greedy_bucketing',
            wq.WORK_QUEUE_ALLOCATION_MODE_GREEDY_BUCKETING,
            max = {},
            min = {},
            expected = e)

# vim: set sts=4 sw=4 ts=4 expandtab ft=python: