osm2city-scripts/worldbuild-worker.py

#! /usr/bin/python3
# Copyright (C) 2018-2020 Merspieler, merspieler _at_ airmail.cc
#
# This program 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.
#
# This program 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 this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.

import os
from subprocess import run, Popen, STDOUT
import sys
import socket
import re
from time import sleep, strftime, time

from common import send_status, get_job, norm, get_south, get_west, get_east, get_north, get_area_name, api_send_status, api_get_job

action = "pending"
host = socket.gethostname()
api = None
api_token = None
port = 12345
prefix = ""
quiet = False
global_db = False

argc = len(sys.argv)
i = 1
first = 1
while i < argc:
	if sys.argv[i] == "--port":
		i += 1
		port = sys.argv[i]
	elif sys.argv[i] == "--host":
		i += 1
		host = sys.argv[i]
	elif sys.argv[i] == "-a" or sys.argv[i] == "--api":
		i += 1
		api = sys.argv[i]
	elif sys.argv[i] == "-t" or sys.argv[i] == "--api-token":
		i += 1
		api_token = sys.argv[i]
	elif sys.argv[i] == "-p" or sys.argv[i] == "--prefix":
		i += 1
		prefix = sys.argv[i]
	elif sys.argv[i] == "-g" or sys.argv[i] == "--global-db":
		global_db = True
	elif sys.argv[i] == "-q" or sys.argv[i] == "--quiet":
		quiet = True
	elif sys.argv[i] == "-a" or sys.argv[i] == "--action":
		i += 1
		if sys.argv[i] == "pending" or sys.argv[i] == "started" or sys.argv[i] == "rebuild" or sys.argv[i] == "skip":
			action = sys.argv[i]
		else:
			print ("ERROR: Unknown action " + sys.argv[i])
			sys.exit(1)
	elif sys.argv[i] == "-h" or sys.argv[i] == "--help":
		print("usage:worldbuild-worker.py [OPTIONS]")
		print("Retrives jobs from the manager and executes them")
		print("")
		print("  -p, --prefix     Database prefix to use")
		print("  -g, --global-db  Use global database")
		print("      --host       Manager host")
		print("      --port       Manager port")
		print("  -a, --api        Manager api url")
		print("  -t, --api-token  Manager api token")
		print("  -q, --quiet      Don't print messages")
		print("  -a, --action     Considered tiles for build. Can be:")
		print("                     pending: Only builds pending tiles. These are always build: default")
		print("                     started: Run tiles marked as started. CAUTION, use with care.")
		print("                     rebuild: Build tiles flaged for rebuild")
		print("                     skip: Force building tile that are marked for being skipped")
		print("  -h, --help       Shows this help and exit")
		sys.exit(0)
	else:
		print("Unknown option " + sys.argv[i])
		sys.exit(1)
	i += 1

if api != None and api_token == None:
	print("Error: API given but no token")
	sys.exit(1)

def cleanup():
	if os.path.isfile("projects/worldbuild-" + name + "/osm2city-exceptions.log"):
		run("mv projects/worldbuild-" + name + "/osm2city-exceptions.log projects/worldbuild/output/error/" + name + "-" + strftime("%Y%m%d-%H%M") + ".exceptions.log", shell=True)

	run("rm -rf projects/worldbuild-" + name, shell=True)

build = None
try:
	running = True
	while running:
		if api != None:
			name = api_get_job(action, api, api_token)
		else:
			name = get_job(action, host, port)
	
			run("mkdir -p projects/worldbuild-" + name, shell=True)
		
			run("cp projects/worldbuild/params.ini projects/worldbuild-" + name + "/", shell=True)
	
			if global_db:
				db_name = prefix + "worldbuild"
			else:
				if get_south(name) >= 80:
					db_name = prefix + "n-pole"
				elif get_south(name) < -80:
					db_name = prefix + "s-pole"
				else:
					match = re.match(r"[0-9]{1,7}", name)
					if match == None:
						print("ERROR: Invalid tile name")
						sys.exit(1)
					else:
						db_name = prefix + get_area_name(get_south(name), get_west(name), major=True)
	
			if get_west(name) < 0:
				bounds = "bounds=*"
			else:
				bounds = "bounds="
		
			bounds += str(get_west(name)) + "_" + str(get_south(name)) + "_" + str(get_east(name)) + "_" + str(get_north(name))
		
			run("sed -i 's/DB_NAME.*/DB_NAME = \"" + db_name + "\"/' projects/worldbuild-" + name + "/params.ini", shell=True)
		
			run("echo '" + bounds + "' > projects/worldbuild-" + name + "/settings", shell=True)
	
			if not quiet:
				print("Building " + name)
	
			logpath = "projects/worldbuild/output/" + get_area_name(get_south(name), get_west(name), major=True) + "/" + get_area_name(get_south(name), get_west(name)) + "/"
			run("mkdir -p " + logpath, shell=True)
			with open(logpath + name + "-" + strftime("%Y%m%d-%H%M") + ".log", "w") as log_file:
				build = Popen("./build -S 10 -t 1 worldbuild-" + name, stdout=log_file, stderr=STDOUT, shell=True, start_new_session=True)
	
			build.wait()
			build = None
	
			cleanup()
		
			if api != None:
				api_send_status(name, "done", api, api_token)
			else:
				send_status(name, "done", host, port)
except KeyboardInterrupt:
	if not quiet:
		print("Graceful shutdown triggered. To force immedate stop, press Ctrl+C again")
	running = False
	try:
		if build != None:
			build.wait()
		cleanup()
		if api != None:
			api_send_status(name, "done", api, api_token)
		else:
			send_status(name, "done", host, port)
	except KeyboardInterrupt:
            #TODO doesn't work
		print("Forcing shutdown...")
		build.terminate()
		sleep(5)
		build.kill()
		sleep(5)
		cleanup()
		sys.exit(0)