<commit_before># Copyright 2020 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Provides a workspace macro to load dependent repositories."""

load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
load("@bazel_tools//tools/build_defs/repo:utils.bzl", "maybe")
load("@rules_cuda//cuda:dependencies.bzl", "rules_cuda_dependencies")
load("@tf_runtime//third_party:repo.bzl", "tfrt_http_archive")
load("@tf_runtime//third_party/cuda:dependencies.bzl", "cuda_dependencies")

def tfrt_dependencies():
    """Loads TFRT external dependencies into WORKSPACE."""

    rules_cuda_dependencies()

    cuda_dependencies()

    # LLVM and MLIR from llvm-project.
    LLVM_COMMIT = "e990fa2170314b179ec025b68fd00fbe9aab398d"
    LLVM_SHA256 = "3d071fad8606faa555518b043e5100c4a03205ac09a503fe30f829fe024eb999"
    LLVM_URLS = [
        "https://storage.googleapis.com/mirror.tensorflow.org/github.com/llvm/llvm-project/archive/{commit}.tar.gz".format(commit = LLVM_COMMIT),
        "https://github.com/llvm/llvm-project/archive/{commit}.tar.gz".format(commit = LLVM_COMMIT),
    ]
    maybe(
        name = "llvm-project",
        repo_rule = tfrt_http_archive,
        sha256 = LLVM_SHA256,
        strip_prefix = "llvm-project-" + LLVM_COMMIT,
        urls = LLVM_URLS,
        additional_build_files = {
            "@tf_runtime//third_party/llvm:llvm.autogenerated.BUILD": "llvm/BUILD",
            "@tf_runtime//third_party/mlir:BUILD": "mlir/BUILD",
            "@tf_runtime//third_party/mlir:test.BUILD": "mlir/test/BUILD",
        },
    )

    # https://github.com/bazelbuild/bazel-skylib/releases
    maybe(
        name = "bazel_skylib",
        repo_rule = http_archive,
        sha256 = "97e70364e9249702246c0e9444bccdc4b847bed1eb03c5a3ece4f83dfe6abc44",
        urls = [
            "https://storage.googleapis.com/mirror.tensorflow.org/github.com/bazelbuild/bazel-skylib/releases/download/1.0.2/bazel-skylib-1.0.2.tar.gz",
            "https://github.com/bazelbuild/bazel-skylib/releases/download/1.0.2/bazel-skylib-1.0.2.tar.gz",
        ],
    )

    # Eigen3
    maybe(
        name = "eigen_archive",
        repo_rule = http_archive,
        build_file = "@tf_runtime//third_party:eigen/BUILD",
        sha256 = "924c7f85d5e2e40beb663489bf3f6908bb4c328dc81ab845c27bb6ce199b8698",  # SHARED_EIGEN_SHA
        strip_prefix = "eigen-3c02fefec53f21d9fad537ff0d62d8dcc8162466",
        urls = [
            "https://storage.googleapis.com/mirror.tensorflow.org/gitlab.com/libeigen/eigen/-/archive/3c02fefec53f21d9fad537ff0d62d8dcc8162466/eigen-3c02fefec53f21d9fad537ff0d62d8dcc8162466.tar.gz",
            "https://gitlab.com/libeigen/eigen/-/archive/3c02fefec53f21d9fad537ff0d62d8dcc8162466/eigen-3c02fefec53f21d9fad537ff0d62d8dcc8162466.tar.gz",
        ],
    )

    maybe(
        name = "dnnl",
        repo_rule = http_archive,
        build_file = "@tf_runtime//third_party/dnnl:BUILD",
        sha256 = "5369f7b2f0b52b40890da50c0632c3a5d1082d98325d0f2bff125d19d0dcaa1d",
        strip_prefix = "oneDNN-1.6.4",
        urls = [
            "https://storage.googleapis.com/mirror.tensorflow.org/github.com/oneapi-src/oneDNN/archive/v1.6.4.tar.gz",
            "https://github.com/oneapi-src/oneDNN/archive/v1.6.4.tar.gz",
        ],
    )

    maybe(
        name = "com_google_googletest",
        repo_rule = tfrt_http_archive,
        sha256 = "ff7a82736e158c077e76188232eac77913a15dac0b22508c390ab3f88e6d6d86",
        strip_prefix = "googletest-b6cd405286ed8635ece71c72f118e659f4ade3fb",
        urls = [
            "https://storage.googleapis.com/mirror.tensorflow.org/github.com/google/googletest/archive/b6cd405286ed8635ece71c72f118e659f4ade3fb.zip",
            "https://github.com/google/googletest/archive/b6cd405286ed8635ece71c72f118e659f4ade3fb.zip",
        ],
    )

    maybe(
        name = "com_github_google_benchmark",
        repo_rule = http_archive,
        strip_prefix = "benchmark-16703ff83c1ae6d53e5155df3bb3ab0bc96083be",
        sha256 = "59f918c8ccd4d74b6ac43484467b500f1d64b40cc1010daa055375b322a43ba3",
        urls = ["https://github.com/google/benchmark/archive/16703ff83c1ae6d53e5155df3bb3ab0bc96083be.zip"],
    )

    maybe(
        name = "com_google_protobuf",
        repo_rule = tfrt_http_archive,
        sha256 = "bf0e5070b4b99240183b29df78155eee335885e53a8af8683964579c214ad301",
        strip_prefix = "protobuf-3.14.0",
        system_build_file = "//third_party/systemlibs:protobuf.BUILD",
        system_link_files = {
            "//third_party/systemlibs:protobuf.bzl": "protobuf.bzl",
        },
        urls = [
            "https://storage.googleapis.com/mirror.tensorflow.org/github.com/protocolbuffers/protobuf/archive/v3.14.0.zip",
            "https://github.com/protocolbuffers/protobuf/archive/v3.14.0.zip",
        ],
    )

    maybe(
        name = "cub_archive",
        repo_rule = tfrt_http_archive,
        build_file = "//third_party:cub/BUILD",
        patch_file = "//third_party:cub/pr170.patch",
        sha256 = "6bfa06ab52a650ae7ee6963143a0bbc667d6504822cbd9670369b598f18c58c3",
        strip_prefix = "cub-1.8.0",
        urls = [
            "https://storage.googleapis.com/mirror.tensorflow.org/github.com/NVlabs/cub/archive/1.8.0.zip",
            "https://github.com/NVlabs/cub/archive/1.8.0.zip",
        ],
    )

    maybe(
        name = "zlib",
        repo_rule = tfrt_http_archive,
        build_file = "//third_party:zlib.BUILD",
        sha256 = "c3e5e9fdd5004dcb542feda5ee4f0ff0744628baf8ed2dd5d66f8ca1197cb1a1",
        strip_prefix = "zlib-1.2.11",
        system_build_file = "//third_party/systemlibs:zlib.BUILD",
        urls = [
            "https://storage.googleapis.com/mirror.tensorflow.org/zlib.net/zlib-1.2.11.tar.gz",
            "https://zlib.net/zlib-1.2.11.tar.gz",
        ],
    )

    maybe(
        name = "py-cpuinfo",
        repo_rule = tfrt_http_archive,
        strip_prefix = "py-cpuinfo-0.2.3",
        sha256 = "f6a016fdbc4e7fadf2d519090fcb4fa9d0831bad4e85245d938e5c2fe7623ca6",
        urls = [
            "https://storage.googleapis.com/mirror.tensorflow.org/pypi.python.org/packages/source/p/py-cpuinfo/py-cpuinfo-0.2.3.tar.gz",
            "https://pypi.python.org/packages/source/p/py-cpuinfo/py-cpuinfo-0.2.3.tar.gz",
        ],
        build_file = "//third_party:py-cpuinfo.BUILD",
    )
<commit_msg>Integrate LLVM at llvm/llvm-project@5657f93e788f<commit_after># Copyright 2020 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Provides a workspace macro to load dependent repositories."""

load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
load("@bazel_tools//tools/build_defs/repo:utils.bzl", "maybe")
load("@rules_cuda//cuda:dependencies.bzl", "rules_cuda_dependencies")
load("@tf_runtime//third_party:repo.bzl", "tfrt_http_archive")
load("@tf_runtime//third_party/cuda:dependencies.bzl", "cuda_dependencies")

def tfrt_dependencies():
    """Loads TFRT external dependencies into WORKSPACE."""

    rules_cuda_dependencies()

    cuda_dependencies()

    # LLVM and MLIR from llvm-project.
    LLVM_COMMIT = "5657f93e788f093c70fb448dd6f9398b149df278"
    LLVM_SHA256 = "47f798395db279dba107a35bffe4c32551d58a5bb443f9cb94e4f309a778f46e"
    LLVM_URLS = [
        "https://storage.googleapis.com/mirror.tensorflow.org/github.com/llvm/llvm-project/archive/{commit}.tar.gz".format(commit = LLVM_COMMIT),
        "https://github.com/llvm/llvm-project/archive/{commit}.tar.gz".format(commit = LLVM_COMMIT),
    ]
    maybe(
        name = "llvm-project",
        repo_rule = tfrt_http_archive,
        sha256 = LLVM_SHA256,
        strip_prefix = "llvm-project-" + LLVM_COMMIT,
        urls = LLVM_URLS,
        additional_build_files = {
            "@tf_runtime//third_party/llvm:llvm.autogenerated.BUILD": "llvm/BUILD",
            "@tf_runtime//third_party/mlir:BUILD": "mlir/BUILD",
            "@tf_runtime//third_party/mlir:test.BUILD": "mlir/test/BUILD",
        },
    )

    # https://github.com/bazelbuild/bazel-skylib/releases
    maybe(
        name = "bazel_skylib",
        repo_rule = http_archive,
        sha256 = "97e70364e9249702246c0e9444bccdc4b847bed1eb03c5a3ece4f83dfe6abc44",
        urls = [
            "https://storage.googleapis.com/mirror.tensorflow.org/github.com/bazelbuild/bazel-skylib/releases/download/1.0.2/bazel-skylib-1.0.2.tar.gz",
            "https://github.com/bazelbuild/bazel-skylib/releases/download/1.0.2/bazel-skylib-1.0.2.tar.gz",
        ],
    )

    # Eigen3
    maybe(
        name = "eigen_archive",
        repo_rule = http_archive,
        build_file = "@tf_runtime//third_party:eigen/BUILD",
        sha256 = "924c7f85d5e2e40beb663489bf3f6908bb4c328dc81ab845c27bb6ce199b8698",  # SHARED_EIGEN_SHA
        strip_prefix = "eigen-3c02fefec53f21d9fad537ff0d62d8dcc8162466",
        urls = [
            "https://storage.googleapis.com/mirror.tensorflow.org/gitlab.com/libeigen/eigen/-/archive/3c02fefec53f21d9fad537ff0d62d8dcc8162466/eigen-3c02fefec53f21d9fad537ff0d62d8dcc8162466.tar.gz",
            "https://gitlab.com/libeigen/eigen/-/archive/3c02fefec53f21d9fad537ff0d62d8dcc8162466/eigen-3c02fefec53f21d9fad537ff0d62d8dcc8162466.tar.gz",
        ],
    )

    maybe(
        name = "dnnl",
        repo_rule = http_archive,
        build_file = "@tf_runtime//third_party/dnnl:BUILD",
        sha256 = "5369f7b2f0b52b40890da50c0632c3a5d1082d98325d0f2bff125d19d0dcaa1d",
        strip_prefix = "oneDNN-1.6.4",
        urls = [
            "https://storage.googleapis.com/mirror.tensorflow.org/github.com/oneapi-src/oneDNN/archive/v1.6.4.tar.gz",
            "https://github.com/oneapi-src/oneDNN/archive/v1.6.4.tar.gz",
        ],
    )

    maybe(
        name = "com_google_googletest",
        repo_rule = tfrt_http_archive,
        sha256 = "ff7a82736e158c077e76188232eac77913a15dac0b22508c390ab3f88e6d6d86",
        strip_prefix = "googletest-b6cd405286ed8635ece71c72f118e659f4ade3fb",
        urls = [
            "https://storage.googleapis.com/mirror.tensorflow.org/github.com/google/googletest/archive/b6cd405286ed8635ece71c72f118e659f4ade3fb.zip",
            "https://github.com/google/googletest/archive/b6cd405286ed8635ece71c72f118e659f4ade3fb.zip",
        ],
    )

    maybe(
        name = "com_github_google_benchmark",
        repo_rule = http_archive,
        strip_prefix = "benchmark-16703ff83c1ae6d53e5155df3bb3ab0bc96083be",
        sha256 = "59f918c8ccd4d74b6ac43484467b500f1d64b40cc1010daa055375b322a43ba3",
        urls = ["https://github.com/google/benchmark/archive/16703ff83c1ae6d53e5155df3bb3ab0bc96083be.zip"],
    )

    maybe(
        name = "com_google_protobuf",
        repo_rule = tfrt_http_archive,
        sha256 = "bf0e5070b4b99240183b29df78155eee335885e53a8af8683964579c214ad301",
        strip_prefix = "protobuf-3.14.0",
        system_build_file = "//third_party/systemlibs:protobuf.BUILD",
        system_link_files = {
            "//third_party/systemlibs:protobuf.bzl": "protobuf.bzl",
        },
        urls = [
            "https://storage.googleapis.com/mirror.tensorflow.org/github.com/protocolbuffers/protobuf/archive/v3.14.0.zip",
            "https://github.com/protocolbuffers/protobuf/archive/v3.14.0.zip",
        ],
    )

    maybe(
        name = "cub_archive",
        repo_rule = tfrt_http_archive,
        build_file = "//third_party:cub/BUILD",
        patch_file = "//third_party:cub/pr170.patch",
        sha256 = "6bfa06ab52a650ae7ee6963143a0bbc667d6504822cbd9670369b598f18c58c3",
        strip_prefix = "cub-1.8.0",
        urls = [
            "https://storage.googleapis.com/mirror.tensorflow.org/github.com/NVlabs/cub/archive/1.8.0.zip",
            "https://github.com/NVlabs/cub/archive/1.8.0.zip",
        ],
    )

    maybe(
        name = "zlib",
        repo_rule = tfrt_http_archive,
        build_file = "//third_party:zlib.BUILD",
        sha256 = "c3e5e9fdd5004dcb542feda5ee4f0ff0744628baf8ed2dd5d66f8ca1197cb1a1",
        strip_prefix = "zlib-1.2.11",
        system_build_file = "//third_party/systemlibs:zlib.BUILD",
        urls = [
            "https://storage.googleapis.com/mirror.tensorflow.org/zlib.net/zlib-1.2.11.tar.gz",
            "https://zlib.net/zlib-1.2.11.tar.gz",
        ],
    )

    maybe(
        name = "py-cpuinfo",
        repo_rule = tfrt_http_archive,
        strip_prefix = "py-cpuinfo-0.2.3",
        sha256 = "f6a016fdbc4e7fadf2d519090fcb4fa9d0831bad4e85245d938e5c2fe7623ca6",
        urls = [
            "https://storage.googleapis.com/mirror.tensorflow.org/pypi.python.org/packages/source/p/py-cpuinfo/py-cpuinfo-0.2.3.tar.gz",
            "https://pypi.python.org/packages/source/p/py-cpuinfo/py-cpuinfo-0.2.3.tar.gz",
        ],
        build_file = "//third_party:py-cpuinfo.BUILD",
    )
<|endoftext|>


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

import os
import time
from contextlib import closing, contextmanager
from xmlrpc.client import ServerProxy

import html2text
import sh
import yaml
from selenium import webdriver
from xvfbwrapper import Xvfb
import schedule

git = sh.git.bake('--no-pager', _cwd="pages")

def main():
    check_pages()
    schedule.every(5).minutes.do(check_pages)
    while True:
        schedule.run_pending()
        time.sleep(60)

def check_pages():
    for conf in pages_conf():
        save_page(conf, download_page(conf))
    if save_changes():
        report = report_diff()
        notify(report)

def download_page(conf):
    with firefox() as browser:
        browser.get(conf['url'])
        if conf.get('delay'):
            time.sleep(conf['delay'])
        tag_name = conf.get('tag')
        if tag_name:
            elem = browser.find_element_by_tag_name(tag_name)
            html = elem.get_attribute('outerHTML')
        else:
            xpath = conf.get('xpath', '//*')
            elem = browser.find_element_by_xpath(xpath)
            html = elem.get_attribute('outerHTML')
        output_format = conf.get('format', 'html')
        if output_format == 'text':
            return sanitize(html)
        elif output_format == 'html':
            return html

def save_page(conf, content):
    with open(os.path.join('pages', conf['name']), 'w') as fp:
        fp.write(content)

def sanitize(html):
    return html2text.html2text(html)

def pages_conf():
    with open('watch.yml') as fp:
        conf = yaml.load(fp)
    return conf['pages']

def notify_conf():
    with open('watch.yml') as fp:
        conf = yaml.load(fp)
    return conf['on_change']

def save_changes():
    git('add', '-A', '.')
    try:
        git.commit('-m', 'Web watch')
        return True
    except sh.ErrorReturnCode_1:
        return False

def report_diff():
    return git.log('-1', '-p', '--no-color').stdout

def notify(report):
    for rule in notify_conf():
        for key, value in rule.items():
            if key == 'slack':
                post_slack(value, report)

def post_slack(channel, text):
    s = ServerProxy('http://localhost:34278', allow_none=True)
    s.post(channel, text)

@contextmanager
def virtual_buffer():
    """
    Try to start xvfb, trying multiple (up to 5) times if a failure
    """
    for i in range(0, 6):
        xvfb_display = Xvfb()
        xvfb_display.start()
        # If Xvfb started, return.
        if xvfb_display.proc is not None:
            try:
                yield
                return
            finally:
                xvfb_display.stop()
    raise Exception("Xvfb could not be started after six attempts.")

@contextmanager
def firefox():
    with virtual_buffer():
        with closing(webdriver.Firefox()) as driver:
            yield driver

if __name__ == '__main__':
    main()
<commit_msg>[wip] independent schedules<commit_after>#!/usr/bin/env python
import json
import os
import time
from contextlib import closing, contextmanager
import functools

import html2text
import requests
import schedule
import sh
import yaml
from selenium import webdriver
from xmlrpc.client import ServerProxy
from xvfbwrapper import Xvfb

git = sh.git.bake('--no-pager', _cwd="pages")

def main():
    conf = pages_conf()
    check_all_pages(conf)
    schedule_checks(conf)
    while True:
        schedule.run_pending()
        time.sleep(60)

def schedule_checks(confs):
    for conf in confs:
        period = conf.get('period', 300)
        schedule.every(period).seconds.do(functools.partial(check_page, conf))

def check_all_pages(confs):
    for conf in confs:
        check_page(conf)

def check_page(conf):
    save_page(conf, download_page(conf))
    if save_changes():
        report = report_diff()
        notify(conf, report)

def download_page(conf):
    url = conf['url']
    output_format = conf.get('format', 'html')
    if output_format in ('asis', 'json'):
        response = requests.get(url)
        if output_format == 'json':
            return json.dumps(response.json, indent=True)
        else:
            return response.text
    else:
        delay = conf.get('delay')
        tag_name = conf.get('tag')
        xpath = conf.get('xpath', '//*')
        with firefox() as browser:
            browser.get(url)
            if delay:
                time.sleep(delay)
            if tag_name:
                elem = browser.find_element_by_tag_name(tag_name)
                html = elem.get_attribute('outerHTML')
            else:
                elem = browser.find_element_by_xpath(xpath)
                html = elem.get_attribute('outerHTML')
            if output_format == 'text':
                return sanitize(html)
            elif output_format == 'html':
                return html

def save_page(conf, content):
    with open(os.path.join('pages', conf['name']), 'w') as fp:
        fp.write(content)

def sanitize(html):
    return html2text.html2text(html)

def pages_conf():
    with open('watch.yml') as fp:
        conf = yaml.load(fp)
    return conf['pages']

def notify_conf():
    with open('watch.yml') as fp:
        conf = yaml.load(fp)
    return conf['on_change']

def save_changes():
    git('add', '-A', '.')
    try:
        git.commit('-m', 'Web watch')
        return True
    except sh.ErrorReturnCode_1:
        return False

def report_diff():
    return git.log('-1', '-p', '--no-color').stdout

def notify(conf, report):
    for rule in notify_conf():
        for key, value in rule.items():
            if key == 'slack':
                post_slack(value, report)

def post_slack(channel, text):
    s = ServerProxy('http://localhost:34278', allow_none=True)
    s.post(channel, text)

@contextmanager
def virtual_buffer():
    """
    Try to start xvfb, trying multiple (up to 5) times if a failure
    """
    for i in range(0, 6):
        xvfb_display = Xvfb()
        xvfb_display.start()
        # If Xvfb started, return.
        if xvfb_display.proc is not None:
            try:
                yield
                return
            finally:
                xvfb_display.stop()
    raise Exception("Xvfb could not be started after six attempts.")

@contextmanager
def firefox():
    with virtual_buffer():
        with closing(webdriver.Firefox()) as driver:
            yield driver

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


<commit_before>from daterange_filter.filter import DateRangeFilter
from django.contrib.auth.models import User
from django.http import HttpResponse
from molo.profiles.admin import ProfileUserAdmin
from import_export import resources
from wagtailmodeladmin.options import ModelAdmin, wagtailmodeladmin_register
from wagtailmodeladmin.views import IndexView

class DateFilter(DateRangeFilter):
    template = 'admin/user_date_range_filter.html'

class FrontendUsersResource(resources.ModelResource):
    class Meta:
        model = User
        exclude = ('id', 'password', 'is_superuser', 'groups',
                   'user_permissions', 'is_staff')

class ModelAdminTemplate(IndexView):
    def post(self, request, *args, **kwargs):

        drf__date_joined__gte = request.GET.get('drf__date_joined__gte')
        drf__date_joined__lte = request.GET.get('drf__date_joined__lte')
        is_active_exact = request.GET.get('is_active_exact')

        filter_list = {
            'date_joined__range': (drf__date_joined__gte,
                                   drf__date_joined__lte) if
            drf__date_joined__gte and drf__date_joined__lte else None,
            'is_active': is_active_exact
        }

        arguments = {}

        for key, value in filter_list.items():
            if value:
                arguments[key] = value

        dataset = FrontendUsersResource().export(
            User.objects.filter(is_staff=False, **arguments))

        response = HttpResponse(dataset.csv, content_type="csv")
        response['Content-Disposition'] = \
            'attachment; filename=frontend_users.csv'
        return response

    def get_template_names(self):
        return 'admin/molo_admin_template.html'

class FrontendUsersModelAdmin(ModelAdmin, ProfileUserAdmin):
    model = User
    menu_label = 'End Users'
    menu_icon = 'user'
    menu_order = 600
    index_view_class = ModelAdminTemplate
    add_to_settings_menu = True
    list_display = ('username', '_alias', '_mobile_number',
                    'email', 'date_joined', 'is_active')

    list_filter = (('date_joined', DateFilter), 'is_active')

    search_fields = ('username',)

    def get_queryset(self, request):
        queryset = User.objects.filter(is_staff=False)
        return queryset

wagtailmodeladmin_register(FrontendUsersModelAdmin)
<commit_msg>cleaned up wagtail_hooks.py<commit_after>from molo.profiles.admin import FrontendUsersModelAdmin
from wagtailmodeladmin.options import wagtailmodeladmin_register

wagtailmodeladmin_register(FrontendUsersModelAdmin)
<|endoftext|>


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

import time
import datetime
import threading
import sys
import os
import signal
from traceback import print_exception

import logbook
import beanstalkc
import anyjson

import tweetsclient

def beanstalk(host='localhost', port=11300, watch=None, use=None):
    beanstalk = beanstalkc.Connection(host=host, port=port)
    if use:
        beanstalk.use(use)
    if watch:
        beanstalk.watch(watch)
    return beanstalk

def configure_log_handler(application_name, loglevel, output):
    if isinstance(loglevel, (str, unicode)):
        loglevel = getattr(logbook, loglevel.upper())

    if not isinstance(loglevel, int):
        raise TypeError("configure_log_handler expects loglevel to be either an integer or a string corresponding to an integer attribute of the logbook module.")

    if output == 'syslog':
        log_handler = logbook.SyslogHandler(
            application_name=application_name,
            facility='user',
            bubble=False,
            level=loglevel)
    elif output == '-' or not output:
        log_handler = logbook.StderrHandler(
            level=loglevel,
            bubble=False)
    else:
        log_handler = logbook.FileHandler(
            filename=output,
            encoding='utf-8',
            level=loglevel,
            bubble=False)

    return log_handler

def run_with_restart(fn, max_restart=0, args=(), kwargs={}):
    restartCounter = 0
    while True:
        try:
            return apply(fn, args, kwargs)
        except AssertionError:
            raise
        except Exception as e:
            logbook.error("Unhandled exception of type {exctype}: {exception}",
                          exctype=type(e),
                          exception=str(e))

            restartCounter += 1
            if max_restart and restartCounter > max_restart:
                logbook.critical("Alreadying restarted {nth} times. Exiting.",
                                 nth=restartCounter)
            else:
                # Sleep longer each time we restart, but no more than 5 minutes
                delay = min(restartCounter * 15, 300)
                logbook.error("Restarting for {nth} time in {sec} seconds.",
                              nth=restartCounter,
                              sec=delay)
                time.sleep(delay)

def restart_process(signum, frame):
    """
    Replaces the current process with a new process invoked
    using the same command line.
    """
    os.execl(sys.executable, sys.executable, *sys.argv)

def start_heartbeat_thread(heart):
    """
    Triggers a regular heartbeat from a background thread
    for scripts that making blocking calls and thus can't
    heartbeat from their main loop.
    """
    def _heartbeat():
        while True:
            heart.sleep()
            heart.beat()

    heartbeat = threading.Thread(target=_heartbeat)
    # This causes the heartbeat thread to die with the main thread
    heartbeat.daemon = True
    heartbeat.start()

def start_watchdog_thread(heart):
    """
    Watch a heartbeat file and restart when the file mtime is either
    too old or too far in the future.
    """

    def _watchdog():
        while True:
            heart.sleep()
            now = datetime.datetime.now()
            try:
                stat = os.stat(heart.filepath)
                mtime = datetime.datetime.fromtimestamp(stat.st_mtime)
            except OSError as e:
                if e.errno == 2: # No such file or directory
                    logbook.warning("Heartbeat file disappeared, restarting via SIGHUP.")
                    os.kill(heart.pid, signal.SIGHUP)
                    return
                else:
                    raise

            if mtime >= now:
                logbook.warning("Heartbeat file mtime is in the future, restarting via SIGHUP.")
                os.kill(heart.pid, signal.SIGHUP)
                return

    watchdog = threading.Thread(target=_watchdog)
    # This causes the watchdog thread to die with the main thread
    watchdog.daemon = True
    watchdog.start()

class Heart(object):
    """
    Updates the access and modification timestamp of a
    file every `interval` seconds.
    """
    def __init__(self):
        self.last_beat = datetime.datetime.now()

        config = tweetsclient.Config().get()
        try:
            self.interval = datetime.timedelta(seconds=float(config.get('tweets-client', 'heartbeat_interval')))
        except:
            logbook.warning("No heartbeat_interval configuration parameter, skipping heartbeat.")
            raise StopIteration

        try:
            directory = config.get('tweets-client', 'heartbeats_directory')
        except:
            logbook.warning("No heartbeats_directory configuration parameter, skipping heartbeat.")
            raise StopIteration

        if not os.path.isdir(directory):
            logbook.warning("The heartbeats_directory parameter ({0}) is not a directory.",
                             directory)
            raise StopIteration

        scriptname = os.path.basename(sys.argv[0])
        self.filepath = os.path.join(directory, scriptname)

        start_time = datetime.datetime.now().isoformat()
        self.pid = os.getpid()
        with file(self.filepath, 'w') as fil:
            fil.write(anyjson.serialize({
                'pid': self.pid,
                'started': start_time
            }))

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        if (exc_type, exc_value, traceback) == (None, None, None):
            os.unlink(self.filepath)
        else:
            with file(self.filepath, 'w') as outf:
                print_exception(exc_type, exc_value, traceback, 1000, outf)

    def sleep(self):
        while True:
            now = datetime.datetime.now()
            since = now - self.last_beat
            if since >= self.interval:
                return
            else:
                time.sleep(self.interval.total_seconds() * 0.10)

    def beat(self):
        now = datetime.datetime.now()
        since = now - self.last_beat
        if since >= self.interval:
            os.utime(self.filepath, None)
            self.last_beat = now
            return True
        else:
            return False

<commit_msg>Moved time calculation to avoid thread-switching skew.<commit_after>#!/user/bin/env python

import time
import datetime
import threading
import sys
import os
import signal
from traceback import print_exception

import logbook
import beanstalkc
import anyjson

import tweetsclient

def beanstalk(host='localhost', port=11300, watch=None, use=None):
    beanstalk = beanstalkc.Connection(host=host, port=port)
    if use:
        beanstalk.use(use)
    if watch:
        beanstalk.watch(watch)
    return beanstalk

def configure_log_handler(application_name, loglevel, output):
    if isinstance(loglevel, (str, unicode)):
        loglevel = getattr(logbook, loglevel.upper())

    if not isinstance(loglevel, int):
        raise TypeError("configure_log_handler expects loglevel to be either an integer or a string corresponding to an integer attribute of the logbook module.")

    if output == 'syslog':
        log_handler = logbook.SyslogHandler(
            application_name=application_name,
            facility='user',
            bubble=False,
            level=loglevel)
    elif output == '-' or not output:
        log_handler = logbook.StderrHandler(
            level=loglevel,
            bubble=False)
    else:
        log_handler = logbook.FileHandler(
            filename=output,
            encoding='utf-8',
            level=loglevel,
            bubble=False)

    return log_handler

def run_with_restart(fn, max_restart=0, args=(), kwargs={}):
    restartCounter = 0
    while True:
        try:
            return apply(fn, args, kwargs)
        except AssertionError:
            raise
        except Exception as e:
            logbook.error("Unhandled exception of type {exctype}: {exception}",
                          exctype=type(e),
                          exception=str(e))

            restartCounter += 1
            if max_restart and restartCounter > max_restart:
                logbook.critical("Alreadying restarted {nth} times. Exiting.",
                                 nth=restartCounter)
            else:
                # Sleep longer each time we restart, but no more than 5 minutes
                delay = min(restartCounter * 15, 300)
                logbook.error("Restarting for {nth} time in {sec} seconds.",
                              nth=restartCounter,
                              sec=delay)
                time.sleep(delay)

def restart_process(signum, frame):
    """
    Replaces the current process with a new process invoked
    using the same command line.
    """
    os.execl(sys.executable, sys.executable, *sys.argv)

def start_heartbeat_thread(heart):
    """
    Triggers a regular heartbeat from a background thread
    for scripts that making blocking calls and thus can't
    heartbeat from their main loop.
    """
    def _heartbeat():
        while True:
            heart.sleep()
            heart.beat()

    heartbeat = threading.Thread(target=_heartbeat)
    # This causes the heartbeat thread to die with the main thread
    heartbeat.daemon = True
    heartbeat.start()

def start_watchdog_thread(heart):
    """
    Watch a heartbeat file and restart when the file mtime is either
    too old or too far in the future.
    """

    def _watchdog():
        while True:
            heart.sleep()
            try:
                stat = os.stat(heart.filepath)
                mtime = datetime.datetime.fromtimestamp(stat.st_mtime)
            except OSError as e:
                if e.errno == 2: # No such file or directory
                    logbook.warning("Heartbeat file disappeared, restarting via SIGHUP.")
                    os.kill(heart.pid, signal.SIGHUP)
                    return
                else:
                    raise

            now = datetime.datetime.now()
            if mtime >= now:
                logbook.warning("Heartbeat file mtime is in the future, restarting via SIGHUP.")
                os.kill(heart.pid, signal.SIGHUP)
                return

    watchdog = threading.Thread(target=_watchdog)
    # This causes the watchdog thread to die with the main thread
    watchdog.daemon = True
    watchdog.start()

class Heart(object):
    """
    Updates the access and modification timestamp of a
    file every `interval` seconds.
    """
    def __init__(self):
        self.last_beat = datetime.datetime.now()

        config = tweetsclient.Config().get()
        try:
            self.interval = datetime.timedelta(seconds=float(config.get('tweets-client', 'heartbeat_interval')))
        except:
            logbook.warning("No heartbeat_interval configuration parameter, skipping heartbeat.")
            raise StopIteration

        try:
            directory = config.get('tweets-client', 'heartbeats_directory')
        except:
            logbook.warning("No heartbeats_directory configuration parameter, skipping heartbeat.")
            raise StopIteration

        if not os.path.isdir(directory):
            logbook.warning("The heartbeats_directory parameter ({0}) is not a directory.",
                             directory)
            raise StopIteration

        scriptname = os.path.basename(sys.argv[0])
        self.filepath = os.path.join(directory, scriptname)

        start_time = datetime.datetime.now().isoformat()
        self.pid = os.getpid()
        with file(self.filepath, 'w') as fil:
            fil.write(anyjson.serialize({
                'pid': self.pid,
                'started': start_time
            }))

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        if (exc_type, exc_value, traceback) == (None, None, None):
            os.unlink(self.filepath)
        else:
            with file(self.filepath, 'w') as outf:
                print_exception(exc_type, exc_value, traceback, 1000, outf)

    def sleep(self):
        while True:
            now = datetime.datetime.now()
            since = now - self.last_beat
            if since >= self.interval:
                return
            else:
                time.sleep(self.interval.total_seconds() * 0.10)

    def beat(self):
        now = datetime.datetime.now()
        since = now - self.last_beat
        if since >= self.interval:
            os.utime(self.filepath, None)
            self.last_beat = now
            return True
        else:
            return False

<|endoftext|>


<commit_before>import mysql.connector
from mysql.connector import InterfaceError

class DBConnection(object):
    '''
    Represents a connection to an RDBMS.
    '''

    def __init__(self, **kwargs):
        self.conn = mysql.connector.connect(**kwargs)

    def _before_db(self):
        """
        MySQL-specific: attempt to reconnect if our connection has timed out
        """
        try:
            self.conn.ping(True)
        except InterfaceError:
            """
            Probably just a stale connection.
            If something worse has gone wrong, we will see it soon anyway.
            """
            pass

    def execute(self, query, return_dictionaries = True, bind_values = None):
        """
        Execute the given query with the given values for placeholders.
        If return_dictionaries is true, each row is a dictionary;
        if it is false, each row is a list of column values.
        In either case, a resultset is a list of rows.
        """
        cursor = self.conn.cursor(dictionary = return_dictionaries, buffered = False)
        # Make this a strong not a weak reference, to prevent premature GC.
        # Works around change:
        # http://bazaar.launchpad.net/~mysql/myconnpy/1.0/revision/201
        cursor._connection = self.conn
        cursor.execute(query, bind_values)
        return cursor

    def callproc(self, procname, return_dictionaries = True, args = []):
        """
        Execute the given-named stored procedure with the given arguments.
        If return_dictionaries is true, each row is a dictionary;
        if it is false, each row is a list of column values.
        In either case, a resultset is a list of rows.
        """
        cursor = self.conn.cursor(dictionary = return_dictionaries, buffered = False)
        cursor.callproc(procname, args)
        return cursor

class CursorIter(object):

    """
    Iterate over a cursor's emitted rows.
    """

    def __init__(self, cursor):
        self.cursor = cursor

    def next(self):
        row = self.cursor.fetchone()
        if row is None:
            raise StopIteration()
        return row

class CursorSliceIter(CursorIter):

    """
    Iterate over a slice of a cursor's emitted rows,
    so that the n'th column of each row is emitted.
    """

    def __init__(self, cursor, index):
        super(CursorSliceIter, self).__init__(cursor)
        self.index = index

    def next(self):
        row = super(CursorSliceIter, self).next()
        return row[self.index]

class ResultSet(object):

    """
    An iterable SQL result set.
    """

    def __init__(self, cursor):
        self.cursor = cursor

    def __iter__(self):
        return CursorIter(self.cursor)

class ResultSetSlice(ResultSet):

    """
    An iterable slice through an SQL result set.
    """

    def __init__(self, cursor, index):
        super(ResultSetSlice, self).__init__(cursor)
        self.index = index

    def __iter__(self):
        return CursorSliceIter(self.cursor, self.index)
<commit_msg>Defeat MySQL's conversion of timestamps to the server's timezone, by setting the per-connection time zone to UTC upon connection.<commit_after>import mysql.connector
from mysql.connector import InterfaceError

class DBConnection(object):
    '''
    Represents a connection to an RDBMS.
    '''

    def __init__(self, **kwargs):
        """
        MySQL stores TIMESTAMP values as UTC, but automatically converts them
        to/from the 'current time zone' on output/input (respectively).
        By default, the 'current time zone' is the server's time zone,
        which is unknown to the client. Thus this is silent data corruption.
        There is no way to disable this conversion, so defeat it by setting
        the 'current time zone' to UTC upon connection.
        """
        if not('time_zone' in kwargs):
            kwargs['time_zone'] = '+00:00'
        self.conn = mysql.connector.connect(**kwargs)

    def _before_db(self):
        """
        MySQL-specific: attempt to reconnect if our connection has timed out
        """
        try:
            self.conn.ping(True)
        except InterfaceError:
            """
            Probably just a stale connection.
            If something worse has gone wrong, we will see it soon anyway.
            """
            pass

    def execute(self, query, return_dictionaries = True, bind_values = None):
        """
        Execute the given query with the given values for placeholders.
        If return_dictionaries is true, each row is a dictionary;
        if it is false, each row is a list of column values.
        In either case, a resultset is a list of rows.
        """
        cursor = self.conn.cursor(dictionary = return_dictionaries, buffered = False)
        # Make this a strong not a weak reference, to prevent premature GC.
        # Works around change:
        # http://bazaar.launchpad.net/~mysql/myconnpy/1.0/revision/201
        cursor._connection = self.conn
        cursor.execute(query, bind_values)
        return cursor

    def callproc(self, procname, return_dictionaries = True, args = []):
        """
        Execute the given-named stored procedure with the given arguments.
        If return_dictionaries is true, each row is a dictionary;
        if it is false, each row is a list of column values.
        In either case, a resultset is a list of rows.
        """
        cursor = self.conn.cursor(dictionary = return_dictionaries, buffered = False)
        cursor.callproc(procname, args)
        return cursor

class CursorIter(object):

    """
    Iterate over a cursor's emitted rows.
    """

    def __init__(self, cursor):
        self.cursor = cursor

    def next(self):
        row = self.cursor.fetchone()
        if row is None:
            raise StopIteration()
        return row

class CursorSliceIter(CursorIter):

    """
    Iterate over a slice of a cursor's emitted rows,
    so that the n'th column of each row is emitted.
    """

    def __init__(self, cursor, index):
        super(CursorSliceIter, self).__init__(cursor)
        self.index = index

    def next(self):
        row = super(CursorSliceIter, self).next()
        return row[self.index]

class ResultSet(object):

    """
    An iterable SQL result set.
    """

    def __init__(self, cursor):
        self.cursor = cursor

    def __iter__(self):
        return CursorIter(self.cursor)

class ResultSetSlice(ResultSet):

    """
    An iterable slice through an SQL result set.
    """

    def __init__(self, cursor, index):
        super(ResultSetSlice, self).__init__(cursor)
        self.index = index

    def __iter__(self):
        return CursorSliceIter(self.cursor, self.index)
<|endoftext|>


<commit_before>#
# Description:
#   This module implements the functions needed to keep the
#   required number of idle glideins
#   plus other miscelaneous functions
#
# Author:
#   Igor Sfiligoi (Sept 19th 2006)
#

import os.path
import sets,string
import condorMonitor,condorExe
import logSupport

class LogFiles:
    def __init__(self,log_dir,max_days,min_days,max_mbs):
        self.log_dir=log_dir
        self.activity_log=logSupport.DayLogFile(os.path.join(log_dir,"frontend"),"info.log")
        self.warning_log=logSupport.DayLogFile(os.path.join(log_dir,"frontend"),"err.log")
        self.debug_log=logSupport.DayLogFile(os.path.join(log_dir,"frontend"),"debug.log")
        self.cleanupObj=logSupport.DirCleanupWSpace(log_dir,"(frontend\.[0-9]*\.info\.log)|(frontend\.[0-9]*\.err\.log)|(frontend\.[0-9]*\.debug\.log)",
                                                    int(max_days*24*3600),int(min_days*24*3600),
                                                    long(max_mbs*(1024.0*1024.0)),
                                                    self.activity_log,self.warning_log)

    def logActivity(self,str):
        try:
            self.activity_log.write(str)
        except:
            # logging must never throw an exception!
            self.logWarning("logActivity failed, was logging: %s"%str,False)

    def logWarning(self,str, log_in_activity=True):
        try:
            self.warning_log.write(str)
        except:
            # logging must throw an exception!
            # silently ignore
            pass
        if log_in_activity:
            self.logActivity("WARNING: %s"%str)

    def logDebug(self,str):
        try:
            self.debug_log.write(str)
        except:
            # logging must never throw an exception!
            # silently ignore
            pass

    def cleanup(self):
        try:
            self.cleanupObj.cleanup()
        except:
            # logging must never throw an exception!
            self.logWarning("log cleanup failed.")

# someone needs to initialize this
# type LogFiles
log_files=None

#############################################################################################

#
# Return a dictionary of schedds containing interesting jobs
# Each element is a condorQ
#
# If not all the jobs of the schedd has to be considered,
# specify the appropriate constraint
#
def getCondorQ(schedd_names,constraint=None,format_list=None):
    if format_list!=None:
        format_list=condorMonitor.complete_format_list(format_list,[('JobStatus','i'),('EnteredCurrentStatus','i'),('ServerTime','i'),('RemoteHost','s')])
    return getCondorQConstrained(schedd_names,"(JobStatus=?=1)||(JobStatus=?=2)",constraint,format_list)

#
# Return a dictionary of schedds containing idle jobs
# Each element is a condorQ
#
# Use the output of getCondorQ
#
def getIdleCondorQ(condorq_dict):
    out={}
    for schedd_name in condorq_dict.keys():
        sq=condorMonitor.SubQuery(condorq_dict[schedd_name],lambda el:(el.has_key('JobStatus') and (el['JobStatus']==1)))
        sq.load()
        out[schedd_name]=sq
    return out

#
# Return a dictionary of schedds containing running jobs
# Each element is a condorQ
#
# Use the output of getCondorQ
#
def getRunningCondorQ(condorq_dict):
    out={}
    for schedd_name in condorq_dict.keys():
        sq=condorMonitor.SubQuery(condorq_dict[schedd_name],lambda el:(el.has_key('JobStatus') and (el['JobStatus']==2)))
        sq.load()
        out[schedd_name]=sq
    return out

def appendRealRunning(condorq_dict, status_dict):
    for schedd_name in condorq_dict:
        condorq = condorq_dict[schedd_name].fetchStored()

        for jid in condorq:
            found = False
            remote_host = condorq[jid]['RemoteHost']

            for collector_name in status_dict:
                condor_status = status_dict[collector_name].fetchStored()
                if remote_host in condor_status:
                    # there is currently no way to get the factory collector from
                    #   condor status so this hack grabs the hostname of the schedd
                    schedd = condor_status[remote_host]['GLIDEIN_Schedd'].split('@')
                    if len(schedd) < 2:
                      break

                    # split by : to remove port number if there
                    fact_pool = schedd[1].split(':')[0]

                    condorq[jid]['RunningOn'] = "%s@%s@%s@%s" % (
                        condor_status[remote_host]['GLIDEIN_Entry_Name'],
                        condor_status[remote_host]['GLIDEIN_Name'],
                        condor_status[remote_host]['GLIDEIN_Factory'],
                        fact_pool)
                    found = True
                    break

            if not found:
                condorq[jid]['RunningOn'] = 'UNKNOWN'

#
# Return a dictionary of schedds containing old jobs
# Each element is a condorQ
#
# Use the output of getCondorQ
#
def getOldCondorQ(condorq_dict,min_age):
    out={}
    for schedd_name in condorq_dict.keys():
        sq=condorMonitor.SubQuery(condorq_dict[schedd_name],lambda el:(el.has_key('ServerTime') and el.has_key('EnteredCurrentStatus') and ((el['ServerTime']-el['EnteredCurrentStatus'])>=min_age)))
        sq.load()
        out[schedd_name]=sq
    return out

#
# Return the number of jobs in the dictionary
# Use the output of getCondorQ
#
def countCondorQ(condorq_dict):
    count=0
    for schedd_name in condorq_dict.keys():
        count+=len(condorq_dict[schedd_name].fetchStored())
    return count

#
# Return a set of users present in the dictionary
# Needs "User" attribute
#

def getCondorQUsers(condorq_dict):
    users_set=sets.Set()
    for schedd_name in condorq_dict.keys():
        condorq_data=condorq_dict[schedd_name].fetchStored()
        for jid in condorq_data.keys():
            job=condorq_data[jid]
            users_set.add(job['User'])

    return users_set

#
# Get the number of jobs that match each glidein
#
# match_obj = compile('(job["MIN_NAME"]<glidein["MIN_NAME"]) && (job["ARCH"]==glidein["ARCH"])',"<string>","eval")
# condorq_dict = output of getidlqCondorQ
# glidein_dict = output of interface.findGlideins
#
# Returns:
#  dictionary of glidein name
#   where elements are number of idle jobs matching
def countMatch(match_obj,condorq_dict,glidein_dict):
    out_glidein_counts={}
    for glidename in glidein_dict:
        glidein=glidein_dict[glidename]
        glidein_count=0
        for schedd in condorq_dict.keys():
            condorq=condorq_dict[schedd]
            condorq_data=condorq.fetchStored()
            schedd_count=0
            for jid in condorq_data.keys():
                job=condorq_data[jid]
                if eval(match_obj):
                    schedd_count+=1
                pass
            glidein_count+=schedd_count
            pass
        out_glidein_counts[glidename]=glidein_count
        pass
    return out_glidein_counts

def countRealRunning(match_obj,condorq_dict,glidein_dict):
    out_glidein_counts={}
    for glidename in glidein_dict:
        # split by : to remove port number if there
        glide_str = "%s@%s" % (glidename[1],glidename[0].split(':')[0])
        glidein=glidein_dict[glidename]
        glidein_count=0
        for schedd in condorq_dict.keys():
            condorq=condorq_dict[schedd]
            condorq_data=condorq.fetchStored()
            schedd_count=0
            for jid in condorq_data.keys():
                job=condorq_data[jid]
                if eval(match_obj) and job['RunningOn'] == glide_str:
                    schedd_count+=1
                pass
            glidein_count+=schedd_count
            pass
        out_glidein_counts[glidename]=glidein_count
        pass
    return out_glidein_counts

#
# Convert frontend param expression in a value
#
# expr_obj = compile('glidein["MaxTimeout"]+frontend["MaxTimeout"]+600',"<string>","eval")
# frontend = the frontend const parameters
# glidein  = glidein factory parameters
#
# Returns:
#  The evaluated value
def evalParamExpr(expr_obj,frontend,glidein):
    return eval(expr_obj)

#
# Return a dictionary of collectors containing interesting classads
# Each element is a condorStatus
#
# If not all the jobs of the schedd has to be considered,
# specify the appropriate constraint
#
def getCondorStatus(collector_names,constraint=None,format_list=None):
    if format_list!=None:
        format_list=condorMonitor.complete_format_list(format_list,[('State','s'),('Activity','s'),('EnteredCurrentState','i'),('EnteredCurrentActivity','i'),('LastHeardFrom','i'),('GLIDEIN_Factory','s'),('GLIDEIN_Name','s'),('GLIDEIN_Entry_Name','s'),('GLIDECLIENT_Name','s'),('GLIDEIN_Schedd','s')])
    return getCondorStatusConstrained(collector_names,'(IS_MONITOR_VM=!=True)&&(GLIDEIN_Factory=!=UNDEFINED)&&(GLIDEIN_Name=!=UNDEFINED)&&(GLIDEIN_Entry_Name=!=UNDEFINED)',constraint,format_list)

#
# Return a dictionary of collectors containing idle(unclaimed) vms
# Each element is a condorStatus
#
# Use the output of getCondorStatus
#
def getIdleCondorStatus(status_dict):
    out={}
    for collector_name in status_dict.keys():
        sq=condorMonitor.SubQuery(status_dict[collector_name],lambda el:(el.has_key('State') and el.has_key('Activity') and (el['State']=="Unclaimed") and (el['Activity']=="Idle")))
        sq.load()
        out[collector_name]=sq
    return out

#
# Return a dictionary of collectors containing running(claimed) vms
# Each element is a condorStatus
#
# Use the output of getCondorStatus
#
def getRunningCondorStatus(status_dict):
    out={}
    for collector_name in status_dict.keys():
        sq=condorMonitor.SubQuery(status_dict[collector_name],lambda el:(el.has_key('State') and el.has_key('Activity') and (el['State']=="Claimed") and (el['Activity'] in ("Busy","Retiring"))))
        sq.load()
        out[collector_name]=sq
    return out

#
# Return a dictionary of collectors containing idle(unclaimed) vms
# Each element is a condorStatus
#
# Use the output of getCondorStatus
#
def getClientCondorStatus(status_dict,frontend_name,group_name,request_name):
    client_name_old="%s@%s.%s"%(request_name,frontend_name,group_name)
    client_name_new="%s.%s"%(frontend_name,group_name)
    out={}
    for collector_name in status_dict.keys():
        sq=condorMonitor.SubQuery(status_dict[collector_name],lambda el:(el.has_key('GLIDECLIENT_Name') and ((el['GLIDECLIENT_Name']==client_name_old) or ((el['GLIDECLIENT_Name']==client_name_new) and (("%s@%s@%s"%(el['GLIDEIN_Entry_Name'],el['GLIDEIN_Name'],el['GLIDEIN_Factory']))==request_name)))))
        sq.load()
        out[collector_name]=sq
    return out

#
# Return the number of vms in the dictionary
# Use the output of getCondorStatus
#
def countCondorStatus(status_dict):
    count=0
    for collector_name in status_dict.keys():
        count+=len(status_dict[collector_name].fetchStored())
    return count

############################################################
#
# I N T E R N A L - Do not use
#
############################################################

#
# Return a dictionary of schedds containing jobs of a certain type
# Each element is a condorQ
#
# If not all the jobs of the schedd has to be considered,
# specify the appropriate additional constraint
#
def getCondorQConstrained(schedd_names,type_constraint,constraint=None,format_list=None):
    out_condorq_dict={}
    for schedd in schedd_names:
        if schedd=='':
            log_files.logWarning("Skipping empty schedd name")
            continue
        condorq=condorMonitor.CondorQ(schedd)
        full_constraint=type_constraint[0:] #make copy
        if constraint!=None:
            full_constraint="(%s) && (%s)"%(full_constraint,constraint)

        try:
            condorq.load(full_constraint,format_list)
        except condorExe.ExeError, e:
            if schedd!=None:
                log_files.logWarning("Failed to talk to schedd %s. See debug log for more details."%schedd)
                log_files.logDebug("Failed to talk to schedd %s: %s"%(schedd, e))
            else:
                log_files.logWarning("Failed to talk to schedd. See debug log for more details.")
                log_files.logDebug("Failed to talk to schedd: %s"%e)
            continue # if schedd not found it is equivalent to no jobs in the queue
        if len(condorq.fetchStored())>0:
            out_condorq_dict[schedd]=condorq
    return out_condorq_dict

#
# Return a dictionary of collectors containing classads of a certain kind
# Each element is a condorStatus
#
# If not all the jobs of the schedd has to be considered,
# specify the appropriate additional constraint
#
def getCondorStatusConstrained(collector_names,type_constraint,constraint=None,format_list=None):
    out_status_dict={}
    for collector in collector_names:
        status=condorMonitor.CondorStatus(pool_name=collector)
        full_constraint=type_constraint[0:] #make copy
        if constraint!=None:
            full_constraint="(%s) && (%s)"%(full_constraint,constraint)

        try:
            status.load(full_constraint,format_list)
        except condorExe.ExeError, e:
            if collector!=None:
                log_files.logWarning("Failed to talk to collector %s. See debug log for more details."%collector)
                log_files.logDebug("Failed to talk to collector %s: %s"%(collector, e))
            else:
                log_files.logWarning("Failed to talk to collector. See debug log for more details.")
                log_files.logDebug("Failed to talk to collector: %s"%e)
            continue # if collector not found it is equivalent to no classads
        if len(status.fetchStored())>0:
            out_status_dict[collector]=status
    return out_status_dict

<commit_msg>Better handle edge case; job claimed to be running, but no RemoteHost attribute - merge from branch_v2_4_2_jeff_fdxml<commit_after>#
# Description:
#   This module implements the functions needed to keep the
#   required number of idle glideins
#   plus other miscelaneous functions
#
# Author:
#   Igor Sfiligoi (Sept 19th 2006)
#

import os.path
import sets,string
import condorMonitor,condorExe
import logSupport

class LogFiles:
    def __init__(self,log_dir,max_days,min_days,max_mbs):
        self.log_dir=log_dir
        self.activity_log=logSupport.DayLogFile(os.path.join(log_dir,"frontend"),"info.log")
        self.warning_log=logSupport.DayLogFile(os.path.join(log_dir,"frontend"),"err.log")
        self.debug_log=logSupport.DayLogFile(os.path.join(log_dir,"frontend"),"debug.log")
        self.cleanupObj=logSupport.DirCleanupWSpace(log_dir,"(frontend\.[0-9]*\.info\.log)|(frontend\.[0-9]*\.err\.log)|(frontend\.[0-9]*\.debug\.log)",
                                                    int(max_days*24*3600),int(min_days*24*3600),
                                                    long(max_mbs*(1024.0*1024.0)),
                                                    self.activity_log,self.warning_log)

    def logActivity(self,str):
        try:
            self.activity_log.write(str)
        except:
            # logging must never throw an exception!
            self.logWarning("logActivity failed, was logging: %s"%str,False)

    def logWarning(self,str, log_in_activity=True):
        try:
            self.warning_log.write(str)
        except:
            # logging must throw an exception!
            # silently ignore
            pass
        if log_in_activity:
            self.logActivity("WARNING: %s"%str)

    def logDebug(self,str):
        try:
            self.debug_log.write(str)
        except:
            # logging must never throw an exception!
            # silently ignore
            pass

    def cleanup(self):
        try:
            self.cleanupObj.cleanup()
        except:
            # logging must never throw an exception!
            self.logWarning("log cleanup failed.")

# someone needs to initialize this
# type LogFiles
log_files=None

#############################################################################################

#
# Return a dictionary of schedds containing interesting jobs
# Each element is a condorQ
#
# If not all the jobs of the schedd has to be considered,
# specify the appropriate constraint
#
def getCondorQ(schedd_names,constraint=None,format_list=None):
    if format_list!=None:
        format_list=condorMonitor.complete_format_list(format_list,[('JobStatus','i'),('EnteredCurrentStatus','i'),('ServerTime','i'),('RemoteHost','s')])
    return getCondorQConstrained(schedd_names,"(JobStatus=?=1)||(JobStatus=?=2)",constraint,format_list)

#
# Return a dictionary of schedds containing idle jobs
# Each element is a condorQ
#
# Use the output of getCondorQ
#
def getIdleCondorQ(condorq_dict):
    out={}
    for schedd_name in condorq_dict.keys():
        sq=condorMonitor.SubQuery(condorq_dict[schedd_name],lambda el:(el.has_key('JobStatus') and (el['JobStatus']==1)))
        sq.load()
        out[schedd_name]=sq
    return out

#
# Return a dictionary of schedds containing running jobs
# Each element is a condorQ
#
# Use the output of getCondorQ
#
def getRunningCondorQ(condorq_dict):
    out={}
    for schedd_name in condorq_dict.keys():
        sq=condorMonitor.SubQuery(condorq_dict[schedd_name],lambda el:(el.has_key('JobStatus') and (el['JobStatus']==2)))
        sq.load()
        out[schedd_name]=sq
    return out

def appendRealRunning(condorq_dict, status_dict):
    for schedd_name in condorq_dict:
        condorq = condorq_dict[schedd_name].fetchStored()

        for jid in condorq:
            found = False

            if condorq[jid].has_key('RemoteHost'):
                remote_host = condorq[jid]['RemoteHost']

                for collector_name in status_dict:
                    condor_status = status_dict[collector_name].fetchStored()
                    if remote_host in condor_status:
                        # there is currently no way to get the factory collector from
                        #   condor status so this hack grabs the hostname of the schedd
                        schedd = condor_status[remote_host]['GLIDEIN_Schedd'].split('@')
                        if len(schedd) < 2:
                          break

                        # split by : to remove port number if there
                        fact_pool = schedd[1].split(':')[0]

                        condorq[jid]['RunningOn'] = "%s@%s@%s@%s" % (
                            condor_status[remote_host]['GLIDEIN_Entry_Name'],
                            condor_status[remote_host]['GLIDEIN_Name'],
                            condor_status[remote_host]['GLIDEIN_Factory'],
                            fact_pool)
                        found = True
                        break

            if not found:
                condorq[jid]['RunningOn'] = 'UNKNOWN'

#
# Return a dictionary of schedds containing old jobs
# Each element is a condorQ
#
# Use the output of getCondorQ
#
def getOldCondorQ(condorq_dict,min_age):
    out={}
    for schedd_name in condorq_dict.keys():
        sq=condorMonitor.SubQuery(condorq_dict[schedd_name],lambda el:(el.has_key('ServerTime') and el.has_key('EnteredCurrentStatus') and ((el['ServerTime']-el['EnteredCurrentStatus'])>=min_age)))
        sq.load()
        out[schedd_name]=sq
    return out

#
# Return the number of jobs in the dictionary
# Use the output of getCondorQ
#
def countCondorQ(condorq_dict):
    count=0
    for schedd_name in condorq_dict.keys():
        count+=len(condorq_dict[schedd_name].fetchStored())
    return count

#
# Return a set of users present in the dictionary
# Needs "User" attribute
#

def getCondorQUsers(condorq_dict):
    users_set=sets.Set()
    for schedd_name in condorq_dict.keys():
        condorq_data=condorq_dict[schedd_name].fetchStored()
        for jid in condorq_data.keys():
            job=condorq_data[jid]
            users_set.add(job['User'])

    return users_set

#
# Get the number of jobs that match each glidein
#
# match_obj = compile('(job["MIN_NAME"]<glidein["MIN_NAME"]) && (job["ARCH"]==glidein["ARCH"])',"<string>","eval")
# condorq_dict = output of getidlqCondorQ
# glidein_dict = output of interface.findGlideins
#
# Returns:
#  dictionary of glidein name
#   where elements are number of idle jobs matching
def countMatch(match_obj,condorq_dict,glidein_dict):
    out_glidein_counts={}
    for glidename in glidein_dict:
        glidein=glidein_dict[glidename]
        glidein_count=0
        for schedd in condorq_dict.keys():
            condorq=condorq_dict[schedd]
            condorq_data=condorq.fetchStored()
            schedd_count=0
            for jid in condorq_data.keys():
                job=condorq_data[jid]
                if eval(match_obj):
                    schedd_count+=1
                pass
            glidein_count+=schedd_count
            pass
        out_glidein_counts[glidename]=glidein_count
        pass
    return out_glidein_counts

def countRealRunning(match_obj,condorq_dict,glidein_dict):
    out_glidein_counts={}
    for glidename in glidein_dict:
        # split by : to remove port number if there
        glide_str = "%s@%s" % (glidename[1],glidename[0].split(':')[0])
        glidein=glidein_dict[glidename]
        glidein_count=0
        for schedd in condorq_dict.keys():
            condorq=condorq_dict[schedd]
            condorq_data=condorq.fetchStored()
            schedd_count=0
            for jid in condorq_data.keys():
                job=condorq_data[jid]
                if eval(match_obj) and job['RunningOn'] == glide_str:
                    schedd_count+=1
                pass
            glidein_count+=schedd_count
            pass
        out_glidein_counts[glidename]=glidein_count
        pass
    return out_glidein_counts

#
# Convert frontend param expression in a value
#
# expr_obj = compile('glidein["MaxTimeout"]+frontend["MaxTimeout"]+600',"<string>","eval")
# frontend = the frontend const parameters
# glidein  = glidein factory parameters
#
# Returns:
#  The evaluated value
def evalParamExpr(expr_obj,frontend,glidein):
    return eval(expr_obj)

#
# Return a dictionary of collectors containing interesting classads
# Each element is a condorStatus
#
# If not all the jobs of the schedd has to be considered,
# specify the appropriate constraint
#
def getCondorStatus(collector_names,constraint=None,format_list=None):
    if format_list!=None:
        format_list=condorMonitor.complete_format_list(format_list,[('State','s'),('Activity','s'),('EnteredCurrentState','i'),('EnteredCurrentActivity','i'),('LastHeardFrom','i'),('GLIDEIN_Factory','s'),('GLIDEIN_Name','s'),('GLIDEIN_Entry_Name','s'),('GLIDECLIENT_Name','s'),('GLIDEIN_Schedd','s')])
    return getCondorStatusConstrained(collector_names,'(IS_MONITOR_VM=!=True)&&(GLIDEIN_Factory=!=UNDEFINED)&&(GLIDEIN_Name=!=UNDEFINED)&&(GLIDEIN_Entry_Name=!=UNDEFINED)',constraint,format_list)

#
# Return a dictionary of collectors containing idle(unclaimed) vms
# Each element is a condorStatus
#
# Use the output of getCondorStatus
#
def getIdleCondorStatus(status_dict):
    out={}
    for collector_name in status_dict.keys():
        sq=condorMonitor.SubQuery(status_dict[collector_name],lambda el:(el.has_key('State') and el.has_key('Activity') and (el['State']=="Unclaimed") and (el['Activity']=="Idle")))
        sq.load()
        out[collector_name]=sq
    return out

#
# Return a dictionary of collectors containing running(claimed) vms
# Each element is a condorStatus
#
# Use the output of getCondorStatus
#
def getRunningCondorStatus(status_dict):
    out={}
    for collector_name in status_dict.keys():
        sq=condorMonitor.SubQuery(status_dict[collector_name],lambda el:(el.has_key('State') and el.has_key('Activity') and (el['State']=="Claimed") and (el['Activity'] in ("Busy","Retiring"))))
        sq.load()
        out[collector_name]=sq
    return out

#
# Return a dictionary of collectors containing idle(unclaimed) vms
# Each element is a condorStatus
#
# Use the output of getCondorStatus
#
def getClientCondorStatus(status_dict,frontend_name,group_name,request_name):
    client_name_old="%s@%s.%s"%(request_name,frontend_name,group_name)
    client_name_new="%s.%s"%(frontend_name,group_name)
    out={}
    for collector_name in status_dict.keys():
        sq=condorMonitor.SubQuery(status_dict[collector_name],lambda el:(el.has_key('GLIDECLIENT_Name') and ((el['GLIDECLIENT_Name']==client_name_old) or ((el['GLIDECLIENT_Name']==client_name_new) and (("%s@%s@%s"%(el['GLIDEIN_Entry_Name'],el['GLIDEIN_Name'],el['GLIDEIN_Factory']))==request_name)))))
        sq.load()
        out[collector_name]=sq
    return out

#
# Return the number of vms in the dictionary
# Use the output of getCondorStatus
#
def countCondorStatus(status_dict):
    count=0
    for collector_name in status_dict.keys():
        count+=len(status_dict[collector_name].fetchStored())
    return count

############################################################
#
# I N T E R N A L - Do not use
#
############################################################

#
# Return a dictionary of schedds containing jobs of a certain type
# Each element is a condorQ
#
# If not all the jobs of the schedd has to be considered,
# specify the appropriate additional constraint
#
def getCondorQConstrained(schedd_names,type_constraint,constraint=None,format_list=None):
    out_condorq_dict={}
    for schedd in schedd_names:
        if schedd=='':
            log_files.logWarning("Skipping empty schedd name")
            continue
        condorq=condorMonitor.CondorQ(schedd)
        full_constraint=type_constraint[0:] #make copy
        if constraint!=None:
            full_constraint="(%s) && (%s)"%(full_constraint,constraint)

        try:
            condorq.load(full_constraint,format_list)
        except condorExe.ExeError, e:
            if schedd!=None:
                log_files.logWarning("Failed to talk to schedd %s. See debug log for more details."%schedd)
                log_files.logDebug("Failed to talk to schedd %s: %s"%(schedd, e))
            else:
                log_files.logWarning("Failed to talk to schedd. See debug log for more details.")
                log_files.logDebug("Failed to talk to schedd: %s"%e)
            continue # if schedd not found it is equivalent to no jobs in the queue
        if len(condorq.fetchStored())>0:
            out_condorq_dict[schedd]=condorq
    return out_condorq_dict

#
# Return a dictionary of collectors containing classads of a certain kind
# Each element is a condorStatus
#
# If not all the jobs of the schedd has to be considered,
# specify the appropriate additional constraint
#
def getCondorStatusConstrained(collector_names,type_constraint,constraint=None,format_list=None):
    out_status_dict={}
    for collector in collector_names:
        status=condorMonitor.CondorStatus(pool_name=collector)
        full_constraint=type_constraint[0:] #make copy
        if constraint!=None:
            full_constraint="(%s) && (%s)"%(full_constraint,constraint)

        try:
            status.load(full_constraint,format_list)
        except condorExe.ExeError, e:
            if collector!=None:
                log_files.logWarning("Failed to talk to collector %s. See debug log for more details."%collector)
                log_files.logDebug("Failed to talk to collector %s: %s"%(collector, e))
            else:
                log_files.logWarning("Failed to talk to collector. See debug log for more details.")
                log_files.logDebug("Failed to talk to collector: %s"%e)
            continue # if collector not found it is equivalent to no classads
        if len(status.fetchStored())>0:
            out_status_dict[collector]=status
    return out_status_dict

<|endoftext|>


<commit_before>import time
import numpy as np
from astropy.table import vstack
from easyquery import Query
from . import cuts as C
from . import build, build2
from .manual_fixes import fixes_to_nsa_v012, fixes_to_nsa_v101
from ..database import FitsTable, Database
from ..hosts import HostCatalog
from ..utils import get_sdss_bands, get_all_colors, add_skycoord, fill_values_by_query

__all__ = ['ObjectCatalog', 'get_unique_objids']

def get_unique_objids(objid_col):
    try:
        objid_col = objid_col.compressed()
    except AttributeError:
        pass
    return np.unique(np.asarray(objid_col, dtype=np.int64))

class ObjectCatalog(object):
    """
    This class provides a high-level interface to access object catalogs
    (also known as "base catalogs")

    Parameters
    ----------
    database : SAGA.Database object

    Returns
    -------
    object_catalog : SAGA.ObjectCatalog object

    Examples
    --------
    >>> import SAGA
    >>> saga_database = SAGA.Database('/path/to/SAGA/Dropbox')
    >>> saga_object_catalog = SAGA.ObjectCatalog(saga_database)
    >>> specs = saga_object_catalog.load(has_spec=True)
    >>> base_anak = saga_object_catalog.load(hosts='AnaK')

    Here specs and base_anak are both astropy tables.
    """

    def __init__(self, database=None):
        self._database = database or Database()
        self._host_catalog = HostCatalog(self._database)

    @staticmethod
    def _annotate_catalog(table, to_add_skycoord=True):
        if 'EXTINCTION_R' in table.colnames:
            for b in get_sdss_bands():
                table['{}_mag'.format(b)] = table[b] - table['EXTINCTION_{}'.format(b.upper())]

        for color in get_all_colors():
            if '{}_mag'.format(color[0]) not in table.colnames or '{}_mag'.format(color[1]) not in table.colnames:
                continue
            table[color] = table['{}_mag'.format(color[0])] - table['{}_mag'.format(color[1])]
            table['{}_err'.format(color)] = np.hypot(table['{}_err'.format(color[0])], table['{}_err'.format(color[1])])

        if to_add_skycoord:
            table = add_skycoord(table)

        return table

    @staticmethod
    def _slice_columns(table, columns, get_coord_later=False):
        if columns is None:
            return table

        if get_coord_later:
            columns_this = list(columns)
            try:
                columns_this.remove('coord')
            except ValueError:
                pass
            if 'RA' not in columns_this:
                columns_this.append('RA')
            if 'DEC' not in columns_this:
                columns_this.append('DEC')

            return table[columns_this]

        return table[columns]

    def load(self, hosts=None, has_spec=None, cuts=None, return_as=None, columns=None, version=None):
        """
        load object catalogs (aka "base catalogs")

        Parameters
        ----------
        hosts : int, str, list, None, optional
            host names/IDs or a list of host names/IDs or short-hand names like
            "paper1" or "paper1_complete"

        has_spec : bool, optional
            If set to True, load only objects that have spectra

        cuts : easyquery.Query, str, tuple, optional
            To apply to the objects when loaded

        return_as : str, optional
            If set to 'list' (default when `has_spec` is None), return a list that contains all tables
            If set to 'stacked' (default when `has_spec` is True), return a stacked table
            If set to 'iter', return an iterator for looping over hosts
            If set to 'dict', return a dictionary with host ids being the keys

        columns : list, optional
            If set, only load a subset of columns

        version : int or str, optional
            Set to 'paper1' for paper1 catalogs

        Returns
        -------
        objects : astropy.table.Table, list, or iterator
            (depending on `return_as`)

        Examples
        --------
        >>> import SAGA
        >>> from SAGA import ObjectCuts as C
        >>> saga_database = SAGA.Database('/path/to/SAGA/Dropbox')
        >>> saga_object_catalog = SAGA.ObjectCatalog(saga_database)

        To load all spectra, with some basic cuts applied:
        >>> specs = saga_object_catalog.load(has_spec=True, cuts=C.basic_cut)

        Load the base catalog for a certain host, with some basic cuts applied:
        >>> specs = saga_object_catalog.load(hosts='AnaK', cuts=C.basic_cut)

        Load base catalog for all paper1 hosts, with some basic cuts applied,
        and stored as a list:
        >>> base_tables = saga_object_catalog.load(hosts='paper1', cuts=C.basic_cut, return_as='list')

        Load base catalog for all paper1 hosts, with some basic cuts applied,
        and stored as one single big table:
        >>> bases_table = saga_object_catalog.load(hosts='paper1', cuts=C.basic_cut, return_as='stacked')
        """

        if return_as is None:
            return_as = 'stacked' if has_spec else 'list'
        return_as = return_as.lower()
        if return_as[0] not in 'slid':
            raise ValueError('`return_as` should be "list", "stacked", "iter", or "dict"')

        if version is None:
            base_key = 'base'
        elif str(version).lower() in ('paper1', 'p1', 'v0p1', '0', '0.1'):
            base_key = 'base_v0p1'
        elif version in (1, 2):
            base_key = 'base_v{}'.format(version)
        else:
            raise ValueError('`version` must be None, \'paper1\', 1 or 2.')

        if has_spec and base_key == 'base_v0p1':
            t = self._database['saga_spectra_May2017'].read()

            if hosts is not None:
                host_ids = self._host_catalog.resolve_id(hosts, 'NSA')
                t = Query((lambda x: np.in1d(x, host_ids), 'HOST_NSAID')).filter(t)

            t = self._annotate_catalog(t)

            if cuts is not None:
                t = Query(cuts).filter(t)

            if return_as[0] != 's':
                if hosts is None:
                    host_ids = np.unique(t['HOST_NSAID'])
                output_iterator = (self._slice_columns(Query('HOST_NSAID == {}'.format(i)).filter(t), columns) for i in host_ids)
                if return_as[0] == 'i':
                    return output_iterator
                if return_as[0] == 'd':
                    return dict(zip(host_ids, output_iterator))
                return list(output_iterator)

            return self._slice_columns(t, columns)

        else:
            q = Query(cuts)
            if has_spec:
                q = q & C.has_spec
            elif has_spec is not None:
                q = q & (~C.has_spec)

            hosts = self._host_catalog.resolve_id(hosts, 'string')

            need_coord = (columns is None or 'coord' in columns)
            to_add_skycoord = (need_coord and return_as[0] != 's') # because skycoord cannot be stacked

            output_iterator = (self._slice_columns(q.filter(self._annotate_catalog(self._database[base_key, host].read(), to_add_skycoord)), columns, (need_coord and not to_add_skycoord)) for host in hosts)

            if return_as[0] == 'i':
                return output_iterator
            if return_as[0] == 's':
                out = vstack(list(output_iterator), 'outer', 'error')
                if out.masked:
                    for name, (dtype, _) in out.dtype.fields.items():
                        if dtype.kind == 'i':
                            out[name].fill_value = -1
                        if dtype.kind == 'b':
                            out[name].fill_value = False
                out = out.filled()
                if need_coord:
                    out = self._slice_columns(add_skycoord(out), columns)
                return out
            if return_as[0] == 'd':
                return dict(zip(hosts, output_iterator))
            return list(output_iterator)

    def load_nsa(self, version='0.1.2'):
        nsa = self._database['nsa_v{}'.format(version)].read()
        if version == '0.1.2':
            nsa = nsa[build.NSA_COLS_USED]
            for nsaid, fixes in fixes_to_nsa_v012.items():
                fill_values_by_query(nsa, 'NSAID == {}'.format(nsaid), fixes)
            # NSA 64408 (127.324917502, 25.75292055) is wrong! For v0.1.2 ONLY!!
            nsa = Query('NSAID != 64408').filter(nsa)
        elif version == '1.0.1':
            nsa = nsa[build2.NSA_COLS_USED]
            for nsaid, fixes in fixes_to_nsa_v101.items():
                fill_values_by_query(nsa, 'NSAID == {}'.format(nsaid), fixes)
        nsa = add_skycoord(nsa)
        return nsa

    def build_and_write_to_database(self, hosts=None, overwrite=False, base_file_path_pattern=None, version=None, return_catalogs=False, raise_exception=False, add_specs_only_before_time=None):
        """
        This function builds the base catalog and writes it to the database.

        !! IMPORTANT !!
        If you want to write the base catalog to an alternative location (not the database)
        Make sure you set the `base_file_path_pattern` option!!

        Parameters
        ----------
        hosts : int, str, list, optional
            host names/IDs or a list of host names/IDs or short-hand names like
            "paper1" or "paper1_complete"

        overwrite : bool, optional (default: False)
            If set to True, overwrite existing base catalog

        base_file_path_pattern : str, optional
        version : int, optional (default: 2)
        return_catalogs : bool, optional (default: False)
        raise_exception : bool, optional (default: False)
        add_specs_only_before_time : astropy.time.Time, optional (default: None)

        Examples
        --------
        >>> saga_database = SAGA.Database('/path/to/SAGA/Dropbox', '/path/to/SAGA/local')
        >>> saga_object_catalog = SAGA.ObjectCatalog(saga_database)

        Overwrite the database (Danger!!)
        >>> saga_object_catalog.build_and_write_to_database('paper1', overwrite=True)

        You can also do
        >>> saga_object_catalog.build_and_write_to_database('paper1', base_file_path_pattern='/other/base/catalog/dir/nsa{}.fits.gz')

        """
        if version not in (None, 1, 2):
            raise ValueError('`version` must be None, 1 or 2.')
        build_module = build if version == 1 else build2

        nsa = self.load_nsa('0.1.2' if version == 1 else '1.0.1')
        spectra = self._database['spectra_raw_all'].read(before_time=add_specs_only_before_time)

        manual_lists = {}
        for survey, col in (('sdss', 'SDSS ID'), ('des', 'DES_OBJID'), ('decals', 'decals_objid')):
            for list_type in ('remove', 'recover'):
                key = '{}_{}'.format(survey, list_type)
                val = get_unique_objids(self._database[key].read()[col])
                if not len(val):
                    val = None
                manual_lists[key] = val

        catalogs_to_return = list()
        host_ids = self._host_catalog.resolve_id(hosts, 'string')

        for i, host_id in enumerate(host_ids):
            if base_file_path_pattern is None:
                base_key = 'base' if version is None else 'base_v{}'.format(version)
                data_obj = self._database[base_key, host_id].remote
            else:
                data_obj = FitsTable(base_file_path_pattern.format(host_id))

            if data_obj.isfile() and not overwrite:
                print(time.strftime('[%m/%d %H:%M:%S]'), 'Base catalog v{} for {} already exists ({}).'.format(version or 2, host_id, data_obj.path), '({}/{})'.format(i+1, len(host_ids)))
                continue

            host = self._host_catalog.load_single(host_id)
            catalogs = ('sdss', 'wise') if version == 1 else ('sdss', 'des', 'decals')

            def get_catalog_or_none(catalog_name):
                try:
                    cat = self._database[catalog_name, host_id].read()
                except OSError:
                    print(time.strftime('[%m/%d %H:%M:%S]'), '[WARNING] Not found: {} catalog for {}.'.format(catalog_name.upper(), host_id))
                    return None
                return cat[build.WISE_COLS_USED] if catalog_name == 'wise' else cat

            catalog_dict = {k: get_catalog_or_none(k) for k in catalogs}

            print(
                time.strftime('[%m/%d %H:%M:%S]'),
                'Use {} to build base catalog v{} for {}'.format(
                    ', '.join((k for k, v in catalog_dict.items() if v is not None)).upper(),
                    (version or 2),
                    host_id
                ),
                '({}/{})'.format(i+1, len(host_ids))
            )

            try:
                base = build_module.build_full_stack(host=host, nsa=nsa, spectra=spectra,
                                                     **manual_lists, **catalog_dict)
            except Exception as e: # pylint: disable=W0703
                print(time.strftime('[%m/%d %H:%M:%S]'), '[ERROR] Fail to build base catalog for {}\n{}'.format(host_id, e))
                base = None
                if raise_exception:
                    raise e
                continue
            finally:
                del catalog_dict
                if return_catalogs:
                    catalogs_to_return.append(base)

            print(time.strftime('[%m/%d %H:%M:%S]'), 'Write base catalog to {}'.format(data_obj.path))
            try:
                data_obj.write(base)
            except (IOError, OSError) as e:
                print(time.strftime('[%m/%d %H:%M:%S]'), '[ERROR] Fail to write base catalog for {}\n{}'.format(host_id, e))
                continue

        if return_catalogs:
            return catalogs_to_return

    def generate_clean_specs(self, version=None):
        """
        gather spectra from all base catalogs
        """
        surveys = ('sdss', 'des', 'decals')
        base_key = 'base' if version is None else 'base_v{}'.format(version)
        out = []

        for host in self._host_catalog.resolve_id('all', 'string'):
            try:
                base = self._database[base_key, host].read()
            except IOError:
                print('[WARNING] cannot find base catalog for', host)

            base = C.has_spec.filter(base)

            if version != 1:
                cols = [col for col in base.colnames if col.startswith('OBJID_') or not any(col.endswith('_'+s) for s in surveys)]
                base = base[cols]
                for s in surveys:
                    key = 'OBJID_' + s
                    if key not in base.colnames:
                        base[key] = -1

            out.append(base)

        return vstack(out, 'exact', 'error')
<commit_msg>expose debug to high-level function<commit_after>import time
import numpy as np
from astropy.table import vstack
from easyquery import Query
from . import cuts as C
from . import build, build2
from .manual_fixes import fixes_to_nsa_v012, fixes_to_nsa_v101
from ..database import FitsTable, Database
from ..hosts import HostCatalog
from ..utils import get_sdss_bands, get_all_colors, add_skycoord, fill_values_by_query

__all__ = ['ObjectCatalog', 'get_unique_objids']

def get_unique_objids(objid_col):
    try:
        objid_col = objid_col.compressed()
    except AttributeError:
        pass
    return np.unique(np.asarray(objid_col, dtype=np.int64))

class ObjectCatalog(object):
    """
    This class provides a high-level interface to access object catalogs
    (also known as "base catalogs")

    Parameters
    ----------
    database : SAGA.Database object

    Returns
    -------
    object_catalog : SAGA.ObjectCatalog object

    Examples
    --------
    >>> import SAGA
    >>> saga_database = SAGA.Database('/path/to/SAGA/Dropbox')
    >>> saga_object_catalog = SAGA.ObjectCatalog(saga_database)
    >>> specs = saga_object_catalog.load(has_spec=True)
    >>> base_anak = saga_object_catalog.load(hosts='AnaK')

    Here specs and base_anak are both astropy tables.
    """

    def __init__(self, database=None):
        self._database = database or Database()
        self._host_catalog = HostCatalog(self._database)

    @staticmethod
    def _annotate_catalog(table, to_add_skycoord=True):
        if 'EXTINCTION_R' in table.colnames:
            for b in get_sdss_bands():
                table['{}_mag'.format(b)] = table[b] - table['EXTINCTION_{}'.format(b.upper())]

        for color in get_all_colors():
            if '{}_mag'.format(color[0]) not in table.colnames or '{}_mag'.format(color[1]) not in table.colnames:
                continue
            table[color] = table['{}_mag'.format(color[0])] - table['{}_mag'.format(color[1])]
            table['{}_err'.format(color)] = np.hypot(table['{}_err'.format(color[0])], table['{}_err'.format(color[1])])

        if to_add_skycoord:
            table = add_skycoord(table)

        return table

    @staticmethod
    def _slice_columns(table, columns, get_coord_later=False):
        if columns is None:
            return table

        if get_coord_later:
            columns_this = list(columns)
            try:
                columns_this.remove('coord')
            except ValueError:
                pass
            if 'RA' not in columns_this:
                columns_this.append('RA')
            if 'DEC' not in columns_this:
                columns_this.append('DEC')

            return table[columns_this]

        return table[columns]

    def load(self, hosts=None, has_spec=None, cuts=None, return_as=None, columns=None, version=None):
        """
        load object catalogs (aka "base catalogs")

        Parameters
        ----------
        hosts : int, str, list, None, optional
            host names/IDs or a list of host names/IDs or short-hand names like
            "paper1" or "paper1_complete"

        has_spec : bool, optional
            If set to True, load only objects that have spectra

        cuts : easyquery.Query, str, tuple, optional
            To apply to the objects when loaded

        return_as : str, optional
            If set to 'list' (default when `has_spec` is None), return a list that contains all tables
            If set to 'stacked' (default when `has_spec` is True), return a stacked table
            If set to 'iter', return an iterator for looping over hosts
            If set to 'dict', return a dictionary with host ids being the keys

        columns : list, optional
            If set, only load a subset of columns

        version : int or str, optional
            Set to 'paper1' for paper1 catalogs

        Returns
        -------
        objects : astropy.table.Table, list, or iterator
            (depending on `return_as`)

        Examples
        --------
        >>> import SAGA
        >>> from SAGA import ObjectCuts as C
        >>> saga_database = SAGA.Database('/path/to/SAGA/Dropbox')
        >>> saga_object_catalog = SAGA.ObjectCatalog(saga_database)

        To load all spectra, with some basic cuts applied:
        >>> specs = saga_object_catalog.load(has_spec=True, cuts=C.basic_cut)

        Load the base catalog for a certain host, with some basic cuts applied:
        >>> specs = saga_object_catalog.load(hosts='AnaK', cuts=C.basic_cut)

        Load base catalog for all paper1 hosts, with some basic cuts applied,
        and stored as a list:
        >>> base_tables = saga_object_catalog.load(hosts='paper1', cuts=C.basic_cut, return_as='list')

        Load base catalog for all paper1 hosts, with some basic cuts applied,
        and stored as one single big table:
        >>> bases_table = saga_object_catalog.load(hosts='paper1', cuts=C.basic_cut, return_as='stacked')
        """

        if return_as is None:
            return_as = 'stacked' if has_spec else 'list'
        return_as = return_as.lower()
        if return_as[0] not in 'slid':
            raise ValueError('`return_as` should be "list", "stacked", "iter", or "dict"')

        if version is None:
            base_key = 'base'
        elif str(version).lower() in ('paper1', 'p1', 'v0p1', '0', '0.1'):
            base_key = 'base_v0p1'
        elif version in (1, 2):
            base_key = 'base_v{}'.format(version)
        else:
            raise ValueError('`version` must be None, \'paper1\', 1 or 2.')

        if has_spec and base_key == 'base_v0p1':
            t = self._database['saga_spectra_May2017'].read()

            if hosts is not None:
                host_ids = self._host_catalog.resolve_id(hosts, 'NSA')
                t = Query((lambda x: np.in1d(x, host_ids), 'HOST_NSAID')).filter(t)

            t = self._annotate_catalog(t)

            if cuts is not None:
                t = Query(cuts).filter(t)

            if return_as[0] != 's':
                if hosts is None:
                    host_ids = np.unique(t['HOST_NSAID'])
                output_iterator = (self._slice_columns(Query('HOST_NSAID == {}'.format(i)).filter(t), columns) for i in host_ids)
                if return_as[0] == 'i':
                    return output_iterator
                if return_as[0] == 'd':
                    return dict(zip(host_ids, output_iterator))
                return list(output_iterator)

            return self._slice_columns(t, columns)

        else:
            q = Query(cuts)
            if has_spec:
                q = q & C.has_spec
            elif has_spec is not None:
                q = q & (~C.has_spec)

            hosts = self._host_catalog.resolve_id(hosts, 'string')

            need_coord = (columns is None or 'coord' in columns)
            to_add_skycoord = (need_coord and return_as[0] != 's') # because skycoord cannot be stacked

            output_iterator = (self._slice_columns(q.filter(self._annotate_catalog(self._database[base_key, host].read(), to_add_skycoord)), columns, (need_coord and not to_add_skycoord)) for host in hosts)

            if return_as[0] == 'i':
                return output_iterator
            if return_as[0] == 's':
                out = vstack(list(output_iterator), 'outer', 'error')
                if out.masked:
                    for name, (dtype, _) in out.dtype.fields.items():
                        if dtype.kind == 'i':
                            out[name].fill_value = -1
                        if dtype.kind == 'b':
                            out[name].fill_value = False
                out = out.filled()
                if need_coord:
                    out = self._slice_columns(add_skycoord(out), columns)
                return out
            if return_as[0] == 'd':
                return dict(zip(hosts, output_iterator))
            return list(output_iterator)

    def load_nsa(self, version='0.1.2'):
        nsa = self._database['nsa_v{}'.format(version)].read()
        if version == '0.1.2':
            nsa = nsa[build.NSA_COLS_USED]
            for nsaid, fixes in fixes_to_nsa_v012.items():
                fill_values_by_query(nsa, 'NSAID == {}'.format(nsaid), fixes)
            # NSA 64408 (127.324917502, 25.75292055) is wrong! For v0.1.2 ONLY!!
            nsa = Query('NSAID != 64408').filter(nsa)
        elif version == '1.0.1':
            nsa = nsa[build2.NSA_COLS_USED]
            for nsaid, fixes in fixes_to_nsa_v101.items():
                fill_values_by_query(nsa, 'NSAID == {}'.format(nsaid), fixes)
        nsa = add_skycoord(nsa)
        return nsa

    def build_and_write_to_database(self, hosts=None, overwrite=False, base_file_path_pattern=None, version=None, return_catalogs=False, raise_exception=False, add_specs_only_before_time=None, debug=None):
        """
        This function builds the base catalog and writes it to the database.

        !! IMPORTANT !!
        If you want to write the base catalog to an alternative location (not the database)
        Make sure you set the `base_file_path_pattern` option!!

        Parameters
        ----------
        hosts : int, str, list, optional
            host names/IDs or a list of host names/IDs or short-hand names like
            "paper1" or "paper1_complete"

        overwrite : bool, optional (default: False)
            If set to True, overwrite existing base catalog

        base_file_path_pattern : str, optional
        version : int, optional (default: 2)
        return_catalogs : bool, optional (default: False)
        raise_exception : bool, optional (default: False)
        add_specs_only_before_time : astropy.time.Time, optional (default: None)

        Examples
        --------
        >>> saga_database = SAGA.Database('/path/to/SAGA/Dropbox', '/path/to/SAGA/local')
        >>> saga_object_catalog = SAGA.ObjectCatalog(saga_database)

        Overwrite the database (Danger!!)
        >>> saga_object_catalog.build_and_write_to_database('paper1', overwrite=True)

        You can also do
        >>> saga_object_catalog.build_and_write_to_database('paper1', base_file_path_pattern='/other/base/catalog/dir/nsa{}.fits.gz')

        """
        if version not in (None, 1, 2):
            raise ValueError('`version` must be None, 1 or 2.')
        build_module = build if version == 1 else build2

        nsa = self.load_nsa('0.1.2' if version == 1 else '1.0.1')
        spectra = self._database['spectra_raw_all'].read(before_time=add_specs_only_before_time)

        manual_lists = {}
        for survey, col in (('sdss', 'SDSS ID'), ('des', 'DES_OBJID'), ('decals', 'decals_objid')):
            for list_type in ('remove', 'recover'):
                key = '{}_{}'.format(survey, list_type)
                val = get_unique_objids(self._database[key].read()[col])
                if not len(val):
                    val = None
                manual_lists[key] = val

        catalogs_to_return = list()
        host_ids = self._host_catalog.resolve_id(hosts, 'string')

        for i, host_id in enumerate(host_ids):
            if base_file_path_pattern is None:
                base_key = 'base' if version is None else 'base_v{}'.format(version)
                data_obj = self._database[base_key, host_id].remote
            else:
                data_obj = FitsTable(base_file_path_pattern.format(host_id))

            if data_obj.isfile() and not overwrite:
                print(time.strftime('[%m/%d %H:%M:%S]'), 'Base catalog v{} for {} already exists ({}).'.format(version or 2, host_id, data_obj.path), '({}/{})'.format(i+1, len(host_ids)))
                continue

            host = self._host_catalog.load_single(host_id)
            catalogs = ('sdss', 'wise') if version == 1 else ('sdss', 'des', 'decals')

            def get_catalog_or_none(catalog_name):
                try:
                    cat = self._database[catalog_name, host_id].read()
                except OSError:
                    print(time.strftime('[%m/%d %H:%M:%S]'), '[WARNING] Not found: {} catalog for {}.'.format(catalog_name.upper(), host_id))
                    return None
                return cat[build.WISE_COLS_USED] if catalog_name == 'wise' else cat

            catalog_dict = {k: get_catalog_or_none(k) for k in catalogs}

            print(
                time.strftime('[%m/%d %H:%M:%S]'),
                'Use {} to build base catalog v{} for {}'.format(
                    ', '.join((k for k, v in catalog_dict.items() if v is not None)).upper(),
                    (version or 2),
                    host_id
                ),
                '({}/{})'.format(i+1, len(host_ids))
            )

            if debug is None:
                debug_this = None
            else:
                debug[host_id] = dict()
                debug_this = debug[host_id]

            try:
                base = build_module.build_full_stack(host=host, nsa=nsa, spectra=spectra, debug=debug_this,
                                                     **manual_lists, **catalog_dict)
            except Exception as e: # pylint: disable=W0703
                print(time.strftime('[%m/%d %H:%M:%S]'), '[ERROR] Fail to build base catalog for {}\n{}'.format(host_id, e))
                base = None
                if raise_exception:
                    raise e
                continue
            finally:
                del catalog_dict
                if return_catalogs:
                    catalogs_to_return.append(base)

            print(time.strftime('[%m/%d %H:%M:%S]'), 'Write base catalog to {}'.format(data_obj.path))
            try:
                data_obj.write(base)
            except (IOError, OSError) as e:
                print(time.strftime('[%m/%d %H:%M:%S]'), '[ERROR] Fail to write base catalog for {}\n{}'.format(host_id, e))
                continue

        if return_catalogs:
            return catalogs_to_return

    def generate_clean_specs(self, version=None):
        """
        gather spectra from all base catalogs
        """
        surveys = ('sdss', 'des', 'decals')
        base_key = 'base' if version is None else 'base_v{}'.format(version)
        out = []

        for host in self._host_catalog.resolve_id('all', 'string'):
            try:
                base = self._database[base_key, host].read()
            except IOError:
                print('[WARNING] cannot find base catalog for', host)

            base = C.has_spec.filter(base)

            if version != 1:
                cols = [col for col in base.colnames if col.startswith('OBJID_') or not any(col.endswith('_'+s) for s in surveys)]
                base = base[cols]
                for s in surveys:
                    key = 'OBJID_' + s
                    if key not in base.colnames:
                        base[key] = -1

            out.append(base)

        return vstack(out, 'exact', 'error')
<|endoftext|>