cmd_queue/client.py at master · sjTaylor/cmd_queue · GitHub

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
import socket
import select
import codes
import funs
import os
import subprocess
import argparse

parser = argparse.ArgumentParser()
parser.add_argument('server_ip', type=str, help='address of the server (e.g. 198.123.1.3)')
parser.add_argument('--server_port', type=int, default=12345, required=False, help='port server is listening on')
args = parser.parse_args()

myid = 0

output_prefix = None
padding = None

try:
    connection = socket.create_connection((args.server_ip, args.server_port))
except:
    print('could not connect to server', flush=True)
    raise SystemError

running = True
cmd_timeout = None

log = funs.get_logger(__name__)

while running:
    readable, foo1, foo2 = select.select([connection], [], [], 2)

    for qq in funs.getinput():
        if 'exit' in qq:
            running = False
            funs.send(connection, funs.encode(codes.disconnecting))

    for s in readable:
        message = funs.recv(s)
        code, data = funs.decode(message)

        json_data = data

        if code == codes.send_config:
            assert 'client_id' in json_data and 'working_directory' in json_data and 'output_prefix' in json_data
            assert 'padding' in json_data and 'timeout' in json_data
            os.chdir(json_data['working_directory'])
            myid = json_data['client_id']
            output_prefix = json_data['output_prefix']
            padding = json_data['padding']
            cmd_timeout = json_data['timeout']

        elif code == codes.send_cmd:
            assert 'command' in json_data and 'cmd_number' in json_data
            command = json_data['command']
            cmdnumber = json_data['cmd_number']
            log.info('Recieved command number : %d' % cmdnumber)
            log.info('--executing : %s' % command)

            log.info('will write out to: |%s|' % funs.do_dir(output_prefix, padding, 'stdout', cmdnumber))
            log.info('will write err to: |%s|' % funs.do_dir(output_prefix, padding, 'stderr', cmdnumber))

            with open(funs.do_dir(output_prefix, padding, 'stdout', cmdnumber), 'w') as sstdout:
                with open(funs.do_dir(output_prefix, padding, 'stderr', cmdnumber), 'w') as sstderr:
                    return_code = subprocess.call(command,
                                                  shell=True,
                                                  stdout=sstdout,
                                                  stderr=sstderr,
                                                  timeout=cmd_timeout)
            if return_code is None:
                log.info('--return_code is None')
                return_code = 1
            # cmd number, client id, return code
            funs.send(connection, funs.encode(codes.finished, (cmdnumber, myid, return_code)))
        if code == codes.exiting:
            log.info('got signal to stop and shut down')
            running = False
        else:
            funs.send(connection, funs.encode(codes.idle))

connection.close()