<commit_before>#  Numenta Platform for Intelligent Computing (NuPIC)
#  Copyright (C) 2019, Numenta, Inc.  Unless you have an agreement
#  with Numenta, Inc., for a separate license for this software code, the
#  following terms and conditions apply:
#
#  This program is free software: you can redistribute it and/or modify
#  it under the terms of the GNU Affero Public License version 3 as
#  published by the Free Software Foundation.
#
#  This program is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
#  See the GNU Affero Public License for more details.
#
#  You should have received a copy of the GNU Affero Public License
#  along with this program.  If not, see http://www.gnu.org/licenses.
#
#  http://numenta.org/licenses/
#
import copy
import os
from pprint import pprint

import ray
import ray.resource_spec
from ray import ray_constants
from ray.tune import Trainable, TuneError, tune
from ray.tune.ray_trial_executor import RESOURCE_REFRESH_PERIOD, RayTrialExecutor
from ray.tune.resources import Resources

from nupic.research.frameworks.pytorch.imagenet import ImagenetExperiment

os.environ["HDF5_USE_FILE_LOCKING"] = "FALSE"


def _get_resources_per_node():
    """
    Maps node id to available resources on that node
    :return: dict with available :class:`Resources` for each node
    """
    def _is_node_key(k):
        return k.startswith(ray.resource_spec.NODE_ID_PREFIX)

    # Only consider active/alive nodes
    nodes = filter(lambda node: node["Alive"], ray.nodes())
    resources = map(lambda node: node["Resources"], nodes)

    # Group resources by node
    resources_by_node = {}
    for item in resources:
        node_id = next(filter(_is_node_key, item.keys()))

        item = item.copy()
        num_cpus = item.pop("CPU", 0)
        num_gpus = item.pop("GPU", 0)
        memory = ray_constants.from_memory_units(item.pop("memory", 0))
        object_store_memory = ray_constants.from_memory_units(
            item.pop("object_store_memory", 0))
        custom_resources = item

        resources_by_node[node_id] = Resources(
            int(num_cpus),
            int(num_gpus),
            memory=int(memory),
            object_store_memory=int(object_store_memory),
            custom_resources=custom_resources)

    return resources_by_node


class AffinityExecutor(RayTrialExecutor):
    """
    Ray Trial Executor used to set node affinity to node trials.
    This trial executor will update the trial configuration with one extra
    parameter (__ray_node_affinity__) that can be used to launch extra workers
    on the same node.
    """

    def __init__(self, queue_trials=False, reuse_actors=False, ray_auto_init=False,
                 refresh_period=RESOURCE_REFRESH_PERIOD):
        self._resources_by_node = {}
        super().__init__(queue_trials, reuse_actors, ray_auto_init, refresh_period)

    def _commit_resources(self, resources):
        resources = copy.deepcopy(resources)
        # FIXME: _commit_resources ignores extra_custom_resources
        # if key is missing from custom_resources
        for k in resources.extra_custom_resources:
            if k not in resources.custom_resources:
                resources.custom_resources[k] = 0.0<commit_msg>Add explanation for updating `resources.custom_resources` keys on `_commit_resources` method and remove FIXME<commit_after>#  Numenta Platform for Intelligent Computing (NuPIC)
#  Copyright (C) 2019, Numenta, Inc.  Unless you have an agreement
#  with Numenta, Inc., for a separate license for this software code, the
#  following terms and conditions apply:
#
#  This program is free software: you can redistribute it and/or modify
#  it under the terms of the GNU Affero Public License version 3 as
#  published by the Free Software Foundation.
#
#  This program is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
#  See the GNU Affero Public License for more details.
#
#  You should have received a copy of the GNU Affero Public License
#  along with this program.  If not, see http://www.gnu.org/licenses.
#
#  http://numenta.org/licenses/
#
import copy
import os
from pprint import pprint

import ray
import ray.resource_spec
from ray import ray_constants
from ray.tune import Trainable, TuneError, tune
from ray.tune.ray_trial_executor import RESOURCE_REFRESH_PERIOD, RayTrialExecutor
from ray.tune.resources import Resources

from nupic.research.frameworks.pytorch.imagenet import ImagenetExperiment

os.environ["HDF5_USE_FILE_LOCKING"] = "FALSE"


def _get_resources_per_node():
    """
    Maps node id to available resources on that node
    :return: dict with available :class:`Resources` for each node
    """
    def _is_node_key(k):
        return k.startswith(ray.resource_spec.NODE_ID_PREFIX)

    # Only consider active/alive nodes
    nodes = filter(lambda node: node["Alive"], ray.nodes())
    resources = map(lambda node: node["Resources"], nodes)

    # Group resources by node
    resources_by_node = {}
    for item in resources:
        node_id = next(filter(_is_node_key, item.keys()))

        item = item.copy()
        num_cpus = item.pop("CPU", 0)
        num_gpus = item.pop("GPU", 0)
        memory = ray_constants.from_memory_units(item.pop("memory", 0))
        object_store_memory = ray_constants.from_memory_units(
            item.pop("object_store_memory", 0))
        custom_resources = item

        resources_by_node[node_id] = Resources(
            int(num_cpus),
            int(num_gpus),
            memory=int(memory),
            object_store_memory=int(object_store_memory),
            custom_resources=custom_resources)

    return resources_by_node


class AffinityExecutor(RayTrialExecutor):
    """
    Ray Trial Executor used to set node affinity to node trials.
    This trial executor will update the trial configuration with one extra
    parameter (__ray_node_affinity__) that can be used to launch extra workers
    on the same node.
    """

    def __init__(self, queue_trials=False, reuse_actors=False, ray_auto_init=False,
                 refresh_period=RESOURCE_REFRESH_PERIOD):
        self._resources_by_node = {}
        super().__init__(queue_trials, reuse_actors, ray_auto_init, refresh_period)

    def _commit_resources(self, resources):
        resources = copy.deepcopy(resources)
        # Make sure "custom_resources" keys match "extra_custom_resources" keys
        # because `RayTrialExecutor._commit_resources` will only commit
        # resources from "extra_custom_resources" if the same resource key
        # is present in the "custom_resources" dict
        for k in resources.extra_custom_resources:
            if k not in resources.custom_resources:
                resources.custom_resources[k] = 0.0
{'fraction_non_alphanumeric': 0.07482407008473359, 'fraction_numerical': 0.0034467901766479965, 'mean_word_length': 3.498708010335917, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, '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': '6886978', 'n_tokens_mistral': 2023, 'n_tokens_neox': 1966, 'n_words': 646}