<commit_before>'''
ARCHES - a program developed to inventory and manage immovable cultural heritage.
Copyright (C) 2013 J. Paul Getty Trust and World Monuments Fund

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU Affero General Public License as
published by the Free Software Foundation, either version 3 of the
License, or (at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
GNU Affero General Public License for more details.

You should have received a copy of the GNU Affero General Public License
along with this program. If not, see <http://www.gnu.org/licenses/>.
'''

from django.conf import settings
from django.http import HttpResponseNotFound
from django.shortcuts import redirect, render
from django.utils.decorators import method_decorator
from django.utils.translation import ugettext as _
from django.views.generic import View
from arches.app.models import models
from arches.app.models.forms import Form
from arches.app.models.card import Card
from arches.app.models.graph import Graph
from arches.app.models.resource import Resource
from arches.app.views.base import BaseManagerView
from arches.app.utils.decorators import group_required
from arches.app.utils.betterJSONSerializer import JSONSerializer, JSONDeserializer
from arches.app.utils.JSONResponse import JSONResponse
from arches.app.search.search_engine_factory import SearchEngineFactory
from arches.app.search.elasticsearch_dsl_builder import Query, Terms
from arches.app.views.concept import Concept
from elasticsearch import Elasticsearch

@method_decorator(group_required('Resource Editor'), name='dispatch')
class ResourceListView(BaseManagerView):
    def get(self, request, graphid=None, resourceid=None):
        context = self.get_context_data(
            main_script='views/resource',
        )

        context['nav']['icon'] = "fa fa-bookmark"
        context['nav']['title'] = "Resource Manager"
        context['nav']['edit_history'] = True
        context['nav']['login'] = True
        context['nav']['help'] = (_('Creating and Editing Resources'),'')

        return render(request, 'views/resource.htm', context)

@method_decorator(group_required('Resource Editor'), name='dispatch')
class ResourceEditorView(BaseManagerView):
    def get(self, request, graphid=None, resourceid=None):
        if graphid is not None:
            # self.graph = Graph.objects.get(graphid=graphid)
            resource_instance = Resource.objects.create(graph_id=graphid)
            resource_instance.index()
            return redirect('resource_editor', resourceid=resource_instance.pk)
        if resourceid is not None:
            resource_instance = models.ResourceInstance.objects.get(pk=resourceid)
            resource_graphs = Graph.objects.exclude(pk=settings.SYSTEM_SETTINGS_RESOURCE_MODEL_ID).exclude(isresource=False).exclude(isactive=False)
            graph = Graph.objects.get(graphid=resource_instance.graph.pk)
            resource_relationship_types = Concept().get_child_concepts('00000000-0000-0000-0000-000000000005', ['member', 'hasTopConcept'], ['prefLabel'], 'prefLabel')
            default_relationshiptype_valueid = None
            for relationship_type in resource_relationship_types:
                if relationship_type[1] == '00000000-0000-0000-0000-000000000007':
                    default_relationshiptype_valueid = relationship_type[5]
            relationship_type_values = {'values':[{'id':str(c[5]), 'text':str(c[3])} for c in resource_relationship_types], 'default': str(default_relationshiptype_valueid)}
            form = Form(resource_instance.pk)
            datatypes = models.DDataType.objects.all()
            widgets = models.Widget.objects.all()
            map_layers = models.MapLayers.objects.all()
            map_sources = models.MapSources.objects.all()
            forms = resource_instance.graph.form_set.filter(visible=True)
            forms_x_cards = models.FormXCard.objects.filter(form__in=forms)
            forms_w_cards = [form_x_card.form for form_x_card in forms_x_cards]
            displayname = Resource.objects.get(pk=resourceid).displayname
            if displayname == 'undefined':
                displayname = 'Unnamed Resource'

            context = self.get_context_data(
                main_script='views/resource/editor',
                resource_type=resource_instance.graph.name,
                relationship_types=relationship_type_values,
                iconclass=resource_instance.graph.iconclass,
                form=JSONSerializer().serialize(form),
                forms=JSONSerializer().serialize(forms_w_cards),
                datatypes_json=JSONSerializer().serialize(datatypes),
                widgets=widgets,
                map_layers=map_layers,
                map_sources=map_sources,
                widgets_json=JSONSerializer().serialize(widgets),
                resourceid=resourceid,
                resource_graphs=resource_graphs,
                graph_json=JSONSerializer().serialize(graph),
                displayname=displayname,
            )

            if graph.iconclass:
                context['nav']['icon'] = graph.iconclass
            context['nav']['title'] = graph.name
            context['nav']['menu'] = True
            context['nav']['edit_history'] = True
            context['nav']['help'] = (_('Creating and Editing Resources'),'')

            return render(request, 'views/resource/editor.htm', context)

        return HttpResponseNotFound()

    def delete(self, request, resourceid=None):
        if resourceid is not None:
            ret = Resource.objects.get(pk=resourceid)
            ret.delete()
            return JSONResponse(ret)
        return HttpResponseNotFound()

@method_decorator(group_required('Resource Editor'), name='dispatch')
class ResourceData(View):
    def get(self, request, resourceid=None, formid=None):
        if formid is not None:
            form = Form(resourceid=resourceid, formid=formid)
            return JSONResponse(form)

        return HttpResponseNotFound()

class ResourceDescriptors(View):
    def get(self, request, resourceid=None):
        if resourceid is not None:
            resource = Resource.objects.get(pk=resourceid)
            return JSONResponse({
                'graphid': resource.graph.pk,
                'graph_name': resource.graph.name,
                'displaydescription': resource.displaydescription,
                'map_popup': resource.map_popup,
                'displayname': resource.displayname,
            })

        return HttpResponseNotFound()

@method_decorator(group_required('Resource Editor'), name='dispatch')
class ResourceReportView(BaseManagerView):
    def get(self, request, resourceid=None):
        resource_instance = models.ResourceInstance.objects.get(pk=resourceid)
        tiles = models.TileModel.objects.filter(resourceinstance=resource_instance)
        try:
           report = models.Report.objects.get(graph=resource_instance.graph, active=True)
        except models.Report.DoesNotExist:
           report = None
        graph = Graph.objects.get(graphid=resource_instance.graph.pk)
        resource_graphs = Graph.objects.exclude(pk=settings.SYSTEM_SETTINGS_RESOURCE_MODEL_ID).exclude(isresource=False).exclude(isactive=False)
        forms = resource_instance.graph.form_set.filter(visible=True)
        forms_x_cards = models.FormXCard.objects.filter(form__in=forms).order_by('sortorder')
        cards = Card.objects.filter(nodegroup__parentnodegroup=None, graph=resource_instance.graph)
        datatypes = models.DDataType.objects.all()
        widgets = models.Widget.objects.all()
        map_layers = models.MapLayers.objects.all()
        map_sources = models.MapSources.objects.all()
        templates = models.ReportTemplate.objects.all()
        context = self.get_context_data(
            main_script='views/resource/report',
            report=JSONSerializer().serialize(report),
            report_templates=templates,
            templates_json=JSONSerializer().serialize(templates),
            forms=JSONSerializer().serialize(forms),
            tiles=JSONSerializer().serialize(tiles),
            forms_x_cards=JSONSerializer().serialize(forms_x_cards),
            cards=JSONSerializer().serialize(cards),
            datatypes_json=JSONSerializer().serialize(datatypes),
            widgets=widgets,
            map_layers=map_layers,
            map_sources=map_sources,
            resource_graphs=resource_graphs,
            graph_id=resource_instance.graph.pk,
            graph_name=resource_instance.graph.name,
            graph_json = JSONSerializer().serialize(graph),
            resourceid=resourceid,
         )

        if graph.iconclass:
            context['nav']['icon'] = graph.iconclass
        context['nav']['title'] = graph.name
        context['nav']['res_edit'] = True
        context['nav']['print'] = True
        context['nav']['print'] = True

        return render(request, 'views/resource/report.htm', context)

@method_decorator(group_required('Resource Editor'), name='dispatch')
class RelatedResourcesView(BaseManagerView):
    def get(self, request, resourceid=None):
        # lang = request.GET.get('lang', settings.LANGUAGE_CODE)
        start = request.GET.get('start', 0)
        resource = Resource.objects.get(pk=resourceid)
        related_resources = resource.get_related_resources(lang="en-US", start=start, limit=15)
        return JSONResponse(related_resources, indent=4)

    def delete(self, request, resourceid=None):
        es = Elasticsearch()
        se = SearchEngineFactory().create()
        req = dict(request.GET)
        ids_to_delete = req['resourcexids[]']
        root_resourceinstanceid = req['root_resourceinstanceid']
        for resourcexid in ids_to_delete:
            try:
                ret = models.ResourceXResource.objects.get(pk=resourcexid).delete()
            except:
                print 'resource relation does not exist'
        start = request.GET.get('start', 0)
        es.indices.refresh(index="resource_relations")
        resource = Resource.objects.get(pk=root_resourceinstanceid[0])
        related_resources = resource.get_related_resources(lang="en-US", start=start, limit=15)
        return JSONResponse(related_resources, indent=4)

    def post(self, request, resourceid=None):
        es = Elasticsearch()
        se = SearchEngineFactory().create()
        res = dict(request.POST)
        relationship_type = res['relationship_properties[relationship_type]'][0]
        datefrom = res['relationship_properties[datefrom]'][0]
        dateto = res['relationship_properties[dateto]'][0]
        dateto = None if dateto == '' else dateto
        datefrom = None if datefrom == '' else datefrom
        notes = res['relationship_properties[notes]'][0]
        root_resourceinstanceid = res['root_resourceinstanceid']
        instances_to_relate = []
        relationships_to_update = []
        if 'instances_to_relate[]' in res:
            instances_to_relate = res['instances_to_relate[]']
        if 'relationship_ids[]' in res:
            relationships_to_update = res['relationship_ids[]']

        def get_relatable_resources(graphid):
            """
            Takes the graphid of a resource, finds the graphs root node, and returns the relatable graphids
            """
            nodes = models.Node.objects.filter(graph_id=graphid)
            top_node = [node for node in nodes if node.istopnode==True][0]
            relatable_resources = [str(node.graph.graphid) for node in top_node.get_relatable_resources()]
            return relatable_resources

        def confirm_relationship_permitted(to_id, from_id):
            resource_instance_to = models.ResourceInstance.objects.filter(resourceinstanceid=to_id)[0]
            resource_instance_from = models.ResourceInstance.objects.filter(resourceinstanceid=from_id)[0]
            relatable_to = get_relatable_resources(resource_instance_to.graph.graphid)
            relatable_from = get_relatable_resources(resource_instance_from.graph.graphid)
            relatable_to_is_valid = str(resource_instance_to.graph.graphid) in relatable_from
            relatable_from_is_valid = str(resource_instance_from.graph.graphid) in relatable_to
            return (relatable_to_is_valid == True and relatable_from_is_valid == True)

        for instanceid in instances_to_relate:
            permitted = confirm_relationship_permitted(instanceid, root_resourceinstanceid[0])
            if permitted == True:
                rr = models.ResourceXResource(
                    resourceinstanceidfrom = Resource(root_resourceinstanceid[0]),
                    resourceinstanceidto = Resource(instanceid),
                    notes = notes,
                    relationshiptype = models.Value(relationship_type),
                    datestarted = datefrom,
                    dateended = dateto
                )
                rr.save()
            else:
                print 'relationship not permitted'

        for relationshipid in relationships_to_update:
            rr = models.ResourceXResource.objects.get(pk=relationshipid)
            rr.notes = notes
            rr.relationshiptype = models.Value(relationship_type)
            rr.datestarted = datefrom
            rr.dateended = dateto
            rr.save()

        start = request.GET.get('start', 0)
        es.indices.refresh(index="resource_relations")
        resource = Resource.objects.get(pk=root_resourceinstanceid[0])
        related_resources = resource.get_related_resources(lang="en-US", start=start, limit=15)
        return JSONResponse(related_resources, indent=4)
<commit_msg>type drop down in related resources time filter re: #1656<commit_after>'''
ARCHES - a program developed to inventory and manage immovable cultural heritage.
Copyright (C) 2013 J. Paul Getty Trust and World Monuments Fund

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU Affero General Public License as
published by the Free Software Foundation, either version 3 of the
License, or (at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
GNU Affero General Public License for more details.

You should have received a copy of the GNU Affero General Public License
along with this program. If not, see <http://www.gnu.org/licenses/>.
'''

from django.conf import settings
from django.http import HttpResponseNotFound
from django.shortcuts import redirect, render
from django.utils.decorators import method_decorator
from django.utils.translation import ugettext as _
from django.views.generic import View
from arches.app.models import models
from arches.app.models.forms import Form
from arches.app.models.card import Card
from arches.app.models.graph import Graph
from arches.app.models.resource import Resource
from arches.app.views.base import BaseManagerView
from arches.app.utils.decorators import group_required
from arches.app.utils.betterJSONSerializer import JSONSerializer, JSONDeserializer
from arches.app.utils.JSONResponse import JSONResponse
from arches.app.search.search_engine_factory import SearchEngineFactory
from arches.app.search.elasticsearch_dsl_builder import Query, Terms
from arches.app.views.concept import Concept
from elasticsearch import Elasticsearch

@method_decorator(group_required('Resource Editor'), name='dispatch')
class ResourceListView(BaseManagerView):
    def get(self, request, graphid=None, resourceid=None):
        context = self.get_context_data(
            main_script='views/resource',
        )

        context['nav']['icon'] = "fa fa-bookmark"
        context['nav']['title'] = "Resource Manager"
        context['nav']['edit_history'] = True
        context['nav']['login'] = True
        context['nav']['help'] = (_('Creating and Editing Resources'),'')

        return render(request, 'views/resource.htm', context)

@method_decorator(group_required('Resource Editor'), name='dispatch')
class ResourceEditorView(BaseManagerView):
    def get(self, request, graphid=None, resourceid=None):
        if graphid is not None:
            # self.graph = Graph.objects.get(graphid=graphid)
            resource_instance = Resource.objects.create(graph_id=graphid)
            resource_instance.index()
            return redirect('resource_editor', resourceid=resource_instance.pk)
        if resourceid is not None:
            resource_instance = models.ResourceInstance.objects.get(pk=resourceid)
            resource_graphs = Graph.objects.exclude(pk=settings.SYSTEM_SETTINGS_RESOURCE_MODEL_ID).exclude(isresource=False).exclude(isactive=False)
            graph = Graph.objects.get(graphid=resource_instance.graph.pk)
            resource_relationship_types = Concept().get_child_concepts('00000000-0000-0000-0000-000000000005', ['member', 'hasTopConcept'], ['prefLabel'], 'prefLabel')
            default_relationshiptype_valueid = None
            for relationship_type in resource_relationship_types:
                if relationship_type[1] == '00000000-0000-0000-0000-000000000007':
                    default_relationshiptype_valueid = relationship_type[5]
            relationship_type_values = {'values':[{'id':str(c[5]), 'text':str(c[3])} for c in resource_relationship_types], 'default': str(default_relationshiptype_valueid)}
            form = Form(resource_instance.pk)
            datatypes = models.DDataType.objects.all()
            widgets = models.Widget.objects.all()
            map_layers = models.MapLayers.objects.all()
            map_sources = models.MapSources.objects.all()
            forms = resource_instance.graph.form_set.filter(visible=True)
            forms_x_cards = models.FormXCard.objects.filter(form__in=forms)
            forms_w_cards = [form_x_card.form for form_x_card in forms_x_cards]
            displayname = Resource.objects.get(pk=resourceid).displayname
            if displayname == 'undefined':
                displayname = 'Unnamed Resource'
            date_nodes = models.Node.objects.filter(datatype='date', graph__isresource=True, graph__isactive=True)
            context = self.get_context_data(
                main_script='views/resource/editor',
                resource_type=resource_instance.graph.name,
                relationship_types=relationship_type_values,
                iconclass=resource_instance.graph.iconclass,
                form=JSONSerializer().serialize(form),
                forms=JSONSerializer().serialize(forms_w_cards),
                datatypes_json=JSONSerializer().serialize(datatypes),
                widgets=widgets,
                date_nodes=date_nodes,
                map_layers=map_layers,
                map_sources=map_sources,
                widgets_json=JSONSerializer().serialize(widgets),
                resourceid=resourceid,
                resource_graphs=resource_graphs,
                graph_json=JSONSerializer().serialize(graph),
                displayname=displayname,
            )

            if graph.iconclass:
                context['nav']['icon'] = graph.iconclass
            context['nav']['title'] = graph.name
            context['nav']['menu'] = True
            context['nav']['edit_history'] = True
            context['nav']['help'] = (_('Creating and Editing Resources'),'')

            return render(request, 'views/resource/editor.htm', context)

        return HttpResponseNotFound()

    def delete(self, request, resourceid=None):
        if resourceid is not None:
            ret = Resource.objects.get(pk=resourceid)
            ret.delete()
            return JSONResponse(ret)
        return HttpResponseNotFound()

@method_decorator(group_required('Resource Editor'), name='dispatch')
class ResourceData(View):
    def get(self, request, resourceid=None, formid=None):
        if formid is not None:
            form = Form(resourceid=resourceid, formid=formid)
            return JSONResponse(form)

        return HttpResponseNotFound()

class ResourceDescriptors(View):
    def get(self, request, resourceid=None):
        if resourceid is not None:
            resource = Resource.objects.get(pk=resourceid)
            return JSONResponse({
                'graphid': resource.graph.pk,
                'graph_name': resource.graph.name,
                'displaydescription': resource.displaydescription,
                'map_popup': resource.map_popup,
                'displayname': resource.displayname,
            })

        return HttpResponseNotFound()

@method_decorator(group_required('Resource Editor'), name='dispatch')
class ResourceReportView(BaseManagerView):
    def get(self, request, resourceid=None):
        resource_instance = models.ResourceInstance.objects.get(pk=resourceid)
        tiles = models.TileModel.objects.filter(resourceinstance=resource_instance)
        try:
           report = models.Report.objects.get(graph=resource_instance.graph, active=True)
        except models.Report.DoesNotExist:
           report = None
        graph = Graph.objects.get(graphid=resource_instance.graph.pk)
        resource_graphs = Graph.objects.exclude(pk=settings.SYSTEM_SETTINGS_RESOURCE_MODEL_ID).exclude(isresource=False).exclude(isactive=False)
        forms = resource_instance.graph.form_set.filter(visible=True)
        forms_x_cards = models.FormXCard.objects.filter(form__in=forms).order_by('sortorder')
        cards = Card.objects.filter(nodegroup__parentnodegroup=None, graph=resource_instance.graph)
        datatypes = models.DDataType.objects.all()
        widgets = models.Widget.objects.all()
        map_layers = models.MapLayers.objects.all()
        map_sources = models.MapSources.objects.all()
        templates = models.ReportTemplate.objects.all()
        context = self.get_context_data(
            main_script='views/resource/report',
            report=JSONSerializer().serialize(report),
            report_templates=templates,
            templates_json=JSONSerializer().serialize(templates),
            forms=JSONSerializer().serialize(forms),
            tiles=JSONSerializer().serialize(tiles),
            forms_x_cards=JSONSerializer().serialize(forms_x_cards),
            cards=JSONSerializer().serialize(cards),
            datatypes_json=JSONSerializer().serialize(datatypes),
            widgets=widgets,
            map_layers=map_layers,
            map_sources=map_sources,
            resource_graphs=resource_graphs,
            graph_id=resource_instance.graph.pk,
            graph_name=resource_instance.graph.name,
            graph_json = JSONSerializer().serialize(graph),
            resourceid=resourceid,
         )

        if graph.iconclass:
            context['nav']['icon'] = graph.iconclass
        context['nav']['title'] = graph.name
        context['nav']['res_edit'] = True
        context['nav']['print'] = True
        context['nav']['print'] = True

        return render(request, 'views/resource/report.htm', context)

@method_decorator(group_required('Resource Editor'), name='dispatch')
class RelatedResourcesView(BaseManagerView):
    def get(self, request, resourceid=None):
        # lang = request.GET.get('lang', settings.LANGUAGE_CODE)
        start = request.GET.get('start', 0)
        resource = Resource.objects.get(pk=resourceid)
        related_resources = resource.get_related_resources(lang="en-US", start=start, limit=15)
        return JSONResponse(related_resources, indent=4)

    def delete(self, request, resourceid=None):
        es = Elasticsearch()
        se = SearchEngineFactory().create()
        req = dict(request.GET)
        ids_to_delete = req['resourcexids[]']
        root_resourceinstanceid = req['root_resourceinstanceid']
        for resourcexid in ids_to_delete:
            try:
                ret = models.ResourceXResource.objects.get(pk=resourcexid).delete()
            except:
                print 'resource relation does not exist'
        start = request.GET.get('start', 0)
        es.indices.refresh(index="resource_relations")
        resource = Resource.objects.get(pk=root_resourceinstanceid[0])
        related_resources = resource.get_related_resources(lang="en-US", start=start, limit=15)
        return JSONResponse(related_resources, indent=4)

    def post(self, request, resourceid=None):
        es = Elasticsearch()
        se = SearchEngineFactory().create()
        res = dict(request.POST)
        relationship_type = res['relationship_properties[relationship_type]'][0]
        datefrom = res['relationship_properties[datefrom]'][0]
        dateto = res['relationship_properties[dateto]'][0]
        dateto = None if dateto == '' else dateto
        datefrom = None if datefrom == '' else datefrom
        notes = res['relationship_properties[notes]'][0]
        root_resourceinstanceid = res['root_resourceinstanceid']
        instances_to_relate = []
        relationships_to_update = []
        if 'instances_to_relate[]' in res:
            instances_to_relate = res['instances_to_relate[]']
        if 'relationship_ids[]' in res:
            relationships_to_update = res['relationship_ids[]']

        def get_relatable_resources(graphid):
            """
            Takes the graphid of a resource, finds the graphs root node, and returns the relatable graphids
            """
            nodes = models.Node.objects.filter(graph_id=graphid)
            top_node = [node for node in nodes if node.istopnode==True][0]
            relatable_resources = [str(node.graph.graphid) for node in top_node.get_relatable_resources()]
            return relatable_resources

        def confirm_relationship_permitted(to_id, from_id):
            resource_instance_to = models.ResourceInstance.objects.filter(resourceinstanceid=to_id)[0]
            resource_instance_from = models.ResourceInstance.objects.filter(resourceinstanceid=from_id)[0]
            relatable_to = get_relatable_resources(resource_instance_to.graph.graphid)
            relatable_from = get_relatable_resources(resource_instance_from.graph.graphid)
            relatable_to_is_valid = str(resource_instance_to.graph.graphid) in relatable_from
            relatable_from_is_valid = str(resource_instance_from.graph.graphid) in relatable_to
            return (relatable_to_is_valid == True and relatable_from_is_valid == True)

        for instanceid in instances_to_relate:
            permitted = confirm_relationship_permitted(instanceid, root_resourceinstanceid[0])
            if permitted == True:
                rr = models.ResourceXResource(
                    resourceinstanceidfrom = Resource(root_resourceinstanceid[0]),
                    resourceinstanceidto = Resource(instanceid),
                    notes = notes,
                    relationshiptype = models.Value(relationship_type),
                    datestarted = datefrom,
                    dateended = dateto
                )
                rr.save()
            else:
                print 'relationship not permitted'

        for relationshipid in relationships_to_update:
            rr = models.ResourceXResource.objects.get(pk=relationshipid)
            rr.notes = notes
            rr.relationshiptype = models.Value(relationship_type)
            rr.datestarted = datefrom
            rr.dateended = dateto
            rr.save()

        start = request.GET.get('start', 0)
        es.indices.refresh(index="resource_relations")
        resource = Resource.objects.get(pk=root_resourceinstanceid[0])
        related_resources = resource.get_related_resources(lang="en-US", start=start, limit=15)
        return JSONResponse(related_resources, indent=4)
<|endoftext|>


<commit_before># Copyright 2014-2015 Ivan Kravets <me@ikravets.com>
#
# 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.

import json
import logging
from urllib import unquote

from bottle import Bottle, request, response

from platformio_api import api, config
from platformio_api.database import db_session
from platformio_api.exception import APIBadRequest, APINotFound

app = Bottle()
logger = logging.getLogger(__name__)

@app.hook("after_request")
def db_disconnect():
    db_session.close()

def finalize_json_response(handler, kwargs):
    assert issubclass(handler, api.APIBase)
    response.set_header("Access-Control-Allow-Origin",
                        config['API_CORS_ORIGIN'])
    response.set_header("Access-Control-Allow-Methods",
                        "GET, POST, PUT, DELETE, OPTIONS")
    response.set_header("Access-Control-Allow-Headers",
                        "Content-Type, Access-Control-Allow-Headers")
    response.set_header("content-type", "application/json")

    status = 200
    error = None
    result = None
    try:
        obj = handler(**kwargs)
        result = obj.get_result()
    except APIBadRequest as error:
        status = 400
    except APINotFound as error:
        status = 404
    except Exception as e:
        status = 500
        error = "Internal server error"
        logger.exception(e)

    if error:
        item = dict(
            status=status,
            title=str(error)
        )
        result = dict(errors=[item])

    response.status = status
    return json.dumps(result)

@app.route("/", method="OPTIONS")
def cors(request):
    """ Preflighted request """
    response.set_header("Access-Control-Allow-Origin",
                        config['API_CORS_ORIGIN'])
    response.set_header("Access-Control-Allow-Methods",
                        "GET, POST, PUT, DELETE, OPTIONS")
    response.set_header("Access-Control-Allow-Headers",
                        "Content-Type, Access-Control-Allow-Headers")
    return None

@app.route("/boards")
def boards():
    return finalize_json_response(api.BoardsAPI, {})

@app.route("/frameworks")
def frameworks():
    return finalize_json_response(api.FrameworksAPI, {})

@app.route("/packages")
def packages():
    if "PlatformIO/" in request.headers.get("User-Agent", ""):
        return finalize_json_response(api.PackagesManifestAPI, {})
    else:
        return finalize_json_response(api.PackagesAPI, {})

@app.route("/packages/manifest")
def packages_manifest():
    return finalize_json_response(api.PackagesManifestAPI, {})

@app.route("/platforms")
def platforms():
    return finalize_json_response(api.PlatformsAPI, {})

@app.route("/lib/search")
def lib_search():
    args = dict(
        query=unquote(request.query.query[:255]),
        page=int(request.query.page) if request.query.page else 0,
        # perpage=int(request.query.perpage) if request.query.perpage else 0
    )
    return finalize_json_response(api.LibSearchAPI, args)

# Uncomment the line below in order to enable the Solr search
# @app.route("/lib/search_v2")
def lib_search_solr():
    strict = request.query.strict
    if strict.lower() in ['0', 'false', 'off']:
        strict = False
    args = dict(
        query=unquote(request.query.query[:255]),
        page=int(request.query.page) if request.query.page else 0,
        strict=bool(strict),
        # perpage=int(request.query.perpage) if request.query.perpage else 0
    )
    return finalize_json_response(api.LibSearchSolrAPI, args)

@app.route("/lib/examples")
def lib_examples():
    args = dict(
        query=unquote(request.query.query[:255]),
        page=int(request.query.page) if request.query.page else 0,
        # perpage=int(request.query.perpage) if request.query.perpage else 0
    )
    return finalize_json_response(api.LibExamplesAPI, args)

@app.route("/lib/info/<id_>")
def lib_info(id_):
    return finalize_json_response(api.LibInfoAPI, dict(id_=id_))

@app.route("/lib/download/<id_:int>")
def lib_download(id_):
    args = dict(
        id_=id_,
        ip=request.remote_addr,
        version=request.query.version,
        ci="CI/1" in request.headers.get("User-Agent", "")
    )
    return finalize_json_response(api.LibDownloadAPI, args)

@app.route("/lib/version/<ids:re:\d+(,\d+)*>")
def lib_version(ids):
    ids = [int(i) for i in ids.split(",")[:50]]
    return finalize_json_response(api.LibVersionAPI, dict(ids=ids))

@app.route("/lib/register", method="POST")
def lib_register():
    return finalize_json_response(
        api.LibRegisterAPI, dict(conf_url=request.forms.get("config_url")))

@app.route("/lib/stats")
def lib_stats():
    return finalize_json_response(api.LibStatsAPI, {})
<commit_msg>Allow to list all versions for the library<commit_after># Copyright 2014-2015 Ivan Kravets <me@ikravets.com>
#
# 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.

import json
import logging
from urllib import unquote

from bottle import Bottle, request, response

from platformio_api import api, config
from platformio_api.database import db_session
from platformio_api.exception import APIBadRequest, APINotFound

app = Bottle()
logger = logging.getLogger(__name__)

@app.hook("after_request")
def db_disconnect():
    db_session.close()

def finalize_json_response(handler, kwargs):
    assert issubclass(handler, api.APIBase)
    response.set_header("Access-Control-Allow-Origin",
                        config['API_CORS_ORIGIN'])
    response.set_header("Access-Control-Allow-Methods",
                        "GET, POST, PUT, DELETE, OPTIONS")
    response.set_header("Access-Control-Allow-Headers",
                        "Content-Type, Access-Control-Allow-Headers")
    response.set_header("content-type", "application/json")

    status = 200
    error = None
    result = None
    try:
        obj = handler(**kwargs)
        result = obj.get_result()
    except APIBadRequest as error:
        status = 400
    except APINotFound as error:
        status = 404
    except Exception as e:
        status = 500
        error = "Internal server error"
        logger.exception(e)

    if error:
        item = dict(status=status, title=str(error))
        result = dict(errors=[item])

    response.status = status
    return json.dumps(result)

@app.route("/", method="OPTIONS")
def cors(request):
    """ Preflighted request """
    response.set_header("Access-Control-Allow-Origin",
                        config['API_CORS_ORIGIN'])
    response.set_header("Access-Control-Allow-Methods",
                        "GET, POST, PUT, DELETE, OPTIONS")
    response.set_header("Access-Control-Allow-Headers",
                        "Content-Type, Access-Control-Allow-Headers")
    return None

@app.route("/boards")
def boards():
    return finalize_json_response(api.BoardsAPI, {})

@app.route("/frameworks")
def frameworks():
    return finalize_json_response(api.FrameworksAPI, {})

@app.route("/packages")
def packages():
    if "PlatformIO/" in request.headers.get("User-Agent", ""):
        return finalize_json_response(api.PackagesManifestAPI, {})
    else:
        return finalize_json_response(api.PackagesAPI, {})

@app.route("/packages/manifest")
def packages_manifest():
    return finalize_json_response(api.PackagesManifestAPI, {})

@app.route("/platforms")
def platforms():
    return finalize_json_response(api.PlatformsAPI, {})

@app.route("/lib/search")
def lib_search():
    args = dict(
        query=unquote(request.query.query[:255]),
        page=int(request.query.page) if request.query.page else 0,
        # perpage=int(request.query.perpage) if request.query.perpage else 0
    )
    return finalize_json_response(api.LibSearchAPI, args)

# Uncomment the line below in order to enable the Solr search
# @app.route("/lib/search_v2")
def lib_search_solr():
    strict = request.query.strict
    if strict.lower() in ['0', 'false', 'off']:
        strict = False
    args = dict(
        query=unquote(request.query.query[:255]),
        page=int(request.query.page) if request.query.page else 0,
        strict=bool(strict),
        # perpage=int(request.query.perpage) if request.query.perpage else 0
    )
    return finalize_json_response(api.LibSearchSolrAPI, args)

@app.route("/lib/examples")
def lib_examples():
    args = dict(
        query=unquote(request.query.query[:255]),
        page=int(request.query.page) if request.query.page else 0,
        # perpage=int(request.query.perpage) if request.query.perpage else 0
    )
    return finalize_json_response(api.LibExamplesAPI, args)

@app.route("/lib/info/<id_>")
def lib_info(id_):
    return finalize_json_response(api.LibInfoAPI, dict(id_=id_))

@app.route("/lib/download/<id_:int>")
def lib_download(id_):
    args = dict(
        id_=id_,
        ip=request.remote_addr,
        version=request.query.version,
        ci="CI/1" in request.headers.get("User-Agent", ""))
    return finalize_json_response(api.LibDownloadAPI, args)

@app.route("/lib/versions/<id_:int>")
def lib_versions(id_):
    return finalize_json_response(api.LibVersionsAPI, dict(id_=id_))

# PlatformIO 2.0
@app.route("/lib/version/<ids:re:\d+(,\d+)*>")
def lib_version(ids):
    ids = [int(i) for i in ids.split(",")[:50]]
    return finalize_json_response(api.LibVersionAPI, dict(ids=ids))

@app.route("/lib/register", method="POST")
def lib_register():
    return finalize_json_response(
        api.LibRegisterAPI, dict(conf_url=request.forms.get("config_url")))

@app.route("/lib/stats")
def lib_stats():
    return finalize_json_response(api.LibStatsAPI, {})
<|endoftext|>


<commit_before>"""Plot posterior densities."""
from typing import Optional
from numbers import Number
import numpy as np
from scipy.stats import mode

from ..data import convert_to_dataset
from ..stats import hpd
from .kdeplot import plot_kde, _fast_kde
from .plot_utils import (
    xarray_var_iter,
    _scale_fig_size,
    make_label,
    default_grid,
    _create_axes_grid,
    get_coords,
    filter_plotters_list,
)
from ..utils import _var_names, format_sig_figs

def plot_posterior(
    data,
    var_names=None,
    coords=None,
    figsize=None,
    textsize=None,
    credible_interval=0.94,
    round_to: Optional[int] = None,
    point_estimate="mean",
    group="posterior",
    rope=None,
    ref_val=None,
    kind="kde",
    bw=4.5,
    bins=None,
    ax=None,
    **kwargs
):
    """Plot Posterior densities in the style of John K. Kruschke's book.

    Parameters
    ----------
    data : obj
        Any object that can be converted to an az.InferenceData object
        Refer to documentation of az.convert_to_dataset for details
    var_names : list of variable names
        Variables to be plotted, two variables are required.
    coords : mapping, optional
        Coordinates of var_names to be plotted. Passed to `Dataset.sel`
    figsize : tuple
        Figure size. If None it will be defined automatically.
    textsize: float
        Text size scaling factor for labels, titles and lines. If None it will be autoscaled based
        on figsize.
    credible_interval : float, optional
        Credible intervals. Defaults to 0.94. Use None to hide the credible interval
    round_to : int, optional
        Controls formatting of floats. Defaults to 2 or the integer part, whichever is bigger.
    point_estimate: str
        Must be in ('mode', 'mean', 'median', None)
    group : str, optional
        Specifies which InferenceData group should be plotted. Defaults to ‘posterior’.
    rope: tuple or dictionary of tuples
        Lower and upper values of the Region Of Practical Equivalence. If a list is provided, its
        length should match the number of variables.
    ref_val: float or dictionary of floats
        display the percentage below and above the values in ref_val. Must be None (default),
        a constant, a list or a dictionary like see an example below. If a list is provided, its
        length should match the number of variables.
    kind: str
        Type of plot to display (kde or hist) For discrete variables this argument is ignored and
        a histogram is always used.
    bw : float
        Bandwidth scaling factor for the KDE. Should be larger than 0. The higher this number the
        smoother the KDE will be. Defaults to 4.5 which is essentially the same as the Scott's rule
        of thumb (the default rule used by SciPy). Only works if `kind == kde`.
    bins : integer or sequence or 'auto', optional
        Controls the number of bins, accepts the same keywords `matplotlib.hist()` does. Only works
        if `kind == hist`. If None (default) it will use `auto` for continuous variables and
        `range(xmin, xmax + 1)` for discrete variables.
    ax : axes
        Matplotlib axes. Defaults to None.
    **kwargs
        Passed as-is to plt.hist() or plt.plot() function depending on the value of `kind`.

    Returns
    -------
    ax : matplotlib axes

    Examples
    --------
    Show a default kernel density plot following style of John Kruschke

    .. plot::
        :context: close-figs

        >>> import arviz as az
        >>> data = az.load_arviz_data('centered_eight')
        >>> az.plot_posterior(data)

    Plot subset variables by specifying variable name exactly

    .. plot::
        :context: close-figs

        >>> az.plot_posterior(data, var_names=['mu'])

    Plot Region of Practical Equivalence (rope) for all distributions

    .. plot::
        :context: close-figs

        >>> az.plot_posterior(data, var_names=['mu', 'theta'], rope=(-1, 1))

    Plot Region of Practical Equivalence for selected distributions

    .. plot::
        :context: close-figs

        >>> rope = {'mu': [{'rope': (-2, 2)}], 'theta': [{'school': 'Choate', 'rope': (2, 4)}]}
        >>> az.plot_posterior(data, var_names=['mu', 'theta'], rope=rope)

    Add reference lines

    .. plot::
        :context: close-figs

        >>> az.plot_posterior(data, var_names=['mu', 'theta'], ref_val=0)

    Show point estimate of distribution

    .. plot::
        :context: close-figs

        >>> az.plot_posterior(data, var_names=['mu', 'theta'], point_estimate='mode')

    Show reference values using variable names and coordinates

    .. plot::
        :context: close-figs

        >>> az.plot_posterior(data, ref_val= {"theta": [{"school": "Deerfield", "ref_val": 4},
                                                        {"school": "Choate", "ref_val": 3}]})

    Show reference values using a list

    .. plot::
        :context: close-figs

        >>> az.plot_posterior(data, ref_val=[1] + [5] * 8 + [1])

    Plot posterior as a histogram

    .. plot::
        :context: close-figs

        >>> az.plot_posterior(data, var_names=['mu'], kind='hist')

    Change size of credible interval

    .. plot::
        :context: close-figs

        >>> az.plot_posterior(data, var_names=['mu'], credible_interval=.75)
    """
    data = convert_to_dataset(data, group=group)
    var_names = _var_names(var_names, data)

    if coords is None:
        coords = {}

    plotters = filter_plotters_list(
        list(xarray_var_iter(get_coords(data, coords), var_names=var_names, combined=True)),
        "plot_posterior",
    )
    length_plotters = len(plotters)
    rows, cols = default_grid(length_plotters)

    (figsize, ax_labelsize, titlesize, xt_labelsize, _linewidth, _) = _scale_fig_size(
        figsize, textsize, rows, cols
    )
    kwargs.setdefault("linewidth", _linewidth)

    if ax is None:
        _, ax = _create_axes_grid(
            length_plotters, rows, cols, figsize=figsize, squeeze=False, constrained_layout=True
        )
    idx = 0
    for (var_name, selection, x), ax_ in zip(plotters, np.ravel(ax)):
        _plot_posterior_op(
            idx,
            x.flatten(),
            var_name,
            selection,
            ax=ax_,
            bw=bw,
            bins=bins,
            kind=kind,
            point_estimate=point_estimate,
            round_to=round_to,
            credible_interval=credible_interval,
            ref_val=ref_val,
            rope=rope,
            ax_labelsize=ax_labelsize,
            xt_labelsize=xt_labelsize,
            **kwargs
        )
        idx += 1
        ax_.set_title(make_label(var_name, selection), fontsize=titlesize, wrap=True)

    return ax

def _plot_posterior_op(
    idx,
    values,
    var_name,
    selection,
    ax,
    bw,
    linewidth,
    bins,
    kind,
    point_estimate,
    credible_interval,
    ref_val,
    rope,
    ax_labelsize,
    xt_labelsize,
    round_to: Optional[int] = None,
    **kwargs
):  # noqa: D202
    """Artist to draw posterior."""

    significant_fig_func = lambda v: format_sig_figs(v, default=round_to)

    def format_as_percent(x, round_to=0):
        return "{0:.{1:d}f}%".format(100 * x, round_to)

    def display_ref_val():
        if ref_val is None:
            return
        elif isinstance(ref_val, dict):
            val = None
            for sel in ref_val.get(var_name, []):
                if all(
                    k in selection and selection[k] == v for k, v in sel.items() if k != "ref_val"
                ):
                    val = sel["ref_val"]
                    break
            if val is None:
                return
        elif isinstance(ref_val, list):
            val = ref_val[idx]
        elif isinstance(ref_val, Number):
            val = ref_val
        else:
            raise ValueError(
                "Argument `ref_val` must be None, a constant, a list or a "
                'dictionary like {"var_name": [{"ref_val": ref_val}]}'
            )
        less_than_ref_probability = (values < val).mean()
        greater_than_ref_probability = (values >= val).mean()
        ref_in_posterior = "{} <{:g}< {}".format(
            format_as_percent(less_than_ref_probability, 1),
            val,
            format_as_percent(greater_than_ref_probability, 1),
        )
        ax.axvline(val, ymin=0.05, ymax=0.75, color="C1", lw=linewidth, alpha=0.65)
        ax.text(
            values.mean(),
            plot_height * 0.6,
            ref_in_posterior,
            size=ax_labelsize,
            color="C1",
            weight="semibold",
            horizontalalignment="center",
        )

    def display_rope():
        if rope is None:
            return
        elif isinstance(rope, dict):
            vals = None
            for sel in rope.get(var_name, []):
                # pylint: disable=line-too-long
                if all(k in selection and selection[k] == v for k, v in sel.items() if k != "rope"):
                    vals = sel["rope"]
                    break
            if vals is None:
                return
        elif len(rope) == 2:
            vals = rope
        else:
            raise ValueError(
                "Argument `rope` must be None, a dictionary like"
                '{"var_name": {"rope": (lo, hi)}}, or an'
                "iterable of length 2"
            )

        ax.plot(
            vals,
            (plot_height * 0.02, plot_height * 0.02),
            lw=linewidth * 5,
            color="C2",
            solid_capstyle="round",
            zorder=0,
            alpha=0.7,
        )
        text_props = {"size": ax_labelsize, "horizontalalignment": "center", "color": "C2"}
        ax.text(vals[0], plot_height * 0.2, vals[0], weight="semibold", **text_props)
        ax.text(vals[1], plot_height * 0.2, vals[1], weight="semibold", **text_props)

    def display_point_estimate():
        if not point_estimate:
            return
        if point_estimate not in ("mode", "mean", "median"):
            raise ValueError("Point Estimate should be in ('mode','mean','median')")
        if point_estimate == "mean":
            point_value = values.mean()
        elif point_estimate == "mode":
            if isinstance(values[0], float):
                density, lower, upper = _fast_kde(values, bw=bw)
                x = np.linspace(lower, upper, len(density))
                point_value = x[np.argmax(density)]
            else:
                point_value = mode(values)[0][0]
        elif point_estimate == "median":
            point_value = np.median(values)
        sig_figs = significant_fig_func(point_value)
        point_text = "{point_estimate}={point_value:.{sig_figs}g}".format(
            point_estimate=point_estimate, point_value=point_value, sig_figs=sig_figs
        )
        ax.text(
            point_value,
            plot_height * 0.8,
            point_text,
            size=ax_labelsize,
            horizontalalignment="center",
        )

    def display_hpd():
        # np.ndarray with 2 entries, min and max
        # pylint: disable=line-too-long
        hpd_intervals = hpd(values, credible_interval=credible_interval)  # type: np.ndarray

        def round_num(n: float) -> str:
            sig_figs = significant_fig_func(n)
            return "{n:.{sig_figs}g}".format(n=n, sig_figs=sig_figs)

        ax.plot(
            hpd_intervals,
            (plot_height * 0.02, plot_height * 0.02),
            lw=linewidth * 2,
            color="k",
            solid_capstyle="round",
        )
        ax.text(
            hpd_intervals[0],
            plot_height * 0.07,
            round_num(hpd_intervals[0]),
            size=ax_labelsize,
            horizontalalignment="center",
        )
        ax.text(
            hpd_intervals[1],
            plot_height * 0.07,
            round_num(hpd_intervals[1]),
            size=ax_labelsize,
            horizontalalignment="center",
        )
        ax.text(
            (hpd_intervals[0] + hpd_intervals[1]) / 2,
            plot_height * 0.3,
            format_as_percent(credible_interval) + " HPD",
            size=ax_labelsize,
            horizontalalignment="center",
        )

    def format_axes():
        ax.yaxis.set_ticks([])
        ax.spines["top"].set_visible(False)
        ax.spines["right"].set_visible(False)
        ax.spines["left"].set_visible(False)
        ax.spines["bottom"].set_visible(True)
        ax.xaxis.set_ticks_position("bottom")
        ax.tick_params(
            axis="x", direction="out", width=1, length=3, color="0.5", labelsize=xt_labelsize
        )
        ax.spines["bottom"].set_color("0.5")

    if kind == "kde" and values.dtype.kind == "f":
        kwargs.setdefault("linewidth", linewidth)
        plot_kde(
            values,
            bw=bw,
            fill_kwargs={"alpha": kwargs.pop("fill_alpha", 0)},
            plot_kwargs=kwargs,
            ax=ax,
            rug=False,
        )
    else:
        if bins is None:
            if values.dtype.kind == "i":
                xmin = values.min()
                xmax = values.max()
                bins = range(xmin, xmax + 2)
                ax.set_xlim(xmin - 0.5, xmax + 0.5)
            else:
                bins = "auto"
        kwargs.setdefault("align", "left")
        kwargs.setdefault("color", "C0")
        ax.hist(values, bins=bins, alpha=0.35, **kwargs)

    plot_height = ax.get_ylim()[1]

    format_axes()
    if credible_interval is not None:
        display_hpd()
    display_point_estimate()
    display_ref_val()
    display_rope()
<commit_msg>Fix warning in ref_val example of plot_posterior (#836)<commit_after>"""Plot posterior densities."""
from typing import Optional
from numbers import Number
import numpy as np
from scipy.stats import mode

from ..data import convert_to_dataset
from ..stats import hpd
from .kdeplot import plot_kde, _fast_kde
from .plot_utils import (
    xarray_var_iter,
    _scale_fig_size,
    make_label,
    default_grid,
    _create_axes_grid,
    get_coords,
    filter_plotters_list,
)
from ..utils import _var_names, format_sig_figs

def plot_posterior(
    data,
    var_names=None,
    coords=None,
    figsize=None,
    textsize=None,
    credible_interval=0.94,
    round_to: Optional[int] = None,
    point_estimate="mean",
    group="posterior",
    rope=None,
    ref_val=None,
    kind="kde",
    bw=4.5,
    bins=None,
    ax=None,
    **kwargs
):
    """Plot Posterior densities in the style of John K. Kruschke's book.

    Parameters
    ----------
    data : obj
        Any object that can be converted to an az.InferenceData object
        Refer to documentation of az.convert_to_dataset for details
    var_names : list of variable names
        Variables to be plotted, two variables are required.
    coords : mapping, optional
        Coordinates of var_names to be plotted. Passed to `Dataset.sel`
    figsize : tuple
        Figure size. If None it will be defined automatically.
    textsize: float
        Text size scaling factor for labels, titles and lines. If None it will be autoscaled based
        on figsize.
    credible_interval : float, optional
        Credible intervals. Defaults to 0.94. Use None to hide the credible interval
    round_to : int, optional
        Controls formatting of floats. Defaults to 2 or the integer part, whichever is bigger.
    point_estimate: str
        Must be in ('mode', 'mean', 'median', None)
    group : str, optional
        Specifies which InferenceData group should be plotted. Defaults to ‘posterior’.
    rope: tuple or dictionary of tuples
        Lower and upper values of the Region Of Practical Equivalence. If a list is provided, its
        length should match the number of variables.
    ref_val: float or dictionary of floats
        display the percentage below and above the values in ref_val. Must be None (default),
        a constant, a list or a dictionary like see an example below. If a list is provided, its
        length should match the number of variables.
    kind: str
        Type of plot to display (kde or hist) For discrete variables this argument is ignored and
        a histogram is always used.
    bw : float
        Bandwidth scaling factor for the KDE. Should be larger than 0. The higher this number the
        smoother the KDE will be. Defaults to 4.5 which is essentially the same as the Scott's rule
        of thumb (the default rule used by SciPy). Only works if `kind == kde`.
    bins : integer or sequence or 'auto', optional
        Controls the number of bins, accepts the same keywords `matplotlib.hist()` does. Only works
        if `kind == hist`. If None (default) it will use `auto` for continuous variables and
        `range(xmin, xmax + 1)` for discrete variables.
    ax : axes
        Matplotlib axes. Defaults to None.
    **kwargs
        Passed as-is to plt.hist() or plt.plot() function depending on the value of `kind`.

    Returns
    -------
    ax : matplotlib axes

    Examples
    --------
    Show a default kernel density plot following style of John Kruschke

    .. plot::
        :context: close-figs

        >>> import arviz as az
        >>> data = az.load_arviz_data('centered_eight')
        >>> az.plot_posterior(data)

    Plot subset variables by specifying variable name exactly

    .. plot::
        :context: close-figs

        >>> az.plot_posterior(data, var_names=['mu'])

    Plot Region of Practical Equivalence (rope) for all distributions

    .. plot::
        :context: close-figs

        >>> az.plot_posterior(data, var_names=['mu', 'theta'], rope=(-1, 1))

    Plot Region of Practical Equivalence for selected distributions

    .. plot::
        :context: close-figs

        >>> rope = {'mu': [{'rope': (-2, 2)}], 'theta': [{'school': 'Choate', 'rope': (2, 4)}]}
        >>> az.plot_posterior(data, var_names=['mu', 'theta'], rope=rope)

    Add reference lines

    .. plot::
        :context: close-figs

        >>> az.plot_posterior(data, var_names=['mu', 'theta'], ref_val=0)

    Show point estimate of distribution

    .. plot::
        :context: close-figs

        >>> az.plot_posterior(data, var_names=['mu', 'theta'], point_estimate='mode')

    Show reference values using variable names and coordinates

    .. plot::
        :context: close-figs

        >>> az.plot_posterior(data, ref_val= {"theta": [{"school": "Deerfield", "ref_val": 4},
        ...                                             {"school": "Choate", "ref_val": 3}]})

    Show reference values using a list

    .. plot::
        :context: close-figs

        >>> az.plot_posterior(data, ref_val=[1] + [5] * 8 + [1])

    Plot posterior as a histogram

    .. plot::
        :context: close-figs

        >>> az.plot_posterior(data, var_names=['mu'], kind='hist')

    Change size of credible interval

    .. plot::
        :context: close-figs

        >>> az.plot_posterior(data, var_names=['mu'], credible_interval=.75)
    """
    data = convert_to_dataset(data, group=group)
    var_names = _var_names(var_names, data)

    if coords is None:
        coords = {}

    plotters = filter_plotters_list(
        list(xarray_var_iter(get_coords(data, coords), var_names=var_names, combined=True)),
        "plot_posterior",
    )
    length_plotters = len(plotters)
    rows, cols = default_grid(length_plotters)

    (figsize, ax_labelsize, titlesize, xt_labelsize, _linewidth, _) = _scale_fig_size(
        figsize, textsize, rows, cols
    )
    kwargs.setdefault("linewidth", _linewidth)

    if ax is None:
        _, ax = _create_axes_grid(
            length_plotters, rows, cols, figsize=figsize, squeeze=False, constrained_layout=True
        )
    idx = 0
    for (var_name, selection, x), ax_ in zip(plotters, np.ravel(ax)):
        _plot_posterior_op(
            idx,
            x.flatten(),
            var_name,
            selection,
            ax=ax_,
            bw=bw,
            bins=bins,
            kind=kind,
            point_estimate=point_estimate,
            round_to=round_to,
            credible_interval=credible_interval,
            ref_val=ref_val,
            rope=rope,
            ax_labelsize=ax_labelsize,
            xt_labelsize=xt_labelsize,
            **kwargs
        )
        idx += 1
        ax_.set_title(make_label(var_name, selection), fontsize=titlesize, wrap=True)

    return ax

def _plot_posterior_op(
    idx,
    values,
    var_name,
    selection,
    ax,
    bw,
    linewidth,
    bins,
    kind,
    point_estimate,
    credible_interval,
    ref_val,
    rope,
    ax_labelsize,
    xt_labelsize,
    round_to: Optional[int] = None,
    **kwargs
):  # noqa: D202
    """Artist to draw posterior."""

    significant_fig_func = lambda v: format_sig_figs(v, default=round_to)

    def format_as_percent(x, round_to=0):
        return "{0:.{1:d}f}%".format(100 * x, round_to)

    def display_ref_val():
        if ref_val is None:
            return
        elif isinstance(ref_val, dict):
            val = None
            for sel in ref_val.get(var_name, []):
                if all(
                    k in selection and selection[k] == v for k, v in sel.items() if k != "ref_val"
                ):
                    val = sel["ref_val"]
                    break
            if val is None:
                return
        elif isinstance(ref_val, list):
            val = ref_val[idx]
        elif isinstance(ref_val, Number):
            val = ref_val
        else:
            raise ValueError(
                "Argument `ref_val` must be None, a constant, a list or a "
                'dictionary like {"var_name": [{"ref_val": ref_val}]}'
            )
        less_than_ref_probability = (values < val).mean()
        greater_than_ref_probability = (values >= val).mean()
        ref_in_posterior = "{} <{:g}< {}".format(
            format_as_percent(less_than_ref_probability, 1),
            val,
            format_as_percent(greater_than_ref_probability, 1),
        )
        ax.axvline(val, ymin=0.05, ymax=0.75, color="C1", lw=linewidth, alpha=0.65)
        ax.text(
            values.mean(),
            plot_height * 0.6,
            ref_in_posterior,
            size=ax_labelsize,
            color="C1",
            weight="semibold",
            horizontalalignment="center",
        )

    def display_rope():
        if rope is None:
            return
        elif isinstance(rope, dict):
            vals = None
            for sel in rope.get(var_name, []):
                # pylint: disable=line-too-long
                if all(k in selection and selection[k] == v for k, v in sel.items() if k != "rope"):
                    vals = sel["rope"]
                    break
            if vals is None:
                return
        elif len(rope) == 2:
            vals = rope
        else:
            raise ValueError(
                "Argument `rope` must be None, a dictionary like"
                '{"var_name": {"rope": (lo, hi)}}, or an'
                "iterable of length 2"
            )

        ax.plot(
            vals,
            (plot_height * 0.02, plot_height * 0.02),
            lw=linewidth * 5,
            color="C2",
            solid_capstyle="round",
            zorder=0,
            alpha=0.7,
        )
        text_props = {"size": ax_labelsize, "horizontalalignment": "center", "color": "C2"}
        ax.text(vals[0], plot_height * 0.2, vals[0], weight="semibold", **text_props)
        ax.text(vals[1], plot_height * 0.2, vals[1], weight="semibold", **text_props)

    def display_point_estimate():
        if not point_estimate:
            return
        if point_estimate not in ("mode", "mean", "median"):
            raise ValueError("Point Estimate should be in ('mode','mean','median')")
        if point_estimate == "mean":
            point_value = values.mean()
        elif point_estimate == "mode":
            if isinstance(values[0], float):
                density, lower, upper = _fast_kde(values, bw=bw)
                x = np.linspace(lower, upper, len(density))
                point_value = x[np.argmax(density)]
            else:
                point_value = mode(values)[0][0]
        elif point_estimate == "median":
            point_value = np.median(values)
        sig_figs = significant_fig_func(point_value)
        point_text = "{point_estimate}={point_value:.{sig_figs}g}".format(
            point_estimate=point_estimate, point_value=point_value, sig_figs=sig_figs
        )
        ax.text(
            point_value,
            plot_height * 0.8,
            point_text,
            size=ax_labelsize,
            horizontalalignment="center",
        )

    def display_hpd():
        # np.ndarray with 2 entries, min and max
        # pylint: disable=line-too-long
        hpd_intervals = hpd(values, credible_interval=credible_interval)  # type: np.ndarray

        def round_num(n: float) -> str:
            sig_figs = significant_fig_func(n)
            return "{n:.{sig_figs}g}".format(n=n, sig_figs=sig_figs)

        ax.plot(
            hpd_intervals,
            (plot_height * 0.02, plot_height * 0.02),
            lw=linewidth * 2,
            color="k",
            solid_capstyle="round",
        )
        ax.text(
            hpd_intervals[0],
            plot_height * 0.07,
            round_num(hpd_intervals[0]),
            size=ax_labelsize,
            horizontalalignment="center",
        )
        ax.text(
            hpd_intervals[1],
            plot_height * 0.07,
            round_num(hpd_intervals[1]),
            size=ax_labelsize,
            horizontalalignment="center",
        )
        ax.text(
            (hpd_intervals[0] + hpd_intervals[1]) / 2,
            plot_height * 0.3,
            format_as_percent(credible_interval) + " HPD",
            size=ax_labelsize,
            horizontalalignment="center",
        )

    def format_axes():
        ax.yaxis.set_ticks([])
        ax.spines["top"].set_visible(False)
        ax.spines["right"].set_visible(False)
        ax.spines["left"].set_visible(False)
        ax.spines["bottom"].set_visible(True)
        ax.xaxis.set_ticks_position("bottom")
        ax.tick_params(
            axis="x", direction="out", width=1, length=3, color="0.5", labelsize=xt_labelsize
        )
        ax.spines["bottom"].set_color("0.5")

    if kind == "kde" and values.dtype.kind == "f":
        kwargs.setdefault("linewidth", linewidth)
        plot_kde(
            values,
            bw=bw,
            fill_kwargs={"alpha": kwargs.pop("fill_alpha", 0)},
            plot_kwargs=kwargs,
            ax=ax,
            rug=False,
        )
    else:
        if bins is None:
            if values.dtype.kind == "i":
                xmin = values.min()
                xmax = values.max()
                bins = range(xmin, xmax + 2)
                ax.set_xlim(xmin - 0.5, xmax + 0.5)
            else:
                bins = "auto"
        kwargs.setdefault("align", "left")
        kwargs.setdefault("color", "C0")
        ax.hist(values, bins=bins, alpha=0.35, **kwargs)

    plot_height = ax.get_ylim()[1]

    format_axes()
    if credible_interval is not None:
        display_hpd()
    display_point_estimate()
    display_ref_val()
    display_rope()
<|endoftext|>


<commit_before>import os
import sys
import time

from six.moves import range
import argparse
import numpy
import posix_ipc
import zmq

try:
    from mpi4py import MPI
except ImportError:
    MPI = None

from util import (PlatoonError, PlatoonFail, mmap, launch_process,
                  op_to_mpi, dtype_to_mpi)

class Controller(object):
    """
    Abstract multi-process controller

    This class provides the necessary features to dispatch data mini-batches
    to workers and handle control requests. Using this class should be done
    by having another class inherit from it and override the method
    `handle_control()`.

    .. warning::

        Due to the underlying implementation it is a bad idea to
        attempt to do both in the same process, even on different
        threads.  This will suffer from interlock problems and may
        negate any speedup you could get from using multiple Workers.

        Because of this issue, the class may be split in the future.

    Parameters
    ----------
    port : int
        The port number to communicate over
    control_port : int
        The control port number.
    device_list : list of strings
        Contains device names in clique order (prefer ring topology)
    hwm : int
        High water mark (see pyzmq docs).

    """

    def __init__(self, control_port, port=None, hwm=10, experiment_name="",
                 local_size=0, device_list=list(), multinode=False):

        self._should_stop = False
        self._workers = set()
        self._need_init = True

        self._device_list = device_list
        self._local_size = local_size
        self._count_workers = 0

        # New interface: Multi node
        self._multinode = multinode
        if self._multinode:
            try:
                self._init_global_comm()
            except PlatoonError as exc:
                # TODO log error
                pass

        # New interface
        if experiment_name:
            logs_folder = os.path.join("PLATOON_LOGS", experiment_name, time.strftime("%Y-%m-%d_%H-%M"))
            os.makedirs(logs_folder)
            for i in range(self._local_size):
                p = launch_process(logs_folder, experiment_name, None, self.device_list[i], "worker")
                self._workers.add(p.pid)

        if port:
            self.init_data(port, hwm)

        self._init_control_socket(control_port)

        # Cleanup and init global lock and job_uid name ##
        self._job_uid = "platoon_{0}_{1}".format(
            os.path.basename(os.path.expanduser('~')), control_port)

        self._lock_name = "{}lock".format(self._job_uid)
        # The ExistentialError is apparently the only way to verify if
        # the semaphore/shared_memory exists.
        try:
            posix_ipc.unlink_semaphore(self._lock_name)
        except posix_ipc.ExistentialError:
            pass
        # Initializing lock
        self._lock = posix_ipc.Semaphore(self._lock_name, posix_ipc.O_CREAT,
                                         initial_value=1)

        self._shmrefs = dict()
        self.shared_buffers = dict()

    def init_data(self, port, hwm=10):
        """
        Initialize the mini-batch socket.

        This must be called before using :meth:`send_mb`.

        Parameters
        ----------
        port : int
            The port to listen on.
        hwm : int
            High water mark, see the pyzmq docs.

        """
        self.acontext = zmq.Context()
        self.asocket = self.acontext.socket(zmq.PUSH)
        self.asocket.set_hwm(hwm)
        self.asocket.bind('tcp://*:{}'.format(port))

    def _init_control_socket(self, port):
        """
        Initialize the control socket.

        This must be called before using :meth:`serve`.

        Parameters
        ----------
        port : int
            The port to listen on.

        """
        self.ccontext = zmq.Context()
        self.csocket = self.ccontext.socket(zmq.REP)
        self.csocket.bind('tcp://*:{}'.format(port))

    def send_mb(self, arrays):
        """
        Send a mini-batch over the socket.

        This function may block if arrays are being sent faster than
        the clients can handle.

        Parameters
        ----------
        arrays : list of ndarrays
            List of numpy.ndarray to send.  All arrays should be
            contiguous for better performance.

        """
        # The buffer protocol only works on contiguous arrays
        arrays = [numpy.ascontiguousarray(array) for array in arrays]
        headers = [numpy.lib.format.header_data_from_array_1_0(array)
                   for array in arrays]
        self.asocket.send_json(headers, zmq.SNDMORE)
        for array in arrays[:-1]:
            self.asocket.send(array, zmq.SNDMORE)
        self.asocket.send(arrays[-1])

    def handle_control(self, req, worker_id, req_info):
        """
        Re-implement or assign a handler to this function to do
        something with control messages.

        The replacement get one parameter which is the request and
        should return the response which must be a json-encodable
        object. Other code is responsible for handling decoding,
        encoding and the network.

        """
        raise NotImplementedError("Classes that "
                                  "inherit from Controller should probably override "
                                  "the method `handle_control()`")

    def _handle_base_control(self, req, worker_id, req_info):
        """
        This method handle base control commands.
        Those commands should not be used in the handle_control method.
        All base control commands should start with "platoon-".
        """
        response = None
        if req == "platoon-get_job_uid":
            response = self._job_uid

        elif req == "platoon-get_device":
            response = self._device_list[req_info['local_rank']]

        elif req == "platoon-need_init":
            response = self._need_init
            self._need_init = False

        elif req == "platoon-get_region_info":
            first = self.is_worker_first()  # See :ref:is_worker_first
            if first:
                self._region_id = b"platoon-" + req_info['region_id']
            response = dict()
            response['region_id'] = self._region_id
            response['region_size'] = self._local_size
            response['regional_rank'] = self._device_list.index(req_info['device'])
            response['multinode'] = self._multinode

        elif req == "platoon-init_new_shmem":
            first = self.is_worker_first()  # See :ref:is_worker_first
            if first:
                self._last_shmem_name = "platoon-{0}_{1}_buffer".format(self._job_uid,
                                                                        len(self.shared_buffers))
                try:
                    posix_ipc.unlink_shared_memory(self._last_shmem_name)
                except posix_ipc.ExistentialError:
                    pass

                size = req_info['size']
                try:
                    shmref = posix_ipc.SharedMemory(self._last_shmem_name,
                                                    posix_ipc.O_CREAT,
                                                    size=size)
                    shm = mmap(fd=self._last_shmref.fd, length=size)
                    shmref.close_fd()
                except Exception as exc:
                    try:
                        shm.unlink()
                    except posix_ipc.ExistentialError:
                        pass
                    raise PlatoonFail()
                # We want every worker to get the same shared memory name that is
                # was declared in the first call of a mass request to this
                # controller for initializing a new shared memory.
                self._shmrefs[self._last_shmem_name] = shmref
                # Keep for unlinking when closing
                self.shared_buffers[self._last_shmem_name] = shm
            response = self._last_shmem_name

        elif req == "platoon-am_i_first":
            response = self.is_worker_first()  # See :ref:is_worker_first

        elif req == "platoon-all_reduce":
            if not self._multinode:
                raise PlatoonFail()
            if MPI is None:
                raise PlatoonFail()
            dtype = req_info['dtype']
            op = req_info['op']
            array = self.shared_buffers[req_info['shmem']]
            try:
                mpi_op = op_to_mpi(op)
                mpi_dtype = dtype_to_mpi(dtype)
                self._global_comm.Allreduce([array, mpi_dtype], [array, mpi_dtype],
                                            op=mpi_op)
            except Exception as exc:
                raise PlatoonFail()

        return response

    def is_worker_first(self):
        """Returns True, if in a mass request in a local platoon (workers in a
        single host) a worker's request reaches first its controller

        This will work only if every single worker participates successfully each
        time in a concurrent request of the same type to their controller.
        """
        self._count_workers = (self._count_workers + 1) % self._local_size
        if self._count_workers == 1:
            return True
        return False

    #  def worker_is_done(self, worker_id):
    #      self._workers.discard(worker_id)
    #      self._should_stop = True

    def serve(self):
        """
        This method will handle control messages until the should_stop flag
        has been raised and that all the known worker are done.
        """
        # spin spin spin
        try:
            #  while (not self._should_stop) or self._workers:
            while self._workers:  # spin while we have still children to watch for
                try:
                    pid, status = os.waitpid(-1, os.WNOHANG)
                except OSError as exc:
                    raise PlatoonFail()
                if pid != 0:  # If a status change has happened at a child
                    if os.WIFEXITED(status):
                        self._workers.discard(pid)
                        self._success = os.WEXITSTATUS(status)
                        if self._success == 0:
                            # A worker has terminated normally. Other workers
                            # are expected to terminate normally too, so
                            # continue.
                            continue
                        else:
                            # A worker has not terminated normally due to an
                            # error or an irrecovable fault
                            raise PlatoonFail()
                    else:  # other status changes are not desirable
                        raise PlatoonFail()
                try:
                    query = self.csocket.recv_json(flags=zmq.NOBLOCK)
                except zmq.Again:  # if a query has not happened, try again
                    continue

                # try default interface, it may raise PlatoonFail
                response = self._handle_base_control(query['req'],
                                                     query['worker_id'],
                                                     query['req_info'])
                if response is None:
                    response = self.handle_control(query['req'],
                                                   query['worker_id'],
                                                   query['req_info'])

                self.csocket.send_json(response)
        except PlatoonFail:  # if platoon fails kill all children workers
            self._kill_workers()
        finally:  # Close sockets and unlink for shared memory
            self._close()
        return self._success

    def _init_global_comm(self):
        if MPI is None:
            raise PlatoonError("mpi4py is not imported")
        self._global_comm = MPI.COMM_WORLD
        self._global_size = MPI.COMM_WORLD.Get_size()
        self._global_rank = MPI.COMM_WORLD.Get_rank()

    def _kill_workers(self):
        import signal
        while self._workers:
            pid = self._workers.pop()
            os.kill(pid, signal.SIGINT)
            pid, status = os.waitpid(pid, 0)

    def _close(self):
        self.csocket.close()
        self.ccontext.term()
        if hasattr(self, 'asocket'):
            self.asocket.close()
            self.acontext.term()
        self._lock.close()
        try:
            self._lock.unlink()
        except posix_ipc.ExistentialError:
            pass
        for shmref in self._shmrefs.values():
            try:
                shmref.unlink()
            except posix_ipc.ExistentialError:
                pass

def parse_arguments():
    parser = argparse.ArgumentParser(
        description="Base Platoon Controller process. Reigns over a computer node.")
    parser.add_argument('experiment_name', help='The name of your experiment. The launcher will expect to find the files <experiment_name>_worker.py and optionally <experiment_name>_controller.py.')
    single_or_multi = parser.add_mutually_exclusive_group(required=True)
    single_or_multi.add_argument('--single', action='store_true',
                                 help='Indicates that this Controller participates in a single-node platoon.')
    single_or_multi.add_argument('--multi', action='store_false',
                                 help='Indicates that this Controller participates in a multi-node platoon. Requires mpi4py')
    parser.add_argument('-D', '--devices', nargs='+', type=str, metavar='devname',
                        required=False, help='List of Theano device names (e.g. gpu0 or cuda1). Each device will be assigned to a separate worker. If this option is specified, experiment will be run in a single node.')
    parser.add_argument('-nw', '--workers', metavar='num_of_workers',
                        required=True, help='Number of workers spawned by this controller for this host.')

    return parser.parse_args()

def spawn_controller():
    args = parse_arguments()
    if args.single and args.devices is not None:
        devices = args.devices
    else:
        from pygpu import gpuarray as ga
        devcount = ga.get_device_count("cuda", 0)
        devices = ["cuda" + str(i) for i in range(devcount)]
        # TODO search for platoonrc or PLATOON_FLAGS
    if args.workers > len(devices):
        print("\nWARNING! Given {0} workers but {1} given devices. Using {1} workers.".format(args.workers, len(devices)))
        workers = len(devices)
    else:
        workers = args.workers

    controller = Controller(control_port=5567,
                            experiment_name=args.experiment_name,
                            local_size=workers,
                            device_list=devices,
                            multinode=not args.single)
    return controller.serve()

if __name__ == '__main__':
    spawn_controller()
<commit_msg>Support killing gracefully controllers in MPI and their workers<commit_after>import os
import sys
import signal
import time

from six.moves import range
import argparse
import numpy
import posix_ipc
import zmq

try:
    from mpi4py import MPI
except ImportError:
    MPI = None

from util import (PlatoonError, PlatoonFail, mmap, launch_process,
                  op_to_mpi, dtype_to_mpi)

class Controller(object):
    """
    Abstract multi-process controller

    This class provides the necessary features to dispatch data mini-batches
    to workers and handle control requests. Using this class should be done
    by having another class inherit from it and override the method
    `handle_control()`.

    .. warning::

        Due to the underlying implementation it is a bad idea to
        attempt to do both in the same process, even on different
        threads.  This will suffer from interlock problems and may
        negate any speedup you could get from using multiple Workers.

        Because of this issue, the class may be split in the future.

    Parameters
    ----------
    port : int
        The port number to communicate over
    control_port : int
        The control port number.
    device_list : list of strings
        Contains device names in clique order (prefer ring topology)
    hwm : int
        High water mark (see pyzmq docs).

    """

    def __init__(self, control_port, port=None, hwm=10, experiment_name="",
                 local_size=0, device_list=list(), multinode=False):
        self._should_stop = False
        self._workers = set()
        self._need_init = True

        self._device_list = device_list
        self._local_size = local_size
        self._count_workers = 0

        # New interface: Multi node
        self._multinode = multinode
        if self._multinode:
            try:
                self._init_global_comm()
                signal.signal(signal.SIGTERM, self._handle_force_close)
            except PlatoonError as exc:
                # TODO log error
                pass

        # New interface
        if experiment_name:
            logs_folder = os.path.join("PLATOON_LOGS", experiment_name, time.strftime("%Y-%m-%d_%H-%M"))
            os.makedirs(logs_folder)
            for i in range(self._local_size):
                p = launch_process(logs_folder, experiment_name, None, self.device_list[i], "worker")
                self._workers.add(p.pid)

        if port:
            self.init_data(port, hwm)

        self._init_control_socket(control_port)

        # Cleanup and init global lock and job_uid name ##
        self._job_uid = "platoon_{0}_{1}".format(
            os.path.basename(os.path.expanduser('~')), control_port)

        self._lock_name = "{}lock".format(self._job_uid)
        # The ExistentialError is apparently the only way to verify if
        # the semaphore/shared_memory exists.
        try:
            posix_ipc.unlink_semaphore(self._lock_name)
        except posix_ipc.ExistentialError:
            pass
        # Initializing lock
        self._lock = posix_ipc.Semaphore(self._lock_name, posix_ipc.O_CREAT,
                                         initial_value=1)

        self._shmrefs = dict()
        self.shared_buffers = dict()

    def init_data(self, port, hwm=10):
        """
        Initialize the mini-batch socket.

        This must be called before using :meth:`send_mb`.

        Parameters
        ----------
        port : int
            The port to listen on.
        hwm : int
            High water mark, see the pyzmq docs.

        """
        self.acontext = zmq.Context()
        self.asocket = self.acontext.socket(zmq.PUSH)
        self.asocket.set_hwm(hwm)
        self.asocket.bind('tcp://*:{}'.format(port))

    def _init_control_socket(self, port):
        """
        Initialize the control socket.

        This must be called before using :meth:`serve`.

        Parameters
        ----------
        port : int
            The port to listen on.

        """
        self.ccontext = zmq.Context()
        self.csocket = self.ccontext.socket(zmq.REP)
        self.csocket.bind('tcp://*:{}'.format(port))

    def send_mb(self, arrays):
        """
        Send a mini-batch over the socket.

        This function may block if arrays are being sent faster than
        the clients can handle.

        Parameters
        ----------
        arrays : list of ndarrays
            List of numpy.ndarray to send.  All arrays should be
            contiguous for better performance.

        """
        # The buffer protocol only works on contiguous arrays
        arrays = [numpy.ascontiguousarray(array) for array in arrays]
        headers = [numpy.lib.format.header_data_from_array_1_0(array)
                   for array in arrays]
        self.asocket.send_json(headers, zmq.SNDMORE)
        for array in arrays[:-1]:
            self.asocket.send(array, zmq.SNDMORE)
        self.asocket.send(arrays[-1])

    def handle_control(self, req, worker_id, req_info):
        """
        Re-implement or assign a handler to this function to do
        something with control messages.

        The replacement get one parameter which is the request and
        should return the response which must be a json-encodable
        object. Other code is responsible for handling decoding,
        encoding and the network.

        """
        raise NotImplementedError("Classes that "
                                  "inherit from Controller should probably override "
                                  "the method `handle_control()`")

    def _handle_base_control(self, req, worker_id, req_info):
        """
        This method handle base control commands.
        Those commands should not be used in the handle_control method.
        All base control commands should start with "platoon-".
        """
        response = None
        if req == "platoon-get_job_uid":
            response = self._job_uid

        elif req == "platoon-get_device":
            response = self._device_list[req_info['local_rank']]

        elif req == "platoon-need_init":
            response = self._need_init
            self._need_init = False

        elif req == "platoon-get_region_info":
            first = self.is_worker_first()  # See :ref:is_worker_first
            if first:
                self._region_id = b"platoon-" + req_info['region_id']
            response = dict()
            response['region_id'] = self._region_id
            response['region_size'] = self._local_size
            response['regional_rank'] = self._device_list.index(req_info['device'])
            response['multinode'] = self._multinode

        elif req == "platoon-init_new_shmem":
            first = self.is_worker_first()  # See :ref:is_worker_first
            if first:
                self._last_shmem_name = "platoon-{0}_{1}_buffer".format(self._job_uid,
                                                                        len(self.shared_buffers))
                try:
                    posix_ipc.unlink_shared_memory(self._last_shmem_name)
                except posix_ipc.ExistentialError:
                    pass

                size = req_info['size']
                try:
                    shmref = posix_ipc.SharedMemory(self._last_shmem_name,
                                                    posix_ipc.O_CREAT,
                                                    size=size)
                    shm = mmap(fd=self._last_shmref.fd, length=size)
                    shmref.close_fd()
                except Exception as exc:
                    try:
                        shm.unlink()
                    except posix_ipc.ExistentialError:
                        pass
                    raise PlatoonFail()
                # We want every worker to get the same shared memory name that is
                # was declared in the first call of a mass request to this
                # controller for initializing a new shared memory.
                self._shmrefs[self._last_shmem_name] = shmref
                # Keep for unlinking when closing
                self.shared_buffers[self._last_shmem_name] = shm
            response = self._last_shmem_name

        elif req == "platoon-am_i_first":
            response = self.is_worker_first()  # See :ref:is_worker_first

        elif req == "platoon-all_reduce":
            if not self._multinode:
                raise PlatoonFail()
            if MPI is None:
                raise PlatoonFail()
            dtype = req_info['dtype']
            op = req_info['op']
            array = self.shared_buffers[req_info['shmem']]
            try:
                mpi_op = op_to_mpi(op)
                mpi_dtype = dtype_to_mpi(dtype)
                self._global_comm.Allreduce([array, mpi_dtype], [array, mpi_dtype],
                                            op=mpi_op)
            except Exception as exc:
                raise PlatoonFail()

        return response

    def is_worker_first(self):
        """Returns True, if in a mass request in a local platoon (workers in a
        single host) a worker's request reaches first its controller

        This will work only if every single worker participates successfully each
        time in a concurrent request of the same type to their controller.
        """
        self._count_workers = (self._count_workers + 1) % self._local_size
        if self._count_workers == 1:
            return True
        return False

    #  def worker_is_done(self, worker_id):
    #      self._workers.discard(worker_id)
    #      self._should_stop = True

    def serve(self):
        """
        This method will handle control messages until the should_stop flag
        has been raised and that all the known worker are done.
        """
        # spin spin spin
        try:
            #  while (not self._should_stop) or self._workers:
            while self._workers:  # spin while we have still children to watch for
                try:
                    pid, status = os.waitpid(-1, os.WNOHANG)
                except OSError as exc:
                    raise PlatoonFail()
                if pid != 0:  # If a status change has happened at a child
                    if os.WIFEXITED(status):
                        self._workers.discard(pid)
                        self._success = os.WEXITSTATUS(status)
                        if self._success == 0:
                            # A worker has terminated normally. Other workers
                            # are expected to terminate normally too, so
                            # continue.
                            continue
                        else:
                            # A worker has not terminated normally due to an
                            # error or an irrecovable fault
                            raise PlatoonFail()
                    else:  # other status changes are not desirable
                        raise PlatoonFail()
                try:
                    query = self.csocket.recv_json(flags=zmq.NOBLOCK)
                except zmq.Again:  # if a query has not happened, try again
                    continue
                except zmq.ZMQError as exc:
                    raise PlatoonFail()

                # try default interface, it may raise PlatoonFail
                response = self._handle_base_control(query['req'],
                                                     query['worker_id'],
                                                     query['req_info'])
                if response is None:
                    response = self.handle_control(query['req'],
                                                   query['worker_id'],
                                                   query['req_info'])

                self.csocket.send_json(response)
        except PlatoonFail as exc:  # if platoon fails kill all children workers
            self._kill_workers()
            if self._multinode:
                self._global_comm.Abort(errorcode=-1)
        finally:  # Close sockets and unlink for shared memory
            self._close()
        return self._success

    def _init_global_comm(self):
        if MPI is None:
            raise PlatoonError("mpi4py is not imported")
        self._global_comm = MPI.COMM_WORLD
        self._global_size = MPI.COMM_WORLD.Get_size()
        self._global_rank = MPI.COMM_WORLD.Get_rank()

    def _kill_workers(self):
        while self._workers:
            pid = self._workers.pop()
            os.kill(pid, signal.SIGINT)
            pid, status = os.waitpid(pid, 0)

    def _close(self):
        self.csocket.close()
        self.ccontext.term()
        if hasattr(self, 'asocket'):
            self.asocket.close()
            self.acontext.term()
        self._lock.close()
        try:
            self._lock.unlink()
        except posix_ipc.ExistentialError:
            pass
        for shmref in self._shmrefs.values():
            try:
                shmref.unlink()
            except posix_ipc.ExistentialError:
                pass

    def _handle_force_close(self, signum, frame):
        """Handle SIGTERM signals from MPI.Abort

        This is expected to happen when something abnormal has happened in other
        controllers over MPI.COMM_WORLD across host which participate in the
        multi-node training.

        """
        self._kill_workers()
        self._close()
        sys.exit(-1)

def parse_arguments():
    parser = argparse.ArgumentParser(
        description="Base Platoon Controller process. Reigns over a computer node.")
    parser.add_argument('experiment_name', help='The name of your experiment. The launcher will expect to find the files <experiment_name>_worker.py and optionally <experiment_name>_controller.py.')
    single_or_multi = parser.add_mutually_exclusive_group(required=True)
    single_or_multi.add_argument('--single', action='store_true',
                                 help='Indicates that this Controller participates in a single-node platoon.')
    single_or_multi.add_argument('--multi', action='store_false',
                                 help='Indicates that this Controller participates in a multi-node platoon. Requires mpi4py')
    parser.add_argument('-D', '--devices', nargs='+', type=str, metavar='devname',
                        required=False, help='List of Theano device names (e.g. gpu0 or cuda1). Each device will be assigned to a separate worker. If this option is specified, experiment will be run in a single node.')
    parser.add_argument('-nw', '--workers', metavar='num_of_workers',
                        required=True, help='Number of workers spawned by this controller for this host.')

    return parser.parse_args()

def spawn_controller():
    args = parse_arguments()
    if args.single and args.devices is not None:
        devices = args.devices
    else:
        from pygpu import gpuarray as ga
        devcount = ga.get_device_count("cuda", 0)
        devices = ["cuda" + str(i) for i in range(devcount)]
        # TODO search for platoonrc or PLATOON_FLAGS
    if args.workers > len(devices):
        print("\nWARNING! Given {0} workers but {1} given devices. Using {1} workers.".format(args.workers, len(devices)))
        workers = len(devices)
    else:
        workers = args.workers

    controller = Controller(control_port=5567,
                            experiment_name=args.experiment_name,
                            local_size=workers,
                            device_list=devices,
                            multinode=not args.single)
    return controller.serve()

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