def select(self):
'''
select all data from table
:return:
'''
try:
nfields = self.table.nfields()
sql, models = "select %s from %s;" % (",".join(quotes(
nfields, '`')), self.name), []
cursor = self.dbc.cursor()
cursor.execute(sql)
results = cursor.fetchall()
for result in results:
model = {}
for idx in range(0, len(result)):
nfield = nfields[idx]
vfield = result[idx]
if isinstance(vfield, str):
vfield = unescapes(vfield)
model[nfield] = vfield
models.append(model)
Logger.info("select from table %s...success", self.name)
return models
except Exception as e:
Logger.error("select from table %s...failed. error %s", self.name,
str(e))
raise e
def create(self, dbc, table):
'''
create table
:return self
'''
try:
#initialize table parameters
self.dbc = dbc
self.table = table
self.name = table.name
#check if table has exists
if self._exists_table():
#exists table
old_table = self.desc()
if old_table != self.table:
if is_subset(old_table.nfields(), self.table.nfields()):
#upgrade table
self._upgrade_table()
else:
#replace table
self._replace_table
else:
#table is the same as in database
pass
else:
#create new table
self._create_table()
Logger.info("create table %s...success.", self.name)
return self
except Exception as e:
Logger.error("create table %s...failed. error: %s", self.name,
str(e))
raise e
def process(self, frame):
time_elapsed = TimeElapsed()
self.lastProcessedFrame = frame
Logger.info('Processing frame - PCL: ' + frame.getFilePath())
cloud = pcl.load_XYZRGB(frame.getFilePath())
self.algorithmSegmentation.segment_image(cloud, True)
time_elapsed.printTimeElapsed()
def test_print_result(self):
device = FakeDevice(SourceType.IMAGE, 'dataset/png/')
# device = FakeDevice(SourceType.IMAGE, '/media/zis/Dados/dev/datasets/putkk.poznan/Dataset_1_Kin_2')
algorithms = []
algorithms.append(
RGBDSegmentation(
Parameter(Segmentation.GRAPH_CANNY,
os.getcwd() + '/results/graph_canny/',
resize=(640, 360),
fix_proportion=False)))
#algorithms.append(RGBDSegmentation(Parameter(Segmentation.RGBD_SALIENCY, os.getcwd() + '/results/rgbd_saliency/', resize=(324, 324))))
#algorithms.append(RGBDSegmentation(Parameter(Segmentation.FCN_TENSORFLOW, os.getcwd() + '/results/fcn_tensorflow/', resize=(224, 224))))
#algorithms.append(RGBDSegmentation(Parameter(Segmentation.FUSENET, os.getcwd() + '/results/fusenet/', resize=(224, 224))))
directory_rgb = device.datasetPath + 'rgb/'
directory_depth = device.datasetPath + 'depth/'
for algorithm in algorithms:
Logger.info('############# Testing algorithm ' +
str(algorithm.parameter.segmentation))
for i in range(0, 10):
frame = RGBDFrame(
RGBFrame(directory_rgb, 'rgb_' + format(i, '05') + '.png'),
DepthFrame(directory_depth,
'depth_' + format(i, '05') + '.png'))
algorithm.process(frame)
algorithm.write_results()
algorithm.finish()
#break
algorithm.release()
break
def open(self, path):
'''
open store or create it if not exist
:return:
'''
try:
with Lock(self.lock):
#init store path
self.path = path
if not path_exists(self.path):
#create database
self._create()
else:
# load database
self._load()
self._rebuild_tindex()
return self
Logger.info("open store %s...success. %d tables.", self.path,
len(self.tables))
except Exception as e:
Logger.error("open store %s...failed. error: %s", self.path,
str(e))
raise e
def get_request_token(self):
Logger.info('Starting to fetch request token for Kite API')
selenium = SeleniumDispatcher(headless=True)
driver = selenium.get_driver()
driver.get(self.kite.login_url())
time.sleep(4)
username_field = driver.find_element_by_xpath("//input[@type='text']")
username_field.send_keys(self.username)
password_field = driver.find_element_by_xpath("//input[@type='password']")
password_field.send_keys(self.password)
password_field.send_keys(Keys.ENTER)
time.sleep(2)
pin_field = driver.find_element_by_xpath("//input[@type='password']")
pin_field.send_keys(self.pin)
pin_field.send_keys(Keys.ENTER)
time.sleep(2)
url = driver.current_url
parsed = urlparse.urlparse(url)
token = parse_qs(parsed.query)['request_token'][0]
Logger.info('Request token received!')
selenium.destroy_driver()
return token
def create_table(self, table):
'''
create table in current database
:param table:
:return:
'''
with Lock(self.lock):
# test if the table has loaded
for t in self.tables:
if t.table == table:
Logger.info("create table %s...exists.", table.name)
return
#create new table
table = FSTable().create(self.path, table)
for i in range(0, len(self.tables)):
t = self.tables[i]
if t.table.name == table.name:
self.tables.pop(i)
break
self.tables.append(table)
self._rebuild_tindex()
def open(self, host, user, pwd, dbn, port=3306):
'''
open database or create it if not exist
:return:
'''
#init store path
try:
self.host, self.port, self.user, self.pwd = host, port, user, pwd
self.dbn = dbn
self.dbc = pymysql.connect(host=host, user=user, passwd=pwd, port=port)
if not self._exists():
#create database
self._create()
self._use()
else:
# load database
self._use()
self._load()
self._rebuild_tindex()
return self
Logger.info("open store mysql://%s:%s@%s:%d/%s...success. %d tables.", user, pwd, host, port, self.dbn, len(self.tables))
except Exception as e:
Logger.error("open store mysql://%s:%s@%s:%d/%s...failed. error: %s", user, pwd, host, port, self.dbn, str(e))
raise e
def indian_business_news(self):
res = requests.get(self.business_url + '&pageSize=100&apiKey=' +
self.api_key).json()
total_results = res['totalResults']
self.total_articles = int(total_results)
# 100 is the max pagesize for newsapi.org
if total_results <= 100:
self.insert_data_into_db(res['articles'])
else:
# break them in chunks of 10
number_of_pages = math.ceil(total_results) / 10
collection = self.client[Config.MONGO_DB][Config.NEWS_COLLECTION]
for page in range(1, number_of_pages):
res = requests.get(
self.business_url + \
'&pageSize=10&page={}&apiKey='.format(page) + \
self.api_key).json()
self.insert_data_into_db(res['articles'])
if self.inserted_articles > 0:
Logger.info(
'News data stored successfully!\nTotal articles: {}\nInserted articles: {}'
.format(self.total_articles, self.inserted_articles))
else:
Logger.info('Total articles: {}\nNot inserted articles: {}'.format(
self.total_articles, self.not_inserted_articles))
def load(self, dbpath, name):
'''
load table
:return: self
'''
try:
#initialize table parameters
self.name = name
self.path = join_paths(dbpath, name)
self.table_file = join_paths(self.path, "table")
self.data_file = join_paths(self.path, "data")
self.table = self.desc()
#load data file
if not is_file(self.data_file):
#create data file if not exists
self._create_data_file()
else:
#replace old data file if needed
with open(self.data_file) as fdata:
nfields = strips(fdata.readline().split(","))
if self.table.nfields() != nfields:
if is_subset(nfields, self.table.nfields()):
self._upgrade_data_file()
else:
self._replace_data_file()
Logger.info("loading table %s...success.", self.name)
return self
except Exception as e:
Logger.info("loading table %s...failed. error: %s", self.name, str(e))
raise e
def select(self):
'''
select all data from table
:return:
'''
try:
with Lock(self.lock):
with open(self.data_file, "r") as fdata:
models = []
#read field names
nfields = strips(fdata.readline().strip().split(","))
#read data records
data = fdata.readline()
while data:
data = data.strip()
vfields = strips(data.split(","))
model = {}
for idx in range(0, len(nfields)):
model[nfields[idx]] = str2obj(vfields[idx], ',')
models.append(model)
data = fdata.readline()
return models
except Exception as e:
Logger.info("select data from table %s...failed. error: %s", self.name, str(e))
raise e
def send(self, data):
# Importing it here to avoid error due to circular dependency
from util.log import Logger
'''Data can be dict, str, int, list'''
Logger.info('Sending slack message to: ' + self.webhook_url)
data = json.dumps({'text': data})
res = requests.post(self.webhook_url, data)
return res
def process(self, frame):
time_elapsed = TimeElapsed()
self.lastProcessedFrame = frame
Logger.info('Processing frame - RGB: ' + frame.rgbFrame.getFilePath() + ', Depth: '+frame.depthFrame.getFilePath())
self.results = self.algorithmSegmentation.segment_image(self.get_image(frame.rgbFrame),
self.get_image(frame.depthFrame))
Logger.info('Objects segmented: ' + str(self.algorithmSegmentation.get_num_objects()))
time_elapsed.printTimeElapsed('Total segmentation - ')
def write_results(self):
time_elapsed = TimeElapsed()
img = self.write_objects()
Logger.info('Saving result to ' + self.parameter.outputDir+self.lastProcessedFrame.rgbFrame.fileName)
if not os.path.exists(self.parameter.outputDir):
os.makedirs(self.parameter.outputDir)
cv2.imwrite(self.parameter.outputDir+self.lastProcessedFrame.rgbFrame.fileName, img)
time_elapsed.printTimeElapsed('Total writing file - ')
def download_data_all(self):
Logger.info('Download started')
if self.expiry_strike_price_map is None:
self.get_expiry_strike_price_map_for_all()
downloaded_files = list()
for expiry in self.expiry_strike_price_map.keys():
for strike_price in self.expiry_strike_price_map[expiry]:
filepath = self.download_data_specific(expiry, strike_price)
downloaded_files.append(filepath)
Logger.info('Download finished')
return downloaded_files
def desc(self):
'''
descrite table from store
:return: Table
'''
try:
with open(self.table_file) as ftable:
table = Table().fromstr(ftable.read())
return table
except Exception as e:
Logger.info("describe table %s...failed. error: %s", self.name, str(e))
raise e
def drop(self):
'''
drop table
:return:
'''
try:
sql = "drop table if exists %s;" % self.name
self.dbc.cursor().execute(sql)
Logger.info("drop table %s...success", self.name)
except Exception as e:
Logger.error("drop table %s...failed. error %s", self.name, str(e))
raise e
def close(self):
'''
close datbase
:return:
'''
#close database connection
try:
if self.dbc is not None:
self.dbc.close()
Logger.info("close store mysql://%s:%s@%s:%d/%s...success.", self.user, self.pwd, self.host, self.port, self.dbn)
except Exception as e:
Logger.info("close store mysql://%s:%s@%s:%d/%s...failed. error: %s", self.user, self.pwd, self.host, self.port, self.dbn, str(e))
raise e
def truncate(self):
'''
truncate table
:return:
'''
try:
sql = "truncate table %s;" % table.name
self.dbc.cursor().execute(sql)
Logger.info("truncate table %s...success", self.name)
except Exception as e:
Logger.error("truncate table %s...failed. error %s", self.name,
str(e))
raise e
def create(self, dbpath, table):
'''
create table
:return self
'''
try:
#initialize table parameters
self.table = table
self.name = table.name
self.path = join_paths(dbpath, table.name)
self.table_file = join_paths(self.path, "table")
self.data_file = join_paths(self.path, "data")
#create table directory if it is not exists
make_dirs(self.path)
#create or replace table file
if is_file(self.table_file):
#replace old table file if needed
old_table = self.desc()
if self.table != old_table:
#replace table file
self._replace_table_file()
else:
#new table is same as exists table
pass
else:
#create new table file
self._create_table_file()
#create or upgrade or replace data file
if is_file(self.data_file):
#replace old data file if needed
with open(self.data_file) as fdata:
nfields = strips(fdata.readline().split(","))
if self.table.nfields() != nfields:
if is_subset(nfields, self.table.nfields()):
self._upgrade_data_file()
else:
self._replace_data_file()
else:
#create new data file
self._create_data_file()
Logger.info("create table %s...success.", self.name)
return self
except Exception as e:
Logger.error("create table %s...failed. error: %s", self.name, str(e))
raise e
def calculate_and_store_pnl(self):
quantitiy = np.array(self.kite_state.numberOfStocksPerCompany)
# Subtracting 1 from all companies as a failsafe for fund exhaustion, in case the stocks goes really high
normalised_quantity = quantitiy - 1
# Rounding negative values to 0
normalised_quantity = normalised_quantity.clip(min=0)
buy = np.array(self.kite_state.buyPrice)
sell = np.array(self.kite_state.sellPrice)
pnl_per_company = np.multiply(sell - buy, normalised_quantity)
self.kite_state.pnlPerCompany = pnl_per_company.tolist()
self.kite_state.pnl = float(np.sum(pnl_per_company))
self.kite_state.save()
Logger.info("Today's PNL: {}".format(self.kite_state.pnl), push_to_slack=True)
class Worker(object):
def __init__(self,
queue_name,
route,
persister,
wnum,
port,
logfile=sys.stdout):
self.queue_name = queue_name
self.route = route
self.persister = persister
self.stop = False
self.pid = os.getpid()
self.worker_id = '-'.join(
['worker', str(wnum), queue_name, route,
str(self.pid)])
self.log = Logger(self.worker_id,
logfile=logfile,
loglevel=logging.DEBUG)
self.log.info("starting")
self.host = socket.gethostbyname(socket.gethostname())
self.port = port
self.register()
self.todo = None
self.stop = False
def register(self):
self.persister.add_worker(self.worker_id, self.host, self.port)
def unregister(self):
self.persister.delete_worker(self.worker_id)
def start_worker(self):
while not self.stop:
if self.todo:
yield self.work(self.todo)
self.todo = None
def stop_worker(self):
self.log.info("shutting down")
self.unregister()
self.stop = True
def work(self, f):
self.persister.set_working()
ret = f()
self.persister.unset_working()
return ret
def run(self):
# Check if we have the list in DB
gl_objects_list = list(self.get_stored_securities_from_db())
if len(gl_objects_list) == 0:
gainers_info_nse = NSEIndiaGLScraper(info_type='Gainers',
view_type='All Securities')
gainers_list_nse = gainers_info_nse.get_instruments(
complete_info=False)
gainers_info_rediff = RediffMoneyGLScraper(view_type='All')
gainers_list_rediff = gainers_info_rediff.get_instruments(
limit_number_of_instruments=self.number_of_rediff_instruments)
gainers_list = list(
set(gainers_list_nse).union(set(gainers_list_rediff)))
percent_delivered_dict = self.equity_scraper.get_info_all(
gainers_list, specific_info_key='deliveryToTradedQuantity')
percent_delivered = list(percent_delivered_dict.values())
previous_close_prices = list(
self.equity_scraper.get_info_all(
gainers_list, specific_info_key='previousClose').values())
low_prices = list(
self.equity_scraper.get_info_all(
gainers_list, specific_info_key='dayLow').values())
# Storing the list for future use
self.store_securities_in_db(gainers_list)
self.store_percent_delivered_in_db(percent_delivered)
self.store_previous_close_price_in_db(previous_close_prices)
self.store_low_price_in_db(low_prices)
else:
# Getting only the first object
percent_delivered = gl_objects_list[-1].percentDelivered
gainers_list = gl_objects_list[-1].listOfCompanies
last_price = self.equity_scraper.get_info_all(
gainers_list, specific_info_key='lastPrice')
self.store_last_price_in_db(list(last_price.values()))
table = PrettyTable(['Security', 'TradedToDelivered', 'LastPrice'])
for security, percent_delivered in zip(gainers_list,
percent_delivered):
table.add_row([security, percent_delivered, last_price[security]])
Logger.info(message=table.get_string(), push_to_slack=self.slack_notif)
Logger.info(message=table.get_html_string(),
push_to_sendgrid=self.sendgrid_notif,
sendgrid_subject=self.class_name)
def print_results(self):
for i in range(self.algorithmSegmentation.get_num_objects()):
obj = self.results[i]
if self.algorithmSegmentation.python_segmentation:
Logger.info('Object id: ' + str(obj.id) + ', número de pontos: ' + str(len(obj.pointsList)))
for j in range(len(obj.pointsList)):
Logger.info('Ponto ' + str(j) + ' - x: ' + str(obj.pointsList[j].x) + ' y: ' + str(obj.pointsList[j].y) + ' z: ' + str(obj.pointsList[j].z))
else:
Logger.info('Object id: ' + str(obj.id) + ', número de pontos: ' + str(obj.pointsLength))
for j in range(obj.pointsLength):
Logger.info('Ponto ' + str(j) + ' - x: ' + str(obj.points[j].x) + ' y: ' + str(obj.points[j].y) + ' z: ' + str(obj.points[j].z))
def insert(self, models):
'''
insert data to table
:param models:
:return:
'''
try:
with Lock(self.lock):
with open(self.data_file, "a") as fdata:
lines = []
for model in models:
vfields = []
for nfield in self.table.nfields():
vfields.append(objtostr(model.get(nfield), ','))
lines.append("%s\n" % ",".join(vfields))
fdata.writelines(lines)
except Exception as e:
Logger.info("insert data to table %s...failed. error: %s", self.name, str(e))
raise e
def desc(self):
'''
descrite table from store
:return: Table
'''
try:
sql = "show create table %s;" % self.name
cursor = self.dbc.cursor()
cursor.execute(sql)
create_sql = cursor.fetchall()[0][1]
table = Table().fromsql(create_sql)
Logger.info("describe table %s...success", self.name)
return table
except Exception as e:
Logger.error("describe table %s...failed. error %s", self.name,
str(e))
raise e
def get_instruments(self, complete_info=False):
# TODO: add other 'view type' and 'info_type' support as well
# @body Rewrite this part of code using nsetools
url = 'https://www1.nseindia.com/live_market/dynaContent/live_analysis/'\
+ self.info_type.lower()\
+ '/'\
+ self.view_type\
+ '.json'
Logger.info(url)
# res = SeleniumDispatcher(headless=False, selenium_wire=True).get_response(url)
res = requests.get(url, headers=Config.NSE_HEADERS)
symbol_info_list = res.json()['data']
# data = res.get_json()
# symbol_info_list = data['data']
symbols = list()
for symbol_info in symbol_info_list:
# Currently considering EQUITY series only
if symbol_info['series'] == 'EQ':
symbols.append(symbol_info['symbol'])
return {True: symbol_info_list, False: symbols}[complete_info]
class Worker(object):
def __init__(self, queue_name, route, persister, wnum, port,
logfile=sys.stdout):
self.queue_name = queue_name
self.route = route
self.persister = persister
self.stop = False
self.pid = os.getpid()
self.worker_id = '-'.join(['worker', str(wnum), queue_name, route, str(self.pid)])
self.log = Logger(self.worker_id,logfile=logfile, loglevel=logging.DEBUG)
self.log.info("starting")
self.host = socket.gethostbyname(socket.gethostname())
self.port = port
self.register()
self.todo = None
self.stop = False
def register(self):
self.persister.add_worker(self.worker_id, self.host, self.port)
def unregister(self):
self.persister.delete_worker(self.worker_id)
def start_worker(self):
while not self.stop:
if self.todo:
yield self.work(self.todo)
self.todo = None
def stop_worker(self):
self.log.info("shutting down")
self.unregister()
self.stop = True
def work(self, f):
self.persister.set_working()
ret = f()
self.persister.unset_working()
return ret
def run(self):
open_price = self.equity_scraper.get_info_all(
self.gl_object.listOfCompanies, specific_info_key='open')
high_price = self.equity_scraper.get_info_all(
self.gl_object.listOfCompanies, specific_info_key='dayHigh')
low_price = self.equity_scraper.get_info_all(
self.gl_object.listOfCompanies, specific_info_key='dayLow')
close_price = self.equity_scraper.get_info_all(
self.gl_object.listOfCompanies, specific_info_key='closePrice')
self.store_open_price_in_db(list(open_price.values()))
self.store_high_price_in_db(list(high_price.values()))
self.store_low_price_in_db(list(low_price.values()))
self.store_close_price_in_db(list(close_price.values()))
report = self.generate_report()
Logger.info(report.get_string(), push_to_slack=self.slack_notif)
Logger.info(report.get_html_string(),
push_to_sendgrid=self.sendgrid_notif,
sendgrid_subject=self.class_name)
def download_data_specific(self, expiry: str, strike_price: str):
url = "https://nseindia.com/live_market/dynaContent/live_watch/get_quote/getFOHistoricalData.jsp"\
+ "?underlying=" + self.symbol_name\
+ '&instrument=' + self.derivative_type_val\
+ '&expiry=' + expiry\
+ '&type=' + self.option_type\
+ '&strike=' + strike_price\
+ "&fromDate=" + 'undefined'\
+ "&toDate=" + 'undefined'\
+ "&datePeriod=" + '3months'\
+ "&fileDnld=" + 'undefined'
download_dir_path = os.path.join(Config.DOWNLOAD_DIRECTORY,
self.derivative_type_val,
self.option_type)
if not os.path.isdir(download_dir_path):
os.makedirs(download_dir_path)
download_file_path = os.path.join(download_dir_path,
expiry + '_' + strike_price + '.csv')
# filename = wget.download(url=url, out=download_file_path)
res = requests.get(url, allow_redirects=True)
Logger.info(download_file_path)
open(download_file_path, 'wb').write(res.content)
return download_file_path
def load(self, dbc, name):
'''
load table from database
:return: self or None
'''
try:
#initialize table parameters
self.dbc = dbc
self.name = name
#create table structure
sql = "show create table %s;" % self.name
cursor = self.dbc.cursor()
cursor.execute(sql)
create_sql = cursor.fetchall()[0][1]
self.table = Table().fromsql(create_sql)
Logger.info("loading table %s...success.", self.name)
return self
except Exception as e:
Logger.info("loading table %s...failed. error: %s", self.name,
str(e))
raise e
class WorkerPool(object):
def __init__(self, queue_name, routing_keys=None,
backend='mongodb', conn_url='localhost:27017',
dbname='skunkqueue', logfile=sys.stdout,
pidfile=None):
"""
routing_keys are a required parameter to specify an n-length list
of routing keys, which will each be assigned to one worker
"""
self.stop = False
self.queue_name = queue_name
self.persister = get_backend(backend)(conn_url=conn_url, dbname=dbname)
self.log = Logger('pool-'+queue_name, logfile=logfile)
self.host = socket.gethostbyname(socket.gethostname())
self.workers = {}
if pidfile:
self.log.info("writing to pidfile %s" % pidfile)
with open(pidfile) as f:
f.write(str(self.pid))
f.close()
# TODO this needs to be in shared memory
wnums = {}
self.s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
self.s.bind((self.host, 0))
self.s.listen(5)
self.port = self.s.getsockname()[1]
for key in routing_keys:
if key not in wnums:
wnums[key] = 0
wnums[key] += 1
worker = Worker(queue_name, key,
self.persister, wnums[key], self.port, logfile)
thread = KillableThread(target=worker.begin_execution)
thread.start()
self.workers[worker.worker_id] = (worker, thread)
def __enter__(self, *args, **kwargs):
self.log.info("starting")
def gentle(signum, frame):
self.log.info("Received gentle shutdown signal %d" % signum)
self.shutdown()
sys.exit(0)
def rough(signum, frame):
self.log.warn("Received non-gentle kill signal %d" % signum)
self.die()
sys.exit(0)
signal.signal(signal.SIGINT, rough )
signal.signal(signal.SIGHUP, gentle)
signal.signal(signal.SIGTERM, gentle)
signal.signal(signal.SIGALRM, gentle)
signal.signal(signal.SIGQUIT, gentle)
return self
def __exit__(self, *args, **kwargs):
self.shutdown()
def shutdown(self):
for _id in self.workers.keys():
worker = self.workers[_id]
worker[0].stop_worker()
del self.workers[_id]
self.s.close()
self.stop = True
def die(self):
for _id in self.workers.keys():
worker = self.workers[_id]
try:
worker[1].raise_exc(TriggeredInterrupt)
self.log.warn("raised an exception in %s" % str(_id))
del self.workers[_id]
except ValueError: # it was dead already
self.log.debug("ignored killing thread %s" % str(_id))
continue
self.s.close()
self.stop = True
def run_cmd(self, terminate=None, kill=None, suicide=None):
# requests a termination
# terminate looks like: <worker_id>
if terminate:
self.workers[terminate][0].stop_worker()
del self.workers[terminate]
# performs a hard kill
# kill looks like: <worker_id>
if kill:
try:
self.workers[kill][1].raise_exc(TriggeredInterrupt)
del self.workers[kill]
except ValueError: # it was already dead
self.log.warn("tried to kill a thread when it was dead already")
# shuts down the entire pool
# suicide looks like: 1
if suicide == 1:
self.stop = True
def listen(self):
while not self.stop:
conn, addr = self.s.accept()
datalen = conn.recv(2)
data = ""
while len(data) < int(datalen):
data += conn.recv(1024)
data = json.loads(data)
self.run_cmd(**data)
# no need to run the pool if it has no workers
if len(self.workers) == 0:
self.shutdown()
class Worker(object):
def __init__(self, queue_name, route, persister, wnum, port,
logfile=sys.stdout):
self.queue_name = queue_name
self.route = route
self.persister = persister
self.stop = False
self.pid = os.getpid()
self.worker_id = '-'.join(['worker', str(wnum), queue_name, route, str(self.pid)])
self.log = Logger(self.worker_id,logfile=logfile, loglevel=logging.DEBUG)
self.log.info("starting")
self.host = socket.gethostbyname(socket.gethostname())
self.port = port
self.register()
def begin_execution(self, *args):
while not self.stop:
try:
job = self.persister.get_job_from_queue(self.queue_name, self.worker_id, self.route)
if job:
self.do_job(job)
sleep(0.1)
except TriggeredInterrupt:
# we were cut off by an interrupt trigger
self.log.warn("received interrupt request; stopping current job")
self.log.warn("no result will be committed and this job will not be restarted")
self.stop_worker()
def register(self):
self.persister.add_worker(self.worker_id, self.host, self.port)
def unregister(self):
self.persister.delete_worker(self.worker_id)
def do_job(self, job):
# depickle
body = pickle.loads(job['body'])
directory = body['dir']
# FIXME a horrible hack where we add ourselves to the pythonpath
sys.path.append(directory)
mod = __import__(body['mod'])
self.log.debug("successfully imported module "+str(mod))
if job['fn_type'] == 'method':
parent = dill.loads(body['parent'])
fn = getattr(parent, body['fn'])
else:
fn = getattr(mod, body['fn'])
args = body['args']
kwargs = body['kwargs']
# call it
self.persister.set_working(self.worker_id)
try:
ret = fn(*args, **kwargs)
self.persister.save_result(job['job_id'], ret, 'complete')
self.log.info(ret)
except Exception as e:
self.persister.save_result(job['job_id'], None, 'error')
self.log.error(str(e))
exc_t, exc_v, exc_tr = sys.exc_info()
self.log.error(str(
'\n'.join(traceback.format_exception(exc_t, exc_v, exc_tr))))
self.log.debug("python path is %s" % str(sys.path))
finally:
a = sys.path.pop()
self.persister.unset_working(self.worker_id)
def stop_worker(self):
self.log.info("shutting down")
self.unregister()
self.stop = True
class WorkerPool(object):
def __init__(self, queue_name, routing_keys=None,
backend='mongodb', conn_url='localhost:27017',
dbname='fwbots', logfile=sys.stdout,
pidfile=None):
"""
routing_keys are a required parameter to specify an n-length list
of routing keys, which will each be assigned to one worker
FWBOTS: we are using routing_keys as the account names to load
"""
self.stop = False
self.name = queue_name
self.log = Logger('pool-'+queue_name, logfile=logfile)
self.persister = get_backend(backend)(conn_url,dbname)
self.port = -1 # too lazy to actually remove this
self.workers = {}
self.manual = {}
self.auto = {}
self.insta = {}
if pidfile:
self.log.info("writing to pidfile %s" % pidfile)
with open(pidfile) as f:
f.write(str(self.pid))
f.close()
# TODO this needs to be in shared memory
wnums = {}
self.ircbot = IrcListener(self,"#fwbots",self.name,"irc.freenode.net")
for key in routing_keys:
errnum = 0
try:
acc = get_twitter_account(key)
if acc.ty == 'auto':
self.log.info('found auto for %s' % key)
self.auto[acc.name] = acc
# if there are multiple manual accs defined, pick only the last one
elif acc.ty == 'manual':
self.log.info('found manual for %s' % key)
self.manual[acc.name] = acc
except NameError:
errnum += 1
try:
acc = get_instagram_account(key)
self.insta[acc.name] = acc
except NameError:
errnum += 1
if errnum > 1:
self.log.warn("Could not find any account called %s"%key)
if key not in wnums:
wnums[key] = 0
wnums[key] += 1
worker = Worker(queue_name, key,
self.persister, wnums[key], self.port, logfile)
thread = KillableThread(target=worker.start_worker)
thread.start()
self.workers[worker.worker_id] = (worker, thread)
# lastly, register yourself
self.persister.add_pool(self)
def __enter__(self, *args, **kwargs):
self.log.info("starting")
def gentle(signum, frame):
self.log.info("Received gentle shutdown signal %d" % signum)
self.shutdown()
sys.exit(0)
def rough(signum, frame):
self.log.warn("Received non-gentle kill signal %d" % signum)
self.die()
sys.exit(0)
signal.signal(signal.SIGINT, rough )
signal.signal(signal.SIGHUP, gentle)
signal.signal(signal.SIGTERM, gentle)
signal.signal(signal.SIGALRM, gentle)
signal.signal(signal.SIGQUIT, gentle)
return self
def __exit__(self, *args, **kwargs):
self.shutdown()
def shutdown(self):
for _id in self.workers.keys():
worker = self.workers[_id]
worker[0].stop_worker()
del self.workers[_id]
self.stop = True
self.persister.delete_pool(self.name)
def die(self):
for _id in self.workers.keys():
worker = self.workers[_id]
try:
worker[1].raise_exc(TriggeredInterrupt)
self.log.warn("raised an exception in %s" % str(_id))
del self.workers[_id]
except ValueError: # it was dead already
self.log.debug("ignored killing thread %s" % str(_id))
continue
self.stop = True
def work(self, f):
w = self.persister.get_avail_workers()[0]
w.todo = f
def run_cmd(self, terminate=None, kill=None, suicide=None,
tweet=None,rtp=None):
# requests a termination
# terminate looks like: <worker_id>
if terminate:
self.workers[terminate][0].stop_worker()
del self.workers[terminate]
# performs a hard kill
# kill looks like: <worker_id>
if kill:
try:
self.workers[kill][1].raise_exc(TriggeredInterrupt)
del self.workers[kill]
except ValueError: # it was already dead
self.log.warn("tried to kill a thread when it was dead already")
# shuts down the entire pool
# suicide looks like: 1
if suicide == 1:
self.stop = True
# sends a tweet from self.manual
# tweet looks like: <msg>
if tweet:
self.log.info("tweeting with manual account %s: '%s'"%(self.manual.name,tweet))
self.work(lambda: self.manual.make_tweet(tweet))
# if rtp is set, have dummies retweet it with probability rtp
if rtp:
def f():
for a in self.auto:
if rtp < random.random():
a.retweet_last(self.manual.name)
self.work(f)
def listen(self):
self.ircbot.start()
