<commit_before>import os
from xml.etree.ElementTree import fromstring

from .vcs_base import find_executable, VcsClientBase

class SvnClient(VcsClientBase):

    type = 'svn'
    _executable = None

    @staticmethod
    def is_repository(path):
        return os.path.isdir(os.path.join(path, '.svn'))

    def __init__(self, path):
        super(SvnClient, self).__init__(path)

    def branch(self, _command):
        cmd_info = [SvnClient._executable, 'info', '--xml']
        result_info = self._run_command(cmd_info)
        if result_info['returncode']:
            result_info['output'] = 'Could not determine url: %s' % result_info['output']
            return result_info
        info = result_info['output']

        try:
            root = fromstring(info)
            entry = root.find('entry')
            url = entry.findtext('url')
            repository = entry.find('repository')
            root_url = repository.findtext('root')
        except Exception as e:
            return {
                'cmd': '',
                'cwd': self.path,
                'output': "Could not determine url from xml: %s" % e,
                'returncode': 1
            }

        if not url.startswith(root_url):
            return {
                'cmd': '',
                'cwd': self.path,
                'output': "Could not determine url suffix. The root url '%s' is not a prefix of the url '%s'" % (root_url, url),
                'returncode': 1
            }

        return {
            'cmd': ' '.join(cmd_info),
            'cwd': self.path,
            'output': url[len(root_url):],
            'returncode': 0,
        }

    def diff(self, command):
        cmd = [SvnClient._executable, 'diff']
        if command.context:
            cmd += ['--unified=%d' % command.context]
        return self._run_command(cmd)

    def export(self, _command):
        cmd_info = [SvnClient._executable, 'info', '--xml']
        result_info = self._run_command(cmd_info)
        if result_info['returncode']:
            result_info['output'] = 'Could not determine url: %s' % result_info['output']
            return result_info
        info = result_info['output']

        try:
            root = fromstring(info)
            entry = root.find('entry')
            url = entry.findtext('url')
        except Exception as e:
            return {
                'cmd': '',
                'cwd': self.path,
                'output': "Could not determine url from xml: %s" % e,
                'returncode': 1
            }

        return {
            'cmd': ' '.join(cmd_info),
            'cwd': self.path,
            'output': url,
            'returncode': 0,
            'export_data': {'url': url}
        }

    def import_(self, command):
        if not command.url:
            return {
                'cmd': '',
                'cwd': self.path,
                'output': "Repository data lacks the 'url' value",
                'returncode': 1
            }

        not_exist = self._create_path()
        if not_exist:
            return not_exist

        cmd_checkout = [SvnClient._executable, '--non-interactive', 'checkout', command.url, '.']
        result_checkout = self._run_command(cmd_checkout)
        if result_checkout['returncode']:
            result_checkout['output'] = "Could not checkout repository '%s': %s" % (command.url, result_checkout['output'])
            return result_checkout

        return {
            'cmd': ' '.join(cmd_checkout),
            'cwd': self.path,
            'output': result_checkout['output'],
            'returncode': 0
        }

    def log(self, command):
        cmd = [SvnClient._executable, 'log', '--limit', '%d' % command.limit]
        return self._run_command(cmd)

    def pull(self, _command):
        cmd = [SvnClient._executable, '--non-interactive', 'update']
        return cmd

    def remotes(self, _command):
        cmd_info = [SvnClient._executable, 'info', '--xml']
        result_info = self._run_command(cmd_info)
        if result_info['returncode']:
            result_info['output'] = 'Could not determine url: %s' % result_info['output']
            return result_info
        info = result_info['output']

        try:
            root = fromstring(info)
            entry = root.find('entry')
            url = entry.findtext('url')
        except Exception as e:
            return {
                'cmd': '',
                'cwd': self.path,
                'output': "Could not determine url from xml: %s" % e,
                'returncode': 1
            }

        return {
            'cmd': ' '.join(cmd_info),
            'cwd': self.path,
            'output': url,
            'returncode': 0,
        }

    def status(self, command):
        cmd = [SvnClient._executable, 'status']
        if command.quiet:
            cmd += ['--quiet']
        return self._run_command(cmd)

if not SvnClient._executable:
    SvnClient._executable = find_executable('svn')
<commit_msg>fix svn pull<commit_after>import os
from xml.etree.ElementTree import fromstring

from .vcs_base import find_executable, VcsClientBase

class SvnClient(VcsClientBase):

    type = 'svn'
    _executable = None

    @staticmethod
    def is_repository(path):
        return os.path.isdir(os.path.join(path, '.svn'))

    def __init__(self, path):
        super(SvnClient, self).__init__(path)

    def branch(self, _command):
        cmd_info = [SvnClient._executable, 'info', '--xml']
        result_info = self._run_command(cmd_info)
        if result_info['returncode']:
            result_info['output'] = 'Could not determine url: %s' % result_info['output']
            return result_info
        info = result_info['output']

        try:
            root = fromstring(info)
            entry = root.find('entry')
            url = entry.findtext('url')
            repository = entry.find('repository')
            root_url = repository.findtext('root')
        except Exception as e:
            return {
                'cmd': '',
                'cwd': self.path,
                'output': "Could not determine url from xml: %s" % e,
                'returncode': 1
            }

        if not url.startswith(root_url):
            return {
                'cmd': '',
                'cwd': self.path,
                'output': "Could not determine url suffix. The root url '%s' is not a prefix of the url '%s'" % (root_url, url),
                'returncode': 1
            }

        return {
            'cmd': ' '.join(cmd_info),
            'cwd': self.path,
            'output': url[len(root_url):],
            'returncode': 0,
        }

    def diff(self, command):
        cmd = [SvnClient._executable, 'diff']
        if command.context:
            cmd += ['--unified=%d' % command.context]
        return self._run_command(cmd)

    def export(self, _command):
        cmd_info = [SvnClient._executable, 'info', '--xml']
        result_info = self._run_command(cmd_info)
        if result_info['returncode']:
            result_info['output'] = 'Could not determine url: %s' % result_info['output']
            return result_info
        info = result_info['output']

        try:
            root = fromstring(info)
            entry = root.find('entry')
            url = entry.findtext('url')
        except Exception as e:
            return {
                'cmd': '',
                'cwd': self.path,
                'output': "Could not determine url from xml: %s" % e,
                'returncode': 1
            }

        return {
            'cmd': ' '.join(cmd_info),
            'cwd': self.path,
            'output': url,
            'returncode': 0,
            'export_data': {'url': url}
        }

    def import_(self, command):
        if not command.url:
            return {
                'cmd': '',
                'cwd': self.path,
                'output': "Repository data lacks the 'url' value",
                'returncode': 1
            }

        not_exist = self._create_path()
        if not_exist:
            return not_exist

        cmd_checkout = [SvnClient._executable, '--non-interactive', 'checkout', command.url, '.']
        result_checkout = self._run_command(cmd_checkout)
        if result_checkout['returncode']:
            result_checkout['output'] = "Could not checkout repository '%s': %s" % (command.url, result_checkout['output'])
            return result_checkout

        return {
            'cmd': ' '.join(cmd_checkout),
            'cwd': self.path,
            'output': result_checkout['output'],
            'returncode': 0
        }

    def log(self, command):
        cmd = [SvnClient._executable, 'log', '--limit', '%d' % command.limit]
        return self._run_command(cmd)

    def pull(self, _command):
        cmd = [SvnClient._executable, '--non-interactive', 'update']
        return self._run_command(cmd)

    def remotes(self, _command):
        cmd_info = [SvnClient._executable, 'info', '--xml']
        result_info = self._run_command(cmd_info)
        if result_info['returncode']:
            result_info['output'] = 'Could not determine url: %s' % result_info['output']
            return result_info
        info = result_info['output']

        try:
            root = fromstring(info)
            entry = root.find('entry')
            url = entry.findtext('url')
        except Exception as e:
            return {
                'cmd': '',
                'cwd': self.path,
                'output': "Could not determine url from xml: %s" % e,
                'returncode': 1
            }

        return {
            'cmd': ' '.join(cmd_info),
            'cwd': self.path,
            'output': url,
            'returncode': 0,
        }

    def status(self, command):
        cmd = [SvnClient._executable, 'status']
        if command.quiet:
            cmd += ['--quiet']
        return self._run_command(cmd)

if not SvnClient._executable:
    SvnClient._executable = find_executable('svn')
<|endoftext|>


<commit_before>"""
Database models of the Dashboard application
"""
import hashlib

from django import core
from django.contrib.auth.models import (User, Group)
from django.contrib.contenttypes import generic
from django.contrib.contenttypes.models import ContentType
from django.db import models
from django.utils.translation import ugettext as _
from django.utils.translation import ungettext

from dashboard_app import managers

def _help_max_length(max_length):
    return ungettext(
            u"Maximum length: {0} character",
            u"Maximum length: {0} characters",
            max_length).format(max_length)

class SoftwarePackage(models.Model):
    """
    Model for software packages.

    This class mirrors launch_control.models.SoftwarePackage.
    """
    name = models.CharField(
            max_length = 64,
            verbose_name = _(u"Package name"),
            help_text = _help_max_length(64))

    version = models.CharField(
            max_length = 32,
            verbose_name = _(u"Package version"),
            help_text = _help_max_length(32))

    class Meta:
        unique_together = (('name', 'version'))

    def __unicode__(self):
        return _(u"{name} {version}").format(
                name = self.name,
                version = self.version)

    @models.permalink
    def get_absolute_url(self):
        return ("dashboard_app.sw-package.detail", [self.name, self.version])

class NamedAttribute(models.Model):
    """
    Model for adding generic named attributes
    to arbitrary other model instances.

    This class mirrors launch_control.models.HardwareDevice.

    Example:
        class Foo(Model):
            attributes = generic.GenericRelation(NamedAttribute)
    """
    name = models.CharField(
            help_text = _help_max_length(32),
            max_length = 32)

    value = models.CharField(
            help_text = _help_max_length(256),
            max_length = 256)

    # Content type plumbing
    content_type = models.ForeignKey(ContentType)
    object_id = models.PositiveIntegerField()
    content_object = generic.GenericForeignKey('content_type', 'object_id')

    def __unicode__(self):
        return _(u"Attribute {name}: {value}").format(
                name = self.name,
                value = self.value)

    class Meta:
        unique_together = (('object_id', 'name'))

class HardwareDevice(models.Model):
    """
    Model for hardware devices

    All devices are simplified into an instance of pre-defined class
    with arbitrary key-value attributes.
    """
    device_type = models.CharField(
            choices = (
                (u"device.cpu", _(u"CPU")),
                (u"device.mem", _(u"Memory")),
                (u"device.usb", _(u"USB device")),
                (u"device.pci", _(u"PCI device")),
                (u"device.board", _(u"Board/Motherboard"))),
            help_text = _(u"One of pre-defined device types"),
            max_length = 32,
            verbose_name = _(u"Device Type"),
            )

    description = models.CharField(
            help_text = _(u"Human readable device summary.") + " " + _help_max_length(256),
            max_length = 256,
            verbose_name = _(u"Description."),
            )

    attributes = generic.GenericRelation(NamedAttribute)

    def __unicode__(self):
        return self.description

    @models.permalink
    def get_absolute_url(self):
        return ("dashboard_app.hw-device.detail", [self.pk])

class BundleStream(models.Model):
    """
    Model for "streams" of bundles.

    Basically it's a named collection of bundles, like directory just
    without the nesting. A simple ACL scheme is also supported,
    a bundle may be uploaded by:
        - specific user when user field is set
        - users of a specific group when group field is set
        - anyone when neither user nor group is set
    """
    user = models.ForeignKey(User,
            blank = True,
            help_text = _("User owning this stream (do not set when group is also set)"),
            null = True,
            verbose_name = _(u"User"),
            )

    group = models.ForeignKey(Group,
            blank = True,
            help_text = _("Group owning this stream (do not set when user is also set)"),
            null = True,
            verbose_name = _(u"Group"),
            )

    slug = models.CharField(
            blank = True,
            help_text = (_(u"Name that you will use when uploading bundles.")
                + " " + _help_max_length(64)),
            max_length = 64,
            verbose_name = _(u"Slug"),
            )

    name = models.CharField(
            blank = True,
            help_text = _help_max_length(64),
            max_length = 64,
            verbose_name = _(u"Name"),
            )

    pathname = models.CharField(
            max_length = 128,
            editable = False,
            unique = True,
            )

    objects = managers.BundleStreamManager()

    def __unicode__(self):
        return _(u"Bundle stream {pathname}").format(
                pathname = self.pathname)

    @models.permalink
    def get_absolute_url(self):
        return ("dashboard_app.bundle_stream_detail", [self.pathname])

    def save(self, *args, **kwargs):
        """
        Save this instance.

        Calls self.clean() to ensure that constraints are met.
        Updates pathname to reflect user/group/slug changes.
        """
        self.pathname = self._calc_pathname()
        self.clean()
        return super(BundleStream, self).save(*args, **kwargs)

    def clean(self):
        """
        Validate instance.

        Makes sure that user and name are not set at the same time
        """
        if self.user is not None and self.group is not None:
            raise core.exceptions.ValidationError('BundleStream cannot '
                    'have both user and name set at the same time')

    def can_access(self, user):
        """
        Returns true if given user can access the contents of this this
        stream.
        """
        if user is None:
            return self.user is None and self.group is None
        else:
            if self.user is not None:
                return self.user.username == user.username
            elif self.group is not None:
                return self.group in user.groups.all()
            else:
                return True

    def _calc_pathname(self):
        """
        Pseudo pathname-like ID of this stream.

        This pathname is user visible and will be presented to users
        when they want to interact with this bundle stream. The
        pathnames are unique and this is enforced at database level (the
        user and name are unique together).
        """
        if self.user is not None:
            if self.slug == "":
                return u"/personal/{user}/".format(
                        user = self.user.username)
            else:
                return u"/personal/{user}/{slug}/".format(
                        user = self.user.username,
                        slug = self.slug)
        elif self.group is not None:
            if self.slug == "":
                return u"/team/{group}/".format(
                        group = self.group.name)
            else:
                return u"/team/{group}/{slug}/".format(
                        group = self.group.name,
                        slug = self.slug)
        else:
            if self.slug == "":
                return u"/anonymous/"
            else:
                return u"/anonymous/{slug}/".format(
                        slug = self.slug)

class Bundle(models.Model):
    """
    Model for "Dashboard Bundles"
    """
    bundle_stream = models.ForeignKey(BundleStream,
            verbose_name = _(u"Stream"),
            related_name = 'bundles')

    uploaded_by = models.ForeignKey(User,
            verbose_name = _(u"Uploaded by"),
            help_text = _(u"The user who submitted this bundle"),
            related_name = 'uploaded_bundles',
            null = True,
            blank = True)

    uploaded_on = models.DateTimeField(
            verbose_name = _(u"Uploaded on"),
            editable = False,
            auto_now_add = True)

    is_deserialized = models.BooleanField(
            verbose_name = _(u"Is deserialized"),
            help_text = _(u"Set when document has been analyzed and loaded"
                " into the database"),
            editable = False)

    content = models.FileField(
            verbose_name = _(u"Content"),
            help_text = _(u"Document in Dashboard Bundle Format 1.0"),
            upload_to = 'bundles',
            null = True)

    content_sha1 = models.CharField(
            editable = False,
            max_length = 40,
            null = True,
            unique = True)

    content_filename = models.CharField(
            verbose_name = _(u"Content file name"),
            help_text = _(u"Name of the originally uploaded bundle"),
            max_length = 256)

    def __unicode__(self):
        return _(u"Bundle {0} ({1})").format(
                self.pk, self.content_filename)

    @models.permalink
    def get_absolute_url(self):
        return ("dashboard_app.bundle.detail", [self.pk])

    def save(self, *args, **kwargs):
        if self.content:
            sha1 = hashlib.sha1()
            for chunk in self.content.chunks():
                sha1.update(chunk)
            self.content_sha1 = sha1.hexdigest()
            self.content.seek(0)
        return super(Bundle, self).save(*args, **kwargs)

class Test(models.Model):
    """
    Model for representing tests.

    Test is a collection of individual test cases.
    """
    test_id = models.CharField(
        max_length = 64,
        verbose_name = _("Test ID"),
        unique = True,
        help_text = _help_max_length(64))

    name = models.CharField(
        blank = True,
        help_text = _help_max_length(64),
        max_length = 64,
        verbose_name = _(u"Name"))

    objects = managers.TestManager()

    @models.permalink
    def get_absolute_url(self):
        return ('dashboard_app.test.detail', [self.test_id])
<commit_msg>Add Test.__unicode__ method<commit_after>"""
Database models of the Dashboard application
"""
import hashlib

from django import core
from django.contrib.auth.models import (User, Group)
from django.contrib.contenttypes import generic
from django.contrib.contenttypes.models import ContentType
from django.db import models
from django.utils.translation import ugettext as _
from django.utils.translation import ungettext

from dashboard_app import managers

def _help_max_length(max_length):
    return ungettext(
            u"Maximum length: {0} character",
            u"Maximum length: {0} characters",
            max_length).format(max_length)

class SoftwarePackage(models.Model):
    """
    Model for software packages.

    This class mirrors launch_control.models.SoftwarePackage.
    """
    name = models.CharField(
            max_length = 64,
            verbose_name = _(u"Package name"),
            help_text = _help_max_length(64))

    version = models.CharField(
            max_length = 32,
            verbose_name = _(u"Package version"),
            help_text = _help_max_length(32))

    class Meta:
        unique_together = (('name', 'version'))

    def __unicode__(self):
        return _(u"{name} {version}").format(
                name = self.name,
                version = self.version)

    @models.permalink
    def get_absolute_url(self):
        return ("dashboard_app.sw-package.detail", [self.name, self.version])

class NamedAttribute(models.Model):
    """
    Model for adding generic named attributes
    to arbitrary other model instances.

    This class mirrors launch_control.models.HardwareDevice.

    Example:
        class Foo(Model):
            attributes = generic.GenericRelation(NamedAttribute)
    """
    name = models.CharField(
            help_text = _help_max_length(32),
            max_length = 32)

    value = models.CharField(
            help_text = _help_max_length(256),
            max_length = 256)

    # Content type plumbing
    content_type = models.ForeignKey(ContentType)
    object_id = models.PositiveIntegerField()
    content_object = generic.GenericForeignKey('content_type', 'object_id')

    def __unicode__(self):
        return _(u"Attribute {name}: {value}").format(
                name = self.name,
                value = self.value)

    class Meta:
        unique_together = (('object_id', 'name'))

class HardwareDevice(models.Model):
    """
    Model for hardware devices

    All devices are simplified into an instance of pre-defined class
    with arbitrary key-value attributes.
    """
    device_type = models.CharField(
            choices = (
                (u"device.cpu", _(u"CPU")),
                (u"device.mem", _(u"Memory")),
                (u"device.usb", _(u"USB device")),
                (u"device.pci", _(u"PCI device")),
                (u"device.board", _(u"Board/Motherboard"))),
            help_text = _(u"One of pre-defined device types"),
            max_length = 32,
            verbose_name = _(u"Device Type"),
            )

    description = models.CharField(
            help_text = _(u"Human readable device summary.") + " " + _help_max_length(256),
            max_length = 256,
            verbose_name = _(u"Description."),
            )

    attributes = generic.GenericRelation(NamedAttribute)

    def __unicode__(self):
        return self.description

    @models.permalink
    def get_absolute_url(self):
        return ("dashboard_app.hw-device.detail", [self.pk])

class BundleStream(models.Model):
    """
    Model for "streams" of bundles.

    Basically it's a named collection of bundles, like directory just
    without the nesting. A simple ACL scheme is also supported,
    a bundle may be uploaded by:
        - specific user when user field is set
        - users of a specific group when group field is set
        - anyone when neither user nor group is set
    """
    user = models.ForeignKey(User,
            blank = True,
            help_text = _("User owning this stream (do not set when group is also set)"),
            null = True,
            verbose_name = _(u"User"),
            )

    group = models.ForeignKey(Group,
            blank = True,
            help_text = _("Group owning this stream (do not set when user is also set)"),
            null = True,
            verbose_name = _(u"Group"),
            )

    slug = models.CharField(
            blank = True,
            help_text = (_(u"Name that you will use when uploading bundles.")
                + " " + _help_max_length(64)),
            max_length = 64,
            verbose_name = _(u"Slug"),
            )

    name = models.CharField(
            blank = True,
            help_text = _help_max_length(64),
            max_length = 64,
            verbose_name = _(u"Name"),
            )

    pathname = models.CharField(
            max_length = 128,
            editable = False,
            unique = True,
            )

    objects = managers.BundleStreamManager()

    def __unicode__(self):
        return _(u"Bundle stream {pathname}").format(
                pathname = self.pathname)

    @models.permalink
    def get_absolute_url(self):
        return ("dashboard_app.bundle_stream_detail", [self.pathname])

    def save(self, *args, **kwargs):
        """
        Save this instance.

        Calls self.clean() to ensure that constraints are met.
        Updates pathname to reflect user/group/slug changes.
        """
        self.pathname = self._calc_pathname()
        self.clean()
        return super(BundleStream, self).save(*args, **kwargs)

    def clean(self):
        """
        Validate instance.

        Makes sure that user and name are not set at the same time
        """
        if self.user is not None and self.group is not None:
            raise core.exceptions.ValidationError('BundleStream cannot '
                    'have both user and name set at the same time')

    def can_access(self, user):
        """
        Returns true if given user can access the contents of this this
        stream.
        """
        if user is None:
            return self.user is None and self.group is None
        else:
            if self.user is not None:
                return self.user.username == user.username
            elif self.group is not None:
                return self.group in user.groups.all()
            else:
                return True

    def _calc_pathname(self):
        """
        Pseudo pathname-like ID of this stream.

        This pathname is user visible and will be presented to users
        when they want to interact with this bundle stream. The
        pathnames are unique and this is enforced at database level (the
        user and name are unique together).
        """
        if self.user is not None:
            if self.slug == "":
                return u"/personal/{user}/".format(
                        user = self.user.username)
            else:
                return u"/personal/{user}/{slug}/".format(
                        user = self.user.username,
                        slug = self.slug)
        elif self.group is not None:
            if self.slug == "":
                return u"/team/{group}/".format(
                        group = self.group.name)
            else:
                return u"/team/{group}/{slug}/".format(
                        group = self.group.name,
                        slug = self.slug)
        else:
            if self.slug == "":
                return u"/anonymous/"
            else:
                return u"/anonymous/{slug}/".format(
                        slug = self.slug)

class Bundle(models.Model):
    """
    Model for "Dashboard Bundles"
    """
    bundle_stream = models.ForeignKey(BundleStream,
            verbose_name = _(u"Stream"),
            related_name = 'bundles')

    uploaded_by = models.ForeignKey(User,
            verbose_name = _(u"Uploaded by"),
            help_text = _(u"The user who submitted this bundle"),
            related_name = 'uploaded_bundles',
            null = True,
            blank = True)

    uploaded_on = models.DateTimeField(
            verbose_name = _(u"Uploaded on"),
            editable = False,
            auto_now_add = True)

    is_deserialized = models.BooleanField(
            verbose_name = _(u"Is deserialized"),
            help_text = _(u"Set when document has been analyzed and loaded"
                " into the database"),
            editable = False)

    content = models.FileField(
            verbose_name = _(u"Content"),
            help_text = _(u"Document in Dashboard Bundle Format 1.0"),
            upload_to = 'bundles',
            null = True)

    content_sha1 = models.CharField(
            editable = False,
            max_length = 40,
            null = True,
            unique = True)

    content_filename = models.CharField(
            verbose_name = _(u"Content file name"),
            help_text = _(u"Name of the originally uploaded bundle"),
            max_length = 256)

    def __unicode__(self):
        return _(u"Bundle {0} ({1})").format(
                self.pk, self.content_filename)

    @models.permalink
    def get_absolute_url(self):
        return ("dashboard_app.bundle.detail", [self.pk])

    def save(self, *args, **kwargs):
        if self.content:
            sha1 = hashlib.sha1()
            for chunk in self.content.chunks():
                sha1.update(chunk)
            self.content_sha1 = sha1.hexdigest()
            self.content.seek(0)
        return super(Bundle, self).save(*args, **kwargs)

class Test(models.Model):
    """
    Model for representing tests.

    Test is a collection of individual test cases.
    """
    test_id = models.CharField(
        max_length = 64,
        verbose_name = _("Test ID"),
        unique = True,
        help_text = _help_max_length(64))

    name = models.CharField(
        blank = True,
        help_text = _help_max_length(64),
        max_length = 64,
        verbose_name = _(u"Name"))

    objects = managers.TestManager()

    def __unicode__(self):
        return _(u"Test {0}").format(self.name or self.test_id)

    @models.permalink
    def get_absolute_url(self):
        return ('dashboard_app.test.detail', [self.test_id])
<|endoftext|>


<commit_before>"""
Make dask workers using condor
"""
from __future__ import division, print_function

import atexit
import logging
import time

import distributed
import tornado

import htcondor
import classad

logger = logging.getLogger(__name__)

JOB_TEMPLATE = dict(
    Executable="/usr/bin/dask-worker",
    Universe="vanilla",
    Output="worker-$(ClusterId).$(ProcId).out",
    Error="worker-$(ClusterId).$(ProcId).err",
    Log="worker-$(ClusterId).$(ProcId).log",
)

JOB_STATUS_IDLE = 1
JOB_STATUS_RUNNING = 2
JOB_STATUS_HELD = 5

_global_schedulers = [] # (scheduler_id, schedd)

@atexit.register
def global_killall():
    for sid, schedd in _global_schedulers:
        condor_rm(schedd, 'DaskSchedulerId == "%s"' % sid)

def worker_constraint(jobid):
    clusterid, procid = jobid.split('.', 1)
    return '(ClusterId == %s && ProcId == %s)' % (clusterid, procid)

def or_constraints(constraints):
    return '(' + ' || '.join(constraints) + ')'

def workers_constraint(jobids):
    return or_constraints([worker_constraint(jid) for jid in jobids])

def condor_rm(schedd, job_spec):
    return schedd.act(htcondor.JobAction.Remove, job_spec)

class HTCondorCluster(object):
    def __init__(self,
                 memory_per_worker=1024,
                 procs_per_worker=1,
                 pool=None,
                 schedd_name=None,
                 threads_per_worker=1,
                 cleanup_interval=1000,
                 worker_timeout=(24 * 60 * 60),
                 **kwargs):

        global _global_schedulers

        if schedd_name is None:
            self.schedd = htcondor.Schedd()
        else:
            collector = htcondor.Collector(pool)
            self.schedd = htcondor.Schedd(
                collector.locate(
                    htcondor.DaemonTypes.Schedd,
                    schedd_name))

        self.local_cluster = distributed.LocalCluster(ip='', n_workers=0,
                                                      **kwargs)

        _global_schedulers.append((self.scheduler.id, self.schedd))

        self.jobs = {}  # {jobid: CLASSAD}
        if int(cleanup_interval) < 1:
            raise ValueError("cleanup_interval must be >= 1")
        self._cleanup_callback = tornado.ioloop.PeriodicCallback(
            callback=self.cleanup_jobs,
            callback_time=cleanup_interval,
            io_loop=self.scheduler.loop)
        self._cleanup_callback.start()

        self.memory_per_worker = memory_per_worker
        self.procs_per_worker = procs_per_worker
        self.threads_per_worker = threads_per_worker
        self.worker_timeout = worker_timeout

    @tornado.gen.coroutine
    def _start(self):
        pass

    @property
    def scheduler(self):
        return self.local_cluster.scheduler

    @property
    def scheduler_address(self):
        return self.scheduler.address

    @property
    def jobids(self):
        return self.jobs.keys()

    @property
    def scheduler_constraint(self):
        return '(DaskSchedulerId == "%s")' % self.scheduler.id

    def start_workers(self,
                      n=1,
                      memory_per_worker=None,
                      procs_per_worker=None,
                      threads_per_worker=None,
                      worker_timeout=None,
                      extra_attribs=None):

        if int(n) < 1:
            raise ValueError("n must be >= 1")

        memory_per_worker = memory_per_worker or self.memory_per_worker
        if int(memory_per_worker) < 1:
            raise ValueError("memory_per_worker must be >= 1 (MB)")
        procs_per_worker = procs_per_worker or self.procs_per_worker
        if int(procs_per_worker) < 1:
            raise ValueError("procs_per_worker must be >= 1")
        threads_per_worker = threads_per_worker or self.threads_per_worker
        if int(threads_per_worker) < 1:
            raise ValueError("threads_per_worker must be >= 1")
        worker_timeout = worker_timeout or self.worker_timeout
        if int(worker_timeout) < 1:
            raise ValueError("worker_timeout must be >= 1 (sec)")

        job = htcondor.Submit(JOB_TEMPLATE)
        args = [self.scheduler_address]
        args.append('--nprocs %d' % procs_per_worker)
        args.append('--nthreads %d' % threads_per_worker)
        request_cpus = procs_per_worker * threads_per_worker
        if procs_per_worker > 1:
            # the nanny takes up a core too
            request_cpus += 1
        else:
            args.append('--no-nanny')
            # can only use --name if --nprocs=1
            worker_name = "htcondor-$(ClusterId).$(ProcId)"
            args.append('--name=' + worker_name)
            # when I tried +DaskWorkerName, then $(ClusterId) and $(ProcId) didn't
            # get expanded (GT #6219)
            job['MY.DaskWorkerName'] = '"' + worker_name + '"'

        args.append('--no-bokeh')

        job['Arguments'] = ' '.join(args)
        job['RequestMemory'] = "%d MB" % memory_per_worker
        job['RequestCpus'] = str(request_cpus)
        job['+DaskSchedulerId'] = '"' + self.scheduler.id + '"'

        job['Periodic_Hold'] = "(time() - JobStartDate) > %d" % (
            worker_timeout)
        job['Periodic_Hold_Reason'] = '"dask-worker max lifetime %d min"' % (
            worker_timeout // 60)

        if extra_attribs:
            job.update(extra_attribs)

        with self.schedd.transaction() as txn:
            classads = []
            clusterid = job.queue(txn, count=n, ad_results=classads)
            logger.info("Started clusterid %s with %d jobs" % (clusterid, n))
            logger.debug(
                "RequestMemory = %s; RequestCpus = %s"
                % (job['RequestMemory'], job['RequestCpus']))
            for ad in classads:
                self.jobs["%s.%s" % (ad['ClusterId'], ad['ProcId'])] = ad

    def killall(self):
        condor_rm(self.schedd, self.scheduler_constraint)

    def submit_worker(self, **kwargs):
        return self.start_workers(n=1, **kwargs)

    def stop_workers(self, worker_ids):
        if isinstance(worker_ids, str):
            worker_ids = [worker_ids]

        constraint = '%s && %s' % (
            self.scheduler_constraint,
            workers_constraint(worker_ids)
            )

        condor_rm(self.schedd, constraint)

    def cleanup_jobs(self):
        active_jobids = \
            ['%s.%s' % (ad['ClusterId'], ad['ProcId'])
             for ad in self.schedd.query(
                self.scheduler_constraint,
                ['ClusterId', 'ProcId', 'JobStatus'])
             if ad['JobStatus'] in (
                JOB_STATUS_IDLE,
                JOB_STATUS_RUNNING,
                JOB_STATUS_HELD)]
        for jobid in self.jobids:
            if jobid not in active_jobids:
                del self.jobs[jobid]

    def close(self):
        self.killall()
        self.local_cluster.close()

    def __del__(self):
        self.close()

    def __enter__(self):
        return self

    def __exit__(self, *args):
        self.close()

    def __str__(self):
        return "<%s: %d workers>" % (self.__class__.__name__, len(self.jobids))

    __repr__ = __str__

<commit_msg>Add --memory-limit using 75% of RequestMemory<commit_after>"""
Make dask workers using condor
"""
from __future__ import division, print_function

import atexit
import logging
import time

import distributed
import tornado

import htcondor
import classad

logger = logging.getLogger(__name__)

JOB_TEMPLATE = dict(
    Executable="/usr/bin/dask-worker",
    Universe="vanilla",
    Output="worker-$(ClusterId).$(ProcId).out",
    Error="worker-$(ClusterId).$(ProcId).err",
    Log="worker-$(ClusterId).$(ProcId).log",
)

JOB_STATUS_IDLE = 1
JOB_STATUS_RUNNING = 2
JOB_STATUS_HELD = 5

_global_schedulers = [] # (scheduler_id, schedd)

@atexit.register
def global_killall():
    for sid, schedd in _global_schedulers:
        condor_rm(schedd, 'DaskSchedulerId == "%s"' % sid)

def worker_constraint(jobid):
    clusterid, procid = jobid.split('.', 1)
    return '(ClusterId == %s && ProcId == %s)' % (clusterid, procid)

def or_constraints(constraints):
    return '(' + ' || '.join(constraints) + ')'

def workers_constraint(jobids):
    return or_constraints([worker_constraint(jid) for jid in jobids])

def condor_rm(schedd, job_spec):
    return schedd.act(htcondor.JobAction.Remove, job_spec)

class HTCondorCluster(object):
    def __init__(self,
                 memory_per_worker=1024,
                 procs_per_worker=1,
                 pool=None,
                 schedd_name=None,
                 threads_per_worker=1,
                 cleanup_interval=1000,
                 worker_timeout=(24 * 60 * 60),
                 **kwargs):

        global _global_schedulers

        if schedd_name is None:
            self.schedd = htcondor.Schedd()
        else:
            collector = htcondor.Collector(pool)
            self.schedd = htcondor.Schedd(
                collector.locate(
                    htcondor.DaemonTypes.Schedd,
                    schedd_name))

        self.local_cluster = distributed.LocalCluster(ip='', n_workers=0,
                                                      **kwargs)

        _global_schedulers.append((self.scheduler.id, self.schedd))

        self.jobs = {}  # {jobid: CLASSAD}
        if int(cleanup_interval) < 1:
            raise ValueError("cleanup_interval must be >= 1")
        self._cleanup_callback = tornado.ioloop.PeriodicCallback(
            callback=self.cleanup_jobs,
            callback_time=cleanup_interval,
            io_loop=self.scheduler.loop)
        self._cleanup_callback.start()

        self.memory_per_worker = memory_per_worker
        self.procs_per_worker = procs_per_worker
        self.threads_per_worker = threads_per_worker
        self.worker_timeout = worker_timeout

    @tornado.gen.coroutine
    def _start(self):
        pass

    @property
    def scheduler(self):
        return self.local_cluster.scheduler

    @property
    def scheduler_address(self):
        return self.scheduler.address

    @property
    def jobids(self):
        return self.jobs.keys()

    @property
    def scheduler_constraint(self):
        return '(DaskSchedulerId == "%s")' % self.scheduler.id

    def start_workers(self,
                      n=1,
                      memory_per_worker=None,
                      procs_per_worker=None,
                      threads_per_worker=None,
                      worker_timeout=None,
                      extra_attribs=None):

        if int(n) < 1:
            raise ValueError("n must be >= 1")

        memory_per_worker = memory_per_worker or self.memory_per_worker
        if int(memory_per_worker) < 1:
            raise ValueError("memory_per_worker must be >= 1 (MB)")
        procs_per_worker = procs_per_worker or self.procs_per_worker
        if int(procs_per_worker) < 1:
            raise ValueError("procs_per_worker must be >= 1")
        threads_per_worker = threads_per_worker or self.threads_per_worker
        if int(threads_per_worker) < 1:
            raise ValueError("threads_per_worker must be >= 1")
        worker_timeout = worker_timeout or self.worker_timeout
        if int(worker_timeout) < 1:
            raise ValueError("worker_timeout must be >= 1 (sec)")

        job = htcondor.Submit(JOB_TEMPLATE)
        args = [self.scheduler_address]
        args.append('--nprocs %d' % procs_per_worker)
        args.append('--nthreads %d' % threads_per_worker)
        # default memory limit is 75% of memory; use 75% of RequestMemory
        # instead
        args.append('--memory-limit %d' %
                    (memory_per_worker * 3 * 1048576 // 4))
        request_cpus = procs_per_worker * threads_per_worker
        if procs_per_worker > 1:
            # the nanny takes up a core too
            request_cpus += 1
        else:
            args.append('--no-nanny')
            # can only use --name if --nprocs=1
            worker_name = "htcondor-$(ClusterId).$(ProcId)"
            args.append('--name=' + worker_name)
            # when I tried +DaskWorkerName, then $(ClusterId) and $(ProcId) didn't
            # get expanded (GT #6219)
            job['MY.DaskWorkerName'] = '"' + worker_name + '"'

        args.append('--no-bokeh')

        job['Arguments'] = ' '.join(args)
        job['RequestMemory'] = "%d MB" % memory_per_worker
        job['RequestCpus'] = str(request_cpus)
        job['+DaskSchedulerId'] = '"' + self.scheduler.id + '"'

        job['Periodic_Hold'] = "(time() - JobStartDate) > %d" % (
            worker_timeout)
        job['Periodic_Hold_Reason'] = '"dask-worker max lifetime %d min"' % (
            worker_timeout // 60)

        if extra_attribs:
            job.update(extra_attribs)

        with self.schedd.transaction() as txn:
            classads = []
            clusterid = job.queue(txn, count=n, ad_results=classads)
            logger.info("Started clusterid %s with %d jobs" % (clusterid, n))
            logger.debug(
                "RequestMemory = %s; RequestCpus = %s"
                % (job['RequestMemory'], job['RequestCpus']))
            for ad in classads:
                self.jobs["%s.%s" % (ad['ClusterId'], ad['ProcId'])] = ad

    def killall(self):
        condor_rm(self.schedd, self.scheduler_constraint)

    def submit_worker(self, **kwargs):
        return self.start_workers(n=1, **kwargs)

    def stop_workers(self, worker_ids):
        if isinstance(worker_ids, str):
            worker_ids = [worker_ids]

        constraint = '%s && %s' % (
            self.scheduler_constraint,
            workers_constraint(worker_ids)
            )

        condor_rm(self.schedd, constraint)

    def cleanup_jobs(self):
        active_jobids = \
            ['%s.%s' % (ad['ClusterId'], ad['ProcId'])
             for ad in self.schedd.query(
                self.scheduler_constraint,
                ['ClusterId', 'ProcId', 'JobStatus'])
             if ad['JobStatus'] in (
                JOB_STATUS_IDLE,
                JOB_STATUS_RUNNING,
                JOB_STATUS_HELD)]
        for jobid in self.jobids:
            if jobid not in active_jobids:
                del self.jobs[jobid]

    def close(self):
        self.killall()
        self.local_cluster.close()

    def __del__(self):
        self.close()

    def __enter__(self):
        return self

    def __exit__(self, *args):
        self.close()

    def __str__(self):
        return "<%s: %d workers>" % (self.__class__.__name__, len(self.jobids))

    __repr__ = __str__

<|endoftext|>


<commit_before>"""
Make dask workers using condor
"""
from __future__ import division, print_function

import atexit
import errno
import logging
import os
import tempfile
import time

import distributed
import tornado

import htcondor
import classad

if not hasattr(htcondor, 'Submit'):
    raise ImportError("htcondor.Submit not found;"
                      " HTCondor 8.6.0 or newer required")

logger = logging.getLogger(__name__)

JOB_TEMPLATE = \
    { 'Executable':           '/usr/bin/dask-worker'
    , 'Universe':             'vanilla'
    # $F(MY.JobId) strips the quotes from MY.JobId
    , 'Output':               'worker-$F(MY.JobId).out'
    , 'Error':                'worker-$F(MY.JobId).err'
    , 'Log':                  'worker-$F(MY.JobId).log'

    # using MY.Arguments instead of Arguments lets the value be a classad
    # expression instead of a string. Thanks TJ!
    , 'MY.Arguments':
        'strcat( MY.DaskSchedulerAddress'
        '      , " --nprocs=", MY.DaskNProcs'
        '      , " --nthreads=", MY.DaskNThreads'
        '      , " --no-bokeh"'
    # default memory limit is 75% of memory; use 75% of RequestMemory instead
        '      , " --memory-limit="'
        '      , floor(RequestMemory * 1048576 * 0.75)'
    # no point in having a nanny if we're only running 1 proc
        '      , ifThenElse( (MY.DaskNProcs < 2)'
        '                  , " --no-nanny "'
        '                  , "")'
    # we can only have a worker name if nprocs == 1
        '      , ifThenElse( isUndefined(MY.DaskWorkerName)'
        '                  , ""'
        '                  , strcat(" --name=", MY.DaskWorkerName))'
        '      )'

    , 'MY.DaskWorkerName':    'ifThenElse( (MY.DaskNProcs < 2)'
                              '          , "htcondor-$F(MY.JobId)"'
                              '          , UNDEFINED)'

    # reserve a CPU for the nanny process if nprocs > 1
    , 'RequestCpus':          'ifThenElse( (MY.DaskNProcs < 2)'
                              '          , MY.DaskNThreads'
                              '          , 1 + MY.DaskNProcs * MY.DaskNThreads)'

    , 'Periodic_Hold':
        '((time() - EnteredCurrentStatus) > MY.DaskWorkerTimeout) &&'
        ' (JobStatus == 2)'
    , 'Periodic_Hold_Reason': 'strcat("dask-worker exceeded max lifetime of ",'
                              '       interval(MY.DaskWorkerTimeout))'
    , 'Transfer_Input_Files': ''
    , 'MY.JobId':             '"$(ClusterId).$(ProcId)"'
    }

JOB_STATUS_IDLE = 1
JOB_STATUS_RUNNING = 2
JOB_STATUS_HELD = 5

SCRIPT_TEMPLATE = """\
#!/bin/bash

if [[ -z $_CONDOR_SCRATCH_DIR ]]; then
    echo '$_CONDOR_SCRATCH_DIR not defined'
    echo 'This script needs to be run as an HTCondor job'
    exit 1
fi

export HOME=$_CONDOR_SCRATCH_DIR

tar xzf ~/%(worker_tarball)s
~/dask_condor_worker/fixpaths.sh
export PATH=~/dask_condor_worker/bin:$PATH
export PYTHONPATH=~:${PYTHONPATH+":$PYTHONPATH"}

args=( "$@" )

# This isn't actually necessary - $TMP is already under $_CONDOR_SCRATCH_DIR
# so that's where mktemp will make its files.
local_directory=$_CONDOR_SCRATCH_DIR/.worker
mkdir -p "$local_directory"
args+=(--local-directory "$local_directory")

exec dask-worker "${args[@]}"
"""

_global_schedulers = [] # (scheduler_id, schedd)

@atexit.register
def global_killall():
    for sid, schedd in _global_schedulers:
        condor_rm(schedd, 'DaskSchedulerId == "%s"' % sid)

def worker_constraint(jobid):
    if '.' in jobid:
        return '(JobId == %s)' % (jobid)
    else:
        return '(ClusterId == %s)' % jobid

def or_constraints(constraints):
    return '(' + ' || '.join(constraints) + ')'

def workers_constraint(jobids):
    return or_constraints([worker_constraint(jid) for jid in jobids])

def condor_rm(schedd, job_spec):
    return schedd.act(htcondor.JobAction.Remove, job_spec)

class Error(Exception):
    pass

class HTCondorCluster(object):
    def __init__(self,
                 memory_per_worker=1024,
                 procs_per_worker=1,
                 pool=None,
                 schedd_name=None,
                 threads_per_worker=1,
                 update_interval=1000,
                 worker_timeout=(24 * 60 * 60),
                 scheduler_port=8786,
                 worker_tarball=None,
                 transfer_files=None,
                 **kwargs):

        self.memory_per_worker = memory_per_worker
        self.procs_per_worker = procs_per_worker
        self.threads_per_worker = threads_per_worker
        if int(update_interval) < 1:
            raise ValueError("update_interval must be >= 1")
        self.worker_timeout = worker_timeout
        self.worker_tarball = worker_tarball
        self.transfer_files = transfer_files

        if schedd_name is None:
            self.schedd = htcondor.Schedd()
        else:
            collector = htcondor.Collector(pool)
            self.schedd = htcondor.Schedd(
                collector.locate(
                    htcondor.DaemonTypes.Schedd,
                    schedd_name))

        self.script = None
        if self.worker_tarball:
            if '://' not in self.worker_tarball:
                self._verify_tarball()
            self.script = tempfile.NamedTemporaryFile(
                suffix='.sh', prefix='dask-worker-wrapper-')
            self.script.write(SCRIPT_TEMPLATE
                % {'worker_tarball': os.path.basename(self.worker_tarball)})
            self.script.flush()

            @atexit.register
            def _erase_script():
                self.script.close()

        self.local_cluster = distributed.LocalCluster(
            ip='', n_workers=0, scheduler_port=scheduler_port, **kwargs)

        global _global_schedulers
        _global_schedulers.append((self.scheduler.id, self.schedd))

        self.jobs = {}  # {jobid: CLASSAD}
        self._update_callback = tornado.ioloop.PeriodicCallback(
            callback=self.update_jobs,
            callback_time=update_interval,
            io_loop=self.scheduler.loop)
        self._update_callback.start()

    @tornado.gen.coroutine
    def _start(self):
        pass

    @property
    def scheduler(self):
        return self.local_cluster.scheduler

    @property
    def scheduler_address(self):
        return self.scheduler.address

    @property
    def jobids(self):
        return self.jobs.keys()

    @property
    def scheduler_constraint(self):
        return '(DaskSchedulerId == "%s")' % self.scheduler.id

    def start_workers(self,
                      n=1,
                      memory_per_worker=None,
                      procs_per_worker=None,
                      threads_per_worker=None,
                      worker_timeout=None,
                      transfer_files=None,
                      extra_attribs=None):
        n = int(n)
        if n < 1:
            raise ValueError("n must be >= 1")
        memory_per_worker = int(memory_per_worker or self.memory_per_worker)
        if memory_per_worker < 1:
            raise ValueError("memory_per_worker must be >= 1 (MB)")
        procs_per_worker = int(procs_per_worker or self.procs_per_worker)
        if procs_per_worker < 1:
            raise ValueError("procs_per_worker must be >= 1")
        threads_per_worker = int(threads_per_worker or self.threads_per_worker)
        if threads_per_worker < 1:
            raise ValueError("threads_per_worker must be >= 1")
        worker_timeout = int(worker_timeout or self.worker_timeout)
        if worker_timeout < 1:
            raise ValueError("worker_timeout must be >= 1 (sec)")
        transfer_files = transfer_files or self.transfer_files
        if transfer_files:
            if not isinstance(transfer_files, str):
                transfer_files = ', '.join(transfer_files)

        job = htcondor.Submit(JOB_TEMPLATE)
        job['MY.DaskSchedulerAddress'] = '"' + self.scheduler_address + '"'
        job['MY.DaskNProcs'] = str(procs_per_worker)
        job['MY.DaskNThreads'] = str(threads_per_worker)
        job['RequestMemory'] = str(memory_per_worker)
        job['MY.DaskSchedulerId'] = '"' + self.scheduler.id + '"'
        job['MY.DaskWorkerTimeout'] = str(worker_timeout)
        if self.script:
            job['Executable'] = self.script.name
            job['Transfer_Input_Files'] = self.worker_tarball \
                + (', ' + transfer_files if transfer_files else '')
        else:
            if transfer_files:
                job['Transfer_Input_Files'] = transfer_files

        if extra_attribs:
            job.update(extra_attribs)

        classads = []
        with self.schedd.transaction() as txn:
            clusterid = job.queue(txn, count=n, ad_results=classads)
        logger.info("%d job(s) submitted to cluster %s." % (n, clusterid))
        for ad in classads:
            self.jobs[ad['JobId']] = ad

    def killall(self):
        condor_rm(self.schedd, self.scheduler_constraint)

    def submit_worker(self, **kwargs):
        return self.start_workers(n=1, **kwargs)

    def stop_workers(self, worker_ids):
        if isinstance(worker_ids, str):
            worker_ids = [worker_ids]

        constraint = '%s && %s' % (
            self.scheduler_constraint,
            workers_constraint(worker_ids)
            )

        condor_rm(self.schedd, constraint)

    def update_jobs(self):
        # Don't want to leave the original in a half-updated state
        new_jobs = self.jobs.copy()
        try:
            ads = self.schedd.xquery(
                self.scheduler_constraint,
                projection=['JobId', 'JobStatus',
                            'EnteredCurrentStatus'])
            active_jobids = []
            for ad in ads:
                jobid = ad['JobId']
                if isinstance(jobid, classad.ExprTree):
                    jobid = jobid.eval()
                jobstatus = ad['JobStatus']
                if isinstance(jobstatus, classad.ExprTree):
                    jobstatus = jobstatus.eval()
                if jobstatus in (
                        JOB_STATUS_IDLE, JOB_STATUS_RUNNING, JOB_STATUS_HELD):
                    for attr in ['JobStatus', 'EnteredCurrentStatus']:
                        new_jobs[jobid][attr] = ad[attr]
                    active_jobids.append(jobid)

            # Evaluate the list of keys now to avoid a RuntimeError when
            # we delete items from the dict mid-iteration
            for jobid in list(new_jobs.keys()):
                if jobid not in active_jobids:
                    del new_jobs[jobid]
        except RuntimeError as err:
            # timeouts happen. Ignore them.
            if 'Timeout when waiting for remote host' in err.message:
                pass
            else:
                raise
        self.jobs = new_jobs

    def close(self):
        self.killall()
        self.local_cluster.close()
        if self.script:
            self.script.close()

    def _verify_tarball(self):
        if not os.path.exists(self.worker_tarball):
            raise IOError(errno.ENOENT, "worker tarball not found",
                          self.worker_tarball)

        import tarfile
        tar = tarfile.open(self.worker_tarball)
        try:
            members = tar.getnames()
            for path in ['dask_condor_worker/fixpaths.sh',
                         'dask_condor_worker/bin/python',
                         'dask_condor_worker/bin/dask-worker']:
                if path not in members:
                    raise Error("Expected file %s not in tarball" % (path))
        finally:
            tar.close()

    def __del__(self):
        self.close()

    def __enter__(self):
        return self

    def __exit__(self, *args):
        self.close()

    def __str__(self):
        total = running = idle = held = 0
        for key in self.jobs:
            status = self.jobs[key]['JobStatus']
            total += 1
            if status == JOB_STATUS_IDLE:
                idle += 1
            elif status == JOB_STATUS_RUNNING:
                running += 1
            elif status == JOB_STATUS_HELD:
                held += 1

        return "<%s: %d workers (%d running, %d idle, %d held)>" \
               % (self.__class__.__name__, total, running, idle, held)

    __repr__ = __str__

<commit_msg>Get rid of second pass through job list<commit_after>"""
Make dask workers using condor
"""
from __future__ import division, print_function

import atexit
import errno
import logging
import os
import tempfile
import time

import distributed
import tornado

import htcondor
import classad

if not hasattr(htcondor, 'Submit'):
    raise ImportError("htcondor.Submit not found;"
                      " HTCondor 8.6.0 or newer required")

logger = logging.getLogger(__name__)

JOB_TEMPLATE = \
    { 'Executable':           '/usr/bin/dask-worker'
    , 'Universe':             'vanilla'
    # $F(MY.JobId) strips the quotes from MY.JobId
    , 'Output':               'worker-$F(MY.JobId).out'
    , 'Error':                'worker-$F(MY.JobId).err'
    , 'Log':                  'worker-$F(MY.JobId).log'

    # using MY.Arguments instead of Arguments lets the value be a classad
    # expression instead of a string. Thanks TJ!
    , 'MY.Arguments':
        'strcat( MY.DaskSchedulerAddress'
        '      , " --nprocs=", MY.DaskNProcs'
        '      , " --nthreads=", MY.DaskNThreads'
        '      , " --no-bokeh"'
    # default memory limit is 75% of memory; use 75% of RequestMemory instead
        '      , " --memory-limit="'
        '      , floor(RequestMemory * 1048576 * 0.75)'
    # no point in having a nanny if we're only running 1 proc
        '      , ifThenElse( (MY.DaskNProcs < 2)'
        '                  , " --no-nanny "'
        '                  , "")'
    # we can only have a worker name if nprocs == 1
        '      , ifThenElse( isUndefined(MY.DaskWorkerName)'
        '                  , ""'
        '                  , strcat(" --name=", MY.DaskWorkerName))'
        '      )'

    , 'MY.DaskWorkerName':    'ifThenElse( (MY.DaskNProcs < 2)'
                              '          , "htcondor-$F(MY.JobId)"'
                              '          , UNDEFINED)'

    # reserve a CPU for the nanny process if nprocs > 1
    , 'RequestCpus':          'ifThenElse( (MY.DaskNProcs < 2)'
                              '          , MY.DaskNThreads'
                              '          , 1 + MY.DaskNProcs * MY.DaskNThreads)'

    , 'Periodic_Hold':
        '((time() - EnteredCurrentStatus) > MY.DaskWorkerTimeout) &&'
        ' (JobStatus == 2)'
    , 'Periodic_Hold_Reason': 'strcat("dask-worker exceeded max lifetime of ",'
                              '       interval(MY.DaskWorkerTimeout))'
    , 'Transfer_Input_Files': ''
    , 'MY.JobId':             '"$(ClusterId).$(ProcId)"'
    }

JOB_STATUS_IDLE = 1
JOB_STATUS_RUNNING = 2
JOB_STATUS_HELD = 5

SCRIPT_TEMPLATE = """\
#!/bin/bash

if [[ -z $_CONDOR_SCRATCH_DIR ]]; then
    echo '$_CONDOR_SCRATCH_DIR not defined'
    echo 'This script needs to be run as an HTCondor job'
    exit 1
fi

export HOME=$_CONDOR_SCRATCH_DIR

tar xzf ~/%(worker_tarball)s
~/dask_condor_worker/fixpaths.sh
export PATH=~/dask_condor_worker/bin:$PATH
export PYTHONPATH=~:${PYTHONPATH+":$PYTHONPATH"}

args=( "$@" )

# This isn't actually necessary - $TMP is already under $_CONDOR_SCRATCH_DIR
# so that's where mktemp will make its files.
local_directory=$_CONDOR_SCRATCH_DIR/.worker
mkdir -p "$local_directory"
args+=(--local-directory "$local_directory")

exec dask-worker "${args[@]}"
"""

_global_schedulers = [] # (scheduler_id, schedd)

@atexit.register
def global_killall():
    for sid, schedd in _global_schedulers:
        condor_rm(schedd, 'DaskSchedulerId == "%s"' % sid)

def worker_constraint(jobid):
    if '.' in jobid:
        return '(JobId == %s)' % (jobid)
    else:
        return '(ClusterId == %s)' % jobid

def or_constraints(constraints):
    return '(' + ' || '.join(constraints) + ')'

def workers_constraint(jobids):
    return or_constraints([worker_constraint(jid) for jid in jobids])

def condor_rm(schedd, job_spec):
    return schedd.act(htcondor.JobAction.Remove, job_spec)

class Error(Exception):
    pass

class HTCondorCluster(object):
    def __init__(self,
                 memory_per_worker=1024,
                 procs_per_worker=1,
                 pool=None,
                 schedd_name=None,
                 threads_per_worker=1,
                 update_interval=1000,
                 worker_timeout=(24 * 60 * 60),
                 scheduler_port=8786,
                 worker_tarball=None,
                 transfer_files=None,
                 **kwargs):

        self.memory_per_worker = memory_per_worker
        self.procs_per_worker = procs_per_worker
        self.threads_per_worker = threads_per_worker
        if int(update_interval) < 1:
            raise ValueError("update_interval must be >= 1")
        self.worker_timeout = worker_timeout
        self.worker_tarball = worker_tarball
        self.transfer_files = transfer_files

        if schedd_name is None:
            self.schedd = htcondor.Schedd()
        else:
            collector = htcondor.Collector(pool)
            self.schedd = htcondor.Schedd(
                collector.locate(
                    htcondor.DaemonTypes.Schedd,
                    schedd_name))

        self.script = None
        if self.worker_tarball:
            if '://' not in self.worker_tarball:
                self._verify_tarball()
            self.script = tempfile.NamedTemporaryFile(
                suffix='.sh', prefix='dask-worker-wrapper-')
            self.script.write(SCRIPT_TEMPLATE
                % {'worker_tarball': os.path.basename(self.worker_tarball)})
            self.script.flush()

            @atexit.register
            def _erase_script():
                self.script.close()

        self.local_cluster = distributed.LocalCluster(
            ip='', n_workers=0, scheduler_port=scheduler_port, **kwargs)

        global _global_schedulers
        _global_schedulers.append((self.scheduler.id, self.schedd))

        self.jobs = {}  # {jobid: CLASSAD}
        self.ignored_jobs = set()  # set of jobids
        self._update_callback = tornado.ioloop.PeriodicCallback(
            callback=self.update_jobs,
            callback_time=update_interval,
            io_loop=self.scheduler.loop)
        self._update_callback.start()

    @tornado.gen.coroutine
    def _start(self):
        pass

    @property
    def scheduler(self):
        return self.local_cluster.scheduler

    @property
    def scheduler_address(self):
        return self.scheduler.address

    @property
    def jobids(self):
        return self.jobs.keys()

    @property
    def scheduler_constraint(self):
        return '(DaskSchedulerId == "%s")' % self.scheduler.id

    def start_workers(self,
                      n=1,
                      memory_per_worker=None,
                      procs_per_worker=None,
                      threads_per_worker=None,
                      worker_timeout=None,
                      transfer_files=None,
                      extra_attribs=None):
        n = int(n)
        if n < 1:
            raise ValueError("n must be >= 1")
        memory_per_worker = int(memory_per_worker or self.memory_per_worker)
        if memory_per_worker < 1:
            raise ValueError("memory_per_worker must be >= 1 (MB)")
        procs_per_worker = int(procs_per_worker or self.procs_per_worker)
        if procs_per_worker < 1:
            raise ValueError("procs_per_worker must be >= 1")
        threads_per_worker = int(threads_per_worker or self.threads_per_worker)
        if threads_per_worker < 1:
            raise ValueError("threads_per_worker must be >= 1")
        worker_timeout = int(worker_timeout or self.worker_timeout)
        if worker_timeout < 1:
            raise ValueError("worker_timeout must be >= 1 (sec)")
        transfer_files = transfer_files or self.transfer_files
        if transfer_files:
            if not isinstance(transfer_files, str):
                transfer_files = ', '.join(transfer_files)

        job = htcondor.Submit(JOB_TEMPLATE)
        job['MY.DaskSchedulerAddress'] = '"' + self.scheduler_address + '"'
        job['MY.DaskNProcs'] = str(procs_per_worker)
        job['MY.DaskNThreads'] = str(threads_per_worker)
        job['RequestMemory'] = str(memory_per_worker)
        job['MY.DaskSchedulerId'] = '"' + self.scheduler.id + '"'
        job['MY.DaskWorkerTimeout'] = str(worker_timeout)
        if self.script:
            job['Executable'] = self.script.name
            job['Transfer_Input_Files'] = self.worker_tarball \
                + (', ' + transfer_files if transfer_files else '')
        else:
            if transfer_files:
                job['Transfer_Input_Files'] = transfer_files

        if extra_attribs:
            job.update(extra_attribs)

        classads = []
        with self.schedd.transaction() as txn:
            clusterid = job.queue(txn, count=n, ad_results=classads)
        logger.info("%d job(s) submitted to cluster %s." % (n, clusterid))
        for ad in classads:
            self.jobs[ad['JobId']] = ad

    def killall(self):
        condor_rm(self.schedd, self.scheduler_constraint)

    def submit_worker(self, **kwargs):
        return self.start_workers(n=1, **kwargs)

    def stop_workers(self, worker_ids):
        if isinstance(worker_ids, str):
            worker_ids = [worker_ids]

        constraint = '%s && %s' % (
            self.scheduler_constraint,
            workers_constraint(worker_ids)
            )

        condor_rm(self.schedd, constraint)

    def update_jobs(self):
        # Don't want to leave the original in a half-updated state
        new_jobs = {}
        try:
            ads = self.schedd.xquery(
                self.scheduler_constraint,
                projection=['JobId', 'JobStatus',
                            'EnteredCurrentStatus'])
            for ad in ads:
                jobid = ad['JobId']
                if isinstance(jobid, classad.ExprTree):
                    jobid = jobid.eval()
                jobstatus = ad['JobStatus']
                if isinstance(jobstatus, classad.ExprTree):
                    jobstatus = jobstatus.eval()
                if jobstatus in (
                        JOB_STATUS_IDLE, JOB_STATUS_RUNNING, JOB_STATUS_HELD):
                    if jobid in self.jobs:
                        new_jobs[jobid] = classad.ClassAd(dict(self.jobs[jobid]))
                        new_jobs[jobid].update(ad)
                    elif jobid not in self.ignored_jobs:
                        logger.warning("Unknown job found: %s, ignoring", jobid)
                        self.ignored_jobs.add(jobid)

            self.jobs = new_jobs
        except RuntimeError as err:
            # timeouts happen. Ignore them.
            if 'Timeout when waiting for remote host' in err.message:
                pass
            else:
                raise

    def close(self):
        self.killall()
        self.local_cluster.close()
        if self.script:
            self.script.close()

    def _verify_tarball(self):
        if not os.path.exists(self.worker_tarball):
            raise IOError(errno.ENOENT, "worker tarball not found",
                          self.worker_tarball)

        import tarfile
        tar = tarfile.open(self.worker_tarball)
        try:
            members = tar.getnames()
            for path in ['dask_condor_worker/fixpaths.sh',
                         'dask_condor_worker/bin/python',
                         'dask_condor_worker/bin/dask-worker']:
                if path not in members:
                    raise Error("Expected file %s not in tarball" % (path))
        finally:
            tar.close()

    def __del__(self):
        self.close()

    def __enter__(self):
        return self

    def __exit__(self, *args):
        self.close()

    def __str__(self):
        total = running = idle = held = 0
        for key in self.jobs:
            status = self.jobs[key]['JobStatus']
            total += 1
            if status == JOB_STATUS_IDLE:
                idle += 1
            elif status == JOB_STATUS_RUNNING:
                running += 1
            elif status == JOB_STATUS_HELD:
                held += 1

        return "<%s: %d workers (%d running, %d idle, %d held)>" \
               % (self.__class__.__name__, total, running, idle, held)

    __repr__ = __str__

<|endoftext|>


<commit_before>from fabric.api import lcd, local, settings

states = [
    ("Alabama", "al"), ("Alaska", "ak"), ("Arizona", "az"), ("Arkansas", "ar"),
    ("California", "ca"), ("Colorado", "co"), ("Connecticut", "ct"),
    ("Delaware", "de"), ("District_of_Columbia", "dc"), ("Florida", "fl"),
    ("Georgia", "ga"), ("Hawaii", "hi"), ("Idaho", "id"), ("Illinois", "il"),
    ("Indiana", "in"), ("Iowa", "ia"), ("Kansas", "ks"), ("Kentucky", "ky"),
    ("Louisiana", "la"), ("Maine", "me"), ("Maryland", "md"),
    ("Massachusetts", "ma"), ("Michigan", "mi"), ("Minnesota", "mn"),
    ("Mississippi", "ms"), ("Missouri", "mo"), ("Montana", "mt"),
    ("Nebraska", "ne"), ("Nevada", "nv"), ("New_Hampshire", "nh"),
    ("New_Jersey", "nj"), ("New_Mexico", "nm"), ("New_York", "ny"),
    ("North_Carolina", "nc"), ("North_Dakota", "nd"), ("Ohio", "oh"),
    ("Oklahoma", "ok"), ("Oregon", "or"), ("Pennsylvania", "pa"),
    ("Rhode_Island", "ri"), ("South_Carolina", "sc"), ("South_Dakota", "sd"),
    ("Tennessee", "tn"), ("Texas", "tx"), ("Utah", "ut"), ("Vermont", "vt"),
    ("Virginia", "va"), ("Washington", "wa"), ("West_Virginia", "wv"),
    ("Wisconsin", "wi"), ("Wyoming", "wy")]

ffiec = "http://www.ffiec.gov/hmdarawdata/OTHER/"

def load_transmittal(working_dir):
    """Download and run the load_transmittal command"""
    with lcd(working_dir):
        local("wget %s2013HMDAInstitutionRecords.zip" % ffiec)
        local("unzip 2013HMDAInstitutionRecords.zip")
        local("rm 2013HMDAInstitutionRecords.zip")
    with lcd("../institutions"):
        local("python manage.py load_transmittal "
              + working_dir + "/2013HMDAInstitutionRecords.txt")
    with lcd(working_dir):
        local("rm 2013HMDAInstitutionRecords.txt")

def load_reporter_panel(working_dir):
    """Download and run the load_reporter_panel command"""
    with lcd(working_dir):
        local("wget %s2013HMDAReporterPanel.zip" % ffiec)
        local("unzip 2013HMDAReporterPanel.zip")
        local("rm 2013HMDAReporterPanel.zip")
    with lcd("../institutions"):
        local("python manage.py load_reporter_panel "
              + working_dir + "/2013HMDAReporterPanel.dat")
    with lcd(working_dir):
        local("rm 2013HMDAReporterPanel.dat")

def load_respondants(working_dir):
    """Load everything for the respondants app"""
    with lcd("../institutions"):
        local("python manage.py loaddata agency")
    load_transmittal(working_dir)
    load_reporter_panel(working_dir)

def load_geos(working_dir):
    """In addition to tracts, we need to load counties, metros, and more."""
    for geo_type in ('county', 'cbsa', 'metdiv'):
        filename = "tl_2014_us_" + geo_type + ".zip"
        with lcd(working_dir):
            local("wget ftp://ftp2.census.gov/geo/tiger/TIGER2014/"
                  + geo_type.upper() + "/" + filename)
            local("unzip " + filename)
        with lcd("../institutions"):
            local("python manage.py load_geos_from " + working_dir + "/"
                  + filename.replace("zip", "shp"))
        with lcd(working_dir):
            local("rm " + filename.replace("zip", "*"))
    with lcd("../institutions"):
        local("python manage.py set_tract_csa_cbsa")

def load_census(working_dir):
    """For all states, download census files and load them into the db"""
    url_tpl = "http://www2.census.gov/census_2010/04-Summary_File_1/%s/"
    url_tpl += "%s2010.sf1.zip"
    for long_name, short_name in states:
        with lcd(working_dir):
            local("wget " + url_tpl % (long_name, short_name))
            local("unzip {0}2010.sf1.zip {0}geo2010.sf1 {0}0000?2010.sf1"
                  .format(short_name))
            local("rm %s2010.sf1.zip" % short_name)
        with lcd("../institutions"):
            local("python manage.py load_summary_one "
                  + working_dir + "/%sgeo2010.sf1" % short_name)
        with lcd(working_dir):
            local("rm %sgeo2010.sf1" % short_name)
            local("rm %s0000?2010.sf1" % short_name)

def load_hmda(working_dir):
    """Download HMDA data and then load it into the db"""
    base_url = "http://www.ffiec.gov/hmdarawdata/LAR/National/"
    filename = "2013HMDALAR - National.zip"
    with lcd(working_dir):
        local("wget '" + base_url + filename + "'")
        local("unzip '" + filename + "'")
        local("rm '" + filename + "'")
        local('split -l 500000 -d "'+ filename.replace("zip", "csv") + '" hmda_csv_')
        local("rm '" + filename.replace("zip", "csv") + "'")
    with lcd("../institutions"):
        local("python manage.py load_hmda '" + working_dir + "' delete_file:true" )
        local("sudo su postgres")
        local("reindex flc_demo)

def precache_hmda():
    """We precache median loan amount per lender x metro"""
    with lcd("../institutions"):
        local("python manage.py calculate_loan_stats")

def load_all(working_dir="/tmp"):
    """Download and import all data for the app"""
    load_respondants(working_dir)
    load_geos(working_dir)
    load_census(working_dir)
    load_hmda(working_dir)
    precache_hmda()
<commit_msg>fixing typo<commit_after>from fabric.api import lcd, local, settings

states = [
    ("Alabama", "al"), ("Alaska", "ak"), ("Arizona", "az"), ("Arkansas", "ar"),
    ("California", "ca"), ("Colorado", "co"), ("Connecticut", "ct"),
    ("Delaware", "de"), ("District_of_Columbia", "dc"), ("Florida", "fl"),
    ("Georgia", "ga"), ("Hawaii", "hi"), ("Idaho", "id"), ("Illinois", "il"),
    ("Indiana", "in"), ("Iowa", "ia"), ("Kansas", "ks"), ("Kentucky", "ky"),
    ("Louisiana", "la"), ("Maine", "me"), ("Maryland", "md"),
    ("Massachusetts", "ma"), ("Michigan", "mi"), ("Minnesota", "mn"),
    ("Mississippi", "ms"), ("Missouri", "mo"), ("Montana", "mt"),
    ("Nebraska", "ne"), ("Nevada", "nv"), ("New_Hampshire", "nh"),
    ("New_Jersey", "nj"), ("New_Mexico", "nm"), ("New_York", "ny"),
    ("North_Carolina", "nc"), ("North_Dakota", "nd"), ("Ohio", "oh"),
    ("Oklahoma", "ok"), ("Oregon", "or"), ("Pennsylvania", "pa"),
    ("Rhode_Island", "ri"), ("South_Carolina", "sc"), ("South_Dakota", "sd"),
    ("Tennessee", "tn"), ("Texas", "tx"), ("Utah", "ut"), ("Vermont", "vt"),
    ("Virginia", "va"), ("Washington", "wa"), ("West_Virginia", "wv"),
    ("Wisconsin", "wi"), ("Wyoming", "wy")]

ffiec = "http://www.ffiec.gov/hmdarawdata/OTHER/"

def load_transmittal(working_dir):
    """Download and run the load_transmittal command"""
    with lcd(working_dir):
        local("wget %s2013HMDAInstitutionRecords.zip" % ffiec)
        local("unzip 2013HMDAInstitutionRecords.zip")
        local("rm 2013HMDAInstitutionRecords.zip")
    with lcd("../institutions"):
        local("python manage.py load_transmittal "
              + working_dir + "/2013HMDAInstitutionRecords.txt")
    with lcd(working_dir):
        local("rm 2013HMDAInstitutionRecords.txt")

def load_reporter_panel(working_dir):
    """Download and run the load_reporter_panel command"""
    with lcd(working_dir):
        local("wget %s2013HMDAReporterPanel.zip" % ffiec)
        local("unzip 2013HMDAReporterPanel.zip")
        local("rm 2013HMDAReporterPanel.zip")
    with lcd("../institutions"):
        local("python manage.py load_reporter_panel "
              + working_dir + "/2013HMDAReporterPanel.dat")
    with lcd(working_dir):
        local("rm 2013HMDAReporterPanel.dat")

def load_respondants(working_dir):
    """Load everything for the respondants app"""
    with lcd("../institutions"):
        local("python manage.py loaddata agency")
    load_transmittal(working_dir)
    load_reporter_panel(working_dir)

def load_geos(working_dir):
    """In addition to tracts, we need to load counties, metros, and more."""
    for geo_type in ('county', 'cbsa', 'metdiv'):
        filename = "tl_2014_us_" + geo_type + ".zip"
        with lcd(working_dir):
            local("wget ftp://ftp2.census.gov/geo/tiger/TIGER2014/"
                  + geo_type.upper() + "/" + filename)
            local("unzip " + filename)
        with lcd("../institutions"):
            local("python manage.py load_geos_from " + working_dir + "/"
                  + filename.replace("zip", "shp"))
        with lcd(working_dir):
            local("rm " + filename.replace("zip", "*"))
    with lcd("../institutions"):
        local("python manage.py set_tract_csa_cbsa")

def load_census(working_dir):
    """For all states, download census files and load them into the db"""
    url_tpl = "http://www2.census.gov/census_2010/04-Summary_File_1/%s/"
    url_tpl += "%s2010.sf1.zip"
    for long_name, short_name in states:
        with lcd(working_dir):
            local("wget " + url_tpl % (long_name, short_name))
            local("unzip {0}2010.sf1.zip {0}geo2010.sf1 {0}0000?2010.sf1"
                  .format(short_name))
            local("rm %s2010.sf1.zip" % short_name)
        with lcd("../institutions"):
            local("python manage.py load_summary_one "
                  + working_dir + "/%sgeo2010.sf1" % short_name)
        with lcd(working_dir):
            local("rm %sgeo2010.sf1" % short_name)
            local("rm %s0000?2010.sf1" % short_name)

def load_hmda(working_dir):
    """Download HMDA data and then load it into the db"""
    base_url = "http://www.ffiec.gov/hmdarawdata/LAR/National/"
    filename = "2013HMDALAR - National.zip"
    with lcd(working_dir):
        local("wget '" + base_url + filename + "'")
        local("unzip '" + filename + "'")
        local("rm '" + filename + "'")
        local('split -l 500000 -d "'+ filename.replace("zip", "csv") + '" hmda_csv_')
        local("rm '" + filename.replace("zip", "csv") + "'")
    with lcd("../institutions"):
        local("python manage.py load_hmda '" + working_dir + "' delete_file:true" )
        local("sudo su postgres")
        local("reindex flc_demo")

def precache_hmda():
    """We precache median loan amount per lender x metro"""
    with lcd("../institutions"):
        local("python manage.py calculate_loan_stats")

def load_all(working_dir="/tmp"):
    """Download and import all data for the app"""
    load_respondants(working_dir)
    load_geos(working_dir)
    load_census(working_dir)
    load_hmda(working_dir)
    precache_hmda()
<|endoftext|>


<commit_before>__version__ = '0.6.1'

__author__ = "Massimiliano Pippi & Federico Frenguelli"

VERSION = __version__  # synonym
<commit_msg>bumped version<commit_after>__version__ = '0.7.0'

__author__ = "Massimiliano Pippi & Federico Frenguelli"

VERSION = __version__  # synonym
<|endoftext|>


<commit_before>"""
Custom Authenticator to use OpenShift OAuth with JupyterHub.

Derived from the GitHub OAuth authenticator.
"""

import json
import os

import requests
from jupyterhub.auth import LocalAuthenticator
from tornado import web
from tornado.auth import OAuth2Mixin
from tornado.httpclient import AsyncHTTPClient, HTTPClient, HTTPRequest
from tornado.httputil import url_concat

from oauthenticator.oauth2 import OAuthenticator
from traitlets import Bool, Set, Unicode, default

class OpenShiftOAuthenticator(OAuthenticator):

    login_service = "OpenShift"

    scope = ['user:info']

    openshift_url = Unicode(
        os.environ.get('OPENSHIFT_URL')
        or 'https://openshift.default.svc.cluster.local',
        config=True,
    )

    validate_cert = Bool(
        True, config=True, help="Set to False to disable certificate validation"
    )

    ca_certs = Unicode(config=True)

    allowed_groups = Set(
        config=True,
        help="Set of OpenShift groups that should be allowed to access the hub.",
    )

    admin_groups = Set(
        config=True,
        help="Set of OpenShift groups that should be given admin access to the hub.",
    )

    @default("ca_certs")
    def _ca_certs_default(self):
        ca_cert_file = "/run/secrets/kubernetes.io/serviceaccount/ca.crt"
        if self.validate_cert and os.path.exists(ca_cert_file):
            return ca_cert_file

        return ''

    openshift_auth_api_url = Unicode(config=True)

    @default("openshift_auth_api_url")
    def _openshift_auth_api_url_default(self):
        auth_info_url = '%s/.well-known/oauth-authorization-server' % self.openshift_url

        resp = requests.get(auth_info_url, verify=self.ca_certs or self.validate_cert)
        resp_json = resp.json()

        return resp_json.get('issuer')

    openshift_rest_api_url = Unicode(
        os.environ.get('OPENSHIFT_REST_API_URL')
        or 'https://openshift.default.svc.cluster.local',
        config=True,
    )

    @default("openshift_rest_api_url")
    def _openshift_rest_api_url_default(self):
        return self.openshift_url

    @default("authorize_url")
    def _authorize_url_default(self):
        return "%s/oauth/authorize" % self.openshift_auth_api_url

    @default("token_url")
    def _token_url_default(self):
        return "%s/oauth/token" % self.openshift_auth_api_url

    @default("userdata_url")
    def _userdata_url_default(self):
        return "%s/apis/user.openshift.io/v1/users/~" % self.openshift_rest_api_url

    @staticmethod
    def user_in_groups(user_groups: set, allowed_groups: set):
        return any(user_groups.intersection(allowed_groups))

    async def authenticate(self, handler, data=None):
        code = handler.get_argument("code")
        # TODO: Configure the curl_httpclient for tornado

        http_client = AsyncHTTPClient()

        # Exchange the OAuth code for a OpenShift Access Token
        #
        # See: https://docs.openshift.org/latest/architecture/additional_concepts/authentication.html#api-authentication

        params = dict(
            client_id=self.client_id,
            client_secret=self.client_secret,
            grant_type="authorization_code",
            code=code,
        )

        url = url_concat(self.token_url, params)

        req = HTTPRequest(
            url,
            method="POST",
            validate_cert=self.validate_cert,
            ca_certs=self.ca_certs,
            headers={"Accept": "application/json"},
            body='',  # Body is required for a POST...
        )

        resp = await http_client.fetch(req)

        resp_json = json.loads(resp.body.decode('utf8', 'replace'))

        access_token = resp_json['access_token']

        # Determine who the logged in user is
        headers = {
            "Accept": "application/json",
            "User-Agent": "JupyterHub",
            "Authorization": "Bearer {}".format(access_token),
        }

        req = HTTPRequest(
            self.userdata_url,
            method="GET",
            validate_cert=self.validate_cert,
            ca_certs=self.ca_certs,
            headers=headers,
        )

        resp = await http_client.fetch(req)
        ocp_user = json.loads(resp.body.decode('utf8', 'replace'))

        username = ocp_user['metadata']['name']

        user_info = {
            'name': username,
            'auth_state': {'access_token': access_token, 'openshift_user': ocp_user},
        }

        if self.allowed_groups or self.admin_groups:
            user_info = await self._add_openshift_group_info(user_info)

        return user_info

    async def _add_openshift_group_info(self, user_info: dict):
        """
        Use the group info stored on the OpenShift User object to determine if a user
        is authenticated based on groups, an admin, or both.
        """
        user_groups = set(user_info['auth_state']['openshift_user']['groups'])

        if self.admin_groups:
            is_admin = self.user_in_groups(user_groups, self.admin_groups)

        user_in_allowed_group = self.user_in_groups(user_groups, self.allowed_groups)

        if self.admin_groups and (is_admin or user_in_allowed_group):
            user_info['admin'] = is_admin
            return user_info
        elif user_in_allowed_group:
            return user_info
        else:
            msg = "username:{username} User not in any of the allowed/admin groups"
            self.log.warning(msg.format(username=user_info['name']))
            return None

class LocalOpenShiftOAuthenticator(LocalAuthenticator, OpenShiftOAuthenticator):

    """A version that mixes in local system user creation"""

    pass
<commit_msg>Update oauthenticator/openshift.py<commit_after>"""
Custom Authenticator to use OpenShift OAuth with JupyterHub.

Derived from the GitHub OAuth authenticator.
"""

import json
import os

import requests
from jupyterhub.auth import LocalAuthenticator
from tornado import web
from tornado.auth import OAuth2Mixin
from tornado.httpclient import AsyncHTTPClient, HTTPClient, HTTPRequest
from tornado.httputil import url_concat

from oauthenticator.oauth2 import OAuthenticator
from traitlets import Bool, Set, Unicode, default

class OpenShiftOAuthenticator(OAuthenticator):

    login_service = "OpenShift"

    scope = ['user:info']

    openshift_url = Unicode(
        os.environ.get('OPENSHIFT_URL')
        or 'https://openshift.default.svc.cluster.local',
        config=True,
    )

    validate_cert = Bool(
        True, config=True, help="Set to False to disable certificate validation"
    )

    ca_certs = Unicode(config=True)

    allowed_groups = Set(
        config=True,
        help="Set of OpenShift groups that should be allowed to access the hub.",
    )

    admin_groups = Set(
        config=True,
        help="Set of OpenShift groups that should be given admin access to the hub.",
    )

    @default("ca_certs")
    def _ca_certs_default(self):
        ca_cert_file = "/run/secrets/kubernetes.io/serviceaccount/ca.crt"
        if self.validate_cert and os.path.exists(ca_cert_file):
            return ca_cert_file

        return ''

    openshift_auth_api_url = Unicode(config=True)

    @default("openshift_auth_api_url")
    def _openshift_auth_api_url_default(self):
        auth_info_url = '%s/.well-known/oauth-authorization-server' % self.openshift_url

        resp = requests.get(auth_info_url, verify=self.ca_certs or self.validate_cert)
        resp_json = resp.json()

        return resp_json.get('issuer')

    openshift_rest_api_url = Unicode(
        os.environ.get('OPENSHIFT_REST_API_URL')
        or 'https://openshift.default.svc.cluster.local',
        config=True,
    )

    @default("openshift_rest_api_url")
    def _openshift_rest_api_url_default(self):
        return self.openshift_url

    @default("authorize_url")
    def _authorize_url_default(self):
        return "%s/oauth/authorize" % self.openshift_auth_api_url

    @default("token_url")
    def _token_url_default(self):
        return "%s/oauth/token" % self.openshift_auth_api_url

    @default("userdata_url")
    def _userdata_url_default(self):
        return "%s/apis/user.openshift.io/v1/users/~" % self.openshift_rest_api_url

    @staticmethod
    def user_in_groups(user_groups: set, allowed_groups: set):
        return any(user_groups.intersection(allowed_groups))

    async def authenticate(self, handler, data=None):
        code = handler.get_argument("code")
        # TODO: Configure the curl_httpclient for tornado

        http_client = AsyncHTTPClient()

        # Exchange the OAuth code for a OpenShift Access Token
        #
        # See: https://docs.openshift.org/latest/architecture/additional_concepts/authentication.html#api-authentication

        params = dict(
            client_id=self.client_id,
            client_secret=self.client_secret,
            grant_type="authorization_code",
            code=code,
        )

        url = url_concat(self.token_url, params)

        req = HTTPRequest(
            url,
            method="POST",
            validate_cert=self.validate_cert,
            ca_certs=self.ca_certs,
            headers={"Accept": "application/json"},
            body='',  # Body is required for a POST...
        )

        resp = await http_client.fetch(req)

        resp_json = json.loads(resp.body.decode('utf8', 'replace'))

        access_token = resp_json['access_token']

        # Determine who the logged in user is
        headers = {
            "Accept": "application/json",
            "User-Agent": "JupyterHub",
            "Authorization": "Bearer {}".format(access_token),
        }

        req = HTTPRequest(
            self.userdata_url,
            method="GET",
            validate_cert=self.validate_cert,
            ca_certs=self.ca_certs,
            headers=headers,
        )

        resp = await http_client.fetch(req)
        ocp_user = json.loads(resp.body.decode('utf8', 'replace'))

        username = ocp_user['metadata']['name']

        user_info = {
            'name': username,
            'auth_state': {'access_token': access_token, 'openshift_user': ocp_user},
        }

        if self.allowed_groups or self.admin_groups:
            user_info = await self._add_openshift_group_info(user_info)

        return user_info

    async def _add_openshift_group_info(self, user_info: dict):
        """
        Use the group info stored on the OpenShift User object to determine if a user
        is authenticated based on groups, an admin, or both.
        """
        user_groups = set(user_info['auth_state']['openshift_user']['groups'])

        if self.admin_groups:
            is_admin = self.user_in_groups(user_groups, self.admin_groups)

        user_in_allowed_group = self.user_in_groups(user_groups, self.allowed_groups)

        if self.admin_groups and (is_admin or user_in_allowed_group):
            user_info['admin'] = is_admin
            return user_info
        elif user_in_allowed_group:
            return user_info
        else:
            msg = f"username:{username} User not in any of the allowed/admin groups"
            self.log.warning(msg.format(username=user_info['name']))
            return None

class LocalOpenShiftOAuthenticator(LocalAuthenticator, OpenShiftOAuthenticator):

    """A version that mixes in local system user creation"""

    pass
<|endoftext|>


<commit_before>__author__ = 'weijia'

INSTALLED_APPS += (
    'mptt',
    'reversion',
    'django_mptt_admin',
    'tagging',
    'ajax_select',
    'django_extensions',
    'geoposition',
    'obj_sys',
    # "obj_sys.apps.ObjSysConfig",
    # 'endless_pagination',
    'el_pagination',
    'crispy_forms',
)

TEMPLATE_CONTEXT_PROCESSORS += (
    'django.core.context_processors.request',
)

# MIDDLEWARE_CLASSES += (
#     'reversion.middleware.RevisionMiddleware',
# )
REST_FRAMEWORK = {
    'DEFAULT_PAGINATION_CLASS': 'rest_framework.pagination.LimitOffsetPagination',
    'PAGE_SIZE': 20,
    'DEFAULT_FILTER_BACKENDS': ('django_filters.rest_framework.backends.DjangoFilterBackend', )
}<commit_msg>Fix GeoPosition setting.<commit_after>__author__ = 'weijia'

INSTALLED_APPS += (
    'mptt',
    'reversion',
    'django_mptt_admin',
    'tagging',
    'ajax_select',
    'django_extensions',
    'geoposition',
    'obj_sys',
    # "obj_sys.apps.ObjSysConfig",
    # 'endless_pagination',
    'el_pagination',
    'crispy_forms',
)

TEMPLATE_CONTEXT_PROCESSORS += (
    'django.core.context_processors.request',
)

# MIDDLEWARE_CLASSES += (
#     'reversion.middleware.RevisionMiddleware',
# )
REST_FRAMEWORK = {
    'DEFAULT_PAGINATION_CLASS': 'rest_framework.pagination.LimitOffsetPagination',
    'PAGE_SIZE': 20,
    'DEFAULT_FILTER_BACKENDS': ('django_filters.rest_framework.backends.DjangoFilterBackend', )
}

GEOPOSITION_GOOGLE_MAPS_API_KEY = ""
<|endoftext|>