<commit_before>"""Functions to communicate with the server using JSONRPC"""

from __future__ import absolute_import, division, print_function

import sys
import os
import time
from copy import deepcopy
from functools import wraps
from datetime import datetime

import logging

from iceprod.core import constants
from iceprod.core import functions
from iceprod.core import dataclasses
from .serialization import dict_to_dataclasses
from .jsonUtil import json_compressor,json_decode
from .rest_client import Client


class ServerComms:
    """
    Setup JSONRPC communications with the IceProd server.

    Args:
        url (str): address to connect to
        passkey (str): passkey for authorization/authentication
        config (:py:class:`iceprod.server.exe.Config`): Config object
        **kwargs: passed to JSONRPC
    """
    def __init__(self, url, passkey, config, **kwargs):
        self.url = url
        self.cfg = config
        self.rest = Client(address=url,auth_key=passkey,**kwargs)

    async def download_task(self, gridspec, resources={}):
        """
        Download new task(s) from the server.

        Args:
            gridspec (str): gridspec the pilot was submitted from
            resources (dict): resources available in the pilot

        Returns:
            list: list of task configs
        """
        hostname = functions.gethostname()
        domain = '.'.join(hostname.split('.')[-2:])
        try:
            ifaces = functions.getInterfaces()
        except Exception:
            ifaces = None
        resources = deepcopy(resources)
        if 'gpu' in resources and isinstance(resources['gpu'],list):
            resources['gpu'] = len(resources['gpu'])
        os_type = os.environ['OS_ARCH'] if 'OS_ARCH' in os.environ else None
        if os_type:
            resources['os'] = os_type
        task = await self.rest.request('POST', '/task_actions/process',
                {'gridspec': gridspec,
                 'hostname': hostname, 
                 'domain': domain,
                 'ifaces': ifaces,
                 'requirements': resources,
                })
        if not task:
            return None
        if task and not isinstance(task,list):
            task = [task]
        # convert dict to Job
        ret = []
        for t in task:
            if not isinstance(t, dataclasses.Job):
                try:
                    ret.append(dict_to_dataclasses(t))
                except Exception:
                    logging.warning('not a Job: %r',t)
                    raise
            else:
                ret.append(t)
        return ret

    async def processing(self, task_id):
        """<commit_msg>get the dataset config and fill in options<commit_after>"""Functions to communicate with the server using JSONRPC"""

from __future__ import absolute_import, division, print_function

import sys
import os
import time
from copy import deepcopy
from functools import wraps
from datetime import datetime

import logging

from iceprod.core import constants
from iceprod.core import functions
from iceprod.core import dataclasses
from .serialization import dict_to_dataclasses
from .jsonUtil import json_compressor,json_decode
from .rest_client import Client


class ServerComms:
    """
    Setup JSONRPC communications with the IceProd server.

    Args:
        url (str): address to connect to
        passkey (str): passkey for authorization/authentication
        config (:py:class:`iceprod.server.exe.Config`): Config object
        **kwargs: passed to JSONRPC
    """
    def __init__(self, url, passkey, config, **kwargs):
        self.url = url
        self.cfg = config
        self.rest = Client(address=url,auth_key=passkey,**kwargs)

    async def download_task(self, gridspec, resources={}):
        """
        Download new task(s) from the server.

        Args:
            gridspec (str): gridspec the pilot was submitted from
            resources (dict): resources available in the pilot

        Returns:
            list: list of task configs
        """
        hostname = functions.gethostname()
        domain = '.'.join(hostname.split('.')[-2:])
        try:
            ifaces = functions.getInterfaces()
        except Exception:
            ifaces = None
        resources = deepcopy(resources)
        if 'gpu' in resources and isinstance(resources['gpu'],list):
            resources['gpu'] = len(resources['gpu'])
        os_type = os.environ['OS_ARCH'] if 'OS_ARCH' in os.environ else None
        if os_type:
            resources['os'] = os_type
        task = await self.rest.request('POST', '/task_actions/process',
                {'gridspec': gridspec,
                 'hostname': hostname, 
                 'domain': domain,
                 'ifaces': ifaces,
                 'requirements': resources,
                })
        if not task:
            return None

        # get config
        try:
            config = await self.rest.request('GET', '/config/{}'.format(task['dataset_id']))
        except Exception:
            logging.warn('failed to get dataset config for dataset %s', task['dataset_id'])
            await self.task_error(task['task_id'], dataset_id=task['dataset_id'],
                                  reason='failed to download dataset config')
            raise

        # convert dict to Job
        if not isinstance(config, dataclasses.Job):
            try:
                config = dict_to_dataclasses(config)
            except Exception:
                logging.warning('not a Job: %r', config)
                raise

        # fill in options
        if 'options' not in config:
            config['options'] = {}
        config['options']['task_id'] = task['task_id']
        config['options']['job_id'] = task['job_id']
        config['options']['dataset_id'] = task['dataset_id']
        config['options']['task'] = task['task_index']
        try:
            job = await self.rest.request('GET', '/jobs/{}'.format(task['job_id']))
            config['options']['job'] = job['job_index']
        except Exception:
            logging.warn('failed to get dataset config for dataset %s', task['dataset_id'])
            await self.task_error(task['task_id'], dataset_id=task['dataset_id'],
                                  reason='failed to download dataset config')
            raise
        return [config]

    async def processing(self, task_id):
        """
{'fraction_non_alphanumeric': 0.10121457489878542, 'fraction_numerical': 0.0003114294612270321, 'mean_word_length': 2.156265356265356, 'pattern_counts': {'":': 0, '<': 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': 8, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4298852', 'n_tokens_mistral': 1684, 'n_tokens_neox': 1644, 'n_words': 508}