from __future__ import print_function

import boto3

import glob

import json

import logging

import os

import re

import subprocess

import sys

import time

import watchtower

import string

#################################

# CONSTANT PATHS IN THE CONTAINER

#################################

DATA_ROOT = '/home/ubuntu/bucket'

PLUGIN_DIR = '/home/ubuntu/CellProfiler-plugins'

QUEUE_URL = os.environ['SQS_QUEUE_URL']

AWS_BUCKET = os.environ['AWS_BUCKET']

LOG_GROUP_NAME= os.environ['LOG_GROUP_NAME']

EXPECTED_NUMBER_FILES= os.environ['EXPECTED_NUMBER_FILES']

if 'MIN_FILE_SIZE_BYTES' not in os.environ:

MIN_FILE_SIZE_BYTES = 1

else:

MIN_FILE_SIZE_BYTES = int(os.environ['MIN_FILE_SIZE_BYTES'])

if 'NECESSARY_STRING' not in os.environ:

NECESSARY_STRING = False

else:

NECESSARY_STRING = os.environ['NECESSARY_STRING']

SCRIPT_DOWNLOAD_URL = os.environ['SCRIPT_DOWNLOAD_URL']

SCRIPT_NAME = os.path.split(SCRIPT_DOWNLOAD_URL)[1]

#################################

# CLASS TO HANDLE THE SQS QUEUE

#################################

class JobQueue():

def __init__(self, queueURL):

self.client = boto3.client('sqs')

self.queueURL = queueURL

def readMessage(self):

response = self.client.receive_message(QueueUrl=self.queueURL, WaitTimeSeconds=20)

if 'Messages' in response.keys():

data = json.loads(response['Messages'][0]['Body'])

handle = response['Messages'][0]['ReceiptHandle']

return data, handle

else:

return None, None

def deleteMessage(self, handle):

self.client.delete_message(QueueUrl=self.queueURL, ReceiptHandle=handle)

return

def returnMessage(self, handle):

self.client.change_message_visibility(QueueUrl=self.queueURL, ReceiptHandle=handle, VisibilityTimeout=60)

return

#################################

# AUXILIARY FUNCTIONS

#################################

def monitorAndLog(process,logger):

while True:

output= process.stdout.readline()

if output== '' and process.poll() is not None:

break

if output:

print(output.strip())

logger.info(output)

def printandlog(text,logger):

print(text)

logger.info(text)

def stringify_metadata_dict(mdict):

met_string = ""

for eachmeta in mdict.keys():

met_string+=eachmeta+"='"+mdict[eachmeta]+"', "

return met_string[:-2]

#################################

# RUN FIJI PROCESS

#################################

def runFIJI(message):

#List the directories in the bucket- this prevents a strange s3fs error

rootlist=os.listdir(DATA_ROOT)

# Configure the logs

logger = logging.getLogger(__name__)

# Prepare paths and parameters

localOut = 'output'

remoteOut = message['output_file_location']

# Start logging

metadataID = stringify_metadata_dict(message['Metadata'])

metadata_for_log_name=metadataID.replace('*','.')

watchtowerlogger=watchtower.CloudWatchLogHandler(log_group=LOG_GROUP_NAME, stream_name=metadataID,create_log_group=False)

logger.addHandler(watchtowerlogger)

# Build and run CellProfiler command

cmd = ["Fiji.app/ImageJ-linux64","--ij2", "--headless", "--console", "--run",os.path.join("Fiji.app/plugins",SCRIPT_NAME)]

cmd.append(stringify_metadata_dict(message['shared_metadata']) + ', ' + metadataID)

print('Running', cmd)

logger.info(cmd)

subp = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)

monitorAndLog(subp,logger)

# Get the outputs and move them to S3

# Figure out how many output files there were - thanks https://stackoverflow.com/a/29769297

print('Checking output folder size')

filenum = 0

for _, _, filenames in os.walk(localOut):

filenum += len(filenames)

print(localOut,filenum)

if filenum>=int(EXPECTED_NUMBER_FILES):

mvtries=0

while mvtries <3:

try:

printandlog('Move attempt #'+str(mvtries+1),logger)

cmd = 'aws s3 mv ' + localOut + ' s3://' + AWS_BUCKET + '/' + remoteOut + ' --recursive'

subp = subprocess.Popen(cmd.split(), stdout=subprocess.PIPE, stderr=subprocess.PIPE)

out,err = subp.communicate()

printandlog('== OUT \n'+out, logger)

if err == '':

break

else:

printandlog('== ERR \n'+err,logger)

mvtries+=1

except:

printandlog('Move failed',logger)

printandlog('== ERR \n'+err,logger)

time.sleep(30)

mvtries+=1

if mvtries<3:

printandlog('SUCCESS',logger)

logger.removeHandler(watchtowerlogger)

return 'SUCCESS'

else:

printandlog('OUTPUT PROBLEM. Only '+str(filenum)+' files detected. Giving up on '+metadataID,logger)

logger.removeHandler(watchtowerlogger)

return 'OUTPUT_PROBLEM'

#################################

# MAIN WORKER LOOP

#################################

def main():

queue = JobQueue(QUEUE_URL)

# Main loop. Keep reading messages while they are available in SQS

while True:

msg, handle = queue.readMessage()

if msg is not None:

result = runFIJI(msg)

if result == 'SUCCESS':

print('Batch completed successfully.')

queue.deleteMessage(handle)

else:

print('Returning message to the queue.')

queue.returnMessage(handle)

else:

print('No messages in the queue')

break

#################################

# MODULE ENTRY POINT

#################################

if __name__ == '__main__':

logging.basicConfig(level=logging.INFO)

print('Worker started')

main()

print('Worker finished')