#!/usr/bin/env python3
# Copyright (C) 2012-2015 ASTRON (Netherlands Institute for Radio Astronomy)
# P.O. Box 2, 7990 AA Dwingeloo, The Netherlands
#
# This file is part of the LOFAR software suite.
# The LOFAR software suite is free software: you can redistribute it and/or
# modify it under the terms of the GNU General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# The LOFAR software suite is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License along
# with the LOFAR software suite. If not, see <http://www.gnu.org/licenses/>.
import os
from optparse import OptionParser, OptionGroup
import logging
logger = logging.getLogger(__name__)
from lofar.sas.tmss.client.tmssbuslistener import *
from lofar.sas.tmss.client.tmss_http_rest_client import TMSSsession
from subprocess import call
from lofar.common.cep4_utils import *
from lofar.common.subprocess_utils import check_output_returning_strings
class TMSSCopyServiceEventMessageHandler(TMSSEventMessageHandler):
'''
'''
def __init__(self, tmss_client_credentials_id: str="TMSSClient"):
super().__init__()
self.tmss_client = TMSSsession.create_from_dbcreds_for_ldap(tmss_client_credentials_id)
self._last_df_check_timestamp = datetime.min
def start_handling(self):
self.tmss_client.open()
super().start_handling()
def stop_handling(self):
super().stop_handling()
self.tmss_client.close()
def before_receive_message(self):
# use 1-sec event loop to poll queued subtasks (rate limited at 60sec)
super().before_receive_message()
if datetime.utcnow() - self._last_df_check_timestamp > timedelta(seconds=60):
self._last_df_check_timestamp = datetime.utcnow()
self.check_and_run_queued_copy_subtask_if_enough_disk_space()
def onSubTaskStatusChanged(self, id: int, status:str):
if status in ('scheduled', 'queued'):
subtask = self.tmss_client.get_subtask(id)
if subtask['subtask_type'] == 'copy':
if status == 'scheduled':
self.queue_copy_subtask(subtask)
elif status == 'queued':
self.run_copy_subtask_if_enough_disk_space(subtask)
else:
logger.info("skipping subtask id=%s status=%s type=%s", subtask['id'], subtask['state_value'], subtask['subtask_type'])
def queue_copy_subtask(self, subtask):
if subtask['subtask_type'] != 'copy':
return
if subtask['state_value'] != 'scheduled':
return
self.tmss_client.set_subtask_status(subtask['id'], 'queueing')
self.tmss_client.set_subtask_status(subtask['id'], 'queued')
def check_and_run_queued_copy_subtask_if_enough_disk_space(self):
subtasks = self.tmss_client.get_subtasks(subtask_type='copy', state='queued')
for subtask in subtasks:
self.run_copy_subtask_if_enough_disk_space(subtask)
def run_copy_subtask_if_enough_disk_space(self, subtask):
if subtask['subtask_type'] != 'copy':
return
if subtask['state_value'] != 'queued':
return
try:
# determine destination host and root_dir
destination = subtask['specifications_doc']['destination']
dst_host = destination[:destination.find(':')] if ':' in destination else ''
dst_root_dir = '/'+destination.replace(dst_host + ':', '').split('/')[1]
# remove unneeded localhost to prevent unneeded ssh calls
dst_host = dst_host.replace('localhost:','').replace('127.0.0.1:','')
df_cmd = ['df', dst_root_dir]
if dst_host:
df_cmd = wrap_command_in_ssh_call(df_cmd, dst_host)
logger.info("checking free disk space for copy-subtask id=%s, executing: %s", subtask['id'], ' '.join(df_cmd))
# run df cmd, and parse output for total free disk space
df_result = check_output_returning_strings(df_cmd)
df_result_line = df_result.splitlines()[-1]
df_result_line_parts = df_result_line.split()
df_bytes = int(df_result_line_parts[3])
input_dp_sizes = self.tmss_client.get_url_as_json_object(subtask['url'].rstrip('/')+'/input_dataproducts?fields=size')
total_size = sum(x['size'] for x in input_dp_sizes)
if df_bytes > total_size:
logger.info("enough free disk space available for copy-subtask id=%s destination=%s df=%d needed=%d", subtask['id'], destination, df_bytes, total_size)
# # clear previously set "not enough free disk space available" error_reason if set
# if subtask.get('error_reason'):
# self.tmss_client.do_request_and_get_result_as_string('PATCH', subtask['url'], {'error_reason': None})
# run it
self.run_copy_subtask(subtask)
else:
msg = "not enough free disk space available to start copy-subtask id=%s df=%d needed=%d" % (subtask['id'], df_bytes, total_size)
logger.warning(msg)
# self.tmss_client.do_request_and_get_result_as_string('PATCH', subtask['url'], {'error_reason': msg})
except Exception as e:
logger.exception(str(e))
# try to run it anyway, maybe it fails on not enough disk space.
# if it fails while running, then it results in an error status, and the user can take appropriate action.
self.run_copy_subtask(subtask)
def run_copy_subtask(self, subtask):
if subtask['subtask_type'] != 'copy':
return
if subtask['state_value'] != 'queued':
return
try:
self.tmss_client.set_subtask_status(subtask['id'], 'starting')
self.tmss_client.set_subtask_status(subtask['id'], 'started')
# cache to reduced rest-calls
# maps producer_id to tuple of producing subtask id and cluster name
_cache = {}
# cache to reduce (ssh) mkdir calls. Only create parent dirs once.
_created_dir_cache = set()
# ToDo: maybe parallelize this? Are multiple parallel rsync's faster?
for input_dataproduct in self.tmss_client.get_subtask_input_dataproducts(subtask['id']):
# fetch producing subtask id and cluster name for cache if needed
if input_dataproduct['producer'] not in _cache:
producer = self.tmss_client.get_url_as_json_object(input_dataproduct['producer'])
filesystem = self.tmss_client.get_url_as_json_object(producer['filesystem'])
cluster = self.tmss_client.get_url_as_json_object(filesystem['cluster'])
_cache[input_dataproduct['producer']] = {'producing_subtask_id': producer['subtask_id'], 'cluster_name': cluster['name'] }
if subtask['specifications_doc'].get('managed_output', False):
output_dataproduct = self.tmss_client.get_subtask_transformed_output_dataproduct(subtask['id'], input_dataproduct['id'])
output_dp_path = output_dataproduct['filepath']
else:
output_dp_path = os.path.join(subtask['specifications_doc']['destination'].rstrip('/'),
('L' + str(_cache[input_dataproduct['producer']]['producing_subtask_id'])) if subtask['specifications_doc'].get('group_by_id', True) else '',
input_dataproduct['filename'])
# split in host & parent_dir_path
dst_host = output_dp_path[:output_dp_path.find(':')] if ':' in output_dp_path else ''
dst_parent_dir_path = output_dp_path.replace(dst_host + ':','').replace(input_dataproduct['filename'],'')
# replace unneeded localhost to prevent under-the-hood ssh wrapping in rsync, and to prevent unneed ssh calls
dst_host = dst_host.replace('localhost:','').replace('127.0.0.1:', '')
dst_parent_dir_path = dst_parent_dir_path.replace('localhost:','').replace('127.0.0.1:','')
if dst_parent_dir_path not in _created_dir_cache:
# create dst_parent_dir_path directories if needed, prepend them to the cmd
mkdir_cmd = ['mkdir', '-p', dst_parent_dir_path]
if dst_host:
mkdir_cmd = wrap_command_in_ssh_call(mkdir_cmd, dst_host)
logger.info("creating parent destination dir if needed for copy-subtask id=%s, executing: %s",subtask['id'], ' '.join(mkdir_cmd))
if call(mkdir_cmd) == 0:
_created_dir_cache.add(dst_parent_dir_path)
else:
msg = "could not create parent destination dir '%s' for copy-subtask id=%s" % (dst_parent_dir_path, subtask['id'])
logger.error(msg)
self.tmss_client.set_subtask_status(subtask['id'], 'error', error_reason=msg)
return
# prepare the actual copy command
cmd = ['rsync', '-a', input_dataproduct['filepath'].rstrip('/'), dst_parent_dir_path]
# wrap in cep4 ssh call if cep4
cluster_name = _cache[input_dataproduct['producer']]['cluster_name']
if cluster_name.lower() == 'cep4':
cmd = wrap_command_in_cep4_available_node_with_lowest_load_ssh_call(cmd, via_head=True)
logger.info("copying dataproduct id=%s for copy-subtask id=%s, executing: %s", input_dataproduct['id'], subtask['id'], ' '.join(cmd))
if call(cmd) == 0:
logger.info("copied dataproduct id=%s for copy-subtask id=%s to '%s'", input_dataproduct['id'], subtask['id'], dst_parent_dir_path)
else:
msg = "could not copy dataproduct id=%s for copy-subtask id=%s to '%s'" % (input_dataproduct['id'], subtask['id'], dst_parent_dir_path)
logger.error(msg)
self.tmss_client.set_subtask_status(subtask['id'], 'error', error_reason=msg)
return
self.tmss_client.set_subtask_status(subtask['id'], 'finishing')
self.tmss_client.set_subtask_status(subtask['id'], 'finished')
except Exception as e:
logger.error(e)
self.tmss_client.set_subtask_status(subtask['id'], 'error', error_reason=str(e))
def create_copy_service(exchange: str=DEFAULT_BUSNAME, broker: str=DEFAULT_BROKER, tmss_client_credentials_id: str="TMSSClient"):
return TMSSBusListener(handler_type=TMSSCopyServiceEventMessageHandler,
handler_kwargs={'tmss_client_credentials_id': tmss_client_credentials_id},
exchange=exchange, broker=broker)
def main():
# make sure we run in UTC timezone
os.environ['TZ'] = 'UTC'
logging.basicConfig(format='%(asctime)s %(levelname)s %(message)s', level=logging.INFO)
# Check the invocation arguments
parser = OptionParser('%prog [options]',
description='run the tmss_copy_service which runs the copy-pipeline for scheduled copy-subtasks')
group = OptionGroup(parser, 'Messaging options')
group.add_option('-b', '--broker', dest='broker', type='string', default=DEFAULT_BROKER, help='Address of the message broker, default: %default')
group.add_option('-e', "--exchange", dest="exchange", type="string", default=DEFAULT_BUSNAME, help="exchange where the TMSS event messages are published. [default: %default]")
parser.add_option_group(group)
group = OptionGroup(parser, 'Django options')
parser.add_option_group(group)
group.add_option('-R', '--tmss_client_credentials_id', dest='tmss_client_credentials_id', type='string', default='TMSSClient', help='TMSS django REST API credentials name, default: %default')
(options, args) = parser.parse_args()
# check TMSS is up and running via the client
TMSSsession.check_connection_and_exit_on_error(options.tmss_client_credentials_id)
from lofar.common.util import waitForInterrupt
with create_copy_service(options.exchange, options.broker, options.tmss_client_credentials_id):
waitForInterrupt()
if __name__ == '__main__':
main()