<commit_before>from contextlib import contextmanager
import fire
import subprocess
from influxdb import InfluxDBClient, DataFrameClient
import logging

from gpumon.nvidia_dmon import nvidia_run_dmon_poll
from gpumon.influxdb_interface import create_influxdb_writer

logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger(__name__)


class Logger(object):
    def __init__(self, ip_or_url, port, username, password, database, series_name):
        self._client = DataFrameClient(ip_or_url, port, username, password, database)
        self._series_name = series_name

    def __call__(self):
        return self._client.query("select * from {}".format(self._series_name))[self._series_name]


@contextmanager
def log_context(ip_or_url, port, username, password, database, series_name, **tags):
    logger.info('Logging GPU to Database {}'.format(ip_or_url))
    process_args = ["python",
                    "influxdb_gpu_logger.py",
                    ip_or_url,
                    port,
                    username,
                    password,
                    database,
                    series_name]

    if tags:
        process_args.extend(tags.to_list()) # TODO:Check

    with subprocess.Popen(process_args, stdout=subprocess.PIPE) as proc:
        yield Logger(ip_or_url, port, username, password, database, series_name)
        proc.terminate()


def _create_database(influxdb_client, database_name):
    dbs = influxdb_client.get_list_database()
    if database_name in [db['name'] for db in dbs]:
        logger.info('Database {} exists')
        influxdb_client.switch_database(database_name)
    else:
        influxdb_client.create_database(database_name)


def main(ip_or_url, port, username, password, database, series_name='gpu_measurements', **tags):
    """ Starts GPU logger

    Logs GPU measurements form nvidia-smi to influxdb

    Parameters
    ----------
    ip_or_url:
    port:
    username:
    password:
    database:
    series_name:
    tags:
    """
    try:
        logger.info('Trying to connect to {} on port {} with {}:{}'.format(ip_or_url, port, username, password))
        client = InfluxDBClient(ip_or_url, port, username, password)
        logger.info('Connected')

        _create_database(client, database)

        to_db = create_influxdb_writer(client, series_name=series_name, **tags)
        logger.info('Starting logging...')
        nvidia_run_dmon_poll(to_db)
    except KeyboardInterrupt:
        logger.info('Exiting')


if __name__=="__main__":
    fire.Fire(main)<commit_msg>Adds debug switch<commit_after>from contextlib import contextmanager
import fire
import subprocess
from influxdb import InfluxDBClient, DataFrameClient
import logging

from gpumon.nvidia_dmon import nvidia_run_dmon_poll
from gpumon.influxdb_interface import create_influxdb_writer

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class Logger(object):
    def __init__(self, ip_or_url, port, username, password, database, series_name):
        self._client = DataFrameClient(ip_or_url, port, username, password, database)
        self._series_name = series_name

    def __call__(self):
        return self._client.query("select * from {}".format(self._series_name))[self._series_name]


@contextmanager
def log_context(ip_or_url, port, username, password, database, series_name, **tags):
    logger.info('Logging GPU to Database {}'.format(ip_or_url))
    process_args = ["python",
                    "influxdb_gpu_logger.py",
                    ip_or_url,
                    port,
                    username,
                    password,
                    database,
                    series_name]

    if tags:
        process_args.extend(tags.to_list()) # TODO:Check

    with subprocess.Popen(process_args, stdout=subprocess.PIPE) as proc:
        yield Logger(ip_or_url, port, username, password, database, series_name)
        proc.terminate()


def _create_database(influxdb_client, database_name):
    dbs = influxdb_client.get_list_database()
    if database_name in [db['name'] for db in dbs]:
        logger.info('Database {} exists')
        influxdb_client.switch_database(database_name)
    else:
        influxdb_client.create_database(database_name)


def main(ip_or_url, port, username, password, database, series_name='gpu_measurements', debug=False, **tags):
    """ Starts GPU logger

    Logs GPU measurements form nvidia-smi to influxdb

    Parameters
    ----------
    ip_or_url:
    port:
    username:
    password:
    database:
    series_name:
    tags:
    """
    if debug:
        logging.basicConfig(level=logging.DEBUG)

    try:
        logger.info('Trying to connect to {} on port {} with {}:{}'.format(ip_or_url, port, username, password))
        client = InfluxDBClient(ip_or_url, port, username, password)
        logger.info('Connected')

        _create_database(client, database)

        to_db = create_influxdb_writer(client, series_name=series_name, **tags)
        logger.info('Starting logging...')
        nvidia_run_dmon_poll(to_db)
    except KeyboardInterrupt:
        logger.info('Exiting')


if __name__=="__main__":
    fire.Fire(main)
{'fraction_non_alphanumeric': 0.09915123456790123, 'fraction_numerical': 0.0, 'mean_word_length': 3.1347687400318978, 'pattern_counts': {'":': 2, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '335209', 'n_tokens_mistral': 1593, 'n_tokens_neox': 1524, 'n_words': 339}