<commit_before>import logging
import sqlite3
import ruamel.yaml

# create logger
logger = logging.getLogger('terminology')
# cashes sql and root concepts
data = {}

def load_terminology_data():
    global data
    f = open('terminology.yml', 'r')
    data = ruamel.yaml.safe_load(f.read())
    f.close()
    return data

def get_dom_sql(dom):
    global data
    if len(data) == 0:
        load_terminology_data()
    return data['dom_sql'][dom]

def get_sql(query):
    global data
    if len(data) == 0:
        load_terminology_data()
    return data['sql'][query]

def get_root_concept(dom):
    global data
    if len(data) == 0:
        load_terminology_data()
    if data['root_concept'][dom]:
        return data['root_concept'][dom]
    else:
        return data['root_concept_default']

# if code (concept id) is 'root' replace with NCIt root concept
def check_for_root(code, dom):
    if 'root' == code:
        return get_root_concept(dom)['code']
    return code

# get connection to sqlite3 terminology database
def get_connection():
    connection = sqlite3.connect('db/terminology.db')
    connection.row_factory = sqlite3.Row
    return connection

# execute submitted sql that must retrieve a resultset
# with two fields named 'code' and 'name'
# function returns [{code: 'code1', name: 'name1'}, {code: 'code2', name: 'name2'}, ...]
def get_code_name_list_result(sql, querytokens=None):
    connection = get_connection()
    results = []
    try:
        cursor = connection.cursor()
        if querytokens:
            cursor.execute(sql, querytokens)
        else:
            cursor.execute(sql)
        for row in cursor.fetchall():
            results.append({'code': row['code'], 'name': row['name']})
        cursor.close()
    except sqlite3.Error as e:
        logger.error("An error occurred:", e.args)
    return results

# get a single row or a single field as a string
def get_single_string_result(sql, querytokens=None):
    connection = get_connection()
    result = None
    try:
        cursor = connection.cursor()
        if querytokens:
            cursor.execute(sql, querytokens)
        else:
            cursor.execute(sql)
        if cursor.arraysize > 0:
            result = cursor.fetchone()[0]
        cursor.close()
    except sqlite3.Error as e:
        logger.error("An error occurred:", e.args[0])
    return result

def search_domain_by_substring(dom, query):
    return get_code_name_list_result(get_dom_sql(dom), ['%' + query + '%', '%' + query + '%'])

def search_diseases_associated_with_anatomic_site(code):
    return get_code_name_list_result(get_sql('search_diseases_associated_with_anatomic_site'),
                                     [check_for_root(code, 'anantomicsite')])

def search_diseases_associated_with_finding(code):
    return get_code_name_list_result(get_sql('search_diseases_associated_with_finding'),
                                     [check_for_root(code, 'finding')])

def search_diseases_associated_with_gene(code):
    return get_code_name_list_result(get_sql('search_diseases_associated_with_gene'),
                                     [check_for_root(code, 'gene')])

def search_diseases_is_stage(code):
    return get_code_name_list_result(get_sql('search_diseases_is_stage'),
                                     [check_for_root(code, 'disease')])

def search_diseases_is_grade(code):
    return get_code_name_list_result(get_sql('search_diseases_is_grade'),
                                     [check_for_root(code, 'disease')])

def get_subtree_codes(code, dom):
    connection = sqlite3.connect('db/terminology.db')
    connection.row_factory = lambda cursor, row: row[0]
    results = []
    try:
        cursor = connection.cursor()
        cursor.execute(get_sql('get_subtree_codes'), [check_for_root(code, dom)])
        results = cursor.fetchall()
        cursor.close()
    except sqlite3.Error as e:
        logger.error("An error occurred:", e)
    return results

def get_child_codes(code, dom):
    # biomarkers are not marked up as such in ncit so we return all flat
    ccode = check_for_root(code, dom)
    # parent codes are in a pipe (|) delimited list
    # just matching on the first term we get spurious substring matches
    # such as C8461 matching C84615
    return get_code_name_list_result(get_sql('get_child_codes'), ['%' + ccode, '%' + ccode + '|%'])

def get_parent_codes(code, dom):
    ccode = check_for_root(code, dom)
    parent_codes = get_single_string_result(get_sql('get_parent_codes'), [ccode])
    results = []
    # parent codes are in a pipe (|) delimited list
    for parent_code in parent_codes.split('|'):
        results.extend(get_code_name_list_result(get_sql('get_parent_codes_sub'), [parent_code]))
    return results
<commit_msg>further dried up the terminology code<commit_after>import logging
import sqlite3
import ruamel.yaml

# create logger
logger = logging.getLogger('terminology')
# cashes sql and root concepts
data = {}

def load_terminology_data():
    global data
    f = open('terminology.yml', 'r')
    data = ruamel.yaml.safe_load(f.read())
    f.close()
    return data

def get_sql(query, is_dom):
    global data
    if len(data) == 0:
        load_terminology_data()
    if is_dom:
        return data['dom_sql'][query]
    else:
        return data['sql'][query]

def get_root_concept(dom):
    global data
    if len(data) == 0:
        load_terminology_data()
    if data['root_concept'][dom]:
        return data['root_concept'][dom]
    else:
        return data['root_concept_default']

# if code (concept id) is 'root' replace with NCIt root concept
def check_for_root(code, dom):
    if 'root' == code:
        return get_root_concept(dom)['code']
    return code

# get connection to sqlite3 terminology database
def get_connection():
    connection = sqlite3.connect('db/terminology.db')
    connection.row_factory = sqlite3.Row
    return connection

# execute submitted sql that must retrieve a resultset
# with two fields named 'code' and 'name'
# function returns [{code: 'code1', name: 'name1'}, {code: 'code2', name: 'name2'}, ...]
def get_code_name_list_result(sql, querytokens=None):
    connection = get_connection()
    results = []
    try:
        cursor = connection.cursor()
        if querytokens:
            cursor.execute(sql, querytokens)
        else:
            cursor.execute(sql)
        for row in cursor.fetchall():
            results.append({'code': row['code'], 'name': row['name']})
        cursor.close()
    except sqlite3.Error as e:
        logger.error("An error occurred:", e.args)
    return results

# get a single row or a single field as a string
def get_single_string_result(sql, querytokens=None):
    connection = get_connection()
    result = None
    try:
        cursor = connection.cursor()
        if querytokens:
            cursor.execute(sql, querytokens)
        else:
            cursor.execute(sql)
        if cursor.arraysize > 0:
            result = cursor.fetchone()[0]
        cursor.close()
    except sqlite3.Error as e:
        logger.error("An error occurred:", e.args[0])
    return result

def search_domain_by_substring(dom, query):
    return get_code_name_list_result(get_sql(dom, True), ['%' + query + '%', '%' + query + '%'])

def search_diseases_associated_with_anatomic_site(code):
    return get_code_name_list_result(get_sql('search_diseases_associated_with_anatomic_site', False),
                                     [check_for_root(code, 'anantomicsite')])

def search_diseases_associated_with_finding(code):
    return get_code_name_list_result(get_sql('search_diseases_associated_with_finding', False),
                                     [check_for_root(code, 'finding')])

def search_diseases_associated_with_gene(code):
    return get_code_name_list_result(get_sql('search_diseases_associated_with_gene', False),
                                     [check_for_root(code, 'gene')])

def search_diseases_is_stage(code):
    return get_code_name_list_result(get_sql('search_diseases_is_stage', False),
                                     [check_for_root(code, 'disease')])

def search_diseases_is_grade(code):
    return get_code_name_list_result(get_sql('search_diseases_is_grade', False),
                                     [check_for_root(code, 'disease')])

def get_subtree_codes(code, dom):
    connection = sqlite3.connect('db/terminology.db')
    connection.row_factory = lambda cursor, row: row[0]
    results = []
    try:
        cursor = connection.cursor()
        cursor.execute(get_sql('get_subtree_codes', False), [check_for_root(code, dom)])
        results = cursor.fetchall()
        cursor.close()
    except sqlite3.Error as e:
        logger.error("An error occurred:", e)
    return results

def get_child_codes(code, dom):
    # biomarkers are not marked up as such in ncit so we return all flat
    ccode = check_for_root(code, dom)
    # parent codes are in a pipe (|) delimited list
    # just matching on the first term we get spurious substring matches
    # such as C8461 matching C84615
    return get_code_name_list_result(get_sql('get_child_codes', False), ['%' + ccode, '%' + ccode + '|%'])

def get_parent_codes(code, dom):
    ccode = check_for_root(code, dom)
    parent_codes = get_single_string_result(get_sql('get_parent_codes', False), [ccode])
    results = []
    # parent codes are in a pipe (|) delimited list
    for parent_code in parent_codes.split('|'):
        results.extend(get_code_name_list_result(get_sql('get_parent_codes_sub', False), [parent_code]))
    return results
<|endoftext|>


<commit_before>ttk_style_opts_to_config = [
('TButton', {
  'background': styled_button_color,
  'bordercolor': 'green',
  'lightcolor': 'blue',
  'darkcolor': 'red',
}),
('TLabel', {
  'background': styled_button_color,
  'bordercolor': 'hot pink', # this does half of what i need
  'lightcolor': 'blue',
  'darkcolor': 'red',
}),
('TFrame', {
  'bordercolor': 'blue', # this does half of what i need
  'lightcolor': 'green',
  'darkcolor': 'red',
}),
]<commit_msg>hotline miameme theme<commit_after>default_style = 'clam'

default_font = None

# {
#     'size': 9,
#     'weight': 'normal',
#     'underline': 0,
#     'overstrike': 0,
#     'slant': 'roman',
#     # 'family': 'Segoe UI'
#     'family': 'Tamsyn6x12'
# }

default_colors = {
    'foreground': 'white',
    'activeForeground': '#222',
    'disabledForeground': '#000',
    'background': '#131313',
    # 'activeBackground': '', # menus
    'insertBackground': '#367AE3',
    'selectColor': '#367AE3',
    # these are for selecting text w/in entries?
    # but not ttk...
    'selectForeground': '#222',
    'selectBackground': '#367AE3',
    'troughColor': '#232626',
    # 'highlightBackground': '',  # this goes around canvases for some reason
    # 'highlightColor': '',
}

# for consistency
default_colors['highlightBackground'] = default_colors['background']

# http://paletton.com/#uid=75a1G0kl1Wx1x+IcEXDsUWkWEVB
pad_colors = [
   # primary color
   ['#FA97CF',
    '#F6008B'],
   # secondary color 1
   ['#FFE79A',
    '#FFC200'],
   # secondary color 2
   ['#A7A3F7',
    '#0E05F2'],
   # complement color
   ['#DCFD99',
    '#A8FC00']
]

# ttk styled specific colors
styled_button_color = '#1F1F1F'
styled_button_border = '#101010'
styled_button_disabled_color = '#2A292A'
styled_button_disabled_color_darker = '#2B282A'

# list of tuples of style to replace & a dict of all their options
ttk_style_opts_to_config = [
('TButton', {
  'background': styled_button_color,
  'bordercolor': 'hot pink',
  'lightcolor': 'blue',
  'darkcolor': 'red',
}),
('TLabel', {
  'background': styled_button_color,
  'bordercolor': 'hot pink',
  'lightcolor': 'red',
  'darkcolor': 'blue',
}),
('TFrame', {
  'bordercolor': 'hot pink',
  'lightcolor': 'blue',
  'darkcolor': 'red',
}),
('fakesep.TFrame', {
  'bordercolor': 'hot pink',
  'lightcolor': 'blue',
  'darkcolor': 'red',
  'background': 'hot pink',
}),
('TEntry', {
  'fieldbackground': styled_button_color,
  'background': styled_button_color,
  'selectbackground': default_colors['selectColor'],
  'bordercolor': 'hot pink',
  'lightcolor': 'blue',
  'darkcolor': 'red',
}),
('TCheckbutton', {
  'background': default_colors['background'],
}),
('TRadiobutton', {
  'background': default_colors['background'],
}),
('Treeview', {
  'indent': 2,
  'borderwidth':  0,
  'relief': 'flat',
  'background': default_colors['activeForeground'],
  'fieldbackground': default_colors['activeForeground'],
  'selectforeground': default_colors['selectColor'],
  'bordercolor': 'hot pink',
  'lightcolor': 'blue',
  'darkcolor': 'red',
}),
('Treeview.Heading', {
  'background': default_colors['background'],
}),
('TNotebook', {
  'relief': 'flat',
  'background': default_colors['background'],
  'bordercolor': 'hot pink',
  'lightcolor': 'red',
  'darkcolor': 'blue',
}),
('TNotebook.Tab', {
  'relief': 'flat',
  'background': styled_button_disabled_color,
  'bordercolor': 'hot pink',
  'lightcolor': 'blue',
  'darkcolor': 'red',
}),
('TLabelFrame', {
  'relief': 'flat',
  'background': default_colors['background'],
  'bordercolor': 'hot pink',
  'lightcolor': 'red',
  'darkcolor': 'blue',
}),
('Vertical.TScrollbar', {
  'background': styled_button_color,
  'bordercolor': 'hot pink',
  'troughcolor': styled_button_disabled_color_darker,
  'lightcolor': 'blue',
  'darkcolor': 'red',
}),
('Horizontal.TScrollbar', {
  'background': styled_button_color,
  'bordercolor': 'hot pink',
  'troughcolor': styled_button_disabled_color_darker,
  'lightcolor': 'red',
  'darkcolor': 'blue',
}),
('Vertical.TScale', {
  'background': styled_button_color,
  'bordercolor': styled_button_color,
  'troughcolor': styled_button_disabled_color_darker,
  'lightcolor': 'blue',
  'darkcolor': 'red',
}),
('Horizontal.TScale', {
  'background': styled_button_color,
  'bordercolor': styled_button_color,
  'troughcolor': styled_button_disabled_color_darker,
  'lightcolor': 'blue',
  'darkcolor': 'red',
}),
('Combobox.field', {
  'background': styled_button_color,
  'fieldbackground': styled_button_color,
  'bordercolor': 'hot pink',
  'lightcolor': 'blue',
  'darkcolor': 'red',
}),

]

# style options that need to be mapped (for active/disabled states usually)
ttk_style_opts_to_map = [
('TButton', {
    'background': [('disabled', styled_button_disabled_color), ('active', styled_button_disabled_color_darker)],
    'relief': [('pressed', '!disabled', 'sunken')]
}),
('TLabel', {
    'background': [('disabled', styled_button_disabled_color), ('active', styled_button_color)],
}),
('fakebut.TLabel', {
    'background': [('disabled', styled_button_disabled_color_darker), ('active', styled_button_color)],
    'foreground': [('disabled', default_colors['foreground']), ('active', default_colors['foreground'])],
}),
('TCheckbutton', {
    'background': [('disabled', styled_button_disabled_color), ('active', styled_button_color)],
}),
('TRadiobutton', {
    'background': [('disabled', styled_button_disabled_color), ('active', styled_button_color)],
}),
('Treeview', {
   'background': [('selected', 'focus', default_colors['selectColor']), ('selected', '!focus', '#6087C4')],
}),
('Treeview.Heading', {
   'background': [('selected', 'focus', default_colors['selectColor']), ('selected', '!focus', '#6087C4')],
}),
('TNotebook.Tab', {
   'background': [('selected', styled_button_color)],
}),
('Vertical.TScrollbar', {
    'background': [('disabled', styled_button_disabled_color), ('active', default_colors['troughColor'])],
}),
('Horizontal.TScrollbar', {
    'background': [('disabled', styled_button_disabled_color), ('active', default_colors['troughColor'])],
}),
('Vertical.TScale', {
    'background': [('disabled', styled_button_disabled_color), ('active', default_colors['troughColor'])],
}),
('Horizontal.TScale', {
    'background': [('disabled', styled_button_disabled_color), ('active', default_colors['troughColor'])],
}),
('TCombobox', {
    'fieldbackground': [('readonly', 'focus', '#6087C4'), ('readonly', styled_button_disabled_color)],
    'background': [('active', styled_button_color), ('pressed', styled_button_border)],
    'foreground': [('readonly', 'focus', default_colors['selectForeground'])],
}),
]
<|endoftext|>


<commit_before>#!/usr/bin/env python

# hook for virtualenv
# switch to the virtualenv where the executor belongs,
# replace all the path for modules
import sys, os.path
P = 'site-packages'
apath = os.path.abspath(__file__)
if P in apath:
    virltualenv = apath[:apath.index(P)]
    sysp = [p[:-len(P)] for p in sys.path if p.endswith(P)][0]
    if sysp != virltualenv:
        sys.path = [p.replace(sysp, virltualenv) for p in sys.path]

import logging
import os, sys, time
import signal
import os.path
import marshal
import cPickle
import multiprocessing
import threading
import SocketServer
import SimpleHTTPServer
import shutil
import socket
import urllib2
import platform
import gc

import zmq

# ignore INFO and DEBUG log
os.environ['GLOG_logtostderr'] = '1'
os.environ['GLOG_minloglevel'] = '1'
try:
    import mesos
    import mesos_pb2
except ImportError:
    import pymesos as mesos
    import pymesos.mesos_pb2 as mesos_pb2

sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))
from dpark.util import compress, decompress, getproctitle, setproctitle, spawn
from dpark.serialize import marshalable
from dpark.accumulator import Accumulator
from dpark.schedule import Success, OtherFailure
from dpark.env import env
from dpark.shuffle import LocalFileShuffle

logger = logging.getLogger("executor@%s" % socket.gethostname())

TASK_RESULT_LIMIT = 1024 * 256
DEFAULT_WEB_PORT = 5055
MAX_TASKS_PER_WORKER = 50
MAX_IDLE_TIME = 60
MAX_IDLE_WORKERS = 8

Script = ''

def reply_status(driver, task, status, data=None):
    update = mesos_pb2.TaskStatus()
    update.task_id.value = task.task_id.value
    update.state = status
    if data is not None:
        update.data = data
    driver.sendStatusUpdate(update)

def run_task(task_data):
    try:
        gc.disable()
        task, ntry = cPickle.loads(decompress(task_data))
        setproctitle('dpark worker %s: run task %s' % (Script, task))

        Accumulator.clear()
        result = task.run(ntry)
        accUpdate = Accumulator.values()

        if marshalable(result):
            flag, data = 0, marshal.dumps(result)
        else:
            flag, data = 1, cPickle.dumps(result, -1)
        data = compress(data)

        if len(data) > TASK_RESULT_LIMIT:
            path = LocalFileShuffle.getOutputFile(0, ntry, task.id, len(data))
            f = open(path, 'w')
            f.write(data)
            f.close()
            data = '/'.join([LocalFileShuffle.getServerUri()] + path.split('/')[-3:])
            flag += 2

        return mesos_pb2.TASK_FINISHED, cPickle.dumps((task.id, Success(), (flag, data), accUpdate), -1)
    except Exception, e:
        import traceback
        msg = traceback.format_exc()
        return mesos_pb2.TASK_FAILED, cPickle.dumps((task.id, OtherFailure(msg), None, None), -1)
    finally:
        setproctitle('dpark worker: idle')
        gc.collect()
        gc.enable()

def init_env(args):
    setproctitle('dpark worker: idle')
    env.start(False, args)

class LocalizedHTTP(SimpleHTTPServer.SimpleHTTPRequestHandler):
    basedir = None
    def translate_path(self, path):
        out = SimpleHTTPServer.SimpleHTTPRequestHandler.translate_path(self, path)
        return self.basedir + '/' + os.path.relpath(out)

    def log_message(self, format, *args):
        pass

def startWebServer(path):
    # check the default web server
    if not os.path.exists(path):
        os.makedirs(path)
    testpath = os.path.join(path, 'test')
    with open(testpath, 'w') as f:
        f.write(path)
    default_uri = 'http://%s:%d/%s' % (socket.gethostname(), DEFAULT_WEB_PORT,
            os.path.basename(path))
    try:
        data = urllib2.urlopen(default_uri + '/' + 'test').read()
        if data == path:
            return default_uri
    except IOError, e:
        pass

    logger.warning("default webserver at %s not available", DEFAULT_WEB_PORT)
    LocalizedHTTP.basedir = os.path.dirname(path)
    ss = SocketServer.TCPServer(('0.0.0.0', 0), LocalizedHTTP)
    spawn(ss.serve_forever)
    uri = "http://%s:%d/%s" % (socket.gethostname(), ss.server_address[1],
            os.path.basename(path))
    return uri

def forword(fd, addr, prefix=''):
    f = os.fdopen(fd, 'r')
    ctx = zmq.Context()
    out = [None]
    buf = []
    def send(buf):
        if not out[0]:
            out[0] = ctx.socket(zmq.PUSH)
            out[0].connect(addr)
        out[0].send(prefix+''.join(buf))

    while True:
        try:
            line = f.readline()
            if not line: break
            buf.append(line)
            if line.endswith('\n'):
                send(buf)
                buf = []
        except IOError:
            break
    if buf:
        send(buf)
    if out[0]:
        out[0].close()
    f.close()
    ctx.shutdown()

def start_forword(addr, prefix=''):
    rfd, wfd = os.pipe()
    t = spawn(forword, rfd, addr, prefix)
    return t, os.fdopen(wfd, 'w', 0)

def get_pool_memory(pool):
    try:
        import psutil
        p = psutil.Process(pool._pool[0].pid)
        return p.get_memory_info()[0] >> 20
    except Exception:
        return 0

def get_task_memory(task):
    for r in task.resources:
        if r.name == 'mem':
            return r.scalar.value
    logger.error("no memory in resource: %s", task.resources)
    return 100 # 100M

def safe(f):
    def _(self, *a, **kw):
        with self.lock:
            r = f(self, *a, **kw)
        return r
    return _

def setup_cleaner_process(workdir):
    ppid = os.getpid()
    pid = os.fork()
    if pid == 0:
        os.setsid()
        pid = os.fork()
        if pid == 0:
            try:
                setproctitle('dpark cleaner %s wait(%d)' % (workdir, ppid))
                import psutil
                psutil.Process(ppid).wait()
                os.killpg(ppid, signal.SIGKILL) # kill workers
            finally:
                setproctitle('dpark cleaning %s ' % workdir)
                for d in workdir:
                    while os.path.exists(d):
                        try: shutil.rmtree(d, True)
                        except: pass
        sys.exit(0)
    os.wait()

class MyExecutor(mesos.Executor):
    def __init__(self):
        self.workdir = None
        self.idle_workers = []
        self.busy_workers = {}
        self.lock = threading.RLock()

    def check_memory(self, driver):
        try:
            import psutil
        except ImportError:
            logger.error("no psutil module")
            return

        mem_limit = {}

        while True:
            self.lock.acquire()

            for tid, (task, pool) in self.busy_workers.items():
                pid = pool._pool[0].pid
                try:
                    p = psutil.Process(pid)
                    rss = p.get_memory_info()[0] >> 20
                except psutil.error.Error, e:
                    logger.error("worker process %d of task %s is dead: %s", pid, tid, e)
                    reply_status(driver, task, mesos_pb2.TASK_LOST)
                    self.busy_workers.pop(tid)
                    continue

                if p.status == psutil.STATUS_ZOMBIE or not p.is_running():
                    logger.error("worker process %d of task %s is zombie", pid, tid)
                    reply_status(driver, task, mesos_pb2.TASK_LOST)
                    self.busy_workers.pop(tid)
                    continue

                offered = get_task_memory(task)
                if not offered:
                    continue
                if rss > offered * 2:
                    logger.warning("task %s used too much memory: %dMB > %dMB * 2, kill it. "
                            + "use -M argument or taskMemory to request more memory.", tid, rss, offered)
                    reply_status(driver, task, mesos_pb2.TASK_KILLED)
                    self.busy_workers.pop(tid)
                    pool.terminate()
                elif rss > offered * mem_limit.get(tid, 1.0):
                    logger.debug("task %s used too much memory: %dMB > %dMB, "
                            + "use -M to request or taskMemory for more memory", tid, rss, offered)
                    mem_limit[tid] = rss / offered + 0.2

            now = time.time()
            n = len([1 for t, p in self.idle_workers if t + MAX_IDLE_TIME < now])
            if n:
                for _, p in self.idle_workers[:n]:
                    p.terminate()
                self.idle_workers = self.idle_workers[n:]

            self.lock.release()

            time.sleep(1)

    @safe
    def registered(self, driver, executorInfo, frameworkInfo, slaveInfo):
        try:
            global Script
            Script, cwd, python_path, osenv, self.parallel, out_logger, err_logger, logLevel, args = marshal.loads(executorInfo.data)
            self.init_args = args
            try:
                os.chdir(cwd)
            except OSError:
                driver.sendFrameworkMessage("switch cwd failed: %s not exists!" % cwd)
            sys.path = python_path
            os.environ.update(osenv)
            prefix = '[%s] ' % socket.gethostname()
            if out_logger:
                self.outt, sys.stdout = start_forword(out_logger, prefix)
            if err_logger:
                self.errt, sys.stderr = start_forword(err_logger, prefix)
            logging.basicConfig(format='%(asctime)-15s [%(levelname)s] [%(name)-9s] %(message)s', level=logLevel)

            self.workdir = args['WORKDIR']
            root = os.path.dirname(self.workdir[0])
            if not os.path.exists(root):
                os.mkdir(root)
                os.chmod(root, 0777) # because umask
            args['SERVER_URI'] = startWebServer(self.workdir[0])
            if 'MESOS_SLAVE_PID' in os.environ: # make unit test happy
                setup_cleaner_process(self.workdir)

            spawn(self.check_memory, driver)

            # wait background threads to been initialized
            time.sleep(0.1)
            if out_logger:
                os.close(1)
                assert os.dup(sys.stdout.fileno()) == 1, 'redirect stdout failed'
            if err_logger:
                os.close(2)
                assert os.dup(sys.stderr.fileno()) == 2, 'redirect stderr failed'

            logger.debug("executor started at %s", slaveInfo.hostname)

        except Exception, e:
            import traceback
            msg = traceback.format_exc()
            logger.error("init executor failed: %s", msg)
            raise

    def get_idle_worker(self):
        try:
            return self.idle_workers.pop()[1]
        except IndexError:
            p = multiprocessing.Pool(1, init_env, [self.init_args])
            p.done = 0
            return p

    @safe
    def launchTask(self, driver, task):
        try:
            reply_status(driver, task, mesos_pb2.TASK_RUNNING)
            logging.debug("launch task %s", task.task_id.value)

            pool = self.get_idle_worker()
            self.busy_workers[task.task_id.value] = (task, pool)

            def callback((state, data)):
                with self.lock:
                    if task.task_id.value not in self.busy_workers:
                        return
                    _, pool = self.busy_workers.pop(task.task_id.value)
                    pool.done += 1
                    reply_status(driver, task, state, data)
                    if (len(self.idle_workers) + len(self.busy_workers) < self.parallel
                            and len(self.idle_workers) < MAX_IDLE_WORKERS
                            and pool.done < MAX_TASKS_PER_WORKER
                            and get_pool_memory(pool) < get_task_memory(task)): # maybe memory leak in executor
                        self.idle_workers.append((time.time(), pool))
                    else:
                        try: pool.terminate()
                        except: pass

            pool.apply_async(run_task, [task.data], callback=callback)

        except Exception, e:
            import traceback
            msg = traceback.format_exc()
            reply_status(driver, task, mesos_pb2.TASK_LOST, msg)
            return

    @safe
    def killTask(self, driver, taskId):
        if taskId.value in self.busy_workers:
            task, pool = self.busy_workers.pop(taskId.value)
            reply_status(driver, task, mesos_pb2.TASK_KILLED)
            pool.terminate()

    @safe
    def shutdown(self, driver=None):
        def terminate(p):
            try:
                for pi in p._pool:
                    os.kill(pi.pid, signal.SIGKILL)
            except Exception, e:
                pass
        for _, p in self.idle_workers:
            terminate(p)
        for _, p in self.busy_workers.itervalues():
            terminate(p)

        # clean work files
        for d in self.workdir:
            try: shutil.rmtree(d, True)
            except: pass

def run():
    executor = MyExecutor()
    driver = mesos.MesosExecutorDriver(executor)
    driver.run()

if __name__ == '__main__':
    run()
<commit_msg>bugfix: catch exception in cleaner process<commit_after>#!/usr/bin/env python

# hook for virtualenv
# switch to the virtualenv where the executor belongs,
# replace all the path for modules
import sys, os.path
P = 'site-packages'
apath = os.path.abspath(__file__)
if P in apath:
    virltualenv = apath[:apath.index(P)]
    sysp = [p[:-len(P)] for p in sys.path if p.endswith(P)][0]
    if sysp != virltualenv:
        sys.path = [p.replace(sysp, virltualenv) for p in sys.path]

import logging
import os, sys, time
import signal
import os.path
import marshal
import cPickle
import multiprocessing
import threading
import SocketServer
import SimpleHTTPServer
import shutil
import socket
import urllib2
import platform
import gc

import zmq

# ignore INFO and DEBUG log
os.environ['GLOG_logtostderr'] = '1'
os.environ['GLOG_minloglevel'] = '1'
try:
    import mesos
    import mesos_pb2
except ImportError:
    import pymesos as mesos
    import pymesos.mesos_pb2 as mesos_pb2

sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))
from dpark.util import compress, decompress, getproctitle, setproctitle, spawn
from dpark.serialize import marshalable
from dpark.accumulator import Accumulator
from dpark.schedule import Success, OtherFailure
from dpark.env import env
from dpark.shuffle import LocalFileShuffle

logger = logging.getLogger("executor@%s" % socket.gethostname())

TASK_RESULT_LIMIT = 1024 * 256
DEFAULT_WEB_PORT = 5055
MAX_TASKS_PER_WORKER = 50
MAX_IDLE_TIME = 60
MAX_IDLE_WORKERS = 8

Script = ''

def reply_status(driver, task, status, data=None):
    update = mesos_pb2.TaskStatus()
    update.task_id.value = task.task_id.value
    update.state = status
    if data is not None:
        update.data = data
    driver.sendStatusUpdate(update)

def run_task(task_data):
    try:
        gc.disable()
        task, ntry = cPickle.loads(decompress(task_data))
        setproctitle('dpark worker %s: run task %s' % (Script, task))

        Accumulator.clear()
        result = task.run(ntry)
        accUpdate = Accumulator.values()

        if marshalable(result):
            flag, data = 0, marshal.dumps(result)
        else:
            flag, data = 1, cPickle.dumps(result, -1)
        data = compress(data)

        if len(data) > TASK_RESULT_LIMIT:
            path = LocalFileShuffle.getOutputFile(0, ntry, task.id, len(data))
            f = open(path, 'w')
            f.write(data)
            f.close()
            data = '/'.join([LocalFileShuffle.getServerUri()] + path.split('/')[-3:])
            flag += 2

        return mesos_pb2.TASK_FINISHED, cPickle.dumps((task.id, Success(), (flag, data), accUpdate), -1)
    except Exception, e:
        import traceback
        msg = traceback.format_exc()
        return mesos_pb2.TASK_FAILED, cPickle.dumps((task.id, OtherFailure(msg), None, None), -1)
    finally:
        setproctitle('dpark worker: idle')
        gc.collect()
        gc.enable()

def init_env(args):
    setproctitle('dpark worker: idle')
    env.start(False, args)

class LocalizedHTTP(SimpleHTTPServer.SimpleHTTPRequestHandler):
    basedir = None
    def translate_path(self, path):
        out = SimpleHTTPServer.SimpleHTTPRequestHandler.translate_path(self, path)
        return self.basedir + '/' + os.path.relpath(out)

    def log_message(self, format, *args):
        pass

def startWebServer(path):
    # check the default web server
    if not os.path.exists(path):
        os.makedirs(path)
    testpath = os.path.join(path, 'test')
    with open(testpath, 'w') as f:
        f.write(path)
    default_uri = 'http://%s:%d/%s' % (socket.gethostname(), DEFAULT_WEB_PORT,
            os.path.basename(path))
    try:
        data = urllib2.urlopen(default_uri + '/' + 'test').read()
        if data == path:
            return default_uri
    except IOError, e:
        pass

    logger.warning("default webserver at %s not available", DEFAULT_WEB_PORT)
    LocalizedHTTP.basedir = os.path.dirname(path)
    ss = SocketServer.TCPServer(('0.0.0.0', 0), LocalizedHTTP)
    spawn(ss.serve_forever)
    uri = "http://%s:%d/%s" % (socket.gethostname(), ss.server_address[1],
            os.path.basename(path))
    return uri

def forword(fd, addr, prefix=''):
    f = os.fdopen(fd, 'r')
    ctx = zmq.Context()
    out = [None]
    buf = []
    def send(buf):
        if not out[0]:
            out[0] = ctx.socket(zmq.PUSH)
            out[0].connect(addr)
        out[0].send(prefix+''.join(buf))

    while True:
        try:
            line = f.readline()
            if not line: break
            buf.append(line)
            if line.endswith('\n'):
                send(buf)
                buf = []
        except IOError:
            break
    if buf:
        send(buf)
    if out[0]:
        out[0].close()
    f.close()
    ctx.shutdown()

def start_forword(addr, prefix=''):
    rfd, wfd = os.pipe()
    t = spawn(forword, rfd, addr, prefix)
    return t, os.fdopen(wfd, 'w', 0)

def get_pool_memory(pool):
    try:
        import psutil
        p = psutil.Process(pool._pool[0].pid)
        return p.get_memory_info()[0] >> 20
    except Exception:
        return 0

def get_task_memory(task):
    for r in task.resources:
        if r.name == 'mem':
            return r.scalar.value
    logger.error("no memory in resource: %s", task.resources)
    return 100 # 100M

def safe(f):
    def _(self, *a, **kw):
        with self.lock:
            r = f(self, *a, **kw)
        return r
    return _

def setup_cleaner_process(workdir):
    ppid = os.getpid()
    pid = os.fork()
    if pid == 0:
        os.setsid()
        pid = os.fork()
        if pid == 0:
            try:
                import psutil
            except ImportError:
                sys.exit(1)
            try:
                setproctitle('dpark cleaner %s wait(%d)' % (workdir, ppid))
                psutil.Process(ppid).wait()
                os.killpg(ppid, signal.SIGKILL) # kill workers
            except Exception, e:
                pass # make sure to exit
            finally:
                setproctitle('dpark cleaning %s ' % workdir)
                for d in workdir:
                    while os.path.exists(d):
                        try: shutil.rmtree(d, True)
                        except: pass
        sys.exit(0)
    os.wait()

class MyExecutor(mesos.Executor):
    def __init__(self):
        self.workdir = None
        self.idle_workers = []
        self.busy_workers = {}
        self.lock = threading.RLock()

    def check_memory(self, driver):
        try:
            import psutil
        except ImportError:
            logger.error("no psutil module")
            return

        mem_limit = {}

        while True:
            self.lock.acquire()

            for tid, (task, pool) in self.busy_workers.items():
                pid = pool._pool[0].pid
                try:
                    p = psutil.Process(pid)
                    rss = p.get_memory_info()[0] >> 20
                except psutil.error.Error, e:
                    logger.error("worker process %d of task %s is dead: %s", pid, tid, e)
                    reply_status(driver, task, mesos_pb2.TASK_LOST)
                    self.busy_workers.pop(tid)
                    continue

                if p.status == psutil.STATUS_ZOMBIE or not p.is_running():
                    logger.error("worker process %d of task %s is zombie", pid, tid)
                    reply_status(driver, task, mesos_pb2.TASK_LOST)
                    self.busy_workers.pop(tid)
                    continue

                offered = get_task_memory(task)
                if not offered:
                    continue
                if rss > offered * 2:
                    logger.warning("task %s used too much memory: %dMB > %dMB * 2, kill it. "
                            + "use -M argument or taskMemory to request more memory.", tid, rss, offered)
                    reply_status(driver, task, mesos_pb2.TASK_KILLED)
                    self.busy_workers.pop(tid)
                    pool.terminate()
                elif rss > offered * mem_limit.get(tid, 1.0):
                    logger.debug("task %s used too much memory: %dMB > %dMB, "
                            + "use -M to request or taskMemory for more memory", tid, rss, offered)
                    mem_limit[tid] = rss / offered + 0.2

            now = time.time()
            n = len([1 for t, p in self.idle_workers if t + MAX_IDLE_TIME < now])
            if n:
                for _, p in self.idle_workers[:n]:
                    p.terminate()
                self.idle_workers = self.idle_workers[n:]

            self.lock.release()

            time.sleep(1)

    @safe
    def registered(self, driver, executorInfo, frameworkInfo, slaveInfo):
        try:
            global Script
            Script, cwd, python_path, osenv, self.parallel, out_logger, err_logger, logLevel, args = marshal.loads(executorInfo.data)
            self.init_args = args
            try:
                os.chdir(cwd)
            except OSError:
                driver.sendFrameworkMessage("switch cwd failed: %s not exists!" % cwd)
            sys.path = python_path
            os.environ.update(osenv)
            prefix = '[%s] ' % socket.gethostname()
            if out_logger:
                self.outt, sys.stdout = start_forword(out_logger, prefix)
            if err_logger:
                self.errt, sys.stderr = start_forword(err_logger, prefix)
            logging.basicConfig(format='%(asctime)-15s [%(levelname)s] [%(name)-9s] %(message)s', level=logLevel)

            self.workdir = args['WORKDIR']
            root = os.path.dirname(self.workdir[0])
            if not os.path.exists(root):
                os.mkdir(root)
                os.chmod(root, 0777) # because umask
            args['SERVER_URI'] = startWebServer(self.workdir[0])
            if 'MESOS_SLAVE_PID' in os.environ: # make unit test happy
                setup_cleaner_process(self.workdir)

            spawn(self.check_memory, driver)

            # wait background threads to been initialized
            time.sleep(0.1)
            if out_logger:
                os.close(1)
                assert os.dup(sys.stdout.fileno()) == 1, 'redirect stdout failed'
            if err_logger:
                os.close(2)
                assert os.dup(sys.stderr.fileno()) == 2, 'redirect stderr failed'

            logger.debug("executor started at %s", slaveInfo.hostname)

        except Exception, e:
            import traceback
            msg = traceback.format_exc()
            logger.error("init executor failed: %s", msg)
            raise

    def get_idle_worker(self):
        try:
            return self.idle_workers.pop()[1]
        except IndexError:
            p = multiprocessing.Pool(1, init_env, [self.init_args])
            p.done = 0
            return p

    @safe
    def launchTask(self, driver, task):
        try:
            reply_status(driver, task, mesos_pb2.TASK_RUNNING)
            logging.debug("launch task %s", task.task_id.value)

            pool = self.get_idle_worker()
            self.busy_workers[task.task_id.value] = (task, pool)

            def callback((state, data)):
                with self.lock:
                    if task.task_id.value not in self.busy_workers:
                        return
                    _, pool = self.busy_workers.pop(task.task_id.value)
                    pool.done += 1
                    reply_status(driver, task, state, data)
                    if (len(self.idle_workers) + len(self.busy_workers) < self.parallel
                            and len(self.idle_workers) < MAX_IDLE_WORKERS
                            and pool.done < MAX_TASKS_PER_WORKER
                            and get_pool_memory(pool) < get_task_memory(task)): # maybe memory leak in executor
                        self.idle_workers.append((time.time(), pool))
                    else:
                        try: pool.terminate()
                        except: pass

            pool.apply_async(run_task, [task.data], callback=callback)

        except Exception, e:
            import traceback
            msg = traceback.format_exc()
            reply_status(driver, task, mesos_pb2.TASK_LOST, msg)
            return

    @safe
    def killTask(self, driver, taskId):
        if taskId.value in self.busy_workers:
            task, pool = self.busy_workers.pop(taskId.value)
            reply_status(driver, task, mesos_pb2.TASK_KILLED)
            pool.terminate()

    @safe
    def shutdown(self, driver=None):
        def terminate(p):
            try:
                for pi in p._pool:
                    os.kill(pi.pid, signal.SIGKILL)
            except Exception, e:
                pass
        for _, p in self.idle_workers:
            terminate(p)
        for _, p in self.busy_workers.itervalues():
            terminate(p)

        # clean work files
        for d in self.workdir:
            try: shutil.rmtree(d, True)
            except: pass

def run():
    executor = MyExecutor()
    driver = mesos.MesosExecutorDriver(executor)
    driver.run()

if __name__ == '__main__':
    run()
<|endoftext|>


<commit_before>#!/usr/bin/env python

# test-client - Simple client for testing parts of Telepathy connection managers
#
# Copyright (C) 2005 Collabora Limited
# Copyright (C) 2005 Nokia Corporation
#
# This library is free software; you can redistribute it and/or
# modify it under the terms of the GNU Lesser General Public
# License as published by the Free Software Foundation; either
# version 2.1 of the License, or (at your option) any later version.
#
# This library is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# Library General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public
# License along with this library; if not, write to the Free Software
# Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA

import dbus
import dbus.glib

assert(getattr(dbus, 'version', (0,0,0)) >= (0,51,0))

import getpass
import gobject
import signal
import sys
import os

from telepathy import *

import telepathy.client

class ContactListChannel(telepathy.client.Channel):
    def __init__(self, conn, object_path, handle):
        telepathy.client.Channel.__init__(self, conn._service_name, object_path)
        self.get_valid_interfaces().add(CHANNEL_TYPE_CONTACT_LIST)
        self._conn = conn
        self._handle = handle
        self._name = handle

    def got_interfaces(self):
        self._conn[CONN_INTERFACE].InspectHandle(self._handle, reply_handler=self.inspect_handle_reply_cb, error_handler=self.error_cb)
        self[CHANNEL_INTERFACE].GetMembers(reply_handler=self.get_members_reply_cb, error_handler=self.error_cb)
        self[CHANNEL_INTERFACE_GROUP].GetLocalPendingMembers(reply_handler=self.get_local_pending_members_reply_cb, error_handler=self.error_cb)
        self[CHANNEL_INTERFACE_GROUP].GetRemotePendingMembers(reply_handler=self.get_remote_pending_members_reply_cb, error_handler=self.error_cb)
        self[CHANNEL_INTERFACE_GROUP].connect_to_signal('MembersChanged', self.members_changed_signal_cb)

    def inspect_handle_reply_cb(self, type, name):
        print "CLC", self._name, "is", name
        self._name = name
        if name == 'subscribe':
            gobject.idle_add(self.subscribe_list_idle_cb)
        elif name == 'publish':
            gobject.idle_add(self.publish_list_idle_cb)

    def subscribe_list_idle_cb(self):
#        handle = self._conn[CONN_INTERFACE].RequestHandle(CONNECTION_HANDLE_TYPE_CONTACT, 'test2@localhost')
#        self[CHANNEL_INTERFACE_GROUP].AddMembers([handle])
        return False

    def publish_list_idle_cb(self):
        return False

    def members_changed_signal_cb(self, message, added, removed, local_p, remote_p):
        print "MembersChanged on CLC", self._name
        print "Message: ", message
        print "Added: ", added
        print "Removed: ", removed
        print "Local Pending: ", local_p
        print "Remote Pending: ", remote_p
#        if added and self._name == 'subscribe':
#            self[CHANNEL_INTERFACE_GROUP].RemoveMembers(added)
#        if removed and self._name == 'subscribe':
#            self[CHANNEL_INTERFACE_GROUP].AddMembers(added)

#        if added:
#            self[CHANNEL_INTERFACE_GROUP].RemoveMembers(added)
#        if local_p:
#            self[CHANNEL_INTERFACE_GROUP].AddMembers(local_p)

    def get_members_reply_cb(self, members):
        print "got members on CLC %s: %s" % (self._name, members)
#        if members and self._name == 'subscribe':
#            self[CHANNEL_INTERFACE_GROUP].RemoveMembers(members)
#        if members:
#            self[CHANNEL_INTERFACE_GROUP].RemoveMembers(members)

    def get_local_pending_members_reply_cb(self, members):
        print "got local pending members on CLC %s: %s" % (self._name, members)
#        if members:
#            self[CHANNEL_INTERFACE_GROUP].AddMembers(members)
#            print "sending addmembers"

    def get_remote_pending_members_reply_cb(self, members):
        print "got remote pending members on CLC %s: %s" % (self._name, members)

class StreamedMediaChannel(telepathy.client.Channel):
    def __init__(self, conn, object_path, handle):
        telepathy.client.Channel.__init__(self, conn._service_name, object_path)
        self._conn = conn
        self.get_valid_interfaces().add(CHANNEL_TYPE_STREAMED_MEDIA)
        self[CHANNEL_TYPE_STREAMED_MEDIA].connect_to_signal('ReceivedMediaParameters', self.received_media_params_cb)
        self[CHANNEL_INTERFACE].connect_to_signal('Closed', self.closed_cb)

    def got_interfaces(self):
        print "SMC got interfaces"
        self[CHANNEL_INTERFACE].GetMembers(reply_handler=self.get_members_reply_cb, error_handler=self.error_cb)
        self[CHANNEL_INTERFACE_GROUP].GetLocalPendingMembers(reply_handler=self.get_local_pending_members_reply_cb, error_handler=self.error_cb)
        self[CHANNEL_INTERFACE_GROUP].GetRemotePendingMembers(reply_handler=self.get_remote_pending_members_reply_cb, error_handler=self.error_cb)
        self[CHANNEL_INTERFACE_GROUP].connect_to_signal('MembersChanged', self.members_changed_signal_cb)

    def received_media_params_cb(self, id, local, remote):
        print "Received media params for %s.\n   local SDP = %s\n  remote SDP = %s" % (id, local, remote)

    def closed_cb(self):
        print "Channel closed"

    def members_changed_signal_cb(self, reason, added, removed, local_pending, remote_pending):
        print "Members Changed"
        print "  added ",added," removed", removed
        print "  local pending",local_pending," remote pending ", remote_pending

        self[CHANNEL_INTERFACE_DTMF].SendDTMF(1, 160)

    def error_cb(self, exception):
        print "Exception received from asynchronous method call:"
        print exception

    def get_members_reply_cb(self, members):
        print "Channel has members", members

    def get_local_pending_members_reply_cb(self, members):
        print "Channel has local pending members", members
        print self[CHANNEL_INTERFACE].GetInterfaces()
        if members:
           if "SIPDEBUG" in os.environ and os.environ["SIPDEBUG"] == "transfer":
               print "getting handle for test2"
               handle = self._conn[CONN_INTERFACE].RequestHandle(CONNECTION_HANDLE_TYPE_CONTACT, 'sip:test2@192.168.1.101')
               print "got handle %d for %s" % (handle, 'sip:test2@192.168.1.101')
               self[CHANNEL_INTERFACE_TRANSFER].Transfer(members[0],handle)
           else:
              print "adding member"
              self[CHANNEL_INTERFACE_GROUP].AddMembers(members)

    def get_remote_pending_members_reply_cb(self, members):
        print "Channel has remote pending members", members

class TextChannel(telepathy.client.Channel):
    def __init__(self, conn, object_path, handle):
        # FIXME: this is racy and buggy and smells, and is deprecated in
        # favour of the stuff in sdp-client.py
        telepathy.client.Channel.__init__(self, conn._service_name, object_path)
        self.get_valid_interfaces().add(CHANNEL_TYPE_TEXT)
        self[CHANNEL_TYPE_TEXT].connect_to_signal('Received', self.received_callback)
        self.doack = True

        self._handled_pending_message = None
        pending_messages = self[CHANNEL_TYPE_TEXT].ListPendingMessages()
        print pending_messages
        for msg in pending_messages:
            (id, timestamp, sender, type, message) = msg
            print "Handling pending message", id
            self.received_callback(id, timestamp, sender, type, message)
            self._handled_pending_message = id

    def received_callback(self, id, timestamp, sender, type, message):
        if self._handled_pending_message != None:
            if id > self._handled_pending_message:
                print "Now handling messages directly"
                self._handled_pending_message = None
            else:
                print "Skipping already handled message", id
                return

        print "Received", id, timestamp, sender, type, message
        self[CHANNEL_TYPE_TEXT].Send(CHANNEL_TEXT_MESSAGE_TYPE_NORMAL, 'got message ' + str(id) + '(' + message + ')')
        if self.doack:
            print "Acknowledging...", id
            self[CHANNEL_TYPE_TEXT].AcknowledgePendingMessage(id)

class TestConnection(telepathy.client.Connection):
    def __init__(self, service_name, object_path, mainloop):
        telepathy.client.Connection.__init__(self, service_name, object_path)

        self._mainloop = mainloop
        self._service_name = service_name

        self._status = CONNECTION_STATUS_CONNECTING
        self._channels = {}

        print "calling GetStatus async"
        # handle race condition when connecting completes
        # before the signal handlers are registered
        self[CONN_INTERFACE].GetStatus(reply_handler=self.get_status_reply_cb, error_handler=self.error_cb)
        print "called GetStatus"

        print 'connecting to StatusChanged'
        self[CONN_INTERFACE].connect_to_signal('StatusChanged', self.status_changed_signal_cb)
        print 'connecting to NewChannel'
        self[CONN_INTERFACE].connect_to_signal('NewChannel', self.new_channel_signal_cb)

    def error_cb(self, exception):
        print "Exception received from asynchronous method call:"
        print exception

    def get_status_reply_cb(self, status):
        print "get_status_reply_cb", status
        self.status_changed_signal_cb(status, CONNECTION_STATUS_REASON_NONE_SPECIFIED)

    def new_channel_signal_cb(self, obj_path, type, handle, supress_handler):
        if obj_path in self._channels:
            return

        print 'NewChannel', obj_path, type, handle, supress_handler

        channel = None

        print "channel type ", type
        if type == CHANNEL_TYPE_TEXT:
            channel = TextChannel(self, obj_path, handle)
        elif type == CHANNEL_TYPE_CONTACT_LIST:
            channel = ContactListChannel(self, obj_path, handle)
        elif type == CHANNEL_TYPE_STREAMED_MEDIA:
            print "got new streamed media channel"
            channel = StreamedMediaChannel(self, obj_path, handle)

        if channel != None:
            self._channels[obj_path] = channel
        else:
            print 'Unknown channel type', type

    def connected_cb(self):
        if CONN_INTERFACE_STREAMED_MEDIA in self.get_valid_interfaces():
            self[CONN_INTERFACE_STREAMED_MEDIA].SetMediaParameters(
"v=0\r\nm=audio 0 RTP/AVP 3 0 8\r\na=rtpmap:3 GSM/8000\r\na=rtpmap  PCMU/8000\r\na=rtpmap:8 PCMA/8000\r\n")
        if "SIPDEBUG" in os.environ and os.environ["SIPDEBUG"][:5] == "call:":
            handle = self[CONN_INTERFACE].RequestHandle(CONNECTION_HANDLE_TYPE_CONTACT, os.environ["SIPDEBUG"][5:])
            print "got handle %d for %s" % (handle, os.environ["SIPDEBUG"][5:])
            print "calling RequestChannel"
            obj = self[CONN_INTERFACE].RequestChannel(CHANNEL_TYPE_STREAMED_MEDIA, handle, True)
            print "called RequestChannel, got ", obj
        elif "SIPDEBUG" in os.environ and os.environ["SIPDEBUG"][:8] == "forward:":
            handle = self[CONN_INTERFACE].RequestHandle(CONNECTION_HANDLE_TYPE_CONTACT, os.environ["SIPDEBUG"][8:])
            print "got handle %d for %s" % (handle, os.environ["SIPDEBUG"][8:])
            #print self[CONN_INTERFACE_PRESENCE].GetStatuses()
            self[CONN_INTERFACE_FORWARDING].SetForwardingHandle(handle)
        return False

    def presence_update_signal_cb(self, presence):
        print "got presence update:", presence

    def get_interfaces_reply_cb(self, interfaces):
        self.get_valid_interfaces().update(interfaces)
        #self[CONN_INTERFACE_PRESENCE].connect_to_signal('PresenceUpdate', self.presence_update_signal_cb)
        gobject.idle_add(self.connected_cb)

    def status_changed_signal_cb(self, status, reason):
        if self._status == status:
            return
        else:
            self._status = status

        print 'StatusChanged', status, reason

        if status == CONNECTION_STATUS_CONNECTED:
            print "connection connected"
            self[CONN_INTERFACE].GetInterfaces(reply_handler=self.get_interfaces_reply_cb, error_handler=self.error_cb)
        if status == CONNECTION_STATUS_DISCONNECTED:
            print "connection terminated"
            self._mainloop.quit()

if __name__ == '__main__':
    reg = telepathy.client.ManagerRegistry()
    reg.LoadManagers()

    protos=reg.GetProtos()

    if len(protos) == 0:
        print "Sorry, no connection managers found!"
        sys.exit(1)

    if len(sys.argv) > 2:
        protocol = sys.argv[2]
    else:
        protocol=''

    while (protocol not in protos):
        protocol = raw_input('Protocol (one of: %s) [%s]: ' % (' '.join(protos),protos[0]))
        if protocol == '':
            protocol = protos[0]

    if len(sys.argv) > 1:
        manager = sys.argv[1]
    else:
        manager = ''
    managers = reg.GetManagers(protocol)

    while (manager not in managers):
        manager = raw_input('Manager (one of: %s) [%s]: ' % (' '.join(managers),managers[0]))
        if manager == '':
            manager = managers[0]

    mgr_bus_name = reg.GetBusName(manager)
    mgr_object_path = reg.GetObjectPath(manager)

    cmdline_params = dict((p.split('=') for p in sys.argv[3:]))
    params={}

    for (name, (type, default)) in reg.GetParams(manager, protocol)[0].iteritems():
        if name in cmdline_params:
            params[name] = dbus.Variant(cmdline_params[name],type)
        elif name == 'password':
            params[name] = dbus.Variant(getpass.getpass(),type)
        else:
            params[name] = dbus.Variant(raw_input(name+': '),type)

    mgr = telepathy.client.ConnectionManager(mgr_bus_name, mgr_object_path)
    print "using params:",params
    bus_name, object_path = mgr[CONN_MGR_INTERFACE].Connect(protocol, params)
    print "done connect"
    mainloop = gobject.MainLoop()
    connection = TestConnection(bus_name, object_path, mainloop)
    print "created connection"

    def quit_cb():
        connection[CONN_INTERFACE].Disconnect()
        mainloop.quit()

    def sigterm_cb():
        gobject.idle_add(quit_cb)

    signal.signal(signal.SIGTERM, sigterm_cb)

    while mainloop.is_running():
        try:
            mainloop.run()
        except KeyboardInterrupt:
            quit_cb()
<commit_msg>update test-client with spec 0.11 handle type arguments<commit_after>#!/usr/bin/env python

# test-client - Simple client for testing parts of Telepathy connection managers
#
# Copyright (C) 2005 Collabora Limited
# Copyright (C) 2005 Nokia Corporation
#
# This library is free software; you can redistribute it and/or
# modify it under the terms of the GNU Lesser General Public
# License as published by the Free Software Foundation; either
# version 2.1 of the License, or (at your option) any later version.
#
# This library is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# Library General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public
# License along with this library; if not, write to the Free Software
# Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA

import dbus
import dbus.glib

assert(getattr(dbus, 'version', (0,0,0)) >= (0,51,0))

import getpass
import gobject
import signal
import sys
import os

from telepathy import *

import telepathy.client

class ContactListChannel(telepathy.client.Channel):
    def __init__(self, conn, object_path, handle_type, handle):
        telepathy.client.Channel.__init__(self, conn._service_name, object_path)
        self.get_valid_interfaces().add(CHANNEL_TYPE_CONTACT_LIST)
        self._conn = conn
        self._handle_type = handle_type
        self._handle = handle
        self._name = handle

    def got_interfaces(self):
        self._conn[CONN_INTERFACE].InspectHandle(self._handle_type, self._handle, reply_handler=self.inspect_handle_reply_cb, error_handler=self.error_cb)
        self[CHANNEL_INTERFACE_GROUP].GetMembers(reply_handler=self.get_members_reply_cb, error_handler=self.error_cb)
        self[CHANNEL_INTERFACE_GROUP].GetLocalPendingMembers(reply_handler=self.get_local_pending_members_reply_cb, error_handler=self.error_cb)
        self[CHANNEL_INTERFACE_GROUP].GetRemotePendingMembers(reply_handler=self.get_remote_pending_members_reply_cb, error_handler=self.error_cb)
        self[CHANNEL_INTERFACE_GROUP].connect_to_signal('MembersChanged', self.members_changed_signal_cb)

    def inspect_handle_reply_cb(self, name):
        print "CLC", self._name, "is", name
        self._name = name
        if name == 'subscribe':
            gobject.idle_add(self.subscribe_list_idle_cb)
        elif name == 'publish':
            gobject.idle_add(self.publish_list_idle_cb)

    def subscribe_list_idle_cb(self):
#        handle = self._conn[CONN_INTERFACE].RequestHandle(CONNECTION_HANDLE_TYPE_CONTACT, 'test2@localhost')
#        self[CHANNEL_INTERFACE_GROUP].AddMembers([handle])
        return False

    def publish_list_idle_cb(self):
        return False

    def members_changed_signal_cb(self, message, added, removed, local_p, remote_p):
        print "MembersChanged on CLC", self._name
        print "Message: ", message
        print "Added: ", added
        print "Removed: ", removed
        print "Local Pending: ", local_p
        print "Remote Pending: ", remote_p
#        if added and self._name == 'subscribe':
#            self[CHANNEL_INTERFACE_GROUP].RemoveMembers(added)
#        if removed and self._name == 'subscribe':
#            self[CHANNEL_INTERFACE_GROUP].AddMembers(added)

#        if added:
#            self[CHANNEL_INTERFACE_GROUP].RemoveMembers(added)
#        if local_p:
#            self[CHANNEL_INTERFACE_GROUP].AddMembers(local_p)

    def get_members_reply_cb(self, members):
        print "got members on CLC %s: %s" % (self._name, members)
#        if members and self._name == 'subscribe':
#            self[CHANNEL_INTERFACE_GROUP].RemoveMembers(members)
#        if members:
#            self[CHANNEL_INTERFACE_GROUP].RemoveMembers(members)

    def get_local_pending_members_reply_cb(self, members):
        print "got local pending members on CLC %s: %s" % (self._name, members)
#        if members:
#            self[CHANNEL_INTERFACE_GROUP].AddMembers(members)
#            print "sending addmembers"

    def get_remote_pending_members_reply_cb(self, members):
        print "got remote pending members on CLC %s: %s" % (self._name, members)

class StreamedMediaChannel(telepathy.client.Channel):
    def __init__(self, conn, object_path, handle_type, handle):
        telepathy.client.Channel.__init__(self, conn._service_name, object_path)
        self._conn = conn
        self.get_valid_interfaces().add(CHANNEL_TYPE_STREAMED_MEDIA)
        self[CHANNEL_TYPE_STREAMED_MEDIA].connect_to_signal('ReceivedMediaParameters', self.received_media_params_cb)
        self[CHANNEL_INTERFACE].connect_to_signal('Closed', self.closed_cb)

    def got_interfaces(self):
        print "SMC got interfaces"
        self[CHANNEL_INTERFACE_GROUP].GetMembers(reply_handler=self.get_members_reply_cb, error_handler=self.error_cb)
        self[CHANNEL_INTERFACE_GROUP].GetLocalPendingMembers(reply_handler=self.get_local_pending_members_reply_cb, error_handler=self.error_cb)
        self[CHANNEL_INTERFACE_GROUP].GetRemotePendingMembers(reply_handler=self.get_remote_pending_members_reply_cb, error_handler=self.error_cb)
        self[CHANNEL_INTERFACE_GROUP].connect_to_signal('MembersChanged', self.members_changed_signal_cb)

    def received_media_params_cb(self, id, local, remote):
        print "Received media params for %s.\n   local SDP = %s\n  remote SDP = %s" % (id, local, remote)

    def closed_cb(self):
        print "Channel closed"

    def members_changed_signal_cb(self, reason, added, removed, local_pending, remote_pending):
        print "Members Changed"
        print "  added ",added," removed", removed
        print "  local pending",local_pending," remote pending ", remote_pending

        self[CHANNEL_INTERFACE_DTMF].SendDTMF(1, 160)

    def error_cb(self, exception):
        print "Exception received from asynchronous method call:"
        print exception

    def get_members_reply_cb(self, members):
        print "Channel has members", members

    def get_local_pending_members_reply_cb(self, members):
        print "Channel has local pending members", members
        print self[CHANNEL_INTERFACE].GetInterfaces()
        if members:
           if "SIPDEBUG" in os.environ and os.environ["SIPDEBUG"] == "transfer":
               print "getting handle for test2"
               handle = self._conn[CONN_INTERFACE].RequestHandle(CONNECTION_HANDLE_TYPE_CONTACT, 'sip:test2@192.168.1.101')
               print "got handle %d for %s" % (handle, 'sip:test2@192.168.1.101')
               self[CHANNEL_INTERFACE_TRANSFER].Transfer(members[0],handle)
           else:
              print "adding member"
              self[CHANNEL_INTERFACE_GROUP].AddMembers(members)

    def get_remote_pending_members_reply_cb(self, members):
        print "Channel has remote pending members", members

class TextChannel(telepathy.client.Channel):
    def __init__(self, conn, object_path, handle_type, handle):
        # FIXME: this is racy and buggy and smells, and is deprecated in
        # favour of the stuff in sdp-client.py
        telepathy.client.Channel.__init__(self, conn._service_name, object_path)
        self.get_valid_interfaces().add(CHANNEL_TYPE_TEXT)
        self[CHANNEL_TYPE_TEXT].connect_to_signal('Received', self.received_callback)
        self.doack = True

        self._handled_pending_message = None
        pending_messages = self[CHANNEL_TYPE_TEXT].ListPendingMessages()
        print pending_messages
        for msg in pending_messages:
            (id, timestamp, sender, type, message) = msg
            print "Handling pending message", id
            self.received_callback(id, timestamp, sender, type, message)
            self._handled_pending_message = id

    def received_callback(self, id, timestamp, sender, type, message):
        if self._handled_pending_message != None:
            if id > self._handled_pending_message:
                print "Now handling messages directly"
                self._handled_pending_message = None
            else:
                print "Skipping already handled message", id
                return

        print "Received", id, timestamp, sender, type, message
        self[CHANNEL_TYPE_TEXT].Send(CHANNEL_TEXT_MESSAGE_TYPE_NORMAL, 'got message ' + str(id) + '(' + message + ')')
        if self.doack:
            print "Acknowledging...", id
            self[CHANNEL_TYPE_TEXT].AcknowledgePendingMessage(id)

class TestConnection(telepathy.client.Connection):
    def __init__(self, service_name, object_path, mainloop):
        telepathy.client.Connection.__init__(self, service_name, object_path)

        self._mainloop = mainloop
        self._service_name = service_name

        self._status = CONNECTION_STATUS_CONNECTING
        self._channels = {}

        print "calling GetStatus async"
        # handle race condition when connecting completes
        # before the signal handlers are registered
        self[CONN_INTERFACE].GetStatus(reply_handler=self.get_status_reply_cb, error_handler=self.error_cb)
        print "called GetStatus"

        print 'connecting to StatusChanged'
        self[CONN_INTERFACE].connect_to_signal('StatusChanged', self.status_changed_signal_cb)
        print 'connecting to NewChannel'
        self[CONN_INTERFACE].connect_to_signal('NewChannel', self.new_channel_signal_cb)

    def error_cb(self, exception):
        print "Exception received from asynchronous method call:"
        print exception

    def get_status_reply_cb(self, status):
        print "get_status_reply_cb", status
        self.status_changed_signal_cb(status, CONNECTION_STATUS_REASON_NONE_SPECIFIED)

    def new_channel_signal_cb(self, obj_path, chan_type, handle_type, handle, supress_handler):
        if obj_path in self._channels:
            return

        print 'NewChannel', obj_path, chan_type, handle_type, handle, supress_handler

        channel = None

        if chan_type == CHANNEL_TYPE_TEXT:
            channel = TextChannel(self, obj_path, handle_type, handle)
        elif chan_type == CHANNEL_TYPE_CONTACT_LIST:
            channel = ContactListChannel(self, obj_path, handle_type, handle)
        elif chan_type == CHANNEL_TYPE_STREAMED_MEDIA:
            channel = StreamedMediaChannel(self, obj_path, handle_type, handle)

        if channel != None:
            self._channels[obj_path] = channel
        else:
            print 'Unknown channel type', handle_type

    def connected_cb(self):
        if CONN_INTERFACE_STREAMED_MEDIA in self.get_valid_interfaces():
            self[CONN_INTERFACE_STREAMED_MEDIA].SetMediaParameters(
"v=0\r\nm=audio 0 RTP/AVP 3 0 8\r\na=rtpmap:3 GSM/8000\r\na=rtpmap  PCMU/8000\r\na=rtpmap:8 PCMA/8000\r\n")
        if "SIPDEBUG" in os.environ and os.environ["SIPDEBUG"][:5] == "call:":
            handle = self[CONN_INTERFACE].RequestHandle(CONNECTION_HANDLE_TYPE_CONTACT, os.environ["SIPDEBUG"][5:])
            print "got handle %d for %s" % (handle, os.environ["SIPDEBUG"][5:])
            print "calling RequestChannel"
            obj = self[CONN_INTERFACE].RequestChannel(CHANNEL_TYPE_STREAMED_MEDIA, CONNECTION_HANDLE_TYPE_CONTACT, handle, True)
            print "called RequestChannel, got ", obj
        elif "SIPDEBUG" in os.environ and os.environ["SIPDEBUG"][:8] == "forward:":
            handle = self[CONN_INTERFACE].RequestHandle(CONNECTION_HANDLE_TYPE_CONTACT, os.environ["SIPDEBUG"][8:])
            print "got handle %d for %s" % (handle, os.environ["SIPDEBUG"][8:])
            #print self[CONN_INTERFACE_PRESENCE].GetStatuses()
            self[CONN_INTERFACE_FORWARDING].SetForwardingHandle(handle)
        return False

    def presence_update_signal_cb(self, presence):
        print "got presence update:", presence

    def get_interfaces_reply_cb(self, interfaces):
        self.get_valid_interfaces().update(interfaces)
        #self[CONN_INTERFACE_PRESENCE].connect_to_signal('PresenceUpdate', self.presence_update_signal_cb)
        gobject.idle_add(self.connected_cb)

    def status_changed_signal_cb(self, status, reason):
        if self._status == status:
            return
        else:
            self._status = status

        print 'StatusChanged', status, reason

        if status == CONNECTION_STATUS_CONNECTED:
            print "connection connected"
            self[CONN_INTERFACE].GetInterfaces(reply_handler=self.get_interfaces_reply_cb, error_handler=self.error_cb)
        if status == CONNECTION_STATUS_DISCONNECTED:
            print "connection terminated"
            self._mainloop.quit()

if __name__ == '__main__':
    reg = telepathy.client.ManagerRegistry()
    reg.LoadManagers()

    protos=reg.GetProtos()

    if len(protos) == 0:
        print "Sorry, no connection managers found!"
        sys.exit(1)

    if len(sys.argv) > 2:
        protocol = sys.argv[2]
    else:
        protocol=''

    while (protocol not in protos):
        protocol = raw_input('Protocol (one of: %s) [%s]: ' % (' '.join(protos),protos[0]))
        if protocol == '':
            protocol = protos[0]

    if len(sys.argv) > 1:
        manager = sys.argv[1]
    else:
        manager = ''
    managers = reg.GetManagers(protocol)

    while (manager not in managers):
        manager = raw_input('Manager (one of: %s) [%s]: ' % (' '.join(managers),managers[0]))
        if manager == '':
            manager = managers[0]

    mgr_bus_name = reg.GetBusName(manager)
    mgr_object_path = reg.GetObjectPath(manager)

    cmdline_params = dict((p.split('=') for p in sys.argv[3:]))
    params={}

    for (name, (type, default)) in reg.GetParams(manager, protocol)[0].iteritems():
        if name in cmdline_params:
            params[name] = dbus.Variant(cmdline_params[name],type)
        elif name == 'password':
            params[name] = dbus.Variant(getpass.getpass(),type)
        else:
            params[name] = dbus.Variant(raw_input(name+': '),type)

    mgr = telepathy.client.ConnectionManager(mgr_bus_name, mgr_object_path)
    print "using params:",params
    bus_name, object_path = mgr[CONN_MGR_INTERFACE].Connect(protocol, params)
    print "done connect"
    mainloop = gobject.MainLoop()
    connection = TestConnection(bus_name, object_path, mainloop)
    print "created connection"

    def quit_cb():
        connection[CONN_INTERFACE].Disconnect()
        mainloop.quit()

    def sigterm_cb():
        gobject.idle_add(quit_cb)

    signal.signal(signal.SIGTERM, sigterm_cb)

    while mainloop.is_running():
        try:
            mainloop.run()
        except KeyboardInterrupt:
            quit_cb()
<|endoftext|>


<commit_before>
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from PIL import Image
from PIL import ImageFont
from PIL import ImageDraw
import copy
import os
from gym.spaces import Box
import numpy as np

import six
from gym.spaces import Discrete
from gym.utils.play import PlayPlot

from tensor2tensor.bin import t2t_trainer
from tensor2tensor.data_generators import generator_utils
from tensor2tensor.data_generators import gym_problems_specs
from tensor2tensor.layers import discretization
from tensor2tensor.rl import rl_trainer_lib
from tensor2tensor.rl.envs.batch_env_factory import batch_env_factory
from tensor2tensor.rl.envs.utils import get_policy
from tensor2tensor.utils import registry
from tensor2tensor.utils import trainer_lib
from tensor2tensor.rl.trainer_model_based import FLAGS, setup_directories, temporary_flags
from gym.utils import play
import tensorflow as tf

HP_SCOPES = ["loop", "model", "ppo"]

_font = None
FONT_SIZE = 20

def _get_font():
  global _font
  if _font is None:
    #weirdness due to various working dirs
    FONT_PATHS = ["tensor-2-tensor-with-mrunner/tensor-2-tensor-with-mrunner/deepsense_experiments/Xerox Serif Narrow.ttf",
                  "tensor-2-tensor-with-mrunner/deepsense_experiments/Xerox Serif Narrow.ttf",
                  "deepsense_experiments/Xerox Serif Narrow.ttf"]

    for path in FONT_PATHS:
      try:
        _font = ImageFont.truetype(path, FONT_SIZE)
        return _font
      except:
        pass

def _assert_image(img):
  if isinstance(img, np.ndarray):
    img = Image.fromarray(np.ndarray.astype(img, np.uint8))
  return img

def write_on_image(img, text="", position=(0, 0), color=(255, 255, 255)):
  img = _assert_image(img)
  if text=="":
    return img
  draw = ImageDraw.Draw(img)
  font = _get_font()
  draw.text(position, text, color, font=font)

  return img

def concatenate_images(*imgs, axis=1):
  imgs = [_assert_image(img) for img in imgs]
  imgs_np = [np.array(img) for img in imgs]
  concatenated_im_np = np.concatenate(imgs_np, axis=axis)

  return _assert_image(concatenated_im_np)

def train_agent(problem_name, agent_model_dir,
                event_dir, world_model_dir, epoch_data_dir, hparams, epoch=0,
                is_final_epoch=False):
  """Train the PPO agent in the simulated environment."""
  gym_problem = registry.problem(problem_name)
  ppo_hparams = trainer_lib.create_hparams(hparams.ppo_params)
  ppo_params_names = ["epochs_num", "epoch_length",
                      "learning_rate", "num_agents",
                      "optimization_epochs"]

  for param_name in ppo_params_names:
    ppo_param_name = "ppo_" + param_name
    if ppo_param_name in hparams:
      ppo_hparams.set_hparam(param_name, hparams.get(ppo_param_name))

  ppo_hparams.save_models_every_epochs = 10
  ppo_hparams.world_model_dir = world_model_dir
  ppo_hparams.add_hparam("force_beginning_resets", True)

  # Adding model hparams for model specific adjustments
  model_hparams = trainer_lib.create_hparams(hparams.generative_model_params)
  ppo_hparams.add_hparam("model_hparams", model_hparams)

  environment_spec = copy.copy(gym_problem.environment_spec)
  environment_spec.simulation_random_starts = hparams.simulation_random_starts
  environment_spec.simulation_flip_first_random_for_beginning = False
  environment_spec.intrinsic_reward_scale = hparams.intrinsic_reward_scale

  ppo_hparams.add_hparam("environment_spec", environment_spec)
  ppo_hparams.num_agents = 1

  with temporary_flags({
      "problem": problem_name,
      "model": hparams.generative_model,
      "hparams_set": hparams.generative_model_params,
      "output_dir": world_model_dir,
      "data_dir": epoch_data_dir,
  }):

    sess = tf.Session()
    env = DebugBatchEnv(ppo_hparams, sess)
    sess.run(tf.global_variables_initializer())
    env.initialize()

    env_model_loader = tf.train.Saver(
      tf.global_variables("next_frame*"))
    trainer_lib.restore_checkpoint(world_model_dir, env_model_loader, sess,
      must_restore=True)

    model_saver = tf.train.Saver(
      tf.global_variables(".*network_parameters.*"))
    trainer_lib.restore_checkpoint(agent_model_dir, model_saver, sess)

    key_mapping = {(): 100, (ord('q'),):1, (ord('a'),):2,
                   (ord('r'),):101,
                   (ord('p'),):102}

    play.play(env, zoom=2, fps=10, keys_to_action=key_mapping)

from gym.core import Env

class DebugBatchEnv(Env):

  def __init__(self, hparams, sess = None):
    if sess == None:
      self.sess = tf.Session()
    else:
      self.sess = sess

    self.action_space = Discrete(6)
    self.observation_space = Box(low=0, high=255, shape=(210, 160+250, 3), dtype=np.uint8)
    self._tmp = 1
    self.res = None
    self._prepare_networks(hparams, self.sess)

  def _prepare_networks(self, hparams, sess):
    self.action = tf.placeholder(shape=(1,), dtype=tf.int32)
    batch_env = batch_env_factory(hparams)
    self.reward, self.done = batch_env.simulate(self.action)
    self.observation = batch_env.observ
    self.reset_op = batch_env.reset(tf.constant([0], dtype=tf.int32))

    environment_wrappers = hparams.environment_spec.wrappers
    wrappers = copy.copy(environment_wrappers) if environment_wrappers else []

    to_initialize = [batch_env]
    for w in wrappers:
      batch_env = w[0](batch_env, **w[1])
      to_initialize.append(batch_env)

    def initialization_lambda():
      for batch_env in to_initialize:
        batch_env.initialize(sess)

    self.initialize = initialization_lambda

    obs_copy = batch_env.observ + 0

    actor_critic = get_policy(tf.expand_dims(obs_copy, 0), hparams)
    self.policy_probs = actor_critic.policy.probs[0, 0, :]
    self.value = actor_critic.value[0, :]

  def render(self, mode='human'):
    raise NotImplemented()

  def _fake_reset(self):
    self._tmp = 0
    _observ = np.ones(shape=(210, 160, 3), dtype=np.uint8) * 10 * self._tmp
    _observ[0, 0, 0] = 0
    _observ[0, 0, 1] = 255
    self.res = (_observ, 0, False, [0.1, 0.5, 0.5], 1.1)
    observ = self._augment_observation()
    return observ

  def reset(self):
    # observ = self.sess.run(self.reset_op)
    self._tmp = 0
    _observ = np.ones(shape=(210, 160, 3), dtype=np.uint8) * 10 * self._tmp
    _observ[0, 0, 0] = 0
    _observ[0, 0, 1] = 255
    self.res = (_observ, 0, False, [0.1, 0.5, 0.5], 1.1)
    observ = self._augment_observation()
    return observ

  def _step_fake(self, action):

    observ = np.ones(shape=(210, 160, 3), dtype=np.uint8)*10*self._tmp
    observ[0, 0, 0] = 0
    observ[0, 0, 1] = 255

    self._tmp += 1
    if self._tmp>20:
      self._tmp = 0

    rew = 1
    done = False
    probs = np.ones(shape=(6,), dtype=np.float32)/6
    vf = 0.0

    return observ, rew, done, probs, vf

  def _step_env(self, action):
    observ, rew, done, probs, vf = self.sess.\
      run([self.observation, self.reward, self.done, self.policy_probs, self.value],
          feed_dict={self.action: [action]})

    return observ[0, ...], rew[0, ...], done[0, ...], probs, vf

  def _augment_observation(self):
    _observ, rew, done, probs, vf = self.res
    info_pane = np.zeros(shape=(210, 250, 3), dtype=np.uint8)
    probs_str = ""
    for p in probs:
      probs_str += "%.2f" % p +", "

    action = np.argmax(probs)

    info_str = " Policy:{}\n Action:{}\n Value function:{}\n Reward:{}".format(probs_str, action,
                                                                           vf, rew)
    info_pane = write_on_image(info_pane, info_str)

    augmented_observ = concatenate_images(_observ, info_pane)
    augmented_observ = np.array(augmented_observ)
    return augmented_observ

  def step(self, action):
    #Special codes
    if action==100:
      #Skip action
      _, rew, done, _, _ = self.res
      observ = self._augment_observation()
      return observ, rew, done, {}

    if action == 101:
      #reset
      self.reset()
      _, rew, done, _, _ = self.res
      observ = self._augment_observation()
      return observ, rew, done, {}

    if action == 102:
      #play
      raise NotImplemented()

    #standard codes
    _observ, rew, done, probs, vf = self._step_env(action)
    self.res = (_observ, rew, done, probs, vf)

    observ = self._augment_observation()
    return observ, rew, done, {"probs": probs, "vf": vf}

def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
  """Run the main training loop."""

  # Directories
  subdirectories = ["data", "tmp", "world_model", "ppo"]
  using_autoencoder = hparams.autoencoder_train_steps > 0
  if using_autoencoder:
    subdirectories.append("autoencoder")
  directories = setup_directories(output_dir, subdirectories)

  if hparams.game in gym_problems_specs.ATARI_GAMES:
    game_with_mode = hparams.game + "_deterministic-v4"
  else:
    game_with_mode = hparams.game
  # Problems
  if using_autoencoder:
    simulated_problem_name = (
        "gym_simulated_discrete_problem_with_agent_on_%s_autoencoded"
        % game_with_mode)
  else:
    simulated_problem_name = ("gym_simulated_discrete_problem_with_agent_on_%s"
                              % game_with_mode)

  epoch_data_dirs = []
  data_dir = os.path.join(directories["data"], "random")
  epoch_data_dirs.append(data_dir)

  epoch = hparams.epochs-1

  # Combine all previously collected environment data
  epoch_data_dir = os.path.join(directories["data"], str(epoch))

  ppo_event_dir = os.path.join(directories["world_model"],
                               "ppo_summaries", str(epoch))
  ppo_model_dir = directories["ppo"]
  if not hparams.ppo_continue_training:
    ppo_model_dir = ppo_event_dir
  train_agent(simulated_problem_name, ppo_model_dir,
              ppo_event_dir, directories["world_model"], epoch_data_dir,
              hparams, epoch=epoch, is_final_epoch=False)

def create_loop_hparams():
  hparams = registry.hparams(FLAGS.loop_hparams_set)
  hparams.parse(FLAGS.loop_hparams)
  return hparams

def main(_):
  hp = create_loop_hparams()
  output_dir = FLAGS.output_dir
  training_loop(hp, output_dir)

if __name__ == "__main__":
  tf.logging.set_verbosity(tf.logging.INFO)
  tf.app.run()
<commit_msg>fixing resets<commit_after>
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from PIL import Image
from PIL import ImageFont
from PIL import ImageDraw
import copy
import os
from gym.spaces import Box
import numpy as np

import six
from gym.spaces import Discrete
from gym.utils.play import PlayPlot

from tensor2tensor.bin import t2t_trainer
from tensor2tensor.data_generators import generator_utils
from tensor2tensor.data_generators import gym_problems_specs
from tensor2tensor.layers import discretization
from tensor2tensor.rl import rl_trainer_lib
from tensor2tensor.rl.envs.batch_env_factory import batch_env_factory
from tensor2tensor.rl.envs.utils import get_policy
from tensor2tensor.utils import registry
from tensor2tensor.utils import trainer_lib
from tensor2tensor.rl.trainer_model_based import FLAGS, setup_directories, temporary_flags
from gym.utils import play
import tensorflow as tf

HP_SCOPES = ["loop", "model", "ppo"]

_font = None
FONT_SIZE = 20

def _get_font():
  global _font
  if _font is None:
    #weirdness due to various working dirs
    FONT_PATHS = ["tensor-2-tensor-with-mrunner/tensor-2-tensor-with-mrunner/deepsense_experiments/Xerox Serif Narrow.ttf",
                  "tensor-2-tensor-with-mrunner/deepsense_experiments/Xerox Serif Narrow.ttf",
                  "deepsense_experiments/Xerox Serif Narrow.ttf"]

    for path in FONT_PATHS:
      try:
        _font = ImageFont.truetype(path, FONT_SIZE)
        return _font
      except:
        pass

def _assert_image(img):
  if isinstance(img, np.ndarray):
    img = Image.fromarray(np.ndarray.astype(img, np.uint8))
  return img

def write_on_image(img, text="", position=(0, 0), color=(255, 255, 255)):
  img = _assert_image(img)
  if text=="":
    return img
  draw = ImageDraw.Draw(img)
  font = _get_font()
  draw.text(position, text, color, font=font)

  return img

def concatenate_images(*imgs, axis=1):
  imgs = [_assert_image(img) for img in imgs]
  imgs_np = [np.array(img) for img in imgs]
  concatenated_im_np = np.concatenate(imgs_np, axis=axis)

  return _assert_image(concatenated_im_np)

def train_agent(problem_name, agent_model_dir,
                event_dir, world_model_dir, epoch_data_dir, hparams, epoch=0,
                is_final_epoch=False):
  """Train the PPO agent in the simulated environment."""
  gym_problem = registry.problem(problem_name)
  ppo_hparams = trainer_lib.create_hparams(hparams.ppo_params)
  ppo_params_names = ["epochs_num", "epoch_length",
                      "learning_rate", "num_agents",
                      "optimization_epochs"]

  for param_name in ppo_params_names:
    ppo_param_name = "ppo_" + param_name
    if ppo_param_name in hparams:
      ppo_hparams.set_hparam(param_name, hparams.get(ppo_param_name))

  ppo_hparams.save_models_every_epochs = 10
  ppo_hparams.world_model_dir = world_model_dir
  ppo_hparams.add_hparam("force_beginning_resets", True)

  # Adding model hparams for model specific adjustments
  model_hparams = trainer_lib.create_hparams(hparams.generative_model_params)
  ppo_hparams.add_hparam("model_hparams", model_hparams)

  environment_spec = copy.copy(gym_problem.environment_spec)
  environment_spec.simulation_random_starts = hparams.simulation_random_starts
  environment_spec.simulation_flip_first_random_for_beginning = False
  environment_spec.intrinsic_reward_scale = hparams.intrinsic_reward_scale

  ppo_hparams.add_hparam("environment_spec", environment_spec)
  ppo_hparams.num_agents = 1

  with temporary_flags({
      "problem": problem_name,
      "model": hparams.generative_model,
      "hparams_set": hparams.generative_model_params,
      "output_dir": world_model_dir,
      "data_dir": epoch_data_dir,
  }):

    sess = tf.Session()
    env = DebugBatchEnv(ppo_hparams, sess)
    sess.run(tf.global_variables_initializer())
    env.initialize()

    env_model_loader = tf.train.Saver(
      tf.global_variables("next_frame*"))
    trainer_lib.restore_checkpoint(world_model_dir, env_model_loader, sess,
      must_restore=True)

    model_saver = tf.train.Saver(
      tf.global_variables(".*network_parameters.*"))
    trainer_lib.restore_checkpoint(agent_model_dir, model_saver, sess)

    key_mapping = {(): 100, (ord('q'),):1, (ord('a'),):2,
                   (ord('r'),):101,
                   (ord('p'),):102}

    play.play(env, zoom=2, fps=10, keys_to_action=key_mapping)

from gym.core import Env

class DebugBatchEnv(Env):

  def __init__(self, hparams, sess = None):
    if sess == None:
      self.sess = tf.Session()
    else:
      self.sess = sess

    self.action_space = Discrete(6)
    self.observation_space = Box(low=0, high=255, shape=(210, 160+250, 3), dtype=np.uint8)
    self._tmp = 1
    self.res = None
    self._prepare_networks(hparams, self.sess)

  def _prepare_networks(self, hparams, sess):
    self.action = tf.placeholder(shape=(1,), dtype=tf.int32)
    batch_env = batch_env_factory(hparams)
    self.reward, self.done = batch_env.simulate(self.action)
    self.observation = batch_env.observ
    self.reset_op = batch_env.reset(tf.constant([0], dtype=tf.int32))

    environment_wrappers = hparams.environment_spec.wrappers
    wrappers = copy.copy(environment_wrappers) if environment_wrappers else []

    to_initialize = [batch_env]
    for w in wrappers:
      batch_env = w[0](batch_env, **w[1])
      to_initialize.append(batch_env)

    def initialization_lambda():
      for batch_env in to_initialize:
        batch_env.initialize(sess)

    self.initialize = initialization_lambda

    obs_copy = batch_env.observ + 0

    actor_critic = get_policy(tf.expand_dims(obs_copy, 0), hparams)
    self.policy_probs = actor_critic.policy.probs[0, 0, :]
    self.value = actor_critic.value[0, :]

  def render(self, mode='human'):
    raise NotImplemented()

  def _fake_reset(self):
    self._tmp = 0
    _observ = np.ones(shape=(210, 160, 3), dtype=np.uint8) * 10 * self._tmp
    _observ[0, 0, 0] = 0
    _observ[0, 0, 1] = 255
    self.res = (_observ, 0, False, [0.1, 0.5, 0.5], 1.1)

  def _reset_env(self):
    _observ = self.sess.run(self.reset_op)[0, ...]
    _observ[0, 0, 0] = 0
    _observ[0, 0, 1] = 255
    #TODO:(put correct numbers)
    self.res = (_observ, 0, False, [0.1, 0.5, 0.5], 1.1)

  def reset(self):
    self._reset_env()
    observ = self._augment_observation()
    return observ

  def _step_fake(self, action):
    observ = np.ones(shape=(210, 160, 3), dtype=np.uint8)*10*self._tmp
    observ[0, 0, 0] = 0
    observ[0, 0, 1] = 255

    self._tmp += 1
    if self._tmp>20:
      self._tmp = 0

    rew = 1
    done = False
    probs = np.ones(shape=(6,), dtype=np.float32)/6
    vf = 0.0

    return observ, rew, done, probs, vf

  def _step_env(self, action):
    observ, rew, done, probs, vf = self.sess.\
      run([self.observation, self.reward, self.done, self.policy_probs, self.value],
          feed_dict={self.action: [action]})

    return observ[0, ...], rew[0, ...], done[0, ...], probs, vf

  def _augment_observation(self):
    _observ, rew, done, probs, vf = self.res
    info_pane = np.zeros(shape=(210, 250, 3), dtype=np.uint8)
    probs_str = ""
    for p in probs:
      probs_str += "%.2f" % p +", "

    action = np.argmax(probs)

    info_str = " Policy:{}\n Action:{}\n Value function:{}\n Reward:{}".format(probs_str, action,
                                                                           vf, rew)
    info_pane = write_on_image(info_pane, info_str)

    augmented_observ = concatenate_images(_observ, info_pane)
    augmented_observ = np.array(augmented_observ)
    return augmented_observ

  def step(self, action):
    #Special codes
    if action==100:
      #Skip action
      _, rew, done, _, _ = self.res
      observ = self._augment_observation()
      return observ, rew, done, {}

    if action == 101:
      #reset
      self.reset()
      _, rew, done, _, _ = self.res
      observ = self._augment_observation()
      return observ, rew, done, {}

    if action == 102:
      #play
      raise NotImplemented()

    #standard codes
    _observ, rew, done, probs, vf = self._step_env(action)
    self.res = (_observ, rew, done, probs, vf)

    observ = self._augment_observation()
    return observ, rew, done, {"probs": probs, "vf": vf}

def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
  """Run the main training loop."""

  # Directories
  subdirectories = ["data", "tmp", "world_model", "ppo"]
  using_autoencoder = hparams.autoencoder_train_steps > 0
  if using_autoencoder:
    subdirectories.append("autoencoder")
  directories = setup_directories(output_dir, subdirectories)

  if hparams.game in gym_problems_specs.ATARI_GAMES:
    game_with_mode = hparams.game + "_deterministic-v4"
  else:
    game_with_mode = hparams.game
  # Problems
  if using_autoencoder:
    simulated_problem_name = (
        "gym_simulated_discrete_problem_with_agent_on_%s_autoencoded"
        % game_with_mode)
  else:
    simulated_problem_name = ("gym_simulated_discrete_problem_with_agent_on_%s"
                              % game_with_mode)

  epoch_data_dirs = []
  data_dir = os.path.join(directories["data"], "random")
  epoch_data_dirs.append(data_dir)

  epoch = hparams.epochs-1

  # Combine all previously collected environment data
  epoch_data_dir = os.path.join(directories["data"], str(epoch))

  ppo_event_dir = os.path.join(directories["world_model"],
                               "ppo_summaries", str(epoch))
  ppo_model_dir = directories["ppo"]
  if not hparams.ppo_continue_training:
    ppo_model_dir = ppo_event_dir
  train_agent(simulated_problem_name, ppo_model_dir,
              ppo_event_dir, directories["world_model"], epoch_data_dir,
              hparams, epoch=epoch, is_final_epoch=False)

def create_loop_hparams():
  hparams = registry.hparams(FLAGS.loop_hparams_set)
  hparams.parse(FLAGS.loop_hparams)
  return hparams

def main(_):
  hp = create_loop_hparams()
  output_dir = FLAGS.output_dir
  training_loop(hp, output_dir)

if __name__ == "__main__":
  tf.logging.set_verbosity(tf.logging.INFO)
  tf.app.run()
<|endoftext|>