<commit_before>from pathlib import Path
from urllib import request
from urllib.error import URLError, HTTPError
from lxml import html, etree
from flask import Flask, Response, request as flask_request, jsonify, render_template, send_file, abort
from werkzeug.contrib.atom import AtomFeed
from werkzeug.utils import secure_filename
from datetime import datetime, timedelta
import io
import json
import os
import pytz
import logging

def request_res_latest_json():
	# Scrape the latest available measurements from CA DWR CDEC.
	blank_val = '--'.rjust(10)
	data = {}
	# Populate at least what is needed to return if there is an error or exception.
	data['res_elev'] = 0
	data['inflow'] = 0
	data['outflow'] = 0
	page = request_cdec_hourly_page()
	if page is None:
		return data
	# There could be multiple tables in the page.
	# The table with hourly measurements has either 14 or 15 rows and exactly 15 columns.
	# The first two rows are for the title, ignore that.
	# The last row may be incomplete (blank values). Ignore these.
	xpath_predicate = '//div[@class="content_left_column"]/table[count(tr) >= 14 and count(tr) <= 15]/tr[count(td) = 15][position() > 2][td[2][text()!="{}"]]'.format(
			blank_val)
	rows = parse_page_table('res_elev', xpath_predicate, page)
	if rows is None:
		return data
	res_elev_row = rows[-1]
	xpath_predicate = '//div[@class="content_left_column"]/table[count(tr) >= 14 and count(tr) <= 15]/tr[count(td) = 15][position() > 2][td[6][text()!="{}"] and td[8][text()!="{}"]]'.format(
			blank_val, blank_val)
	rows = parse_page_table('flow', xpath_predicate, page)
	if rows is None:
		return data
	flow_row = rows[-1]
	res_elev_datetime = '{:%b %-d %-I%P}'.format(
		datetime.strptime(res_elev_row[0].xpath('text()')[0], '%m/%d/%Y %H:%M'))
	res_elev = float(res_elev_row[1].xpath('text()')[0])
	flow_datetime = '{:%b %-d %-I%P}'.format(
		datetime.strptime(flow_row[0].xpath('text()')[0], '%m/%d/%Y %H:%M'))
	outflow = int(flow_row[5].xpath('./text()')[0])
	inflow = int(flow_row[7].xpath('./text()')[0])
	data['timestamp_str'] = str(datetime.utcnow())
	data['res_elev_datetime'] = res_elev_datetime
	data['res_elev'] = res_elev
	data['flow_datetime'] = flow_datetime
	data['inflow'] = inflow
	data['outflow'] = outflow
	return data

def request_gauges_latest_json():
	# Scrape the latest available measurements from CA DWR CDEC.
	blank_val = '--'.rjust(10)
	data = {}
	# Populate at least what is needed to return if there is an error or exception.
	data['res_elev'] = 0
	data['inflow'] = 0
	data['outflow'] = 0
	PT = pytz.timezone('America/Los_Angeles')  # Pacific Timezone
	page = request_cdec_hourly_page()
	if page is None:
		return data
	# There could be multiple tables in the page.
	# The table with hourly measurements has either 14 or 15 rows and exactly 15 colums.
	# The first two rows are for the title, ignore that.
	# The last row may be incomplete (blank values). Ignore these.
	xpath_predicate = '//div[@class="content_left_column"]/table[count(tr) >= 14 and count(tr) <= 15]/tr[count(td) = 15][position() > 2][td[2][text()!="{}"] and td[6][text()!="{}"] and td[8][text()!="{}"]]'.format(
			blank_val, blank_val, blank_val)
	rows = parse_page_table('gauge', xpath_predicate, page)
	if rows is None:
		return data
	row = rows[-1]
	row_datetime = datetime.strptime(row[0].xpath('text()')[0], '%m/%d/%Y %H:%M')
	row_datetime = PT.localize(row_datetime)
	datetime_str = '{:%B %-d, %Y %-I%P} {}'.format(
		row_datetime, row_datetime.tzname())
	res_elev = float(row[1].xpath('text()')[0])
	outflow = int(row[5].xpath('./text()')[0])
	inflow = int(row[7].xpath('./text()')[0])
	data['timestamp_str'] = str(datetime.utcnow())
	data['datetime'] = datetime_str
	data['res_elev'] = res_elev
	data['inflow'] = inflow
	data['outflow'] = outflow
	return data

def res_latest():
	site_root = Path(__file__).parent
	path = Path(site_root, 'data', 'reslatest.json')
	data = {}
	if path.is_file():
		with path.open('r') as data_file:
			data = json.load(data_file)
	# If more than a minute old, get the latest.
	if 'timestamp_str' not in data or \
		abs((datetime.utcnow() - datetime.strptime(data['timestamp_str'], '%Y-%m-%d %H:%M:%S.%f')).total_seconds()) >= 60:
		data = request_res_latest_json()
		with path.open('w') as data_file:
			json.dump(data, data_file)
	return data

def gauges_latest():
	site_root = Path(__file__).parent
	path = Path(site_root, 'data', 'gaugeslatest.json')
	data = {}
	if path.is_file():
		with path.open('r') as data_file:
			data = json.load(data_file)
	# If more than a minute old, get the latest.
	if 'timestamp_str' not in data or \
		abs((datetime.utcnow() - datetime.strptime(data['timestamp_str'], '%Y-%m-%d %H:%M:%S.%f')).total_seconds()) >= 60:
		data = request_gauges_latest_json()
		with path.open('w') as data_file:
			json.dump(data, data_file)
	return data

def sync_gauges_json(last_datetime_str):
	# If json file in data directory is stale, delete it so will get refreshed.
	data = {}
	site_root = Path(__file__).parent
	path = Path(site_root, 'data', 'gaugeslatest.json')
	if path.is_file():
		with path.open('r') as data_file:
			data = json.load(data_file)
		if last_datetime_str != data.get('datetime'):
			os.remove(str(path))

def save_debug_data(data):
	site_root = Path(__file__).parent
	path = Path(site_root, 'data', 'debug-data-{:%Y%m%d%H%M%S}'.format(datetime.utcnow()))
	with path.open('w') as debug_file:
		debug_file.write(data)
		logging.warning('Wrote debug data file: {}'.format(path))

def request_cdec_hourly_page():
	url = 'http://cdec.water.ca.gov/cgi-progs/queryF?ORO'
	req = request.Request(url=url)
	attempt = 0
	page = None
	while attempt < 3:
		attempt += 1
		try:
			with request.urlopen(req) as response:
				page = response.read()
				break
		except (URLError, HTTPError) as e:
			logging.warning('The following error has occurred: {}'.format(repr(e)))
			break
		except ConnectionResetError as e:
		   logging.warning('The following error has occurred: {}'.format(repr(e)))
	return page

def parse_page_table(label, xpath_predicate, page):
	try:
		tree = html.fromstring(page)
	except etree.XMLSyntaxError as e:
		# Bad markup returned.
		logging.warning('The following error has occurred: {}'.format(repr(e)))
		return None
	rows = tree.xpath(xpath_predicate)
	if len(rows) < 1:
		# Unexpected data somewhere.
		logging.warning('The following error has occurred: Expected {} rows, found only: {} rows'.format(
			label, len(rows)))
		save_debug_data(str(page))
		return None
	return rows

app = Flask(__name__, instance_relative_config=True)
app.config.from_pyfile('config.py', silent=True)

@app.route('/')
def index():
	data = res_latest()
	data['analytics_id'] = app.config['ANALYTICS_ID']
	return render_template('index.html', data=data)

@app.route('/gauges/')
def gauges():
	data = gauges_latest()
	return render_template('gauges.html', data=data)

@app.route('/gaugeslatest/', methods=['GET'])
def gauges_latest_json():
	return jsonify(gauges_latest())

@app.route('/reslatest/', methods=['GET'])
def res_latest_json():
	return jsonify(res_latest())

@app.route('/latest/', methods=['GET'])
@app.route('/latest/<filename>.<file_ext>', methods=['GET'])
@app.route('/latest/<filename>', methods=['GET'])
def latest(filename='latest', file_ext='png'):
	url = 'https://phantomjscloud.com/api/browser/v2/{}/?request={{url:%22https://orovilledam.org/gauges/%22,renderType:%22png%22,renderSettings:{{viewport:{{width:600,height:350}}}}}}'.format(
		app.config['PHANTOMJSCLOUD_API_KEY'])
	output_types = {
		'png': 'image/png',
	}
	if file_ext and file_ext.lower() not in output_types:
	   abort(404)
	req = request.Request(url=url)
	with request.urlopen(req) as response:
		mimetype = response.headers['Content-Type']
		image_data = io.BytesIO(response.read())
	return send_file(
		image_data,
		attachment_filename=secure_filename('{}.{}'.format(filename, file_ext)),
		mimetype=mimetype)

@app.route('/feed/', methods=['GET'])
@app.route('/rss/', methods=['GET'])
def feed():
	blank_val = '--'.rjust(10)
	PT = pytz.timezone('America/Los_Angeles')  # Pacific Timezone
	feed = AtomFeed('Recent entries', feed_url=flask_request.url, url=flask_request.url_root)
	page = request_cdec_hourly_page()
	if page is None:
		return feed.get_response()
	# There could be multiple tables in the page.
	# The table with hourly measurements has either 14 or 15 rows and exactly 15 colums.
	# The first two rows are for the title, ignore that.
	# The last row may be incomplete (blank values). Ignore these.
	last_datetime_str = None
	xpath_predicate = '//div[@class="content_left_column"]/table[count(tr) >= 14 and count(tr) <= 15]/tr[count(td) = 15][position() > 2][td[2][text()!="{}"] and td[6][text()!="{}"] and td[8][text()!="{}"]]'.format(
			blank_val, blank_val, blank_val)
	rows = parse_page_table('feed', xpath_predicate, page)
	if rows is None:
		return feed.get_response()
	for row in rows:
		data = {}
		row_datetime = datetime.strptime(row.xpath('td')[0].xpath('text()')[0], '%m/%d/%Y %H:%M')
		row_datetime = PT.localize(row_datetime)
		last_datetime_str = '{:%B %-d, %Y %-I%P} {}'.format(row_datetime, row_datetime.tzname())
		data['entryTitle'] = last_datetime_str
		data['entryText'] = 'Lake Level: {:,.2f} ft, Inflow: {:,d} cfps, Outflow: {:,d} cfps'.format(
			float(row.xpath('td')[1].xpath('text()')[0]),
			int(row.xpath('td')[7].xpath('text()')[0]),
			int(row.xpath('td')[5].xpath('text()')[0]))
		feed.add(
		   data['entryTitle'],
		   data['entryText'],
		   id='https://cdec.water.ca.gov/cgi-progs/queryF?ORO#{:%Y%m%d%H%M%S}'.format(row_datetime),
		   content_type='text',
		   author='California Dept of Water Resources https://cdec.water.ca.gov/cgi-progs/queryF?ORO',
		   url='https://cdec.water.ca.gov/cgi-progs/queryF?ORO',
		   updated=row_datetime,
		   published=row_datetime)
	# After doing a /feed, if there's a new entry, /gauge will be accessed.
	# But that could return a stale view since the gauge json was recent.
	# So if /feed has a newer reading than the gauge json, purge it so it will be recreated next time, fresh.
	sync_gauges_json(last_datetime_str)
	return feed.get_response()
<commit_msg>Only overwrite 'last' record when data was valid.<commit_after>from pathlib import Path
from urllib import request
from urllib.error import URLError, HTTPError
from lxml import html, etree
from flask import Flask, Response, request as flask_request, jsonify, render_template, send_file, abort
from werkzeug.contrib.atom import AtomFeed
from werkzeug.utils import secure_filename
from datetime import datetime, timedelta
import io
import json
import os
import pytz
import logging

def request_res_latest_json():
	# Scrape the latest available measurements from CA DWR CDEC.
	blank_val = '--'.rjust(10)
	data = {}
	# Populate at least what is needed to return if there is an error or exception.
	data['res_elev'] = 0
	data['inflow'] = 0
	data['outflow'] = 0
	page = request_cdec_hourly_page()
	if page is None:
		return data
	# There could be multiple tables in the page.
	# The table with hourly measurements has either 14 or 15 rows and exactly 15 columns.
	# The first two rows are for the title, ignore that.
	# The last row may be incomplete (blank values). Ignore these.
	xpath_predicate = '//div[@class="content_left_column"]/table[count(tr) >= 14 and count(tr) <= 15]/tr[count(td) = 15][position() > 2][td[2][text()!="{}"]]'.format(
			blank_val)
	rows = parse_page_table('res_elev', xpath_predicate, page)
	if rows is None:
		return data
	res_elev_row = rows[-1]
	xpath_predicate = '//div[@class="content_left_column"]/table[count(tr) >= 14 and count(tr) <= 15]/tr[count(td) = 15][position() > 2][td[6][text()!="{}"] and td[8][text()!="{}"]]'.format(
			blank_val, blank_val)
	rows = parse_page_table('flow', xpath_predicate, page)
	if rows is None:
		return data
	flow_row = rows[-1]
	res_elev_datetime = '{:%b %-d %-I%P}'.format(
		datetime.strptime(res_elev_row[0].xpath('text()')[0], '%m/%d/%Y %H:%M'))
	res_elev = float(res_elev_row[1].xpath('text()')[0])
	flow_datetime = '{:%b %-d %-I%P}'.format(
		datetime.strptime(flow_row[0].xpath('text()')[0], '%m/%d/%Y %H:%M'))
	outflow = int(flow_row[5].xpath('./text()')[0])
	inflow = int(flow_row[7].xpath('./text()')[0])
	data['timestamp_str'] = str(datetime.utcnow())
	data['res_elev_datetime'] = res_elev_datetime
	data['res_elev'] = res_elev
	data['flow_datetime'] = flow_datetime
	data['inflow'] = inflow
	data['outflow'] = outflow
	return data

def request_gauges_latest_json():
	# Scrape the latest available measurements from CA DWR CDEC.
	blank_val = '--'.rjust(10)
	data = {}
	# Populate at least what is needed to return if there is an error or exception.
	data['res_elev'] = 0
	data['inflow'] = 0
	data['outflow'] = 0
	PT = pytz.timezone('America/Los_Angeles')  # Pacific Timezone
	page = request_cdec_hourly_page()
	if page is None:
		return data
	# There could be multiple tables in the page.
	# The table with hourly measurements has either 14 or 15 rows and exactly 15 colums.
	# The first two rows are for the title, ignore that.
	# The last row may be incomplete (blank values). Ignore these.
	xpath_predicate = '//div[@class="content_left_column"]/table[count(tr) >= 14 and count(tr) <= 15]/tr[count(td) = 15][position() > 2][td[2][text()!="{}"] and td[6][text()!="{}"] and td[8][text()!="{}"]]'.format(
			blank_val, blank_val, blank_val)
	rows = parse_page_table('gauge', xpath_predicate, page)
	if rows is None:
		return data
	row = rows[-1]
	row_datetime = datetime.strptime(row[0].xpath('text()')[0], '%m/%d/%Y %H:%M')
	row_datetime = PT.localize(row_datetime)
	datetime_str = '{:%B %-d, %Y %-I%P} {}'.format(
		row_datetime, row_datetime.tzname())
	res_elev = float(row[1].xpath('text()')[0])
	outflow = int(row[5].xpath('./text()')[0])
	inflow = int(row[7].xpath('./text()')[0])
	data['timestamp_str'] = str(datetime.utcnow())
	data['datetime'] = datetime_str
	data['res_elev'] = res_elev
	data['inflow'] = inflow
	data['outflow'] = outflow
	return data

def res_latest():
	site_root = Path(__file__).parent
	path = Path(site_root, 'data', 'reslatest.json')
	data = {}
	if path.is_file():
		with path.open('r') as data_file:
			data = json.load(data_file)
	# If more than a minute old, get the latest.
	if 'timestamp_str' not in data or \
		abs((datetime.utcnow() - datetime.strptime(data['timestamp_str'], '%Y-%m-%d %H:%M:%S.%f')).total_seconds()) >= 60:
		data = request_res_latest_json()
		if 'datetime' in data:
			with path.open('w') as data_file:
				json.dump(data, data_file)
	return data

def gauges_latest():
	site_root = Path(__file__).parent
	path = Path(site_root, 'data', 'gaugeslatest.json')
	data = {}
	if path.is_file():
		with path.open('r') as data_file:
			data = json.load(data_file)
	# If more than a minute old, get the latest.
	if 'timestamp_str' not in data or \
		abs((datetime.utcnow() - datetime.strptime(data['timestamp_str'], '%Y-%m-%d %H:%M:%S.%f')).total_seconds()) >= 60:
		data = request_gauges_latest_json()
		if 'datetime' in data:
			with path.open('w') as data_file:
				json.dump(data, data_file)
	return data

def sync_gauges_json(last_datetime_str):
	# If json file in data directory is stale, delete it so will get refreshed.
	data = {}
	site_root = Path(__file__).parent
	path = Path(site_root, 'data', 'gaugeslatest.json')
	if path.is_file():
		with path.open('r') as data_file:
			data = json.load(data_file)
		if last_datetime_str != data.get('datetime'):
			os.remove(str(path))

def save_debug_data(data):
	site_root = Path(__file__).parent
	path = Path(site_root, 'data', 'debug-data-{:%Y%m%d%H%M%S}'.format(datetime.utcnow()))
	with path.open('w') as debug_file:
		debug_file.write(data)
		logging.warning('Wrote debug data file: {}'.format(path))

def request_cdec_hourly_page():
	url = 'http://cdec.water.ca.gov/cgi-progs/queryF?ORO'
	req = request.Request(url=url)
	attempt = 0
	page = None
	while attempt < 3:
		attempt += 1
		try:
			with request.urlopen(req) as response:
				page = response.read()
				break
		except (URLError, HTTPError) as e:
			logging.warning('The following error has occurred: {}'.format(repr(e)))
			break
		except ConnectionResetError as e:
		   logging.warning('The following error has occurred: {}'.format(repr(e)))
	return page

def parse_page_table(label, xpath_predicate, page):
	try:
		tree = html.fromstring(page)
	except etree.XMLSyntaxError as e:
		# Bad markup returned.
		logging.warning('The following error has occurred: {}'.format(repr(e)))
		return None
	rows = tree.xpath(xpath_predicate)
	if len(rows) < 1:
		# Unexpected data somewhere.
		logging.warning('The following error has occurred: Expected {} rows, found only: {} rows'.format(
			label, len(rows)))
		save_debug_data(str(page))
		return None
	return rows

app = Flask(__name__, instance_relative_config=True)
app.config.from_pyfile('config.py', silent=True)

@app.route('/')
def index():
	data = res_latest()
	data['analytics_id'] = app.config['ANALYTICS_ID']
	return render_template('index.html', data=data)

@app.route('/gauges/')
def gauges():
	data = gauges_latest()
	return render_template('gauges.html', data=data)

@app.route('/gaugeslatest/', methods=['GET'])
def gauges_latest_json():
	return jsonify(gauges_latest())

@app.route('/reslatest/', methods=['GET'])
def res_latest_json():
	return jsonify(res_latest())

@app.route('/latest/', methods=['GET'])
@app.route('/latest/<filename>.<file_ext>', methods=['GET'])
@app.route('/latest/<filename>', methods=['GET'])
def latest(filename='latest', file_ext='png'):
	url = 'https://phantomjscloud.com/api/browser/v2/{}/?request={{url:%22https://orovilledam.org/gauges/%22,renderType:%22png%22,renderSettings:{{viewport:{{width:600,height:350}}}}}}'.format(
		app.config['PHANTOMJSCLOUD_API_KEY'])
	output_types = {
		'png': 'image/png',
	}
	if file_ext and file_ext.lower() not in output_types:
	   abort(404)
	req = request.Request(url=url)
	with request.urlopen(req) as response:
		mimetype = response.headers['Content-Type']
		image_data = io.BytesIO(response.read())
	return send_file(
		image_data,
		attachment_filename=secure_filename('{}.{}'.format(filename, file_ext)),
		mimetype=mimetype)

@app.route('/feed/', methods=['GET'])
@app.route('/rss/', methods=['GET'])
def feed():
	blank_val = '--'.rjust(10)
	PT = pytz.timezone('America/Los_Angeles')  # Pacific Timezone
	feed = AtomFeed('Recent entries', feed_url=flask_request.url, url=flask_request.url_root)
	page = request_cdec_hourly_page()
	if page is None:
		return feed.get_response()
	# There could be multiple tables in the page.
	# The table with hourly measurements has either 14 or 15 rows and exactly 15 colums.
	# The first two rows are for the title, ignore that.
	# The last row may be incomplete (blank values). Ignore these.
	last_datetime_str = None
	xpath_predicate = '//div[@class="content_left_column"]/table[count(tr) >= 14 and count(tr) <= 15]/tr[count(td) = 15][position() > 2][td[2][text()!="{}"] and td[6][text()!="{}"] and td[8][text()!="{}"]]'.format(
			blank_val, blank_val, blank_val)
	rows = parse_page_table('feed', xpath_predicate, page)
	if rows is None:
		return feed.get_response()
	for row in rows:
		data = {}
		row_datetime = datetime.strptime(row.xpath('td')[0].xpath('text()')[0], '%m/%d/%Y %H:%M')
		row_datetime = PT.localize(row_datetime)
		last_datetime_str = '{:%B %-d, %Y %-I%P} {}'.format(row_datetime, row_datetime.tzname())
		data['entryTitle'] = last_datetime_str
		data['entryText'] = 'Lake Level: {:,.2f} ft, Inflow: {:,d} cfps, Outflow: {:,d} cfps'.format(
			float(row.xpath('td')[1].xpath('text()')[0]),
			int(row.xpath('td')[7].xpath('text()')[0]),
			int(row.xpath('td')[5].xpath('text()')[0]))
		feed.add(
		   data['entryTitle'],
		   data['entryText'],
		   id='https://cdec.water.ca.gov/cgi-progs/queryF?ORO#{:%Y%m%d%H%M%S}'.format(row_datetime),
		   content_type='text',
		   author='California Dept of Water Resources https://cdec.water.ca.gov/cgi-progs/queryF?ORO',
		   url='https://cdec.water.ca.gov/cgi-progs/queryF?ORO',
		   updated=row_datetime,
		   published=row_datetime)
	# After doing a /feed, if there's a new entry, /gauge will be accessed.
	# But that could return a stale view since the gauge json was recent.
	# So if /feed has a newer reading than the gauge json, purge it so it will be recreated next time, fresh.
	sync_gauges_json(last_datetime_str)
	return feed.get_response()
<|endoftext|>


<commit_before>from operator import itemgetter
from web_ui import app

class PageManager(object):
    """ Responsible for creating pages. Reads in possible layouts,
        and tries to arrange content in the most beneficial way.
    """

    def __init__(self):
        """ Initializes screen size and  loads all layouts
        for the right context, using the context defined in
        the concrete subclass.
        """

        self.screen_size = (12.0, 8.0)
        self.all_layouts = app.config['LAYOUT_DEFINITIONS']

    def _create_page_section(self, cell, content):
        """ Creates a section of a page consisting of a dict
        with the containing the css_class and the content of
        the section.
        CSS class name format
         col   column at which the css container begins
         row   row at which it begins
         w     width of the container
         h     height of the container
        """
        css_class = "col{} row{} w{} h{}".format(
            cell[0],
            cell[1],
            cell[2],
            cell[3])

        return {'css_class': css_class, 'content': content}

    def _get_layouts_for(self, context):
        """ Returns all layouts appropriate for context """

        return [layout for layout in self.all_layouts if
                context in layout['context']]

    def create_star_layout(self):
        """Returns a page for creating stars."""

        context = 'create_star_page'
        layouts = self._get_layouts_for(context)

        # currently no logic to choose among different layouts
        assert(len(layouts) == 1)
        page = {'create_star_form': []}

        for cell in layouts[0]['create_star_form']:
            page['create_star_form'].append(
                self._create_page_section(cell, None))

        return page

    def star_layout(self, stars):
        """Return the optimal layouted page for the given stars."""

        context = 'star_page'
        layouts = self._get_layouts_for(context)

        # Rank stars by score
        stars_ranked = sorted(stars, key=lambda s: s.hot(), reverse=True)

        # Find best layout by filling each one with stars
        # and determining which one gives the best score
        layout_scores = dict()
        for layout in layouts:
            # print("\nLayout: {}".format(layout['name']))
            layout_scores[layout['name']] = 0

            for i, star_cell in enumerate(layout['stars']):
                if i >= len(stars_ranked):
                    continue
                star = stars_ranked[i]

                cell_score = self._cell_score(star_cell)
                layout_scores[layout['name']] += star.hot() * cell_score
                # print("{}\t{}\t{}".format(star, star.hot() * cell_score, cell_score))
            # print("Score: {}".format(layout_scores[layout['name']]))

        # Select best layout
        selected_layouts = sorted(
            layout_scores.iteritems(),
            key=itemgetter(1),
            reverse=True)

        if len(selected_layouts) == 0:
            # TODO: Throw exception here, if no layout found the PM failed
            app.logger.error("No fitting layout found")
            return

        for layout in layouts:
            if layout['name'] == selected_layouts[0][0]:
                break

        # print("Chosen {}".format(layout))

        # Create list of elements in layout
        page = {'stars': []}
        for i, star_cell in enumerate(layout['stars']):
            if i >= len(stars_ranked):
                break

            star = stars_ranked[i]
            page['stars'].append(self._create_page_section(star_cell, star))

        return page

    def _cell_score(self, cell):
        """ Return a score that describes how valuable a given cell on
        the screen is. Cell on the top left is 1.0, score diminishes to
        the right and bottom. Bigger cells get higher scores, cells off
        the screen get a 0.0 """

        import math

        # position score
        if cell[0] > self.screen_size[0] or cell[1] > self.screen_size[1]:
            pscore = 0.0
        else:
            score_x = 1.0 if cell[0] == 0 else 1.0 / (1.0 + (cell[0] / self.screen_size[0]))
            score_y = 1.0 if cell[1] == 0 else 1.0 / (1.0 + (cell[1] / self.screen_size[1]))
            pscore = (score_x + score_y) / 2.0

        # size score (sigmoid)
        area = cell[2] * cell[3]
        sscore = 1.0 / (1 + pow(math.exp(1), -0.1 * (area - 12.0)))
        return pscore * sscore
<commit_msg>PageManager now reaturns Page instances instead of dicts. OOP for the win!<commit_after>from operator import itemgetter
from web_ui import app

class PageManager(object):
    """ Holds all factory methods for Page creation.
    Those methods read in possible layouts, and create
    and a Page instance arrangeing content in the most
    beneficial way. Also see Page class below.
    """

    def __init__(self):
        """ Initializes screen size and  loads all layouts
        for the right context, using the context defined in
        the concrete subclass.
        """

        self.screen_size = (12.0, 8.0)
        self.all_layouts = app.config['LAYOUT_DEFINITIONS']

    def _get_layouts_for(self, context):
        """ Returns all layouts appropriate for context """

        return [layout for layout in self.all_layouts if
                context in layout['context']]

    def create_star_layout(self):
        """Returns a page for creating stars."""

        # use layouts for create_star_page context
        context = 'create_star_page'
        layouts = self._get_layouts_for(context)

        # currently no logic to choose among different layouts
        assert(len(layouts) == 1)

        # page entry containing the form will be called:
        section = 'create_star_form'

        page = Page()

        for cell in layouts[0][section]:
            page.add_to_section(section, cell, None)

        return page

    def star_layout(self, stars):
        """Return the optimal layouted page for the given stars."""

        context = 'star_page'
        layouts = self._get_layouts_for(context)

        section = 'stars'

        # Rank stars by score
        stars_ranked = sorted(stars, key=lambda s: s.hot(), reverse=True)

        # Find best layout by filling each one with stars
        # and determining which one gives the best score
        layout_scores = dict()
        for layout in layouts:
            # print("\nLayout: {}".format(layout['name']))
            layout_scores[layout['name']] = 0

            for i, star_cell in enumerate(layout[section]):
                if i >= len(stars_ranked):
                    continue
                star = stars_ranked[i]

                cell_score = self._cell_score(star_cell)
                layout_scores[layout['name']] += star.hot() * cell_score
                # print("{}\t{}\t{}".format(star, star.hot() * cell_score, cell_score))
            # print("Score: {}".format(layout_scores[layout['name']]))

        # Select best layout
        selected_layouts = sorted(
            layout_scores.iteritems(),
            key=itemgetter(1),
            reverse=True)

        if len(selected_layouts) == 0:
            # TODO: Throw exception here, if no layout found the PM failed
            app.logger.error("No fitting layout found")
            return

        for layout in layouts:
            if layout['name'] == selected_layouts[0][0]:
                break

        # print("Chosen {}".format(layout))

        page = Page()
        for i, star_cell in enumerate(layout['stars']):
            if i >= len(stars_ranked):
                break

            star = stars_ranked[i]
            page.add_to_section(section, star_cell, star)

        return page

    def _cell_score(self, cell):
        """ Return a score that describes how valuable a given cell on
        the screen is. Cell on the top left is 1.0, score diminishes to
        the right and bottom. Bigger cells get higher scores, cells off
        the screen get a 0.0 """

        import math

        # position score
        if cell[0] > self.screen_size[0] or cell[1] > self.screen_size[1]:
            pscore = 0.0
        else:
            score_x = 1.0 if cell[0] == 0 else 1.0 / (1.0 + (cell[0] / self.screen_size[0]))
            score_y = 1.0 if cell[1] == 0 else 1.0 / (1.0 + (cell[1] / self.screen_size[1]))
            pscore = (score_x + score_y) / 2.0

        # size score (sigmoid)
        area = cell[2] * cell[3]
        sscore = 1.0 / (1 + pow(math.exp(1), -0.1 * (area - 12.0)))
        return pscore * sscore

class Page(object):
    """Responsible for the layout of a page. Instances of Page
    hold variables for each dynamic section of the page. This
    variables hold a list with the entries of the section.
    An entry is a dict with the keys 'css_class' and 'content'.
    Example:
        star_page:
            + header:[
                {css_class:'col1 row1 w3 h1', content: None}
              ]
            + stars: [
                {css_class:'col1 row2 w3 h1', content: star1},
                {css_class:'col1 row3 w3 h1', content: star2}
              ]
    """

    def add_to_section(self, section, entry, content):
        """ Adds a new entry to the page section 'section'
        (and creates it if necessary). """

        section_entry = self._create_entry(entry, content)

        if not hasattr(self, section):
            setattr(self, section, [])

        attr = getattr(self, section)
        attr.append(section_entry)

    def _create_entry(self, cell, content):
        """ Creates a section of a page consisting of a dict
        containing the css_class and the content of the section.
        CSS class name format
         col   column at which the css container begins
         row   row at which it begins
         w     width of the container
         h     height of the container
        """
        css_class = "col{} row{} w{} h{}".format(
            cell[0],
            cell[1],
            cell[2],
            cell[3])

        return {'css_class': css_class, 'content': content}
<|endoftext|>


<commit_before>from os import path, mkdir
from subprocess import call
from time import time
import numpy as np
import pandas as pd
import yaml
from nbodykit.cosmology import Cosmology
from nbodykit.cosmology.power import LinearPower, ZeldovichPower

# I could have this accept kwargs,
# and have main
def submit_sims(lhc_fname, output_dir, powerspectrum, cosmic_var, initial_z, final_z,\
                npart, boxsize, ncores, sim_time):

    assert path.isdir(output_dir), "Invalid output directory %s"%output_dir

    lhc = pd.read_csv(lhc_fname, sep=' ', index_col=None)
    param_names = list(lhc)

    k = np.logspace(-4, 2, 1000)
    # others?
    assert powerspectrum.lower() in {'linear', 'zeldovich'}, "Invalid initial power spectrum"

    if powerspectrum.lower() == 'linear':
        power = LinearPower
    else:
        power = ZeldovichPower

    cosmic_var = str(cosmic_var)
    assert cosmic_var.lower() in {'true', 'false'}

    initial_a, final_a = 1./(1+initial_z), 1./(1+final_z)

    for idx, row in lhc.iterrows():
        # TODO what to do if already exists? likely will throw some kind of error
        boxdir = path.join(output_dir, "Box_%03d/"%idx)
        if not path.isdir(boxdir):
            mkdir(boxdir)

        cosmo = make_cosmo(param_names, row)

        p = power(cosmo, final_z)(k)

        np.savetxt(path.join(boxdir,'powerspec.txt'), np.array((k,p)).T)

        with open(path.join(boxdir, 'nbodykit.lua'), 'w') as f:
           f.write(sim_config_template.format(nc = npart, boxsize = boxsize,
                                              initial_a = initial_a,
                                              final_a = final_a,
                                              final_z = final_z,
                                              cosmic_var = cosmic_var,
                                              omega_m = cosmo.Om0,
                                              h = cosmo.h,
                                              seed = int(time())%1000,
                                              boxdir = boxdir))

        # now, write submission file and submit

        with open(path.join(boxdir, 'submit_sim.sbatch'), 'w') as f:
            f.write(submission_file_template.format(jobname = 'Box_%03d'%idx,
                                                    boxdir = boxdir,
                                                    time = sim_time*60, #TODO in config
                                                    ntasks = ncores))

        #call("sbatch {boxdir}submit_sim.sbatch".format(boxdir=boxdir), shell=True)

def make_cosmo(param_names, param_vals):
    """
    Frequently, the parameters of our input will not align with what we need.
    So, we convert them to the appropriate format.
    :param param_names:
    :param param_vals:
    :return:
    """
    param_dict = dict([(pn, param_vals[pn]) for pn in param_names] )
    # TODO it'd be nice if this could be somehow generalized.
    param_dict['N_ncdm'] = 3.0
    param_dict['N_ur'] = param_dict['Neff']-3.0
    del param_dict['Neff']

    #param_dict['h'] = param_dict['H0']/100
    #del param_dict['H0']
    param_dict['h'] = param_dict['h0']
    del param_dict['h0']

    param_dict['w0_fld'] = param_dict['w']
    del param_dict['w']

    param_dict['Omega0_b'] = param_dict['Omega_b']
    param_dict['Omega0_cdm'] = param_dict['Omega_m'] - param_dict['Omega_b']
    del param_dict['Omega_b']
    del param_dict['Omega_m']

    param_dict['ln_1e10_A_s'] = param_dict['ln(10^{10}A_s)']
    del param_dict['ln(10^{10}A_s)']

    param_dict['Omega0_ncdm_tot'] = param_dict['Omeganuh2']/(param_dict['h']**2)
    del param_dict['Omeganuh2']

    return Cosmology(**param_dict)

sim_config_template = """
nc = {nc:d}
boxsize  = {boxsize:.1f}

time_step = linspace({initial_a:.3f}, {final_a:.3f}, 20)

output_redshifts = {{{final_z:.2f}}} --redshift of output

omega_m = {omega_m: f}
h = {h:f}

read_powerspectrum = "{boxdir}powerspec.txt"
remove_cosmic_variance = {cosmic_var}
random_seed = {seed:d}

pm_nc_factor = 2
np_alloc_factor = 4.0

fof_nmin = 30

write_snapshot = "{boxdir}fastpm"
write_nonlineark= "{boxdir}fastpm"
write_fof = "{boxdir}fastpm"

-- 1d power spectrum (raw), without shotnoise correction
write_powerspectrum = "{boxdir}powerspec-debug"
"""

submission_file_template = """#!/bin/bash
#SBATCH -p regular
#SBATCH --time={time:d}:00
##SBATCH --qos regular
#SBATCH -A cosmosim
#SBATCH --mail-type=ALL
#SBATCH --mail-user=swmclau2@stanford.edu
#SBATCH -N {ntasks:d}
#SBATCH --ntasks={ntasks:d}
#SBATCH --cpus-per-task=16
#SBATCH --mem-per-cpu=MaxMemPerCPU
#SBATCH -J {jobname}
#SBATCH --output={boxdir}{jobname}.out
#SBATCH --error={boxdir}{jobname}.err

srun /global/u2/s/swmclau2/Git/fastpm/src/fastpm {boxdir}nbodykit.lua
"""
#may need to be an mpirun?

if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser(description='Submit an ensemble of simulations.')
    parser.add_argument('config_fname', help = 'Filename of the configfile.')
    # TODO here is where i would add an argument to parse what sims to run
    # TODO add arg to do halos, or should that just be a separate file.
    args = parser.parse_args()
    config_fname = args.config_fname
    # NOTE could use a vars(args) to turn into a dictionary.
    # currently not doing that if some of those end up not being arguments for submit_sims
    try:
        assert path.isfile(config_fname)
    except AssertionError:
        raise AssertionError("%s is not a valid filename." % config_fname)

    with open(config_fname, 'r') as ymlfile:
        cfg = yaml.load(ymlfile)

    try:
        submit_sims(**cfg)
    except TypeError:
        raise TypeError("Not all require arguments specified in yaml config.")
<commit_msg>almost done being at war with nbodykit.<commit_after>from os import path, mkdir
from subprocess import call
from time import time
import numpy as np
import pandas as pd
import yaml
from nbodykit.cosmology import Cosmology
from nbodykit.cosmology.power import LinearPower, ZeldovichPower

# I could have this accept kwargs,
# and have main
def submit_sims(lhc_fname, output_dir, powerspectrum, cosmic_var, initial_z, final_z,\
                npart, boxsize, ncores, sim_time):

    assert path.isdir(output_dir), "Invalid output directory %s"%output_dir

    lhc = pd.read_csv(lhc_fname, sep=' ', index_col=None)
    param_names = list(lhc)

    k = np.logspace(-4, 2, 1000)
    # others?
    assert powerspectrum.lower() in {'linear', 'zeldovich'}, "Invalid initial power spectrum"

    if powerspectrum.lower() == 'linear':
        power = LinearPower
    else:
        power = ZeldovichPower

    cosmic_var = str(cosmic_var)
    assert cosmic_var.lower() in {'true', 'false'}

    initial_a, final_a = 1./(1+initial_z), 1./(1+final_z)

    for idx, row in lhc.iterrows():
        # TODO what to do if already exists? likely will throw some kind of error
        boxdir = path.join(output_dir, "Box_%03d/"%idx)
        if not path.isdir(boxdir):
            mkdir(boxdir)

        cosmo = make_cosmo(param_names, row)
	p = power(cosmo, 0.0)(k) #only works for z=0

        np.savetxt(path.join(boxdir,'powerspec.txt'), np.array((k,p)).T)

        with open(path.join(boxdir, 'nbodykit.lua'), 'w') as f:
           f.write(sim_config_template.format(nc = npart, boxsize = boxsize,
                                              initial_a = initial_a,
                                              final_a = final_a,
                                              final_z = final_z,
                                              cosmic_var = cosmic_var,
                                              omega_m = cosmo.Om0,
                                              h = cosmo.h,
                                              seed = int(time())%1000,
                                              boxdir = boxdir))

        # now, write submission file and submit

        with open(path.join(boxdir, 'submit_sim.sbatch'), 'w') as f:
            f.write(submission_file_template.format(jobname = 'Box_%03d'%idx,
                                                    boxdir = boxdir,
                                                    time = sim_time*60, #TODO in config
                                                    ntasks = ncores))

        #call("sbatch {boxdir}submit_sim.sbatch".format(boxdir=boxdir), shell=True)

def make_cosmo(param_names, param_vals):
    """
    Frequently, the parameters of our input will not align with what we need.
    So, we convert them to the appropriate format.
    :param param_names:
    :param param_vals:
    :return:
    """
    param_dict = dict([(pn, param_vals[pn]) for pn in param_names] )
    # TODO it'd be nice if this could be somehow generalized.
    param_dict['N_ncdm'] = 3.0
    param_dict['N_ur'] = param_dict['Neff']-3.0
    del param_dict['Neff']

    #param_dict['h'] = param_dict['H0']/100
    #del param_dict['H0']
    param_dict['h'] = param_dict['h0']
    del param_dict['h0']

    #param_dict['w0_fld'] = param_dict['w']
    #del param_dict['w']

    param_dict['Omega_cdm'] = param_dict['Omega_m'] - param_dict['Omega_b']
    del param_dict['Omega_b']
    del param_dict['Omega_m']

    param_dict['Omega_ncdm'] = [param_dict['Omeganuh2']/(param_dict['h']**2), 0.0, 0.0]
    #param_dict['m_ncdm']=None
    #param_dict['omch2'] = (param_dict['Omega_m'] - param_dict['Omega_b'] - param_dict['Omega0_ncdm_tot'])*(param_dict['h']**2)
    del param_dict['Omeganuh2']

    #param_dict['A_s'] = 10**(np.log10(np.exp(param_dict['ln_1e10_A_s']))-10.0)
    param_dict['A_s'] = 10**(-10)*np.exp(param_dict['ln(10^{10}A_s)'])
    del param_dict['ln(10^10{10}A_s']

    #print param_dict
    # this is seriously what i have to do here.
    C= Cosmology()
    C2 = C.from_dict(param_dict)
    C3 = C2.clone(w0_fld = param_dict['w'])
    return C3#Cosmology(**param_dict)

sim_config_template = """
nc = {nc:d}
boxsize  = {boxsize:.1f}

time_step = linspace({initial_a:.3f}, 1.0, 20)

output_redshifts = {{{final_z:.2f}, 0.0}} --redshift of output

omega_m = {omega_m: f}
h = {h:f}

read_powerspectrum = "{boxdir}powerspec.txt"
remove_cosmic_variance = {cosmic_var}
random_seed = {seed:d}

pm_nc_factor = 2
np_alloc_factor = 4.0

fof_nmin = 30

write_snapshot = "{boxdir}fastpm"
write_nonlineark= "{boxdir}fastpm"
write_fof = "{boxdir}fastpm"

-- 1d power spectrum (raw), without shotnoise correction
write_powerspectrum = "{boxdir}powerspec-debug"
"""

submission_file_template = """#!/bin/bash
#SBATCH -p regular
#SBATCH --time={time:d}:00
##SBATCH --qos regular
#SBATCH -A cosmosim
#SBATCH --mail-type=ALL
#SBATCH --mail-user=swmclau2@stanford.edu
#SBATCH -N {ntasks:d}
#SBATCH --ntasks={ntasks:d}
#SBATCH --cpus-per-task=16
#SBATCH --mem-per-cpu=MaxMemPerCPU
#SBATCH -J {jobname}
#SBATCH --output={boxdir}{jobname}.out
#SBATCH --error={boxdir}{jobname}.err

srun /global/u2/s/swmclau2/Git/fastpm/src/fastpm {boxdir}nbodykit.lua
"""
#may need to be an mpirun?

if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser(description='Submit an ensemble of simulations.')
    parser.add_argument('config_fname', help = 'Filename of the configfile.')
    # TODO here is where i would add an argument to parse what sims to run
    # TODO add arg to do halos, or should that just be a separate file.
    args = parser.parse_args()
    config_fname = args.config_fname
    # NOTE could use a vars(args) to turn into a dictionary.
    # currently not doing that if some of those end up not being arguments for submit_sims
    try:
        assert path.isfile(config_fname)
    except AssertionError:
        raise AssertionError("%s is not a valid filename." % config_fname)

    with open(config_fname, 'r') as ymlfile:
        cfg = yaml.load(ymlfile)

    try:
        submit_sims(**cfg)
    except TypeError:
        raise TypeError("Not all require arguments specified in yaml config.")
<|endoftext|>


<commit_before>from flask_restful import marshal, marshal_with, fields, reqparse
from flask import abort, g
from flask import Blueprint as FlaskBlueprint

import datetime
import logging
import json

from pebbles.models import db, Blueprint, Instance, InstanceLog, User
from pebbles.forms import InstanceForm, UserIPForm
from pebbles.server import app, restful
from pebbles.utils import requires_admin, memoize
from pebbles.tasks import run_update, update_user_connectivity, fetch_running_instance_logs
from pebbles.views.commons import auth, is_group_manager
from pebbles.rules import apply_rules_instances, get_group_blueprint_ids_for_instances

from pebbles.tasks import send_mails

instances = FlaskBlueprint('instances', __name__)

USER_INSTANCE_LIMIT = 5

instance_fields = {
    'id': fields.String,
    'name': fields.String,
    'provisioned_at': fields.DateTime,
    'lifetime_left': fields.Integer,
    'maximum_lifetime': fields.Integer,
    'runtime': fields.Float,
    'state': fields.String,
    'to_be_deleted': fields.Boolean,
    'error_msg': fields.String,
    'username': fields.String,
    'user_id': fields.String,
    'blueprint': fields.String,
    'blueprint_id': fields.String,
    'cost_multiplier': fields.Float(default=1.0),
    'can_update_connectivity': fields.Boolean(default=False),
    'instance_data': fields.Raw,
    'public_ip': fields.String,
    'client_ip': fields.String(default='not set'),
    'logs': fields.Raw,
}

instance_log_fields = {
    'id': fields.String,
    'instance_id': fields.String,
    'log_type': fields.String,
    'log_level': fields.String,
    'timestamp': fields.Float,
    'message': fields.String
}

def query_blueprint(blueprint_id):
    return Blueprint.query.filter_by(id=blueprint_id).first()

def query_user(user_id):
    return User.query.filter_by(id=user_id).first()

def positive_integer(input_value):
    """Return input_value if valid, raise an exception in other case."""
    try:
        input_int = int(input_value)
    except:
        raise ValueError('{} is not a valid integer'.format(input_value))
    if input_int >= 0:
        return input_int
    else:
        raise ValueError('{} is not a positive integer'.format(input_value))

class InstanceList(restful.Resource):
    parser = reqparse.RequestParser()
    parser.add_argument('show_deleted', type=bool, default=False, location='args')
    parser.add_argument('show_only_mine', type=bool, default=False, location='args')
    parser.add_argument('offset', type=positive_integer, location='args')
    parser.add_argument('limit', type=positive_integer, location='args')

    @auth.login_required
    @marshal_with(instance_fields)
    def get(self):
        user = g.user
        args = self.parser.parse_args()
        q = apply_rules_instances(user, args)
        q = q.order_by(Instance.provisioned_at)
        instances = q.all()

        get_blueprint = memoize(query_blueprint)
        get_user = memoize(query_user)
        for instance in instances:
            instance_logs = get_logs_from_db(instance.id)
            instance.logs = marshal(instance_logs, instance_log_fields)

            user = get_user(instance.user_id)
            if user:
                instance.username = user.eppn

            blueprint = get_blueprint(instance.blueprint_id)
            if not blueprint:
                logging.warn("instance %s has a reference to non-existing blueprint" % instance.id)
                continue

            age = 0
            if instance.provisioned_at:
                age = (datetime.datetime.utcnow() - instance.provisioned_at).total_seconds()
            instance.lifetime_left = max(blueprint.maximum_lifetime - age, 0)
            instance.maximum_lifetime = blueprint.maximum_lifetime
            instance.cost_multiplier = blueprint.cost_multiplier

            if instance.to_be_deleted:
                instance.state = Instance.STATE_DELETING

        return instances

    @auth.login_required
    def post(self):
        user = g.user

        form = InstanceForm()
        if not form.validate_on_submit():
            logging.warn("validation error on user login")
            return form.errors, 422

        blueprint_id = form.blueprint.data
        allowed_blueprint_ids = get_group_blueprint_ids_for_instances(user)
        if not user.is_admin and blueprint_id not in allowed_blueprint_ids:
            abort(403)
        blueprint = Blueprint.query.filter_by(id=blueprint_id, is_enabled=True).first()
        if not blueprint:
            abort(404)
        if user.quota_exceeded():
            return {'error': 'USER_OVER_QUOTA'}, 409

        if blueprint.preallocated_credits:
            preconsumed_amount = blueprint.cost()
            total_credits_spent = preconsumed_amount + user.credits_spent
            if user.credits_quota < total_credits_spent:
                return {'error': 'USER_OVER_QUOTA'}, 409

        instances_for_user = Instance.query.filter_by(
            blueprint_id=blueprint.id,
            user_id=user.id
        ).filter(Instance.state != 'deleted').all()

        user_instance_limit = blueprint.full_config.get('maximum_instances_per_user', USER_INSTANCE_LIMIT)
        if instances_for_user and len(instances_for_user) >= user_instance_limit:
            return {'error': 'BLUEPRINT_INSTANCE_LIMIT_REACHED'}, 409

        instance = Instance(blueprint, user)
        # XXX: Choosing the name should be done in the model's constructor method
        # decide on a name that is not used currently
        existing_names = set(x.name for x in Instance.query.all())
        # Note: the potential race is solved by unique constraint in database
        while True:
            c_name = Instance.generate_name(prefix=app.dynamic_config.get('INSTANCE_NAME_PREFIX'))
            if c_name not in existing_names:
                instance.name = c_name
                break
        db.session.add(instance)
        db.session.commit()

        if not app.dynamic_config.get('SKIP_TASK_QUEUE'):
            run_update.delay(instance.id)

        return marshal(instance, instance_fields), 200

class InstanceView(restful.Resource):
    parser = reqparse.RequestParser()
    parser.add_argument('state', type=str)
    parser.add_argument('public_ip', type=str)
    parser.add_argument('error_msg', type=str)
    parser.add_argument('client_ip', type=str)
    parser.add_argument('instance_data', type=str)
    parser.add_argument('to_be_deleted', type=bool)
    parser.add_argument('send_email', type=bool)

    @auth.login_required
    @marshal_with(instance_fields)
    def get(self, instance_id):
        user = g.user
        args = {'instance_id': instance_id}
        query = apply_rules_instances(user, args)
        instance = query.first()
        if not instance:
            abort(404)

        blueprint = Blueprint.query.filter_by(id=instance.blueprint_id).first()
        instance.blueprint_id = blueprint.id
        instance.username = instance.user
        instance_logs = get_logs_from_db(instance.id)
        instance.logs = marshal(instance_logs, instance_log_fields)

        if 'allow_update_client_connectivity' in blueprint.full_config \
                and blueprint.full_config['allow_update_client_connectivity']:
            instance.can_update_connectivity = True

        age = 0
        if instance.provisioned_at:
            age = (datetime.datetime.utcnow() - instance.provisioned_at).total_seconds()
        instance.lifetime_left = max(blueprint.maximum_lifetime - age, 0)
        instance.maximum_lifetime = blueprint.maximum_lifetime
        instance.cost_multiplier = blueprint.cost_multiplier

        return instance

    @auth.login_required
    @marshal_with(instance_fields)
    def post(self, instance_id):
        args = self.parser.parse_args()
        user = g.user
        query = apply_rules_instances(user, {'instance_id': instance_id})
        instance = query.first()
        if not instance:
            abort(404)
        if args.get('send_email') and instance.state != 'running':
            text = {"subject": " ", "message": " "}
            admin_users = User.query.filter_by(is_admin=True).all()
            for admins in admin_users:
                if admins.eppn != 'worker@pebbles':
                    text['subject'] = "Notebooks.csc.fi:WARNING"
                    text['message'] = instance.name + " is taking more than ten minutes to launch"
                    # send email only through email_id because some eppn bounce back.
                    # Also the email_id will be updated once they login so here it is available
                    send_mails.delay([(admins.email_id, 'None', 't')], text)
        return instance

    @auth.login_required
    def delete(self, instance_id):
        user = g.user
        query = Instance.query.filter_by(id=instance_id)
        if not user.is_admin and not is_group_manager(user):
            query = query.filter_by(user_id=user.id)
        instance = query.first()
        if not instance:
            abort(404)
        group = instance.blueprint.group
        if not user.is_admin and not is_group_manager(user, group) and instance.user_id != user.id:
            abort(403)
        instance.to_be_deleted = True
        instance.state = Instance.STATE_DELETING
        instance.deprovisioned_at = datetime.datetime.utcnow()
        db.session.commit()
        if not app.dynamic_config.get('SKIP_TASK_QUEUE'):
            run_update.delay(instance.id)

    @auth.login_required
    def put(self, instance_id):
        user = g.user
        form = UserIPForm()
        if not form.validate_on_submit():
            logging.warn("validation error on UserIPForm")
            return form.errors, 422

        instance = Instance.query.filter_by(id=instance_id, user_id=user.id).first()
        if not instance:
            abort(404)

        blueprint = Blueprint.query.filter_by(id=instance.blueprint_id).first()
        if 'allow_update_client_connectivity' in blueprint.full_config \
                and blueprint.full_config['allow_update_client_connectivity']:
            instance.client_ip = form.client_ip.data
            if not app.dynamic_config.get('SKIP_TASK_QUEUE'):
                update_user_connectivity.delay(instance.id)
            db.session.commit()

        else:
            abort(400)

    @auth.login_required
    @requires_admin
    def patch(self, instance_id):
        args = self.parser.parse_args()
        instance = Instance.query.filter_by(id=instance_id).first()
        if not instance:
            abort(404)

        if args.get('state'):
            instance.state = args['state']
            if instance.state == Instance.STATE_RUNNING:
                if not instance.provisioned_at:
                    instance.provisioned_at = datetime.datetime.utcnow()
            if args['state'] == Instance.STATE_FAILED:
                instance.errored = True

            db.session.commit()

        if args.get('to_be_deleted'):
            instance.to_be_deleted = args['to_be_deleted']
            db.session.commit()

        if args.get('error_msg'):
            instance.error_msg = args['error_msg']
            db.session.commit()

        if args.get('public_ip'):
            instance.public_ip = args['public_ip']
            db.session.commit()

        if args.get('instance_data'):
            try:
                instance.instance_data = json.loads(args['instance_data'])
            except ValueError:
                logging.warn("invalid instance_data passed to view: %s" % args['instance_data'])
            db.session.commit()

class InstanceLogs(restful.Resource):
    parser = reqparse.RequestParser()
    parser.add_argument('log_record', type=dict)
    parser.add_argument('log_type', type=str)
    parser.add_argument('send_log_fetch_task', type=bool)

    @auth.login_required
    @marshal_with(instance_log_fields)
    def get(self, instance_id):
        args = self.parser.parse_args()
        instance = Instance.query.filter_by(id=instance_id).first()
        if not instance:
            abort(404)
        if args.get('log_type'):
            log_type = args['log_type']
        instance_logs = get_logs_from_db(instance_id, log_type)
        return instance_logs

    @auth.login_required
    @requires_admin
    def patch(self, instance_id):
        args = self.parser.parse_args()
        instance = Instance.query.filter_by(id=instance_id).first()
        if not instance:
            abort(404)
        if args.get('send_log_fetch_task'):
            if not app.dynamic_config.get('SKIP_TASK_QUEUE'):
                fetch_running_instance_logs.delay(instance_id)
        if args.get('log_record'):
            log_record = args['log_record']
            instance_log = process_logs(instance_id, log_record)
            db.session.add(instance_log)
            db.session.commit()

        return 'ok'

    @auth.login_required
    @requires_admin
    def delete(self, instance_id):
        args = self.parser.parse_args()
        log_type = args.get('log_type')
        if log_type:
            instance_logs = get_logs_from_db(instance_id, args.get('log_type'))
        else:
            instance_logs = get_logs_from_db(instance_id)
        if not instance_logs:
            logging.warn("There are no log entries to be deleted")

        for instance_log in instance_logs:
            db.session.delete(instance_log)
        db.session.commit()

def get_logs_from_db(instance_id, log_type=None):
    logs_query = InstanceLog.query\
        .filter_by(instance_id=instance_id)
    if log_type:
        logs_query = logs_query.filter_by(log_type=log_type)
    logs_query = logs_query.order_by(InstanceLog.timestamp)
    logs = logs_query.all()
    return logs

def process_logs(instance_id, log_record):

    check_running_log = get_logs_from_db(instance_id, "running")

    if check_running_log:  # in case of running logs, the whole message is replaced again (thus only 1 running log record with all info)
        instance_log = check_running_log[0]
        if log_record['log_type'] == "running":
            instance_log.timestamp = float(log_record['timestamp'])
            instance_log.message = log_record['message']  # replace the whole text (older text now appended with the new text)
            return instance_log

    instance_log = InstanceLog(instance_id)
    instance_log.log_type = log_record['log_type']
    instance_log.log_level = log_record['log_level']
    instance_log.timestamp = float(log_record['timestamp'])
    instance_log.message = log_record['message']

    return instance_log
<commit_msg>Running logs has been enabled for group owners<commit_after>from flask_restful import marshal, marshal_with, fields, reqparse
from flask import abort, g
from flask import Blueprint as FlaskBlueprint

import datetime
import logging
import json

from pebbles.models import db, Blueprint, Instance, InstanceLog, User
from pebbles.forms import InstanceForm, UserIPForm
from pebbles.server import app, restful
from pebbles.utils import requires_admin, requires_group_owner_or_admin, memoize
from pebbles.tasks import run_update, update_user_connectivity, fetch_running_instance_logs
from pebbles.views.commons import auth, is_group_manager
from pebbles.rules import apply_rules_instances, get_group_blueprint_ids_for_instances

from pebbles.tasks import send_mails

instances = FlaskBlueprint('instances', __name__)

USER_INSTANCE_LIMIT = 5

instance_fields = {
    'id': fields.String,
    'name': fields.String,
    'provisioned_at': fields.DateTime,
    'lifetime_left': fields.Integer,
    'maximum_lifetime': fields.Integer,
    'runtime': fields.Float,
    'state': fields.String,
    'to_be_deleted': fields.Boolean,
    'error_msg': fields.String,
    'username': fields.String,
    'user_id': fields.String,
    'blueprint': fields.String,
    'blueprint_id': fields.String,
    'cost_multiplier': fields.Float(default=1.0),
    'can_update_connectivity': fields.Boolean(default=False),
    'instance_data': fields.Raw,
    'public_ip': fields.String,
    'client_ip': fields.String(default='not set'),
    'logs': fields.Raw,
}

instance_log_fields = {
    'id': fields.String,
    'instance_id': fields.String,
    'log_type': fields.String,
    'log_level': fields.String,
    'timestamp': fields.Float,
    'message': fields.String
}

def query_blueprint(blueprint_id):
    return Blueprint.query.filter_by(id=blueprint_id).first()

def query_user(user_id):
    return User.query.filter_by(id=user_id).first()

def positive_integer(input_value):
    """Return input_value if valid, raise an exception in other case."""
    try:
        input_int = int(input_value)
    except:
        raise ValueError('{} is not a valid integer'.format(input_value))
    if input_int >= 0:
        return input_int
    else:
        raise ValueError('{} is not a positive integer'.format(input_value))

class InstanceList(restful.Resource):
    parser = reqparse.RequestParser()
    parser.add_argument('show_deleted', type=bool, default=False, location='args')
    parser.add_argument('show_only_mine', type=bool, default=False, location='args')
    parser.add_argument('offset', type=positive_integer, location='args')
    parser.add_argument('limit', type=positive_integer, location='args')

    @auth.login_required
    @marshal_with(instance_fields)
    def get(self):
        user = g.user
        args = self.parser.parse_args()
        q = apply_rules_instances(user, args)
        q = q.order_by(Instance.provisioned_at)
        instances = q.all()

        get_blueprint = memoize(query_blueprint)
        get_user = memoize(query_user)
        for instance in instances:
            instance_logs = get_logs_from_db(instance.id)
            instance.logs = marshal(instance_logs, instance_log_fields)

            user = get_user(instance.user_id)
            if user:
                instance.username = user.eppn

            blueprint = get_blueprint(instance.blueprint_id)
            if not blueprint:
                logging.warn("instance %s has a reference to non-existing blueprint" % instance.id)
                continue

            age = 0
            if instance.provisioned_at:
                age = (datetime.datetime.utcnow() - instance.provisioned_at).total_seconds()
            instance.lifetime_left = max(blueprint.maximum_lifetime - age, 0)
            instance.maximum_lifetime = blueprint.maximum_lifetime
            instance.cost_multiplier = blueprint.cost_multiplier

            if instance.to_be_deleted:
                instance.state = Instance.STATE_DELETING

        return instances

    @auth.login_required
    def post(self):
        user = g.user

        form = InstanceForm()
        if not form.validate_on_submit():
            logging.warn("validation error on user login")
            return form.errors, 422

        blueprint_id = form.blueprint.data
        allowed_blueprint_ids = get_group_blueprint_ids_for_instances(user)
        if not user.is_admin and blueprint_id not in allowed_blueprint_ids:
            abort(403)
        blueprint = Blueprint.query.filter_by(id=blueprint_id, is_enabled=True).first()
        if not blueprint:
            abort(404)
        if user.quota_exceeded():
            return {'error': 'USER_OVER_QUOTA'}, 409

        if blueprint.preallocated_credits:
            preconsumed_amount = blueprint.cost()
            total_credits_spent = preconsumed_amount + user.credits_spent
            if user.credits_quota < total_credits_spent:
                return {'error': 'USER_OVER_QUOTA'}, 409

        instances_for_user = Instance.query.filter_by(
            blueprint_id=blueprint.id,
            user_id=user.id
        ).filter(Instance.state != 'deleted').all()

        user_instance_limit = blueprint.full_config.get('maximum_instances_per_user', USER_INSTANCE_LIMIT)
        if instances_for_user and len(instances_for_user) >= user_instance_limit:
            return {'error': 'BLUEPRINT_INSTANCE_LIMIT_REACHED'}, 409

        instance = Instance(blueprint, user)
        # XXX: Choosing the name should be done in the model's constructor method
        # decide on a name that is not used currently
        existing_names = set(x.name for x in Instance.query.all())
        # Note: the potential race is solved by unique constraint in database
        while True:
            c_name = Instance.generate_name(prefix=app.dynamic_config.get('INSTANCE_NAME_PREFIX'))
            if c_name not in existing_names:
                instance.name = c_name
                break
        db.session.add(instance)
        db.session.commit()

        if not app.dynamic_config.get('SKIP_TASK_QUEUE'):
            run_update.delay(instance.id)

        return marshal(instance, instance_fields), 200

class InstanceView(restful.Resource):
    parser = reqparse.RequestParser()
    parser.add_argument('state', type=str)
    parser.add_argument('public_ip', type=str)
    parser.add_argument('error_msg', type=str)
    parser.add_argument('client_ip', type=str)
    parser.add_argument('instance_data', type=str)
    parser.add_argument('to_be_deleted', type=bool)
    parser.add_argument('send_email', type=bool)

    @auth.login_required
    @marshal_with(instance_fields)
    def get(self, instance_id):
        user = g.user
        args = {'instance_id': instance_id}
        query = apply_rules_instances(user, args)
        instance = query.first()
        if not instance:
            abort(404)

        blueprint = Blueprint.query.filter_by(id=instance.blueprint_id).first()
        instance.blueprint_id = blueprint.id
        instance.username = instance.user
        instance_logs = get_logs_from_db(instance.id)
        instance.logs = marshal(instance_logs, instance_log_fields)

        if 'allow_update_client_connectivity' in blueprint.full_config \
                and blueprint.full_config['allow_update_client_connectivity']:
            instance.can_update_connectivity = True

        age = 0
        if instance.provisioned_at:
            age = (datetime.datetime.utcnow() - instance.provisioned_at).total_seconds()
        instance.lifetime_left = max(blueprint.maximum_lifetime - age, 0)
        instance.maximum_lifetime = blueprint.maximum_lifetime
        instance.cost_multiplier = blueprint.cost_multiplier

        return instance

    @auth.login_required
    @marshal_with(instance_fields)
    def post(self, instance_id):
        args = self.parser.parse_args()
        user = g.user
        query = apply_rules_instances(user, {'instance_id': instance_id})
        instance = query.first()
        if not instance:
            abort(404)
        if args.get('send_email') and instance.state != 'running':
            text = {"subject": " ", "message": " "}
            admin_users = User.query.filter_by(is_admin=True).all()
            for admins in admin_users:
                if admins.eppn != 'worker@pebbles':
                    text['subject'] = "Notebooks.csc.fi:WARNING"
                    text['message'] = instance.name + " is taking more than ten minutes to launch"
                    # send email only through email_id because some eppn bounce back.
                    # Also the email_id will be updated once they login so here it is available
                    send_mails.delay([(admins.email_id, 'None', 't')], text)
        return instance

    @auth.login_required
    def delete(self, instance_id):
        user = g.user
        query = Instance.query.filter_by(id=instance_id)
        if not user.is_admin and not is_group_manager(user):
            query = query.filter_by(user_id=user.id)
        instance = query.first()
        if not instance:
            abort(404)
        group = instance.blueprint.group
        if not user.is_admin and not is_group_manager(user, group) and instance.user_id != user.id:
            abort(403)
        instance.to_be_deleted = True
        instance.state = Instance.STATE_DELETING
        instance.deprovisioned_at = datetime.datetime.utcnow()
        db.session.commit()
        if not app.dynamic_config.get('SKIP_TASK_QUEUE'):
            run_update.delay(instance.id)

    @auth.login_required
    def put(self, instance_id):
        user = g.user
        form = UserIPForm()
        if not form.validate_on_submit():
            logging.warn("validation error on UserIPForm")
            return form.errors, 422

        instance = Instance.query.filter_by(id=instance_id, user_id=user.id).first()
        if not instance:
            abort(404)

        blueprint = Blueprint.query.filter_by(id=instance.blueprint_id).first()
        if 'allow_update_client_connectivity' in blueprint.full_config \
                and blueprint.full_config['allow_update_client_connectivity']:
            instance.client_ip = form.client_ip.data
            if not app.dynamic_config.get('SKIP_TASK_QUEUE'):
                update_user_connectivity.delay(instance.id)
            db.session.commit()

        else:
            abort(400)

    @auth.login_required
    @requires_admin
    def patch(self, instance_id):
        args = self.parser.parse_args()
        instance = Instance.query.filter_by(id=instance_id).first()
        if not instance:
            abort(404)

        if args.get('state'):
            instance.state = args['state']
            if instance.state == Instance.STATE_RUNNING:
                if not instance.provisioned_at:
                    instance.provisioned_at = datetime.datetime.utcnow()
            if args['state'] == Instance.STATE_FAILED:
                instance.errored = True

            db.session.commit()

        if args.get('to_be_deleted'):
            instance.to_be_deleted = args['to_be_deleted']
            db.session.commit()

        if args.get('error_msg'):
            instance.error_msg = args['error_msg']
            db.session.commit()

        if args.get('public_ip'):
            instance.public_ip = args['public_ip']
            db.session.commit()

        if args.get('instance_data'):
            try:
                instance.instance_data = json.loads(args['instance_data'])
            except ValueError:
                logging.warn("invalid instance_data passed to view: %s" % args['instance_data'])
            db.session.commit()

class InstanceLogs(restful.Resource):
    parser = reqparse.RequestParser()
    parser.add_argument('log_record', type=dict)
    parser.add_argument('log_type', type=str)
    parser.add_argument('send_log_fetch_task', type=bool)

    @auth.login_required
    @marshal_with(instance_log_fields)
    def get(self, instance_id):
        args = self.parser.parse_args()
        instance = Instance.query.filter_by(id=instance_id).first()
        if not instance:
            abort(404)
        if args.get('log_type'):
            log_type = args['log_type']
        instance_logs = get_logs_from_db(instance_id, log_type)
        return instance_logs

    @auth.login_required
    @requires_group_owner_or_admin
    def patch(self, instance_id):
        args = self.parser.parse_args()
        instance = Instance.query.filter_by(id=instance_id).first()
        if not instance:
            abort(404)
        if args.get('send_log_fetch_task'):
            if not app.dynamic_config.get('SKIP_TASK_QUEUE'):
                fetch_running_instance_logs.delay(instance_id)
        if args.get('log_record'):
            log_record = args['log_record']
            instance_log = process_logs(instance_id, log_record)
            db.session.add(instance_log)
            db.session.commit()

        return 'ok'

    @auth.login_required
    @requires_admin
    def delete(self, instance_id):
        args = self.parser.parse_args()
        log_type = args.get('log_type')
        if log_type:
            instance_logs = get_logs_from_db(instance_id, args.get('log_type'))
        else:
            instance_logs = get_logs_from_db(instance_id)
        if not instance_logs:
            logging.warn("There are no log entries to be deleted")

        for instance_log in instance_logs:
            db.session.delete(instance_log)
        db.session.commit()

def get_logs_from_db(instance_id, log_type=None):
    logs_query = InstanceLog.query\
        .filter_by(instance_id=instance_id)
    if log_type:
        logs_query = logs_query.filter_by(log_type=log_type)
    logs_query = logs_query.order_by(InstanceLog.timestamp)
    logs = logs_query.all()
    return logs

def process_logs(instance_id, log_record):

    check_running_log = get_logs_from_db(instance_id, "running")

    if check_running_log:  # in case of running logs, the whole message is replaced again (thus only 1 running log record with all info)
        instance_log = check_running_log[0]
        if log_record['log_type'] == "running":
            instance_log.timestamp = float(log_record['timestamp'])
            instance_log.message = log_record['message']  # replace the whole text (older text now appended with the new text)
            return instance_log

    instance_log = InstanceLog(instance_id)
    instance_log.log_type = log_record['log_type']
    instance_log.log_level = log_record['log_level']
    instance_log.timestamp = float(log_record['timestamp'])
    instance_log.message = log_record['message']

    return instance_log
<|endoftext|>


<commit_before><commit_msg>simple mute<commit_after><|endoftext|>


<commit_before>import logging
from typing import List, Optional, Union

import torch
from torch import nn
from torch.nn import Module

logger = logging.getLogger(__name__)

class PositionalEmbedding(Module):
    def __init__(self, num_embeddings: int, embedding_dim: int, pad_index: int):
        super().__init__()
        self.embedding = nn.Embedding(num_embeddings, embedding_dim, pad_index)
        self.pad_index = pad_index

    def forward(self, input):
        positions = self._make_positions(input, self.pad_index)
        return self.embedding(positions)

    def max_positions(self):
        if self.pad_index is not None:
            return self.num_embeddings - self.pad_index - 1
        else:
            return self.num_embeddings

    def _make_positions(self, tensor, pad_index: int):
        masked = tensor.ne(pad_index).long()
        return torch.cumsum(masked, dim=1) * masked + pad_index

class TransformerEncoderLayer(Module):
    def __init__(
        self,
        embedding_dim: int,
        num_attention_heads: int,
        ffn_dimension: Optional[int] = None,
        dropout: float = 0.1,
        normalize_before: bool = False,
        scaling: Optional[float] = None,
    ):
        super().__init__()
        # TODO Manually setting scaling is not allowed
        ffn_dimension = ffn_dimension or embedding_dim * 4

        self.better_transformer = torch.nn.TransformerEncoderLayer(
            d_model=embedding_dim,
            nhead=num_attention_heads,
            dim_feedforward=ffn_dimension,
            dropout=dropout,
            batch_first=True,
            activation="gelu",
            norm_first=normalize_before,
        )

    def _load_from_state_dict(
        self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs
    ):
        better_to_old_names = {
            "better_transformer.self_attn.in_proj_weight": "attention.input_projection.weight",
            "better_transformer.self_attn.in_proj_bias": "attention.input_projection.bias",
            "better_transformer.self_attn.out_proj.weight": "attention.output_projection.weight",
            "better_transformer.self_attn.out_proj.bias": "attention.output_projection.bias",
            "better_transformer.linear1.weight": "residual_mlp.mlp.0.weight",
            "better_transformer.linear1.bias": "residual_mlp.mlp.0.bias",
            "better_transformer.linear2.weight": "residual_mlp.mlp.3.weight",
            "better_transformer.linear2.bias": "residual_mlp.mlp.3.bias",
            "better_transformer.norm1.weight": "attention_layer_norm.weight",
            "better_transformer.norm1.bias": "attention_layer_norm.bias",
            "better_transformer.norm2.weight": "final_layer_norm.weight",
            "better_transformer.norm2.bias": "final_layer_norm.bias",
        }
        for better, old in better_to_old_names.items():
            better_name = prefix + better
            old_name = prefix + old
            if old_name in state_dict:
                state_dict[better_name] = state_dict[old_name]
                state_dict.pop(old_name)
            elif better_name in state_dict:
                # Do nothing
                pass
            elif strict:
                missing_keys.append(better_name)

        super(TransformerEncoderLayer, self)._load_from_state_dict(
            state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs
        )

    def forward(self, input: torch.Tensor, key_padding_mask: torch.Tensor, attn_mask: Optional[torch.Tensor] = None):
        # torch.nn.TransformerEncodeLayer's attn_mask and key_padding_mask's
        # order is reversed
        return self.better_transformer(input.transpose(0, 1), attn_mask, key_padding_mask).transpose(0, 1)

class TransformerEncoder(Module):
    def __init__(
        self,
        vocab_size: int,
        embedding_dim: int,
        padding_idx: int,
        max_seq_len: int,
        num_encoder_layers: int,
        num_attention_heads: int,
        ffn_dimension: Optional[int] = None,
        dropout: float = 0.1,
        normalize_before: bool = False,
        scaling: Optional[float] = None,
        return_all_layers: bool = False,
    ):
        super().__init__()
        self.padding_idx = padding_idx
        self.token_embedding = nn.Embedding(vocab_size, embedding_dim, padding_idx)
        self.layers = nn.ModuleList(
            [
                TransformerEncoderLayer(
                    embedding_dim=embedding_dim,
                    num_attention_heads=num_attention_heads,
                    ffn_dimension=ffn_dimension,
                    dropout=dropout,
                    normalize_before=normalize_before,
                    scaling=scaling,
                )
                for _ in range(num_encoder_layers)
            ]
        )
        self.positional_embedding = PositionalEmbedding(max_seq_len, embedding_dim, padding_idx)
        self.embedding_layer_norm = nn.LayerNorm(embedding_dim)
        self.dropout = nn.Dropout(dropout)
        self.normalize_before = normalize_before
        self.return_all_layers = return_all_layers

    def forward(
        self, tokens: torch.Tensor, attn_mask: Optional[torch.Tensor] = None
    ) -> Union[torch.Tensor, List[torch.Tensor]]:
        if attn_mask is not None:
            torch._assert(
                attn_mask.is_floating_point() or attn_mask.dtype == torch.bool,
                f"Only float or bool types are supported for attn_mask not {attn_mask.dtype}",
            )

        padding_mask = tokens.eq(self.padding_idx)

        token_embeddings = self.token_embedding(tokens)
        embedded_positions = self.positional_embedding(tokens)

        embedded = token_embeddings + embedded_positions

        if not hasattr(self, "normalize_before"):
            self.normalize_before = False
        if not self.normalize_before:
            embedded = self.embedding_layer_norm(embedded)
        embedded = self.dropout(embedded)

        padded_embedded = embedded * (1 - padding_mask.unsqueeze(-1).type_as(embedded))

        encoded = padded_embedded.transpose(0, 1)

        if self.return_all_layers:
            states = [encoded]

            for layer in self.layers:
                encoded = layer(encoded, padding_mask, attn_mask)
                states.append(encoded)

            if self.normalize_before:
                for i, state in enumerate(states):
                    states[i] = self.embedding_layer_norm(state)

            # states are returned as T x B x C
            return states
        else:
            for layer in self.layers:
                encoded = layer(encoded, padding_mask, attn_mask)

            if self.normalize_before:
                encoded = self.embedding_layer_norm(encoded)

            # states are returned as T x B x C
            return encoded
<commit_msg>[fbsync] Replace TransformerEncoder in torchtext with better transformer (#1703)<commit_after>import logging
from typing import List, Optional, Union

import torch
from torch import nn
from torch.nn import Module

logger = logging.getLogger(__name__)

class PositionalEmbedding(Module):
    def __init__(self, num_embeddings: int, embedding_dim: int, pad_index: int):
        super().__init__()
        self.embedding = nn.Embedding(num_embeddings, embedding_dim, pad_index)
        self.pad_index = pad_index

    def forward(self, input):
        positions = self._make_positions(input, self.pad_index)
        return self.embedding(positions)

    def max_positions(self):
        if self.pad_index is not None:
            return self.num_embeddings - self.pad_index - 1
        else:
            return self.num_embeddings

    def _make_positions(self, tensor, pad_index: int):
        masked = tensor.ne(pad_index).long()
        return torch.cumsum(masked, dim=1) * masked + pad_index

class TransformerEncoderLayer(Module):
    def __init__(
        self,
        embedding_dim: int,
        num_attention_heads: int,
        ffn_dimension: Optional[int] = None,
        dropout: float = 0.1,
        normalize_before: bool = False,
        scaling: Optional[float] = None,
    ):
        super().__init__()
        # TODO Manually setting scaling is not allowed
        ffn_dimension = ffn_dimension or embedding_dim * 4

        self.better_transformer = torch.nn.TransformerEncoderLayer(
            d_model=embedding_dim,
            nhead=num_attention_heads,
            dim_feedforward=ffn_dimension,
            dropout=dropout,
            batch_first=True,
            activation="gelu",
            norm_first=normalize_before,
        )

    def _load_from_state_dict(
        self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs
    ):
        better_to_old_names = {
            "better_transformer.self_attn.in_proj_weight": "attention.input_projection.weight",
            "better_transformer.self_attn.in_proj_bias": "attention.input_projection.bias",
            "better_transformer.self_attn.out_proj.weight": "attention.output_projection.weight",
            "better_transformer.self_attn.out_proj.bias": "attention.output_projection.bias",
            "better_transformer.linear1.weight": "residual_mlp.mlp.0.weight",
            "better_transformer.linear1.bias": "residual_mlp.mlp.0.bias",
            "better_transformer.linear2.weight": "residual_mlp.mlp.3.weight",
            "better_transformer.linear2.bias": "residual_mlp.mlp.3.bias",
            "better_transformer.norm1.weight": "attention_layer_norm.weight",
            "better_transformer.norm1.bias": "attention_layer_norm.bias",
            "better_transformer.norm2.weight": "final_layer_norm.weight",
            "better_transformer.norm2.bias": "final_layer_norm.bias",
        }
        for better, old in better_to_old_names.items():
            better_name = prefix + better
            old_name = prefix + old
            if old_name in state_dict:
                state_dict[better_name] = state_dict[old_name]
                state_dict.pop(old_name)
            elif better_name in state_dict:
                # Do nothing
                pass
            elif strict:
                missing_keys.append(better_name)

        super(TransformerEncoderLayer, self)._load_from_state_dict(
            state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs
        )

    def forward(self, input: torch.Tensor, key_padding_mask: torch.Tensor, attn_mask: Optional[torch.Tensor] = None):
        # torch.nn.TransformerEncodeLayer's attn_mask and key_padding_mask's
        # order is reversed
        return self.better_transformer(input.transpose(0, 1), attn_mask, key_padding_mask).transpose(0, 1)

class TransformerEncoder(Module):
    def __init__(
        self,
        vocab_size: int,
        embedding_dim: int,
        padding_idx: int,
        max_seq_len: int,
        num_encoder_layers: int,
        num_attention_heads: int,
        ffn_dimension: Optional[int] = None,
        dropout: float = 0.1,
        normalize_before: bool = False,
        scaling: Optional[float] = None,
        return_all_layers: bool = False,
    ):
        super().__init__()
        self.padding_idx = padding_idx
        self.token_embedding = nn.Embedding(vocab_size, embedding_dim, padding_idx)
        ffn_dimension = ffn_dimension or 4 * embedding_dim
        layer = torch.nn.TransformerEncoderLayer(
            d_model=embedding_dim,
            nhead=num_attention_heads,
            dim_feedforward=ffn_dimension,
            dropout=dropout,
            activation="gelu",
            batch_first=True,
            norm_first=normalize_before,
        )
        self.layers = torch.nn.TransformerEncoder(encoder_layer=layer, num_layers=num_encoder_layers)
        self.positional_embedding = PositionalEmbedding(max_seq_len, embedding_dim, padding_idx)
        self.embedding_layer_norm = nn.LayerNorm(embedding_dim)
        self.dropout = nn.Dropout(dropout)
        self.normalize_before = normalize_before
        self.return_all_layers = return_all_layers

    def forward(
        self, tokens: torch.Tensor, attn_mask: Optional[torch.Tensor] = None
    ) -> Union[torch.Tensor, List[torch.Tensor]]:
        if attn_mask is not None:
            torch._assert(
                attn_mask.is_floating_point() or attn_mask.dtype == torch.bool,
                f"Only float or bool types are supported for attn_mask not {attn_mask.dtype}",
            )

        padding_mask = tokens.eq(self.padding_idx)

        token_embeddings = self.token_embedding(tokens)
        embedded_positions = self.positional_embedding(tokens)

        embedded = token_embeddings + embedded_positions

        if not hasattr(self, "normalize_before"):
            self.normalize_before = False
        if not self.normalize_before:
            embedded = self.embedding_layer_norm(embedded)
        embedded = self.dropout(embedded)

        padded_embedded = embedded * (1 - padding_mask.unsqueeze(-1).type_as(embedded))

        if self.return_all_layers:
            encoded = padded_embedded
            # B x T x C
            # Then transpose back to T x B x C
            states = [encoded.transpose(1, 0)]
            for layer in self.layers.layers:
                encoded = layer(encoded, padding_mask, attn_mask)
                encoded_t = encoded.transpose(1, 0)
                states.append(encoded_t)
            if self.normalize_before:
                for i, state in enumerate(states):
                    states[i] = self.embedding_layer_norm(state)
            return states
        else:
            # B x T x C
            # Then transpose back to T x B x C
            encoded = self.layers(padded_embedded).transpose(1, 0)
            if self.normalize_before:
                encoded = self.embedding_layer_norm(encoded)
            return encoded

    def _load_from_state_dict(
        self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs
    ):
        better_to_old_names = {
            "self_attn.in_proj_weight": "attention.input_projection.weight",
            "self_attn.in_proj_bias": "attention.input_projection.bias",
            "self_attn.out_proj.weight": "attention.output_projection.weight",
            "self_attn.out_proj.bias": "attention.output_projection.bias",
            "linear1.weight": "residual_mlp.mlp.0.weight",
            "linear1.bias": "residual_mlp.mlp.0.bias",
            "linear2.weight": "residual_mlp.mlp.3.weight",
            "linear2.bias": "residual_mlp.mlp.3.bias",
            "norm1.weight": "attention_layer_norm.weight",
            "norm1.bias": "attention_layer_norm.bias",
            "norm2.weight": "final_layer_norm.weight",
            "norm2.bias": "final_layer_norm.bias",
        }
        for i in range(self.layers.num_layers):
            for better, old in better_to_old_names.items():
                better_name = prefix + "layers.layers.{}.".format(i) + better
                old_name = prefix + "layers.{}.".format(i) + old
                if old_name in state_dict:
                    state_dict[better_name] = state_dict[old_name]
                    state_dict.pop(old_name)
                elif better_name in state_dict:
                    # Do nothing
                    pass
                elif strict:
                    missing_keys.append(better_name)

        super(TransformerEncoder, self)._load_from_state_dict(
            state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs
        )
<|endoftext|>


<commit_before>import os
from functools import wraps
from os import path
from shutil import copy

from .. import docker_interface
from ..graph.artifacts import File
from ..graph.nodes import Software
from ..helper import Logger
from ..storage import Memory

def _get_cmd(interface, force_exec=False):
    def _get_cmd_decorator(func):
        @wraps(func)
        def func_wrapper(*args):
            assert isinstance(args[0], Software)
            cmd = Software_manager._get_cmd_args(args[0], interface)
            if cmd or force_exec:
                return func(cmd, *args)
        return func_wrapper
    return _get_cmd_decorator

class Software_manager:

    @staticmethod
    @_get_cmd('create', force_exec=True)
    def create(cmd, node):
        Software_manager._copy_files(node)
        if cmd is not None:
            docker_interface.update_container(node.host_container, cmd)

    @staticmethod
    @_get_cmd('configure')
    def configure(cmd, node):
        docker_interface.update_container(node.host_container, cmd)

    @staticmethod
    @_get_cmd('start')
    def start(cmd, node):
        _log = Logger.get(__name__)
        try:
            docker_interface.exec_cmd(node.host_container, cmd)
        except Exception as e:
            _log.debug('is not running!')
            # docker_interface.delete_container(node.host_container)
            docker_interface.create_container(node.host_container,
                                              cmd=cmd,
                                              entrypoint='',
                                              from_saved=True,
                                              force=True)
            docker_interface.start_container(node.host_container)

    @staticmethod
    @_get_cmd('stop')
    def stop(cmd, node):
        _log = Logger.get(__name__)
        if docker_interface.is_running(node.host_container):
            _log.debug('exec stop command!')
            docker_interface.exec_cmd(node.host_container, cmd)

    @staticmethod
    @_get_cmd('delete')
    def delete(cmd, node):
        _log = Logger.get(__name__)
        _log.debug('exec delete command!')
        docker_interface.update_container(node.host_container, cmd)

    @staticmethod
    def _copy_files(node):
        _log = Logger.get(__name__)
        # generate path for the tmp folder
        tmp = path.join(node.tpl.tmp_dir,
                        node.host_container.name,
                        node.name)

        # create the folder for the software
        try:
            os.makedirs(tmp)
        except os.error:
            pass

        # copy all the interfaces scripts
        for key, value in node.interfaces.items():
            _log.debug(value['cmd'].file_path)
            copy(value['cmd'].file_path, tmp)

        # if present copy all the artifacts
        for art in node.artifacts:
            copy(art.file_path, tmp)

    @staticmethod
    def _get_cmd_args(node, interface):
        _log = Logger.get(__name__)

        def get_inside_path(p):
            if isinstance(p, File):
                p = p.file
            return path.join('/tmp/dt/', node.name, p)

        def get_cmd():
            def get_echo(string, where):
                return 'echo {} > {}'.format(string, where)

            def set_in_state():
                if interface == 'start':
                    return get_echo(Memory.STATE.STARTED.value,
                                    get_inside_path('state')) + ';'
                return ''

            def set_out_state():
                if interface == 'create':
                    return ';' + get_echo(Memory.STATE.CREATED.value,
                                          get_inside_path('state'))
                elif interface == 'start':
                    return ';' + get_echo(Memory.STATE.CREATED.value,
                                          get_inside_path('state'))
                elif interface == 'stop':
                    return ';' + get_echo(Memory.STATE.CREATED.value,
                                          get_inside_path('state'))
                elif interface == 'delete':
                    return ';' + get_echo(Memory.STATE.DELETED.value,
                                          get_inside_path('state'))
                return ''

            return '{}sh {}{}'.format(
                set_in_state(),
                get_inside_path(node.interfaces[interface]['cmd']),
                set_out_state()
            )

        if interface not in node.interfaces:
            return None
        _log.debug('interface: {}'
                   ''.format(node.interfaces[interface]))
        args = []
        args_env = []
        res = None
        if 'inputs' in node.interfaces[interface]:
            for key, value in node.interfaces[interface]['inputs'].items():
                if isinstance(value, File):
                    value = get_inside_path(value)
                args.append('--{} {}'.format(key, value))
                args_env.append(
                    'export INPUT_{}={}'.format(key.upper(), value))

            res = "sh -c '{};{}'".format(
                ';'.join(args_env),
                get_cmd()  # ,
                # ' '.join(args)
            )
        else:
            res = "sh -c '{}'".format(get_cmd())

        _log.debug('{} command ({}) on container {}'
                   ''.format(interface, res,
                             node.host_container))
        return res
<commit_msg>Update software_manager.py<commit_after>import os
from functools import wraps
from os import path
from shutil import copy

from .. import docker_interface
from ..graph.artifacts import File
from ..graph.nodes import Software
from ..helper import Logger
from ..storage import Memory

def _get_cmd(interface, force_exec=False):
    def _get_cmd_decorator(func):
        @wraps(func)
        def func_wrapper(*args):
            assert isinstance(args[0], Software)
            cmd = Software_manager._get_cmd_args(args[0], interface)
            if cmd or force_exec:
                return func(cmd, *args)
        return func_wrapper
    return _get_cmd_decorator

class Software_manager:

    @staticmethod
    @_get_cmd('create', force_exec=True)
    def create(cmd, node):
        Software_manager._copy_files(node)
        if cmd is not None:
            docker_interface.update_container(node.host_container, cmd)

    @staticmethod
    @_get_cmd('configure')
    def configure(cmd, node):
        docker_interface.update_container(node.host_container, cmd)

    @staticmethod
    @_get_cmd('start')
    def start(cmd, node):
        _log = Logger.get(__name__)
        try:
            docker_interface.exec_cmd(node.host_container, cmd)
        except Exception as e:
            _log.debug('is not running!')
            # docker_interface.delete_container(node.host_container)
            docker_interface.create_container(node.host_container,
                                              cmd=cmd,
                                              entrypoint='',
                                              from_saved=True,
                                              force=True)
            docker_interface.start_container(node.host_container)

    @staticmethod
    @_get_cmd('stop')
    def stop(cmd, node):
        _log = Logger.get(__name__)
        if docker_interface.is_running(node.host_container):
            _log.debug('exec stop command!')
            docker_interface.exec_cmd(node.host_container, cmd)

    @staticmethod
    @_get_cmd('delete')
    def delete(cmd, node):
        _log = Logger.get(__name__)
        _log.debug('exec delete command!')
        docker_interface.update_container(node.host_container, cmd)

    @staticmethod
    def _copy_files(node):
        # generate path for the tmp folder
        tmp = path.join(node.tpl.tmp_dir,
                        node.host_container.name,
                        node.name)

        # create the folder for the software
        try:
            os.makedirs(tmp)
        except os.error:
            pass

        # copy all the interfaces scripts
        for key, value in node.interfaces.items():
            copy(value['cmd'].file_path, tmp)

        # if present copy all the artifacts
        for art in node.artifacts:
            copy(art.file_path, tmp)

    @staticmethod
    def _get_cmd_args(node, interface):
        _log = Logger.get(__name__)

        def get_inside_path(p):
            if isinstance(p, File):
                p = p.file
            return path.join('/tmp/dt/', node.name, p)

        def get_cmd():
            def get_echo(string, where):
                return 'echo {} > {}'.format(string, where)

            def set_in_state():
                if interface == 'start':
                    return get_echo(Memory.STATE.STARTED.value,
                                    get_inside_path('state')) + ';'
                return ''

            def set_out_state():
                if interface == 'create':
                    return ';' + get_echo(Memory.STATE.CREATED.value,
                                          get_inside_path('state'))
                elif interface == 'start':
                    return ';' + get_echo(Memory.STATE.CREATED.value,
                                          get_inside_path('state'))
                elif interface == 'stop':
                    return ';' + get_echo(Memory.STATE.CREATED.value,
                                          get_inside_path('state'))
                elif interface == 'delete':
                    return ';' + get_echo(Memory.STATE.DELETED.value,
                                          get_inside_path('state'))
                return ''

            return '{}sh {}{}'.format(
                set_in_state(),
                get_inside_path(node.interfaces[interface]['cmd']),
                set_out_state()
            )

        if interface not in node.interfaces:
            return None
        _log.debug('interface: {}'
                   ''.format(node.interfaces[interface]))
        args = []
        args_env = []
        res = None
        if 'inputs' in node.interfaces[interface]:
            for key, value in node.interfaces[interface]['inputs'].items():
                if isinstance(value, File):
                    value = get_inside_path(value)
                args.append('--{} {}'.format(key, value))
                args_env.append(
                    'export INPUT_{}={}'.format(key.upper(), value))

            res = "sh -c '{};{}'".format(
                ';'.join(args_env),
                get_cmd()  # ,
                # ' '.join(args)
            )
        else:
            res = "sh -c '{}'".format(get_cmd())

        _log.debug('{} command ({}) on container {}'
                   ''.format(interface, res,
                             node.host_container))
        return res
<|endoftext|>