<commit_before>from cendr import app
from cendr.models import *
from flask import request, abort
import string
import random
from googleapiclient import discovery
from oauth2client.client import GoogleCredentials
credentials = GoogleCredentials.get_application_default()
compute = discovery.build('compute', 'v1', credentials=credentials)
import time

def id_generator(size = 4, chars = string.ascii_lowercase):
    """
        Generate random 4 letter id.
    """
    return ''.join(random.choice(chars) for _ in range(size))

def create_mapping_instance(gce_name, params):
    # Get latest CeNDR mapping image
    image_response = compute.images().getFromFamily(project='andersen-lab',
                                                     family = 'cendr').execute()
    source_disk_image = image_response['selfLink']

    config = {
        'name': 'cendr-' + gce_name + "-" + id_generator(4),
        'zone': 'projects/andersen-lab/zones/us-central1-a',
        'machineType': "zones/us-central1-a/machineTypes/custom-1-6656",
        'tags': {
            "items": ['cendr-mapping']
        },
        # Specify the boot disk and the image to use as a source.
        'disks': [
            {
                'boot': True,
                'autoDelete': True,
                'initializeParams': {
                    'sourceImage': source_disk_image,
                    'diskType': "projects/andersen-lab/zones/us-central1-a/diskTypes/pd-standard",
                    'diskSizeGb': "10"
                }
            }
        ],

        # Specify a network interface with NAT to access the public
        # internet.
        'networkInterfaces': [{
            'network': 'projects/andersen-lab/global/networks/default',
            'accessConfigs': [
                {
                  "name": "External NAT",
                  "type": "ONE_TO_ONE_NAT"
                }
              ]
        }],

        # Allow the instance to access cloud storage and logging.
        'serviceAccounts': [{
            'email': 'andersen-lab@appspot.gserviceaccount.com',
            'scopes': [
                'https://www.googleapis.com/auth/cloud-platform'
            ]
        }],

        # Metadata is readable from the instance and allows you to
        # pass configuration from deployment scripts to instances.
        'metadata': {
            'items': params
        }
    }

    instance = compute.instances().insert(
        project='andersen-lab',
        zone='us-central1-a',
        body=config).execute()
    return instance

def instance_count():
    instances = compute.instances().list(project='andersen-lab', zone = 'us-central1-a').execute()['items']
    mapping_instances = [x['name'] for x in instances if x['name'].startswith("cendr")]
    return len(mapping_instances)

@app.route("/cronmapping")
def launch_mapping(verify_request = True):
    if instance_count() < 5:
        job_submissions = list(trait.select(trait.id.alias('trait_id'), trait, report)
                  .join(report)
                  .filter(trait.status == "queue")
                  .order_by(trait.submission_date)
                  .limit(5)
                  .dicts()
                  .execute())
        for job in job_submissions:
            job['submission_date'] = job['submission_date'].isoformat()
            if job['submission_complete']:
                job['submission_complete'] = job['submission_complete'].isoformat()
            params = [{'key': k, 'value': v} for k,v in job.items()]
            gce_name = job['report_slug'] + "-" + job['trait_slug']
            startup_script = open('cendr/task/startup-script.sh','r').read()
            run_pipeline = open('cendr/task/run_pipeline.py', 'r').read()
            pipeline_script = open('cendr/task/pipeline.R', 'r').read()
            models = open('cendr/models.py', 'r').read()
            params += [{'key':'startup-script','value': startup_script}]
            params += [{'key':'run_pipeline','value': run_pipeline}]
            params += [{'key':'pipeline','value':pipeline_script}]
            params += [{'key':'models','value':models}]
            create_mapping_instance(gce_name, params)
            # Update queue status
            t = trait.get(trait.id==job['trait_id'])
            t.status = 'Initializing'
            t.save()
    return "success", 200

<commit_msg>add verification back<commit_after>from cendr import app
from cendr.models import *
from flask import request, abort
import string
import random
from googleapiclient import discovery
from oauth2client.client import GoogleCredentials
credentials = GoogleCredentials.get_application_default()
compute = discovery.build('compute', 'v1', credentials=credentials)
import time

def id_generator(size = 4, chars = string.ascii_lowercase):
    """
        Generate random 4 letter id.
    """
    return ''.join(random.choice(chars) for _ in range(size))

def create_mapping_instance(gce_name, params):
    # Get latest CeNDR mapping image
    image_response = compute.images().getFromFamily(project='andersen-lab',
                                                     family = 'cendr').execute()
    source_disk_image = image_response['selfLink']

    config = {
        'name': 'cendr-' + gce_name + "-" + id_generator(4),
        'zone': 'projects/andersen-lab/zones/us-central1-a',
        'machineType': "zones/us-central1-a/machineTypes/custom-1-6656",
        'tags': {
            "items": ['cendr-mapping']
        },
        # Specify the boot disk and the image to use as a source.
        'disks': [
            {
                'boot': True,
                'autoDelete': True,
                'initializeParams': {
                    'sourceImage': source_disk_image,
                    'diskType': "projects/andersen-lab/zones/us-central1-a/diskTypes/pd-standard",
                    'diskSizeGb': "10"
                }
            }
        ],

        # Specify a network interface with NAT to access the public
        # internet.
        'networkInterfaces': [{
            'network': 'projects/andersen-lab/global/networks/default',
            'accessConfigs': [
                {
                  "name": "External NAT",
                  "type": "ONE_TO_ONE_NAT"
                }
              ]
        }],

        # Allow the instance to access cloud storage and logging.
        'serviceAccounts': [{
            'email': 'andersen-lab@appspot.gserviceaccount.com',
            'scopes': [
                'https://www.googleapis.com/auth/cloud-platform'
            ]
        }],

        # Metadata is readable from the instance and allows you to
        # pass configuration from deployment scripts to instances.
        'metadata': {
            'items': params
        }
    }

    instance = compute.instances().insert(
        project='andersen-lab',
        zone='us-central1-a',
        body=config).execute()
    return instance

def instance_count():
    instances = compute.instances().list(project='andersen-lab', zone = 'us-central1-a').execute()['items']
    mapping_instances = [x['name'] for x in instances if x['name'].startswith("cendr")]
    return len(mapping_instances)

@app.route("/cronmapping")
def launch_mapping(verify_request = True):
    # Verify cron submission
    if verify_request:
        gae_header = request.headers.get("X-Appengine-Cron", None)
        if not gae_header:
            abort(404)

    if instance_count() < 5:
        job_submissions = list(trait.select(trait.id.alias('trait_id'), trait, report)
                  .join(report)
                  .filter(trait.status == "queue")
                  .order_by(trait.submission_date)
                  .limit(5)
                  .dicts()
                  .execute())
        for job in job_submissions:
            job['submission_date'] = job['submission_date'].isoformat()
            if job['submission_complete']:
                job['submission_complete'] = job['submission_complete'].isoformat()
            params = [{'key': k, 'value': v} for k,v in job.items()]
            gce_name = job['report_slug'] + "-" + job['trait_slug']
            startup_script = open('cendr/task/startup-script.sh','r').read()
            run_pipeline = open('cendr/task/run_pipeline.py', 'r').read()
            pipeline_script = open('cendr/task/pipeline.R', 'r').read()
            models = open('cendr/models.py', 'r').read()
            params += [{'key':'startup-script','value': startup_script}]
            params += [{'key':'run_pipeline','value': run_pipeline}]
            params += [{'key':'pipeline','value':pipeline_script}]
            params += [{'key':'models','value':models}]
            create_mapping_instance(gce_name, params)
            # Update queue status
            t = trait.get(trait.id==job['trait_id'])
            t.status = 'Initializing'
            t.save()
    return "success", 200

<|endoftext|>


<commit_before>#
#  The OpenDiamond Platform for Interactive Search
#
#  Copyright (c) 2012 Carnegie Mellon University
#  All rights reserved.
#
#  This software is distributed under the terms of the Eclipse Public
#  License, Version 1.0 which can be found in the file named LICENSE.
#  ANY USE, REPRODUCTION OR DISTRIBUTION OF THIS SOFTWARE CONSTITUTES
#  RECIPIENT'S ACCEPTANCE OF THIS AGREEMENT
#

'''JSON Blaster request handlers.'''

from datetime import timedelta
import logging
import simplejson as json
import os
from sockjs.tornado import SockJSConnection
from urlparse import urlparse
from tornado.curl_httpclient import CurlAsyncHTTPClient as AsyncHTTPClient
from tornado import gen
from tornado.ioloop import IOLoop
from tornado.options import define, options
from tornado.web import asynchronous, RequestHandler, HTTPError
import validictory

import opendiamond
from opendiamond.attributes import (StringAttributeCodec,
        IntegerAttributeCodec, DoubleAttributeCodec, PatchesAttributeCodec)
from opendiamond.blaster.cache import SearchCacheLoadError
from opendiamond.blaster.search import (Blob, EmptyBlob, DiamondSearch,
        FilterSpec)
from opendiamond.helpers import sha256
from opendiamond.scope import ScopeCookie, ScopeError

CACHE_URN_SCHEME = 'blob'
STATS_INTERVAL = timedelta(milliseconds=1000)

# HTTP method handlers have specific argument lists rather than
# (self, *args, **kwargs) as in the superclass.
# pylint: disable=W0221

define('enable_testui', default=True,
        help='Enable the example user interface')
define('http_proxy', type=str, default=None,
        metavar='HOST:PORT', help='Use a proxy for HTTP client requests')

_log = logging.getLogger(__name__)

def _load_schema(name):
    with open(os.path.join(os.path.dirname(__file__), name)) as fh:
        return json.load(fh)
EVENT_SCHEMA = _load_schema('schema-event.json')
SEARCH_SCHEMA = _load_schema('schema-search.json')

def _validate_json(schema, obj):
    # required_by_default=False and blank_by_default=True for
    # JSON Schema draft 3 semantics
    validictory.validate(obj, schema, required_by_default=False,
            blank_by_default=True)

def _make_object_json(obj, url_factory):
    '''Convert an object attribute dict into a dict suitable for JSON
    encoding.  url_factory(k) is a function to convert an attribute name
    into a URL.'''
    result = {}
    for k, v in obj.iteritems():
        data = None
        # Inline known attribute types that can be represented in JSON
        if k.endswith('.int'):
            data = IntegerAttributeCodec().decode(v)
        elif k.endswith('.double'):
            data = DoubleAttributeCodec().decode(v)
        elif k.endswith('.patches'):
            distance, patches = PatchesAttributeCodec().decode(v)
            data = {
                'distance': distance,
                'patches': [{
                    'x0': tl[0],
                    'y0': tl[1],
                    'x1': br[0],
                    'y1': br[1]
                } for tl, br in patches],
            }
        else:
            # Treat remaining attributes as strings if: they don't have
            # a suffix representing a known binary type, they can be decoded
            # by the string codec (i.e., their last byte is 0), they are
            # valid UTF-8, and they are not the '' (object data) attribute.
            try:
                _base, suffix = k.rsplit('.', 1)
            except ValueError:
                suffix = None
            try:
                if k != '' and suffix not in ('jpeg', 'rgbimage', 'binary'):
                    data = StringAttributeCodec().decode(v).decode('UTF-8')
            except ValueError:
                pass

        if data is not None:
            result[k] = {
                'data': data,
            }
        else:
            result[k] = {
                'url': url_factory(k),
            }
    return result

class _BlasterRequestHandler(RequestHandler):
    @property
    def blob_cache(self):
        return self.application.blob_cache

    @property
    def search_cache(self):
        return self.application.search_cache

    def write_error(self, code, **kwargs):
        exc_type, exc_value, _exc_tb = kwargs.get('exc_info', [None] * 3)
        if exc_type is not None and issubclass(exc_type, HTTPError):
            self.set_header('Content-Type', 'text/plain')
            if exc_value.log_message:
                self.write(exc_value.log_message + '\n')
        else:
            RequestHandler.write_error(self, code, **kwargs)

class _BlasterBlob(Blob):
    '''Instances with the same URI compare equal and hash to the same value.'''

    def __init__(self, uri, expected_sha256=None):
        Blob.__init__(self)
        self.uri = uri
        self._expected_sha256 = expected_sha256
        self._data = None

    def __str__(self):
        if self._data is None:
            raise RuntimeError('Attempting to read an unfetched blob')
        return self._data

    def __repr__(self):
        return '<_BlasterBlob %s>' % (self.uri)

    def __hash__(self):
        return hash(self.uri)

    def __eq__(self, other):
        return isinstance(other, type(self)) and self.uri == other.uri

    # pylint doesn't understand named tuples
    # pylint: disable=E1101
    @gen.engine
    def fetch(self, blob_cache, callback=None):
        if self._data is None:
            # Fetch data
            parts = urlparse(self.uri)
            if parts.scheme == 'blob':
                try:
                    data = blob_cache[parts.path]
                except KeyError:
                    raise HTTPError(400, 'Blob missing from blob cache')
            elif parts.scheme == 'http' or parts.scheme == 'https':
                client = AsyncHTTPClient()
                if options.http_proxy is not None:
                    proxy_host, proxy_port = options.http_proxy.split(':', 1)
                    proxy_port = int(proxy_port)
                else:
                    proxy_host = proxy_port = None
                response = yield gen.Task(client.fetch, self.uri,
                        user_agent='JSONBlaster/%s' % opendiamond.__version__,
                        proxy_host=proxy_host, proxy_port=proxy_port,
                        validate_cert=False)
                if response.error:
                    raise HTTPError(400, 'Error fetching <%s>: %s' % (
                            self.uri, str(response.error)))
                data = response.body
            else:
                raise HTTPError(400, 'Unacceptable blob URI scheme')

            # Check hash if requested
            if self._expected_sha256 is not None:
                if sha256(data).hexdigest() != self._expected_sha256:
                    raise HTTPError(400, 'SHA-256 mismatch on %s' % self.uri)

            # Commit
            self._data = data

        if callback is not None:
            callback()
    # pylint: enable=E1101

class _SearchSpec(object):
    def __init__(self, data):
        # Load JSON
        try:
            config = json.loads(data)
            _validate_json(SEARCH_SCHEMA, config)
        except ValueError, e:
            raise HTTPError(400, str(e))

        # Build cookies
        # Assume each "cookie" may actually be a megacookie
        try:
            self.cookies = [ScopeCookie.parse(c) for mc in config['cookies']
                    for c in ScopeCookie.split(mc)]
        except ScopeError, e:
            raise HTTPError(400, 'Invalid scope cookie: %s' % e)
        if not self.cookies:
            # No cookies could be parsed out of the client's cookie list
            raise HTTPError(400, 'No scope cookies found')

        # Build filters
        blobs = {}  # blob -> itself  (for deduplication)
        def make_blob(obj):
            if obj is not None:
                blob = _BlasterBlob(obj['uri'], obj.get('sha256'))
                # Intern the blob
                return blobs.setdefault(blob, blob)
            else:
                return EmptyBlob()
        self.filters = [FilterSpec(
                    name=f['name'],
                    code=make_blob(f['code']),
                    arguments=f.get('arguments', []),
                    blob_argument=make_blob(f.get('blob')),
                    dependencies=f.get('dependencies', []),
                    min_score=f.get('min_score', float('-inf')),
                    max_score=f.get('max_score', float('inf'))
                ) for f in config['filters']]
        self.blobs = blobs.values()

    @gen.engine
    def fetch_blobs(self, blob_cache, callback=None):
        yield [gen.Task(blob.fetch, blob_cache) for blob in self.blobs]
        if callback is not None:
            callback()

    def make_search(self, **kwargs):
        return DiamondSearch(self.cookies, self.filters, **kwargs)

class SearchHandler(_BlasterRequestHandler):
    def get(self):
        if options.enable_testui:
            self.render('testui/search.html')
        else:
            raise HTTPError(405, 'Method not allowed')

    @asynchronous
    @gen.engine
    def post(self):
        # Build search spec
        if self.request.headers['Content-Type'] != 'application/json':
            raise HTTPError(415, 'Content type must be application/json')
        spec = _SearchSpec(self.request.body)
        yield gen.Task(spec.fetch_blobs, self.blob_cache)

        # Store it
        search_key = self.search_cache.put_search(spec)

        # Return result
        self.set_status(204)
        self.set_header('Location', '/search')
        self.set_header('X-Search-Key', search_key)
        self.finish()

class PostBlobHandler(_BlasterRequestHandler):
    def post(self):
        sig = self.blob_cache.add(self.request.body)
        self.set_header('Location', '%s:%s' % (CACHE_URN_SCHEME, sig))
        self.set_status(204)

class AttributeHandler(_BlasterRequestHandler):
    def get(self, search_key, object_key, attr_name):
        try:
            data = self.search_cache.get_object_attribute(search_key,
                    object_key, attr_name)
        except KeyError:
            raise HTTPError(404, 'Not found')

        if attr_name.endswith('.jpeg'):
            self.set_header('Content-Type', 'image/jpeg')
        else:
            self.set_header('Content-Type', 'application/octet-stream')

        self.write(data)

class ResultsHandler(_BlasterRequestHandler):
    def get(self):
        if options.enable_testui:
            self.render('testui/results.html')
        else:
            raise HTTPError(403, 'Forbidden')

class _StructuredSocketConnection(SockJSConnection):
    @classmethod
    def event(cls, func):
        '''Decorator specifying that this function is an event handler.'''
        func.event_handler = True
        return func

    # pylint is confused by msg.get()
    # pylint: disable=E1103
    def on_message(self, data):
        try:
            msg = json.loads(data)
            _validate_json(EVENT_SCHEMA, msg)
        except ValueError, e:
            raise HTTPError(400, str(e))
        event = msg['event']
        try:
            handler = getattr(self, event)
        except AttributeError:
            _log.warning('Unknown event type %s' % event)
            return
        if not getattr(handler, 'event_handler', False):
            _log.warning('Event %s has invalid handler' % event)
            return
        handler(**msg.get('data', {}))
    # pylint: enable=E1103

    def emit(self, event, **args):
        self.send(json.dumps({
            'event': event,
            'data': args,
        }))

class SearchConnection(_StructuredSocketConnection):
    def __init__(self, *args, **kwargs):
        _StructuredSocketConnection.__init__(self, *args, **kwargs)
        self._search = None
        self._search_key = None

    @property
    def search_cache(self):
        return self.session.handler.application.search_cache

    def reverse_url(self, *args, **kwargs):
        return self.session.handler.application.reverse_url(*args, **kwargs)

    @_StructuredSocketConnection.event
    @gen.engine
    def start(self, search_key):
        # Sanity checks
        if self._search is not None:
            raise HTTPError(400, 'Search already started')

        # Load the search spec
        try:
            search_spec = self.search_cache.get_search(search_key)
        except KeyError:
            raise HTTPError(400, 'Invalid search key')
        except SearchCacheLoadError:
            raise HTTPError(400, 'Corrupt search key')

        # Start the search
        print 'start', search_key
        self._search_key = search_key
        self._search = search_spec.make_search(
            object_callback=self._result,
            finished_callback=self._finished,
            close_callback=self._closed,
        )
        search_id = yield gen.Task(self._search.start)

        # Return search ID to client
        self.emit('search_started', search_id=search_id)

        # Start statistics coroutine
        self._stats_coroutine()

    @gen.engine
    def _stats_coroutine(self):
        '''Statistics coroutine.'''
        while self._search is not None:
            yield gen.Task(self._send_stats)
            yield gen.Task(IOLoop.instance().add_timeout, STATS_INTERVAL)

    @gen.engine
    def _send_stats(self, callback=None):
        '''Fetch stats and send them to the client.'''
        stats = yield gen.Task(self._search.get_stats)
        self.emit('statistics', **stats)
        if callback is not None:
            callback()

    def _result(self, obj):
        '''Blast channel result.'''
        object_key = self.search_cache.put_search_result(self._search_key,
                obj['_ObjectID'], obj)
        def factory(k):
            return self.reverse_url('attribute', self._search_key,
                    object_key, k)
        result = _make_object_json(obj, factory)
        self.emit('result', **result)

    @gen.engine
    def _finished(self):
        '''Search has completed.'''
        # Send final stats
        yield gen.Task(self._send_stats)
        self.emit('search_complete')
        self._search.close()

    def _closed(self):
        '''Search closed.'''
        # Close the SockJS connection
        self.close()

    @_StructuredSocketConnection.event
    def pause(self):
        if self._search is None:
            raise HTTPError(400, 'Search not yet started')
        self._search.pause()

    @_StructuredSocketConnection.event
    def resume(self):
        if self._search is None:
            raise HTTPError(400, 'Search not yet started')
        self._search.resume()

    def on_close(self):
        '''SockJS connection closed.'''
        # Close the Diamond connection
        if self._search is not None:
            print 'close'
            search = self._search
            self._search = None
            search.close()
<commit_msg>blaster: Autodetect MIME type for binary object attributes<commit_after>#
#  The OpenDiamond Platform for Interactive Search
#
#  Copyright (c) 2012 Carnegie Mellon University
#  All rights reserved.
#
#  This software is distributed under the terms of the Eclipse Public
#  License, Version 1.0 which can be found in the file named LICENSE.
#  ANY USE, REPRODUCTION OR DISTRIBUTION OF THIS SOFTWARE CONSTITUTES
#  RECIPIENT'S ACCEPTANCE OF THIS AGREEMENT
#

'''JSON Blaster request handlers.'''

from datetime import timedelta
import logging
import magic
import simplejson as json
import os
from sockjs.tornado import SockJSConnection
from urlparse import urlparse
from tornado.curl_httpclient import CurlAsyncHTTPClient as AsyncHTTPClient
from tornado import gen
from tornado.ioloop import IOLoop
from tornado.options import define, options
from tornado.web import asynchronous, RequestHandler, HTTPError
import validictory

import opendiamond
from opendiamond.attributes import (StringAttributeCodec,
        IntegerAttributeCodec, DoubleAttributeCodec, PatchesAttributeCodec)
from opendiamond.blaster.cache import SearchCacheLoadError
from opendiamond.blaster.search import (Blob, EmptyBlob, DiamondSearch,
        FilterSpec)
from opendiamond.helpers import sha256
from opendiamond.scope import ScopeCookie, ScopeError

CACHE_URN_SCHEME = 'blob'
STATS_INTERVAL = timedelta(milliseconds=1000)

# HTTP method handlers have specific argument lists rather than
# (self, *args, **kwargs) as in the superclass.
# pylint: disable=W0221

define('enable_testui', default=True,
        help='Enable the example user interface')
define('http_proxy', type=str, default=None,
        metavar='HOST:PORT', help='Use a proxy for HTTP client requests')

_log = logging.getLogger(__name__)

_magic = magic.open(magic.NONE)
_magic.setflags(magic.MIME_TYPE)
_magic.load()

def _load_schema(name):
    with open(os.path.join(os.path.dirname(__file__), name)) as fh:
        return json.load(fh)
EVENT_SCHEMA = _load_schema('schema-event.json')
SEARCH_SCHEMA = _load_schema('schema-search.json')

def _validate_json(schema, obj):
    # required_by_default=False and blank_by_default=True for
    # JSON Schema draft 3 semantics
    validictory.validate(obj, schema, required_by_default=False,
            blank_by_default=True)

def _make_object_json(obj, url_factory):
    '''Convert an object attribute dict into a dict suitable for JSON
    encoding.  url_factory(k) is a function to convert an attribute name
    into a URL.'''
    result = {}
    for k, v in obj.iteritems():
        data = None
        # Inline known attribute types that can be represented in JSON
        if k.endswith('.int'):
            data = IntegerAttributeCodec().decode(v)
        elif k.endswith('.double'):
            data = DoubleAttributeCodec().decode(v)
        elif k.endswith('.patches'):
            distance, patches = PatchesAttributeCodec().decode(v)
            data = {
                'distance': distance,
                'patches': [{
                    'x0': tl[0],
                    'y0': tl[1],
                    'x1': br[0],
                    'y1': br[1]
                } for tl, br in patches],
            }
        else:
            # Treat remaining attributes as strings if: they don't have
            # a suffix representing a known binary type, they can be decoded
            # by the string codec (i.e., their last byte is 0), they are
            # valid UTF-8, and they are not the '' (object data) attribute.
            try:
                _base, suffix = k.rsplit('.', 1)
            except ValueError:
                suffix = None
            try:
                if k != '' and suffix not in ('jpeg', 'rgbimage', 'binary'):
                    data = StringAttributeCodec().decode(v).decode('UTF-8')
            except ValueError:
                pass

        if data is not None:
            result[k] = {
                'data': data,
            }
        else:
            result[k] = {
                'url': url_factory(k),
            }
    return result

class _BlasterRequestHandler(RequestHandler):
    @property
    def blob_cache(self):
        return self.application.blob_cache

    @property
    def search_cache(self):
        return self.application.search_cache

    def write_error(self, code, **kwargs):
        exc_type, exc_value, _exc_tb = kwargs.get('exc_info', [None] * 3)
        if exc_type is not None and issubclass(exc_type, HTTPError):
            self.set_header('Content-Type', 'text/plain')
            if exc_value.log_message:
                self.write(exc_value.log_message + '\n')
        else:
            RequestHandler.write_error(self, code, **kwargs)

class _BlasterBlob(Blob):
    '''Instances with the same URI compare equal and hash to the same value.'''

    def __init__(self, uri, expected_sha256=None):
        Blob.__init__(self)
        self.uri = uri
        self._expected_sha256 = expected_sha256
        self._data = None

    def __str__(self):
        if self._data is None:
            raise RuntimeError('Attempting to read an unfetched blob')
        return self._data

    def __repr__(self):
        return '<_BlasterBlob %s>' % (self.uri)

    def __hash__(self):
        return hash(self.uri)

    def __eq__(self, other):
        return isinstance(other, type(self)) and self.uri == other.uri

    # pylint doesn't understand named tuples
    # pylint: disable=E1101
    @gen.engine
    def fetch(self, blob_cache, callback=None):
        if self._data is None:
            # Fetch data
            parts = urlparse(self.uri)
            if parts.scheme == 'blob':
                try:
                    data = blob_cache[parts.path]
                except KeyError:
                    raise HTTPError(400, 'Blob missing from blob cache')
            elif parts.scheme == 'http' or parts.scheme == 'https':
                client = AsyncHTTPClient()
                if options.http_proxy is not None:
                    proxy_host, proxy_port = options.http_proxy.split(':', 1)
                    proxy_port = int(proxy_port)
                else:
                    proxy_host = proxy_port = None
                response = yield gen.Task(client.fetch, self.uri,
                        user_agent='JSONBlaster/%s' % opendiamond.__version__,
                        proxy_host=proxy_host, proxy_port=proxy_port,
                        validate_cert=False)
                if response.error:
                    raise HTTPError(400, 'Error fetching <%s>: %s' % (
                            self.uri, str(response.error)))
                data = response.body
            else:
                raise HTTPError(400, 'Unacceptable blob URI scheme')

            # Check hash if requested
            if self._expected_sha256 is not None:
                if sha256(data).hexdigest() != self._expected_sha256:
                    raise HTTPError(400, 'SHA-256 mismatch on %s' % self.uri)

            # Commit
            self._data = data

        if callback is not None:
            callback()
    # pylint: enable=E1101

class _SearchSpec(object):
    def __init__(self, data):
        # Load JSON
        try:
            config = json.loads(data)
            _validate_json(SEARCH_SCHEMA, config)
        except ValueError, e:
            raise HTTPError(400, str(e))

        # Build cookies
        # Assume each "cookie" may actually be a megacookie
        try:
            self.cookies = [ScopeCookie.parse(c) for mc in config['cookies']
                    for c in ScopeCookie.split(mc)]
        except ScopeError, e:
            raise HTTPError(400, 'Invalid scope cookie: %s' % e)
        if not self.cookies:
            # No cookies could be parsed out of the client's cookie list
            raise HTTPError(400, 'No scope cookies found')

        # Build filters
        blobs = {}  # blob -> itself  (for deduplication)
        def make_blob(obj):
            if obj is not None:
                blob = _BlasterBlob(obj['uri'], obj.get('sha256'))
                # Intern the blob
                return blobs.setdefault(blob, blob)
            else:
                return EmptyBlob()
        self.filters = [FilterSpec(
                    name=f['name'],
                    code=make_blob(f['code']),
                    arguments=f.get('arguments', []),
                    blob_argument=make_blob(f.get('blob')),
                    dependencies=f.get('dependencies', []),
                    min_score=f.get('min_score', float('-inf')),
                    max_score=f.get('max_score', float('inf'))
                ) for f in config['filters']]
        self.blobs = blobs.values()

    @gen.engine
    def fetch_blobs(self, blob_cache, callback=None):
        yield [gen.Task(blob.fetch, blob_cache) for blob in self.blobs]
        if callback is not None:
            callback()

    def make_search(self, **kwargs):
        return DiamondSearch(self.cookies, self.filters, **kwargs)

class SearchHandler(_BlasterRequestHandler):
    def get(self):
        if options.enable_testui:
            self.render('testui/search.html')
        else:
            raise HTTPError(405, 'Method not allowed')

    @asynchronous
    @gen.engine
    def post(self):
        # Build search spec
        if self.request.headers['Content-Type'] != 'application/json':
            raise HTTPError(415, 'Content type must be application/json')
        spec = _SearchSpec(self.request.body)
        yield gen.Task(spec.fetch_blobs, self.blob_cache)

        # Store it
        search_key = self.search_cache.put_search(spec)

        # Return result
        self.set_status(204)
        self.set_header('Location', '/search')
        self.set_header('X-Search-Key', search_key)
        self.finish()

class PostBlobHandler(_BlasterRequestHandler):
    def post(self):
        sig = self.blob_cache.add(self.request.body)
        self.set_header('Location', '%s:%s' % (CACHE_URN_SCHEME, sig))
        self.set_status(204)

class AttributeHandler(_BlasterRequestHandler):
    def get(self, search_key, object_key, attr_name):
        try:
            data = self.search_cache.get_object_attribute(search_key,
                    object_key, attr_name)
        except KeyError:
            raise HTTPError(404, 'Not found')

        if attr_name.endswith('.jpeg'):
            mime = 'image/jpeg'
        else:
            mime = _magic.buffer(data)

        self.set_header('Content-Type', mime)
        self.write(data)

class ResultsHandler(_BlasterRequestHandler):
    def get(self):
        if options.enable_testui:
            self.render('testui/results.html')
        else:
            raise HTTPError(403, 'Forbidden')

class _StructuredSocketConnection(SockJSConnection):
    @classmethod
    def event(cls, func):
        '''Decorator specifying that this function is an event handler.'''
        func.event_handler = True
        return func

    # pylint is confused by msg.get()
    # pylint: disable=E1103
    def on_message(self, data):
        try:
            msg = json.loads(data)
            _validate_json(EVENT_SCHEMA, msg)
        except ValueError, e:
            raise HTTPError(400, str(e))
        event = msg['event']
        try:
            handler = getattr(self, event)
        except AttributeError:
            _log.warning('Unknown event type %s' % event)
            return
        if not getattr(handler, 'event_handler', False):
            _log.warning('Event %s has invalid handler' % event)
            return
        handler(**msg.get('data', {}))
    # pylint: enable=E1103

    def emit(self, event, **args):
        self.send(json.dumps({
            'event': event,
            'data': args,
        }))

class SearchConnection(_StructuredSocketConnection):
    def __init__(self, *args, **kwargs):
        _StructuredSocketConnection.__init__(self, *args, **kwargs)
        self._search = None
        self._search_key = None

    @property
    def search_cache(self):
        return self.session.handler.application.search_cache

    def reverse_url(self, *args, **kwargs):
        return self.session.handler.application.reverse_url(*args, **kwargs)

    @_StructuredSocketConnection.event
    @gen.engine
    def start(self, search_key):
        # Sanity checks
        if self._search is not None:
            raise HTTPError(400, 'Search already started')

        # Load the search spec
        try:
            search_spec = self.search_cache.get_search(search_key)
        except KeyError:
            raise HTTPError(400, 'Invalid search key')
        except SearchCacheLoadError:
            raise HTTPError(400, 'Corrupt search key')

        # Start the search
        print 'start', search_key
        self._search_key = search_key
        self._search = search_spec.make_search(
            object_callback=self._result,
            finished_callback=self._finished,
            close_callback=self._closed,
        )
        search_id = yield gen.Task(self._search.start)

        # Return search ID to client
        self.emit('search_started', search_id=search_id)

        # Start statistics coroutine
        self._stats_coroutine()

    @gen.engine
    def _stats_coroutine(self):
        '''Statistics coroutine.'''
        while self._search is not None:
            yield gen.Task(self._send_stats)
            yield gen.Task(IOLoop.instance().add_timeout, STATS_INTERVAL)

    @gen.engine
    def _send_stats(self, callback=None):
        '''Fetch stats and send them to the client.'''
        stats = yield gen.Task(self._search.get_stats)
        self.emit('statistics', **stats)
        if callback is not None:
            callback()

    def _result(self, obj):
        '''Blast channel result.'''
        object_key = self.search_cache.put_search_result(self._search_key,
                obj['_ObjectID'], obj)
        def factory(k):
            return self.reverse_url('attribute', self._search_key,
                    object_key, k)
        result = _make_object_json(obj, factory)
        self.emit('result', **result)

    @gen.engine
    def _finished(self):
        '''Search has completed.'''
        # Send final stats
        yield gen.Task(self._send_stats)
        self.emit('search_complete')
        self._search.close()

    def _closed(self):
        '''Search closed.'''
        # Close the SockJS connection
        self.close()

    @_StructuredSocketConnection.event
    def pause(self):
        if self._search is None:
            raise HTTPError(400, 'Search not yet started')
        self._search.pause()

    @_StructuredSocketConnection.event
    def resume(self):
        if self._search is None:
            raise HTTPError(400, 'Search not yet started')
        self._search.resume()

    def on_close(self):
        '''SockJS connection closed.'''
        # Close the Diamond connection
        if self._search is not None:
            print 'close'
            search = self._search
            self._search = None
            search.close()
<|endoftext|>


<commit_before># Copyright 2014 Google Inc. All Rights Reserved.

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

#     http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""User-facing interface to an ADB device.

This module provides the user-facing interface to controlling ADB devices with
this library.  See adb_protocol.py for a breakdown of the various layers.
Essentially, this module provides the AdbDevice class, which takes a UsbHandle
and uses the layers in between to provide a high level interface that is
designed to be similar to the 'adb' commandline interface.

Some additional useful mechanisms are provided for more common programmatic
needs, see shell_service.py for a description of how to use asynchronous shell
commands.

Note we don't provide a device listing facility, as this is not actually an ADB
function, but rather a USB function - listing devices with a specific interface
class, subclass, and protocol.
"""

try:
  import cStringIO as io
except ImportError:
  import io
import logging
import os.path

try:
  from M2Crypto import RSA
except ImportError:
  logging.error('Failed to import M2Crypto, did you pip install '
                'openhtf[usb_plugs]?')
  raise

from openhtf.plugs.usb import adb_protocol
from openhtf.plugs.usb import filesync_service
from openhtf.plugs.usb import shell_service
from openhtf.plugs.usb import usb_exceptions

from openhtf.util import timeouts

# USB interface class, subclass, and protocol for matching against.
CLASS = 0xFF
SUBCLASS = 0x42
PROTOCOL = 0x01

class M2CryptoSigner(adb_protocol.AuthSigner):
  """AuthSigner using M2Crypto."""

  def __init__(self, rsa_key_path):
    with open(rsa_key_path + '.pub') as rsa_pub_file:
      self.public_key = rsa_pub_file.read()

    self.rsa_key = RSA.load_key(rsa_key_path)

  def sign(self, data):
    return self.rsa_key.sign(data, 'sha1')

  def get_public_key(self):
    """Return the public key."""
    return self.public_key

class AdbDevice(object):
  """Exposes adb-like methods for use.

  Some methods are more-pythonic and/or have more options.

  Attributes:
    filesync_service: Direct access to this device's FilesyncService.
    shell_service: Direct access to this device's ShellService.
  """

  def __init__(self, adb_connection):
    self._adb_connection = adb_connection
    self.filesync_service = filesync_service.FilesyncService.using_connection(
        adb_connection)
    self.shell_service = shell_service.ShellService.using_connection(
        adb_connection)

  def __str__(self):
    return '<%s: %s(%s) @%s>' % (type(self).__name__,
                                 self._adb_connection.serial,
                                 self._adb_connection.systemtype,
                                 self._adb_connection.transport)
  __repr__ = __str__

  def get_system_type(self):
    """Return the system type."""
    return self._adb_connection.systemtype

  def get_serial(self):
    """Return the device serial."""
    return self._adb_connection.serial

  def close(self):
    """Close the ADB connection."""
    self.filesync_service.close()
    self._adb_connection.close()

  def install(self, apk_path, destination_dir=None, timeout_ms=None):
    """Install apk to device.

    Doesn't support verifier file, instead allows destination directory to be
    overridden.

    Arguments:
      apk_path: Local path to apk to install.
      destination_dir: Optional destination directory. Use /system/app/ for
        persistent applications.
      timeout_ms: Expected timeout for pushing and installing.

    Returns:
      The pm install output.
    """
    if not destination_dir:
      destination_dir = '/data/local/tmp/'
    basename = os.path.basename(apk_path)
    destination_path = destination_dir + basename
    self.push(apk_path, destination_path, timeout_ms=timeout_ms)
    return self.Shell('pm install -r "%s"' % destination_path,
                      timeout_ms=timeout_ms)

  def push(self, source_file, device_filename, timeout_ms=None):
    """Push source_file to file on device.

    Arguments:
      source_file: Either a filename or file-like object to push to the device.
        If a filename, will set the remote mtime to match the local mtime,
        otherwise will use the current time.
      device_filename: The filename on the device to write to.
      timeout_ms: Expected timeout for any part of the push.
    """
    mtime = 0
    if isinstance(source_file, str):
      mtime = os.path.getmtime(source_file)
      source_file = open(source_file)

    self.filesync_service.send(
        source_file, device_filename, mtime=mtime,
        timeout=timeouts.PolledTimeout.from_millis(timeout_ms))

  def pull(self, device_filename, dest_file=None, timeout_ms=None):
    """Pull file from device.

    Arguments:
      device_filename: The filename on the device to pull.
      dest_file: If set, a filename or writable file-like object.
      timeout_ms: Expected timeout for the pull.

    Returns:
      The file data if dest_file is not set, None otherwise.
    """
    if isinstance(dest_file, str):
      dest_file = open(dest_file, 'w')
    elif not dest_file:
      dest_file = io.StringIO()
    self.filesync_service.recv(device_filename, dest_file,
                               timeouts.PolledTimeout.from_millis(timeout_ms))
    # An empty call to cStringIO.StringIO returns an instance of
    # cStringIO.OutputType.
    if isinstance(dest_file, io.OutputType):
      return dest_file.getvalue()

  def list(self, device_path, timeout_ms=None):
    """Yield filesync_service.DeviceFileStat objects for directory contents."""
    return self.filesync_service.list(
        device_path, timeouts.PolledTimeout.from_millis(timeout_ms))

  def command(self, command, raw=False, timeout_ms=None):
    """Run command on the device, returning the output."""
    return self.shell_service.command(command, raw=raw, timeout_ms=timeout_ms)
  Shell = command  #pylint: disable=invalid-name

  def async_command(self, command, raw=False, timeout_ms=None):
    """See shell_service.ShellService.async_command()."""
    return self.shell_service.async_command(command, raw=raw,
                                            timeout_ms=timeout_ms)

  def _check_remote_command(self, destination, timeout_ms, success_msgs=None):
    """Open a stream to destination, check for remote errors.

    Used for reboot, remount, and root services.  If this method returns, the
    command was successful, otherwise an appropriate error will have been
    raised.

    Args:
      destination: Stream destination to open.
      timeout_ms: Timeout in milliseconds for the operation.
      success_msgs: If provided, a list of messages that, if returned from the
        device, indicate success, so don't treat them as errors.

    Raises:
      AdbRemoteError: If the remote command fails, will contain any message we
        got back from the device.
      AdbStreamUnavailableError: The service requested isn't supported.
    """
    timeout = timeouts.PolledTimeout.from_millis(timeout_ms)
    stream = self._adb_connection.open_stream(destination, timeout)
    if not stream:
      raise usb_exceptions.AdbStreamUnavailableError(
          'Service %s not supported', destination)
    try:
      message = stream.read(timeout_ms=timeout)
      # Some commands report success messages, ignore them.
      if any([m in message for m in success_msgs]):
        return
    except usb_exceptions.CommonUsbError:
      if destination.startswith('reboot:'):
        # We expect this if the device is rebooting.
        return
      raise
    raise usb_exceptions.AdbRemoteError('Device message: %s', message)

  def reboot(self, destination='', timeout_ms=None):
    """Reboot device, specify 'bootloader' for fastboot."""
    self._check_remote_command('reboot:%s' % destination, timeout_ms)

  def remount(self, timeout_ms=None):
    """Remount / as read-write."""
    self._check_remote_command('remount:', timeout_ms, ['remount succeeded'])

  def root(self, timeout_ms=None):
    """Restart adbd as root on device."""
    self._check_remote_command('root:', timeout_ms,
                               ['already running as root',
                                'restarting adbd as root'])

  @classmethod
  def connect(cls, usb_handle, **kwargs):
    """Connect to the device.

    Args:
      usb_handle: UsbHandle instance to use.
      **kwargs: See AdbConnection.connect for kwargs. Includes rsa_keys, and
        auth_timeout_ms.

    Returns:
      An instance of this class if the device connected successfully.
    """
    adb_connection = adb_protocol.AdbConnection.connect(usb_handle, **kwargs)
    return cls(adb_connection)
<commit_msg>Cleaner fix for python 2/3 in adb_device.py (#738)<commit_after># Copyright 2014 Google Inc. All Rights Reserved.

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

#     http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""User-facing interface to an ADB device.

This module provides the user-facing interface to controlling ADB devices with
this library.  See adb_protocol.py for a breakdown of the various layers.
Essentially, this module provides the AdbDevice class, which takes a UsbHandle
and uses the layers in between to provide a high level interface that is
designed to be similar to the 'adb' commandline interface.

Some additional useful mechanisms are provided for more common programmatic
needs, see shell_service.py for a description of how to use asynchronous shell
commands.

Note we don't provide a device listing facility, as this is not actually an ADB
function, but rather a USB function - listing devices with a specific interface
class, subclass, and protocol.
"""

import io
import logging
import os.path

try:
  from M2Crypto import RSA
except ImportError:
  logging.error('Failed to import M2Crypto, did you pip install '
                'openhtf[usb_plugs]?')
  raise

from openhtf.plugs.usb import adb_protocol
from openhtf.plugs.usb import filesync_service
from openhtf.plugs.usb import shell_service
from openhtf.plugs.usb import usb_exceptions

from openhtf.util import timeouts

# USB interface class, subclass, and protocol for matching against.
CLASS = 0xFF
SUBCLASS = 0x42
PROTOCOL = 0x01

class M2CryptoSigner(adb_protocol.AuthSigner):
  """AuthSigner using M2Crypto."""

  def __init__(self, rsa_key_path):
    with open(rsa_key_path + '.pub') as rsa_pub_file:
      self.public_key = rsa_pub_file.read()

    self.rsa_key = RSA.load_key(rsa_key_path)

  def sign(self, data):
    return self.rsa_key.sign(data, 'sha1')

  def get_public_key(self):
    """Return the public key."""
    return self.public_key

class AdbDevice(object):
  """Exposes adb-like methods for use.

  Some methods are more-pythonic and/or have more options.

  Attributes:
    filesync_service: Direct access to this device's FilesyncService.
    shell_service: Direct access to this device's ShellService.
  """

  def __init__(self, adb_connection):
    self._adb_connection = adb_connection
    self.filesync_service = filesync_service.FilesyncService.using_connection(
        adb_connection)
    self.shell_service = shell_service.ShellService.using_connection(
        adb_connection)

  def __str__(self):
    return '<%s: %s(%s) @%s>' % (type(self).__name__,
                                 self._adb_connection.serial,
                                 self._adb_connection.systemtype,
                                 self._adb_connection.transport)
  __repr__ = __str__

  def get_system_type(self):
    """Return the system type."""
    return self._adb_connection.systemtype

  def get_serial(self):
    """Return the device serial."""
    return self._adb_connection.serial

  def close(self):
    """Close the ADB connection."""
    self.filesync_service.close()
    self._adb_connection.close()

  def install(self, apk_path, destination_dir=None, timeout_ms=None):
    """Install apk to device.

    Doesn't support verifier file, instead allows destination directory to be
    overridden.

    Arguments:
      apk_path: Local path to apk to install.
      destination_dir: Optional destination directory. Use /system/app/ for
        persistent applications.
      timeout_ms: Expected timeout for pushing and installing.

    Returns:
      The pm install output.
    """
    if not destination_dir:
      destination_dir = '/data/local/tmp/'
    basename = os.path.basename(apk_path)
    destination_path = destination_dir + basename
    self.push(apk_path, destination_path, timeout_ms=timeout_ms)
    return self.Shell('pm install -r "%s"' % destination_path,
                      timeout_ms=timeout_ms)

  def push(self, source_file, device_filename, timeout_ms=None):
    """Push source_file to file on device.

    Arguments:
      source_file: Either a filename or file-like object to push to the device.
        If a filename, will set the remote mtime to match the local mtime,
        otherwise will use the current time.
      device_filename: The filename on the device to write to.
      timeout_ms: Expected timeout for any part of the push.
    """
    mtime = 0
    if isinstance(source_file, str):
      mtime = os.path.getmtime(source_file)
      source_file = open(source_file)

    self.filesync_service.send(
        source_file, device_filename, mtime=mtime,
        timeout=timeouts.PolledTimeout.from_millis(timeout_ms))

  def pull(self, device_filename, dest_file=None, timeout_ms=None):
    """Pull file from device.

    Arguments:
      device_filename: The filename on the device to pull.
      dest_file: If set, a filename or writable file-like object.
      timeout_ms: Expected timeout for the pull.

    Returns:
      The file data if dest_file is not set, None otherwise.
    """
    should_return_data = dest_file is None
    if isinstance(dest_file, str):
      dest_file = open(dest_file, 'w')
    elif dest_file is None:
      dest_file = io.StringIO()
    self.filesync_service.recv(device_filename, dest_file,
                               timeouts.PolledTimeout.from_millis(timeout_ms))
    if should_return_data:
      return dest_file.getvalue()

  def list(self, device_path, timeout_ms=None):
    """Yield filesync_service.DeviceFileStat objects for directory contents."""
    return self.filesync_service.list(
        device_path, timeouts.PolledTimeout.from_millis(timeout_ms))

  def command(self, command, raw=False, timeout_ms=None):
    """Run command on the device, returning the output."""
    return self.shell_service.command(command, raw=raw, timeout_ms=timeout_ms)
  Shell = command  #pylint: disable=invalid-name

  def async_command(self, command, raw=False, timeout_ms=None):
    """See shell_service.ShellService.async_command()."""
    return self.shell_service.async_command(command, raw=raw,
                                            timeout_ms=timeout_ms)

  def _check_remote_command(self, destination, timeout_ms, success_msgs=None):
    """Open a stream to destination, check for remote errors.

    Used for reboot, remount, and root services.  If this method returns, the
    command was successful, otherwise an appropriate error will have been
    raised.

    Args:
      destination: Stream destination to open.
      timeout_ms: Timeout in milliseconds for the operation.
      success_msgs: If provided, a list of messages that, if returned from the
        device, indicate success, so don't treat them as errors.

    Raises:
      AdbRemoteError: If the remote command fails, will contain any message we
        got back from the device.
      AdbStreamUnavailableError: The service requested isn't supported.
    """
    timeout = timeouts.PolledTimeout.from_millis(timeout_ms)
    stream = self._adb_connection.open_stream(destination, timeout)
    if not stream:
      raise usb_exceptions.AdbStreamUnavailableError(
          'Service %s not supported', destination)
    try:
      message = stream.read(timeout_ms=timeout)
      # Some commands report success messages, ignore them.
      if any([m in message for m in success_msgs]):
        return
    except usb_exceptions.CommonUsbError:
      if destination.startswith('reboot:'):
        # We expect this if the device is rebooting.
        return
      raise
    raise usb_exceptions.AdbRemoteError('Device message: %s', message)

  def reboot(self, destination='', timeout_ms=None):
    """Reboot device, specify 'bootloader' for fastboot."""
    self._check_remote_command('reboot:%s' % destination, timeout_ms)

  def remount(self, timeout_ms=None):
    """Remount / as read-write."""
    self._check_remote_command('remount:', timeout_ms, ['remount succeeded'])

  def root(self, timeout_ms=None):
    """Restart adbd as root on device."""
    self._check_remote_command('root:', timeout_ms,
                               ['already running as root',
                                'restarting adbd as root'])

  @classmethod
  def connect(cls, usb_handle, **kwargs):
    """Connect to the device.

    Args:
      usb_handle: UsbHandle instance to use.
      **kwargs: See AdbConnection.connect for kwargs. Includes rsa_keys, and
        auth_timeout_ms.

    Returns:
      An instance of this class if the device connected successfully.
    """
    adb_connection = adb_protocol.AdbConnection.connect(usb_handle, **kwargs)
    return cls(adb_connection)
<|endoftext|>


<commit_before>from math import floor
from random import random, randint, seed, shuffle
from mad import MadFile
from datetime import datetime
import os
import magic
from urllib import unquote_plus
import syslog
import ConfigParser

from constantina_shared import GlobalConfig, BaseFiles, BaseState
from medusa_cards import MedusaState

syslog.openlog(ident='constantina_state')

class ConstantinaState(BaseState):
    """
    Constantina State Object, an aggregate of the following:
       MedusaState: blog cards. images, quotes, news, songs
       ZooState: forum cards. threads, images, songs
       DraculaState: mail cards. ??

    Aggregate all other states here, and do counting for things like total
    cards, the number of pages, etc. Lots of helper functions for deciding
    things about the Constantina page are tracked here, since these conditions
    depend on properties of the various states
    """
    def __init__(self, in_state=None):
        BaseState.__init__(self, in_state, None)

        # Getting defaults from the other states requires us to first import
        # any random seed value. Then, we can finish setting the imported state
        self.__import_random_seed()
        self.__set_state_defaults()
        self.__import_state()

    def __set_state_defaults(self):
        """
        Set default values for special_state properties and normal content-card
        properties, as well as upper limits on how many cards can exist on a
        single Constantina page for this type of application. Doing so means we
        must import all of the sub-application states here, rather than in the
        import function like you'd do intuitively :(
        """
        self.max_items = 0   # Aggregate max items per page
        self.filtered = 0    # Aggregate filtered card count

        # For the values in the special_states config, create variables, i.e.
        #   state.appearance, state.page
        for spctype, spcfield in GlobalConfig.items("special_states"):
            setattr(self, spcfield, None)       # All state vals are expected to exist

        # Subapplication states are held in this object too
        self.medusa = None
        self.zoo = None
        # self.dracula = None

        # Based on modes, enable Medusa/Zoo/other states
        if GlobalConfig.get("authentication", "mode") == "blog":
            self.medusa = MedusaState(self.in_state)
            self.max_items += self.medusa.max_items
            self.filtered += self.medusa.filtered

        if GlobalConfig.get("authentication", "mode") == "zoo":
            self.zoo = ZooState(self.in_state)
            self.max_items += self.zoo.max_items
            self.filtered += self.zoo.filtered

        if GlobalConfig.get("authentication", "mode") == "combined":
            self.medusa = MedusaState(self.in_state)
            self.max_items += self.medusa.max_items
            self.filtered += self.medusa.filtered
            self.zoo = ZooState(self.in_state)
            self.max_items += self.zoo.max_items
            self.filtered += self.zoo.filtered

    def get(self, application, value):
        """
        Allow retrieving a medusa/zoo state value or function, or returning a
        default None value.
        """
        app_state = getattr(self, application, None)
        if app_state == None:
            return None

        item = getattr(app_state, value)
        if callable(item):
            return item()
        else:
            return item

    def all(self, value, mode="append"):
        """
        For all available applications, get the value or function asked for.
        The mode can be a sum of available values, or logic like 'and' or 'or'
        """
        items = []
        for application in GlobalConfig.get("applications", "list"):
            items.append = self.get(application, value)
        if mode == "append":
            items = filter(None, items)
            if items == []:
                return None
            return [item for sublist in items for item in sublist]   # Flatten
        elif mode == "sum":
            return sum(items)
        elif mode == "or":
            return any(items)   # Logical OR across the list
        elif mode == "and":
            return all(items)   # Logical AND across the list
        else:
            return None

    def __import_page_count_state(self):
        """
        For all subsequent "infinite-scroll AJAX" content after the initial page
        load, we track the current page number of content.

        Output must be an integer.
        """
        self.page = BaseState._find_state_variable(self, 'xp')

        # If page was read in as a special state variable, use that (for search results)
        if self.page is not None:
            self.page = BaseState._int_translate(self, self.page, 1, 0)
        else:
            self.page = 0

    def __import_random_seed(self):
        """
        Set the return seed based on a 14-digit string from the state variable.
        As an input to seed(), this has to be a float between zero and one.

        This seed is used to consistently seed the shuffle function, so that
        between page loads, we know the shuffled card functions give the same
        shuffle ordering.
        """
        self.seed = BaseState._find_state_variable(self, "seed")
        if self.seed is None:
            self.seed = round(random(), 14)
        else:
            self.seed = float(str("0." + self.seed))
        seed(self.seed)   # Now the RNG is seeded with our consistent value

    def _import_theme_state(self):
        """
        In the top-level state object, we track an appearance variable, which
        corresponds to the exact state variable imported (and exported) for the
        appearance of the entire Constantina site.

        The appearance value lets us look up which theme we display for the user.
        This theme value is a path fragment to a theme's images and stylesheets.
        """
        appearance_state = BaseState._find_state_variable('xa')

        if appearance_state is not None:
            # Read in single char of theme state value
            self.appearance = BaseState._int_translate(appearance_state, 1, 0)

        theme_count = len(GlobalConfig.items("themes")) - 1
        self.theme = None
        if self.appearance is None:
            self.theme = GlobalConfig.get("themes", "default")
        elif self.appearance >= theme_count:
            self.theme = GlobalConfig.get("themes", str(self.appearance % theme_count))
        else:
            self.theme = GlobalConfig.get("themes", str(self.appearance))

        # If the configuration supports a random theme, and we didn't have a
        # theme provided in the initial state, let's choose one randomly
        if (appearance_state is None) and (self.theme == "random"):
            seed()   # Enable non-seeded choice
            choice = randint(0, theme_count - 1)
            self.theme = GlobalConfig.get("themes", str(choice))
            if self.seed:   # Re-enable seeded nonrandom choice
                seed(self.seed)

    def __import_state(self):
        """
        When setting defaults, we imported the sub application states. But now
        we need to import the relevant states for the entire application, like
        the page state and the chosen theme.
        """
        self.__import_page_count_state()
        self.__import_theme_state()

    def __export_page_count_state(self):
        """
        If we had search results and used a page number, write an incremented page
        number into the next search state for loading
        """
        page_string = None
        if self.search_mode() is True:
            export_page = int(self.page) + 1
            page_string = "xp" + str(export_page)
        return page_string

    def __export_random_seed(self):
        """Export the random seed for adding to the state variable"""
        return str(self.seed).replace("0.", "")

    def _export_theme_state(self):
        """If tracking an appearance or theme, include it in state links"""
        appearance_string = None
        if self.appearance is not None:
            appearance_string = "xa" + str(self.appearance)
        return appearance_string

    def export_state(self):
        """Export all medusa/other states, as well as appearance/page here"""
        export_parts = [ self.export_random_seed(),
                         self.export_page_count_state(),
                         self.get("medusa", "export_state"),
                         self.get("zoo", "export_state") ]

        export_parts = filter(None, export_parts)
        export_string = ':'.join(export_parts)
        return export_string

    def fresh_mode(self):
        """Either an empty state, or just an empty state and a theme is set"""
        if (((self.in_state is None) or (self.all("configured_states") == ['appearance'])) and
             (self.page == 0) and
             (self.all("reshuffle", "or") is False)):
            return True
        else:
            return False

    def reshuffle_mode(self):
        """An empty search was given, so reshuffle the page"""
        if ((self.all("search") is not None) and
            (self.all("reshuffle", "or") is True) and
            (self.all("card_filter") is None)):
            return True
        else:
            return False

    def permalink_mode(self):
        """Is one of the permalink modes on?"""
	# TODO: check other states for active permalinks!
        if ((self.medusa.news_permalink is not None) or
            (self.medusa.features_permalink is not None) or
            (self.medusa.topics_permalink is not None)):
            return True
        else:
            return False

    def exclude_cardtype(self, application, ctype):
        """
        Is a card filter in place, and if so, is the given card type being filtered?
        If this returns true, it means the user either wants cards of this type, or
        that no card filtering is currently in place.
        """
        app_ctype = self.get(application, ctype)
        if app_ctype == None:   # No app or ctype, so no cards of this type
            return False

        if ((self.all("card_filter") is not None) and
            (app_ctype.filtertype is False)):
            return True
        else:
            return False

    def filter_processed_mode(self):
        """
        Is it a search state, and did we already convert #hashtag strings
        into filter queries?
        """
        states = self.all("configured_states")
        if (('search' in states) or
            ('card_filter' in states)):
            return True
        else:
            return False

    def filter_only_mode(self):
        """All search queries converted into card filters"""
        if self.all("configured_states") == ['card_filter']:
            return True
        else:
            return False

    def search_only_mode(self):
        """There is a search query, but no terms converted into card filters"""
        if self.all("configured_states") == ['search']:
            return True
        else:
            return False

    def search_mode(self):
        """Any valid state from a search mode will trigger this mode"""
        if ((self.all("search") is not None) or
            (self.all("card_filter") is not None) or
            (self.filtered != 0)):
            return True
        else:
            return False

    def out_of_content(self, card_count):
        """
        Count the number of cards that are part of our page counting. If we've
        already displayed this number of cards, we are out of content.
        """
        card_limit = 0
        for application in GlobalConfig.get("application", "list"):
            app_state = getattr(self, application)
            for ctype in app_state.config.get("card_properties", "pagecount").replace(" ", "").split(","):
                card_limit = card_limit + getattr(self, ctype).file_count
        # syslog.syslog("card_limit: " + str(card_limit) + "   card_count: " + str(card_count))
        if card_count >= card_limit:
            return True
        else:
            return False
<commit_msg>think constantina_state is cleaned up, even if not tested yet<commit_after>from math import floor
from random import random, randint, seed, shuffle
from mad import MadFile
from datetime import datetime
import os
import magic
from urllib import unquote_plus
import syslog
import ConfigParser

from constantina_shared import GlobalConfig, BaseFiles, BaseState
from medusa_cards import MedusaState

syslog.openlog(ident='constantina_state')

class ConstantinaState(BaseState):
    """
    Constantina State Object, an aggregate of the following:
       MedusaState: blog cards. images, quotes, news, songs
       ZooState: forum cards. threads, images, songs
       DraculaState: mail cards. ??

    Aggregate all other states here, and do counting for things like total
    cards, the number of pages, etc. Lots of helper functions for deciding
    things about the Constantina page are tracked here, since these conditions
    depend on properties of the various states
    """
    def __init__(self, in_state=None):
        BaseState.__init__(self, in_state, None)

        # Getting defaults from the other states requires us to first import
        # any random seed value. Then, we can finish setting the imported state
        self.__import_random_seed()
        self.__set_state_defaults()
        self.__import_state()

    def __set_state_defaults(self):
        """
        Set default values for special_state properties and normal content-card
        properties, as well as upper limits on how many cards can exist on a
        single Constantina page for this type of application. Doing so means we
        must import all of the sub-application states here, rather than in the
        import function like you'd do intuitively :(
        """
        self.max_items = 0   # Aggregate max items per page
        self.filtered = 0    # Aggregate filtered card count

        # For the values in the special_states config, create variables, i.e.
        #   state.appearance, state.page
        for spctype, spcfield in GlobalConfig.items("special_states"):
            setattr(self, spcfield, None)       # All state vals are expected to exist

        # Subapplication states are held in this object too
        self.medusa = None
        self.zoo = None
        # self.dracula = None

        # Based on modes, enable Medusa/Zoo/other states
        if GlobalConfig.get("authentication", "mode") == "blog":
            self.medusa = MedusaState(self.in_state)
            self.max_items += self.medusa.max_items
            self.filtered += self.medusa.filtered

        if GlobalConfig.get("authentication", "mode") == "zoo":
            self.zoo = ZooState(self.in_state)
            self.max_items += self.zoo.max_items
            self.filtered += self.zoo.filtered

        if GlobalConfig.get("authentication", "mode") == "combined":
            self.medusa = MedusaState(self.in_state)
            self.max_items += self.medusa.max_items
            self.filtered += self.medusa.filtered
            self.zoo = ZooState(self.in_state)
            self.max_items += self.zoo.max_items
            self.filtered += self.zoo.filtered

    def get(self, application, value):
        """
        Allow retrieving a medusa/zoo state value or function, or returning a
        default None value.
        """
        app_state = getattr(self, application, None)
        if app_state == None:
            return None

        item = getattr(app_state, value)
        if callable(item):
            return item()
        else:
            return item

    def all(self, value, mode="append"):
        """
        For all available applications, get the value or function asked for.
        The mode can be a sum of available values, or logic like 'and' or 'or'
        """
        items = []
        for application in GlobalConfig.get("applications", "list"):
            items.append(self.get(application, value))
        if mode == "append":
            items = filter(None, items)
            if items == []:
                return None
            return [item for sublist in items for item in sublist]   # Flatten
        elif mode == "sum":
            return sum(items)
        elif mode == "or":
            return any(items)   # Logical OR across the list
        elif mode == "and":
            return all(items)   # Logical AND across the list
        else:
            return None

    def __import_page_count_state(self):
        """
        For all subsequent "infinite-scroll AJAX" content after the initial page
        load, we track the current page number of content.

        Output must be an integer.
        """
        self.page = BaseState._find_state_variable(self, 'xp')

        # If page was read in as a special state variable, use that (for search results)
        if self.page is not None:
            self.page = BaseState._int_translate(self, self.page, 1, 0)
        else:
            self.page = 0

    def __import_random_seed(self):
        """
        Set the return seed based on a 14-digit string from the state variable.
        As an input to seed(), this has to be a float between zero and one.

        This seed is used to consistently seed the shuffle function, so that
        between page loads, we know the shuffled card functions give the same
        shuffle ordering.
        """
        self.seed = BaseState._find_state_variable(self, "seed")
        if self.seed is None:
            self.seed = round(random(), 14)
        else:
            self.seed = float(str("0." + self.seed))
        seed(self.seed)   # Now the RNG is seeded with our consistent value

    def __import_theme_state(self):
        """
        In the top-level state object, we track an appearance variable, which
        corresponds to the exact state variable imported (and exported) for the
        appearance of the entire Constantina site.

        The appearance value lets us look up which theme we display for the user.
        This theme value is a path fragment to a theme's images and stylesheets.
        """
        appearance_state = BaseState._find_state_variable(self, 'xa')

        if appearance_state is not None:
            # Read in single char of theme state value
            self.appearance = BaseState._int_translate(self, appearance_state, 1, 0)

        theme_count = len(GlobalConfig.items("themes")) - 1
        self.theme = None
        if self.appearance is None:
            self.theme = GlobalConfig.get("themes", "default")
        elif self.appearance >= theme_count:
            self.theme = GlobalConfig.get("themes", str(self.appearance % theme_count))
        else:
            self.theme = GlobalConfig.get("themes", str(self.appearance))

        # If the configuration supports a random theme, and we didn't have a
        # theme provided in the initial state, let's choose one randomly
        if (appearance_state is None) and (self.theme == "random"):
            seed()   # Enable non-seeded choice
            choice = randint(0, theme_count - 1)
            self.theme = GlobalConfig.get("themes", str(choice))
            if self.seed:   # Re-enable seeded nonrandom choice
                seed(self.seed)

    def __import_state(self):
        """
        When setting defaults, we imported the sub application states. But now
        we need to import the relevant states for the entire application, like
        the page state and the chosen theme.
        """
        self.__import_page_count_state()
        self.__import_theme_state()

    def __export_page_count_state(self):
        """
        If we had search results and used a page number, write an incremented page
        number into the next search state for loading
        """
        page_string = None
        if self.search_mode() is True:
            export_page = int(self.page) + 1
            page_string = "xp" + str(export_page)
        return page_string

    def __export_random_seed(self):
        """Export the random seed for adding to the state variable"""
        return str(self.seed).replace("0.", "")

    def _export_theme_state(self):
        """If tracking an appearance or theme, include it in state links"""
        appearance_string = None
        if self.appearance is not None:
            appearance_string = "xa" + str(self.appearance)
        return appearance_string

    def export_state(self):
        """Export all medusa/other states, as well as appearance/page here"""
        export_parts = [ self.export_random_seed(),
                         self.export_page_count_state(),
                         self.get("medusa", "export_state"),
                         self.get("zoo", "export_state") ]

        export_parts = filter(None, export_parts)
        export_string = ':'.join(export_parts)
        return export_string

    def fresh_mode(self):
        """Either an empty state, or just an empty state and a theme is set"""
        if (((self.in_state is None) or (self.all("configured_states") == ['appearance'])) and
             (self.page == 0) and
             (self.all("reshuffle", "or") is False)):
            return True
        else:
            return False

    def reshuffle_mode(self):
        """An empty search was given, so reshuffle the page"""
        if ((self.all("search") is not None) and
            (self.all("reshuffle", "or") is True) and
            (self.all("card_filter") is None)):
            return True
        else:
            return False

    def permalink_mode(self):
        """Is one of the permalink modes on?"""
	# TODO: check other states for active permalinks!
        if ((self.medusa.news_permalink is not None) or
            (self.medusa.features_permalink is not None) or
            (self.medusa.topics_permalink is not None)):
            return True
        else:
            return False

    def exclude_cardtype(self, application, ctype):
        """
        Is a card filter in place, and if so, is the given card type being filtered?
        If this returns true, it means the user either wants cards of this type, or
        that no card filtering is currently in place.
        """
        app_ctype = self.get(application, ctype)
        if app_ctype == None:   # No app or ctype, so no cards of this type
            return False

        if ((self.all("card_filter") is not None) and
            (app_ctype.filtertype is False)):
            return True
        else:
            return False

    def filter_processed_mode(self):
        """
        Is it a search state, and did we already convert #hashtag strings
        into filter queries?
        """
        states = self.all("configured_states")
        if (('search' in states) or
            ('card_filter' in states)):
            return True
        else:
            return False

    def filter_only_mode(self):
        """All search queries converted into card filters"""
        if self.all("configured_states") == ['card_filter']:
            return True
        else:
            return False

    def search_only_mode(self):
        """There is a search query, but no terms converted into card filters"""
        if self.all("configured_states") == ['search']:
            return True
        else:
            return False

    def search_mode(self):
        """Any valid state from a search mode will trigger this mode"""
        if ((self.all("search") is not None) or
            (self.all("card_filter") is not None) or
            (self.filtered != 0)):
            return True
        else:
            return False

    def out_of_content(self, card_count):
        """
        Count the number of cards that are part of our page counting. If we've
        already displayed this number of cards, we are out of content.
        """
        card_limit = 0
        for application in GlobalConfig.get("application", "list"):
            app_state = getattr(self, application)
            for ctype in app_state.config.get("card_properties", "pagecount").replace(" ", "").split(","):
                card_limit = card_limit + getattr(self, ctype).file_count
        # syslog.syslog("card_limit: " + str(card_limit) + "   card_count: " + str(card_count))
        if card_count >= card_limit:
            return True
        else:
            return False
<|endoftext|>


<commit_before>import os
import sys
import signal
import socket
import subprocess
import multiprocessing
from openquake.baselib import zeromq as z, parallel as p

def streamer(task_in_url, task_out_url):
    """A streamer thread for the zmq workers"""
    try:
        z.zmq.proxy(z.bind(task_in_url, z.zmq.PULL),
                    z.bind(task_out_url, z.zmq.PUSH))
    except (KeyboardInterrupt, z.zmq.ZMQError):
        pass

def _starmap(func, iterargs, task_in_url, receiver_url):
    # called by parallel.Starmap; should not be used directly
    with z.Socket(receiver_url, z.zmq.PULL, 'bind') as receiver, \
            z.Socket(task_in_url, z.zmq.PUSH, 'connect') as sender:
        n = 0
        for args in iterargs:
            # args[-1] is a Monitor instance
            args[-1].backurl = receiver.true_end_point
            sender.send((func, args))
            n += 1
        yield n
        for _ in range(n):
            # receive n responses for the n requests sent
            yield receiver.zsocket.recv_pyobj()

class WorkerMaster(object):
    """
    :param frontend_url: url where to send the tasks
    :param task_out_url: url with a range of ports to receive the results
    :param ctrl_port: port on which the worker pools listen
    :param host_cores: names of the remote hosts and number of cores
    :param remote_python: path of the Python executable on the remote hosts
    """
    def __init__(self, task_in_url, task_out_url,
                 ctrl_port, host_cores, remote_python=None):
        self.task_in_url = task_in_url
        self.task_out_url = task_out_url
        self.ctrl_port = int(ctrl_port)
        self.host_cores = [hc.split() for hc in host_cores.split(',')]
        self.remote_python = remote_python or sys.executable

    def status(self, host=None):
        """
        :returns: a list of pairs (hostname, 'running'|'not-running')
        """
        if host is None:
            host_cores = self.host_cores
        else:
            host_cores = [hc for hc in self.host_cores if hc[0] == host]
        lst = []
        for host, _ in host_cores:
            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            try:
                err = sock.connect_ex((host, self.ctrl_port))
            finally:
                sock.close()
            lst.append((host, 'not-running' if err else 'running'))
        return lst

    def start(self):
        """
        Start multiple workerpools, possibly on remote servers via ssh
        """
        for host, cores in self.host_cores:
            if self.status(host)[0][1] == 'running':
                print('%s already running' % host)
                continue
            ctrl_url = 'tcp://%s:%s' % (host, self.ctrl_port)

            if host == '127.0.0.1':  # localhost
                args = [sys.executable]
            else:
                args = ['ssh', host, self.remote_python]
            args += ['-m', 'openquake.baselib.workerpool',
                     ctrl_url, self.task_out_url, cores]
            print('starting ' + ' '.join(args))
            subprocess.Popen(args)
        return '%s started' % ctrl_url

    def stop(self):
        """
        Send a "stop" command to all worker pools
        """
        for host, _ in self.host_cores:
            if self.status(host)[0][1] == 'not-running':
                print('%s not running' % host)
                continue
            ctrl_url = 'tcp://%s:%s' % (host, self.ctrl_port)
            with z.Socket(ctrl_url, z.zmq.REQ, 'connect') as sock:
                print(sock.send('stop'))
        return 'stopped'

    def kill(self):
        """
        Send a "kill" command to all worker pools
        """
        for host, _ in self.host_cores:
            if self.status(host)[0][1] == 'not-running':
                print('%s not running' % host)
                continue
            ctrl_url = 'tcp://%s:%s' % (host, self.ctrl_port)
            with z.Socket(ctrl_url, z.zmq.REQ, 'connect') as sock:
                print(sock.send('kill'))
        return 'killed'

class WorkerPool(object):
    """
    A pool of workers accepting the command 'stop' and 'kill' and reading
    tasks to perform from the task_out_url.

    :param ctrl_url: zmq address of the control socket
    :param task_out_url: zmq address of the task streamer
    :param num_workers: a string with the number of workers (or '-1')
    """
    def __init__(self, ctrl_url, task_out_url, num_workers='-1'):
        self.ctrl_url = ctrl_url
        self.task_out_url = task_out_url
        self.num_workers = (multiprocessing.cpu_count()
                            if num_workers == '-1' else int(num_workers))
        self.pid = os.getpid()

    def worker(self, sock):
        p.setproctitle('oq worker')
        for cmd, args in sock:
            backurl = args[-1].backurl  # attached to the monitor
            with z.Socket(backurl, z.zmq.PUSH, 'connect') as s:
                s.send(p.safely_call(cmd, args))

    def start(self):
        p.setproctitle('oq workerpool')
        # start workers
        self.workers = []
        for _ in range(self.num_workers):
            sock = z.Socket(self.task_out_url, z.zmq.PULL, 'connect')
            proc = multiprocessing.Process(target=self.worker, args=(sock,))
            proc.start()
            sock.pid = proc.pid
            self.workers.append(sock)

        # start control loop accepting the commands stop and kill
        ctrlsock = z.Socket(self.ctrl_url, z.zmq.REP, 'bind')
        for cmd in ctrlsock:
            if cmd in ('stop', 'kill'):
                msg = getattr(self, cmd)()
                ctrlsock.send(msg)
                break
            elif cmd == 'getpid':
                ctrlsock.send(self.pid)

    def stop(self):
        for sock in self.workers:
            os.kill(sock.pid, signal.SIGINT)
        return 'WorkerPool %s stopped' % self.ctrl_url

    def kill(self):
        for sock in self.workers:
            os.kill(sock.pid, signal.SIGTERM)
        return 'WorkerPool %s killed' % self.ctrl_url

if __name__ == '__main__':
    ctrl_url, task_out_url, num_workers = sys.argv[1:]
    WorkerPool(ctrl_url, task_out_url, num_workers).start()
<commit_msg>Improved process title [skip CI]<commit_after>import os
import sys
import signal
import socket
import subprocess
import multiprocessing
from openquake.baselib import zeromq as z, parallel as p

def streamer(task_in_url, task_out_url):
    """A streamer thread for the zmq workers"""
    try:
        z.zmq.proxy(z.bind(task_in_url, z.zmq.PULL),
                    z.bind(task_out_url, z.zmq.PUSH))
    except (KeyboardInterrupt, z.zmq.ZMQError):
        pass

def _starmap(func, iterargs, task_in_url, receiver_url):
    # called by parallel.Starmap; should not be used directly
    with z.Socket(receiver_url, z.zmq.PULL, 'bind') as receiver, \
            z.Socket(task_in_url, z.zmq.PUSH, 'connect') as sender:
        n = 0
        for args in iterargs:
            # args[-1] is a Monitor instance
            args[-1].backurl = receiver.true_end_point
            sender.send((func, args))
            n += 1
        yield n
        for _ in range(n):
            # receive n responses for the n requests sent
            yield receiver.zsocket.recv_pyobj()

class WorkerMaster(object):
    """
    :param frontend_url: url where to send the tasks
    :param task_out_url: url with a range of ports to receive the results
    :param ctrl_port: port on which the worker pools listen
    :param host_cores: names of the remote hosts and number of cores
    :param remote_python: path of the Python executable on the remote hosts
    """
    def __init__(self, task_in_url, task_out_url,
                 ctrl_port, host_cores, remote_python=None):
        self.task_in_url = task_in_url
        self.task_out_url = task_out_url
        self.ctrl_port = int(ctrl_port)
        self.host_cores = [hc.split() for hc in host_cores.split(',')]
        self.remote_python = remote_python or sys.executable

    def status(self, host=None):
        """
        :returns: a list of pairs (hostname, 'running'|'not-running')
        """
        if host is None:
            host_cores = self.host_cores
        else:
            host_cores = [hc for hc in self.host_cores if hc[0] == host]
        lst = []
        for host, _ in host_cores:
            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            try:
                err = sock.connect_ex((host, self.ctrl_port))
            finally:
                sock.close()
            lst.append((host, 'not-running' if err else 'running'))
        return lst

    def start(self):
        """
        Start multiple workerpools, possibly on remote servers via ssh
        """
        for host, cores in self.host_cores:
            if self.status(host)[0][1] == 'running':
                print('%s already running' % host)
                continue
            ctrl_url = 'tcp://%s:%s' % (host, self.ctrl_port)

            if host == '127.0.0.1':  # localhost
                args = [sys.executable]
            else:
                args = ['ssh', host, self.remote_python]
            args += ['-m', 'openquake.baselib.workerpool',
                     ctrl_url, self.task_out_url, cores]
            print('starting ' + ' '.join(args))
            subprocess.Popen(args)
        return '%s started' % ctrl_url

    def stop(self):
        """
        Send a "stop" command to all worker pools
        """
        for host, _ in self.host_cores:
            if self.status(host)[0][1] == 'not-running':
                print('%s not running' % host)
                continue
            ctrl_url = 'tcp://%s:%s' % (host, self.ctrl_port)
            with z.Socket(ctrl_url, z.zmq.REQ, 'connect') as sock:
                print(sock.send('stop'))
        return 'stopped'

    def kill(self):
        """
        Send a "kill" command to all worker pools
        """
        for host, _ in self.host_cores:
            if self.status(host)[0][1] == 'not-running':
                print('%s not running' % host)
                continue
            ctrl_url = 'tcp://%s:%s' % (host, self.ctrl_port)
            with z.Socket(ctrl_url, z.zmq.REQ, 'connect') as sock:
                print(sock.send('kill'))
        return 'killed'

class WorkerPool(object):
    """
    A pool of workers accepting the command 'stop' and 'kill' and reading
    tasks to perform from the task_out_url.

    :param ctrl_url: zmq address of the control socket
    :param task_out_url: zmq address of the task streamer
    :param num_workers: a string with the number of workers (or '-1')
    """
    def __init__(self, ctrl_url, task_out_url, num_workers='-1'):
        self.ctrl_url = ctrl_url
        self.task_out_url = task_out_url
        self.num_workers = (multiprocessing.cpu_count()
                            if num_workers == '-1' else int(num_workers))
        self.pid = os.getpid()

    def worker(self, sock):
        p.setproctitle('oq worker')
        for cmd, args in sock:
            backurl = args[-1].backurl  # attached to the monitor
            with z.Socket(backurl, z.zmq.PUSH, 'connect') as s:
                s.send(p.safely_call(cmd, args))

    def start(self):
        p.setproctitle('oq workerpool %s' % self.ctrl_url[6:])  # strip tcp://
        # start workers
        self.workers = []
        for _ in range(self.num_workers):
            sock = z.Socket(self.task_out_url, z.zmq.PULL, 'connect')
            proc = multiprocessing.Process(target=self.worker, args=(sock,))
            proc.start()
            sock.pid = proc.pid
            self.workers.append(sock)

        # start control loop accepting the commands stop and kill
        ctrlsock = z.Socket(self.ctrl_url, z.zmq.REP, 'bind')
        for cmd in ctrlsock:
            if cmd in ('stop', 'kill'):
                msg = getattr(self, cmd)()
                ctrlsock.send(msg)
                break
            elif cmd == 'getpid':
                ctrlsock.send(self.pid)

    def stop(self):
        for sock in self.workers:
            os.kill(sock.pid, signal.SIGINT)
        return 'WorkerPool %s stopped' % self.ctrl_url

    def kill(self):
        for sock in self.workers:
            os.kill(sock.pid, signal.SIGTERM)
        return 'WorkerPool %s killed' % self.ctrl_url

if __name__ == '__main__':
    ctrl_url, task_out_url, num_workers = sys.argv[1:]
    WorkerPool(ctrl_url, task_out_url, num_workers).start()
<|endoftext|>