-
Notifications
You must be signed in to change notification settings - Fork 0
/
startup-signal.py
67 lines (56 loc) · 2 KB
/
startup-signal.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
# This is a startup script running on trainer VMs
# This script is invoked by an init.sh script in background
# Upon startup, here is the sequence of the work
# 1. init subscriber on TODO topic
# 2. recieve task
# 3. lay down event marker
# 4. process task by invoking module
# 5. if response is non-zero => completed and (copy ouput artifact done by task script)
# 5.1 remove event marker
# 5.2 notify the DONE queue
from google.cloud import pubsub_v1
from pathlib import Path
from pubsub.publisher import Publisher, getDonePublisher, getTodoPublisher
from pubsub.subscriber import Subscriber, getTodoSubscriber
import logging
import sys
import time
import const as c
import util
from subprocess import Popen, PIPE, STDOUT
def log_subprocess_output(pipe):
for line in iter(pipe.readline, b''): # b'\n'-separated lines
logging.info('subprocess: %r', line)
######## MAIN #########
logging.basicConfig(filename='/tmp/cflow/cflow.log', level=logging.DEBUG)
logger = logging.getLogger(__name__)
sub = getTodoSubscriber()
while True:
# retrieve the first data
logger.info("pulling message from queue")
message = sub.pull()
if message is not None:
logger.info("received message %s" % message)
#lay down marker with complete "message"
util.placeMarkerFile(message)
#Use subprocess to invoke the real script and wait for the return value
#Script should pass the entire script
#create a command
cmd = c.TASK_MODULE_CMD + " " + message
logger.debug("executing command [%s]" % cmd)
process = Popen(cmd, shell=True, stdout=PIPE, stderr=STDOUT)
with process.stdout:
log_subprocess_output(process.stdout)
exitcode = process.wait() # 0 means success
#remove marker
util.removeMarkerFile()
if exitcode == 0:
logger.info("successfully executed subprocess")
#notify completion with original message
getDonePublisher().publish(message)
else:
logger.info("failed to execute subprocess")
getTodoPublisher().publish(message)
else:
logger.info("no message to process; wait for next task")
time.sleep(5)