<commit_before>#!/usr/bin/python
# -*- coding: utf-8 -*-

'''sabd-cli

Because I need gurbani on the cli.

Usage:
  sabd-cli [options] -q <query>
  sabd-cli [options] -s <sabd_id>
  sabd-cli (-h | --help)

Examples:
  sabd-cli -q hjkk <to find "har jio kirpa karo">
  sabd-cli -s 2289 <to bring all lines of "har jio kirpa karo" back>

Options:
  -q <query>, --first-letter-search <query>     To do a first letter search
  -s <sabad_id>, --sabd <sabad_id>              To bring a specific sabd back
  -o, --output-as html|txt|md                   Output sabd as html|txt|md
  --html-template line-by-line|presentation     If html then output as line-by-line[default] mode or presentation
  -h, --help                                    Show this screen
'''

import docopt
from clint.textui import colored
import tempfile
# for html output
from jinja2 import Environment, PackageLoader
import webbrowser

import application.bootstrap as bootstrap
from application.models.IsgGurbaniDb import IsgGurbaniDb

def main():
    '''Main entry point for sabd-cli.'''
    args = docopt.docopt(__doc__, help=True, version=None, options_first=False)

    if args['--output-as'] and args['--output-as'] not in ['md', 'txt', 'html']:
        bootstrap.logger.info("we don't support that file format yet")
        quit()

    model_gurbani = IsgGurbaniDb(bootstrap.config)
    if args['--first-letter-search']:
        query = args['--first-letter-search'].strip()
        bootstrap.logger.info("Doing Gurmukhi first letter search with " + query)

        try:
            data = model_gurbani.first_letter_search(query)
        except Exception as e:
            bootstrap.logger.critical("There was an issue searching for " + query, e.message)

        if data:
            print (getattr(colored, 'red')("sabd_id") + "\t")
            print ("Gurmukhi" + "\t")
            for line in data:
                print (getattr(colored, 'red')(line[1]) + "\t")
                print (line[2] + "\t")
        else:
            bootstrap.logger.warn("couldn't find any sabds with first letter search: " + query)

    # get a sabad
    if args['--sabd']:
        sabd_id = int(args['--sabd'])
        bootstrap.logger.info("getting sabad by id " + str(sabd_id))
        bootstrap.logger.info("Searching for sabd by " + str(sabd_id))

        try:
            data = model_gurbani.get_sabad_by_sabad_id(sabd_id)
        except Exception as e:
            bootstrap.logger.critical("There was an issue searching for " + str(sabd_id), e.message)

        if data:
            if args['--output-as'] == 'txt':
                for line in data:
                    # print (getattr(colored, 'cyan')(line[0])) + "\t", #ID
                    #print (getattr(colored, 'cyan')(line[1])) + "\t", #ANG
                    #print (getattr(colored, 'cyan')(line[2])) + "\t", #source
                    #print (getattr(colored, 'cyan')(line[3])) + "\t", # line
                    ##print (getattr(colored, 'red')(line[4])) + "\t", # sabd_id
                    print (line[5])  # gurmukhi
                    print (line[6])  #transliteration
                    print (getattr(colored, 'cyan')(line[7]))  #english translation
                    print ("\n")
                    #print " " + (getattr(colored, 'magenta')line[1].ljust(40)
            elif args['--output-as'] == 'md':
                bootstrap.logger.warn("md output no yet implemented")
            else:  # default is presentation mode html
                html_mode = 'line-by-line'
                if args['--html-template']:
                    if args['--html-template'].strip() == 'presentation':
                        html_mode = 'presentation'

                output_file = tempfile.NamedTemporaryFile(mode='w+b', buffering=-1, encoding=None, newline=None, suffix='.html', prefix='sabd-cli-', dir=None, delete=False)
                bootstrap.logger.info('output filename has not been specified, will write to ' + str(output_file.name))
                outputSabadAsFile(output_file, 'html', data, html_mode)

        else:
            bootstrap.logger.warn("couldn't find a sabad by " + str(sabd_id))

def outputSabadAsFile(file, type, data, html_template_type='line-by-line'):
    """
    Output sabd as certain filetypes
    @param file the file object
    @param string md|html|txt
    @param dictionary data
    @param string html template type
    """
    if type == 'md':
        for line in data:
            file.write("**" + line[5].encode("UTF-8") + "**" + "\n")
            file.write("*" + line[6].encode("UTF-8") + "*" + "\n")
            file.write(line[7].encode("UTF-8") + "\n")
            file.write("\n")
        file.flush()
        file.close()

    if type == 'html':
        # use jinja templates
        env = Environment(loader=PackageLoader('application', 'views'))
        template = env.get_template(html_template_type + '.html')
        template_data = {'title': 'Sabad CLI output [' + html_template_type + ']', 'gurbani': data}
        output = template.render(template_data)
        file.write(output.encode("UTF-8"))
        file.flush()
        file.close()
        webbrowser.open(file.name, 2)  #open file in new browser 2=new tab
        quit()

    if type == 'txt':
        for line in data:
            file.write(line[5].encode("UTF-8") + "\n")
            file.write(line[6].encode("UTF-8") + "\n")
            file.write(line[7].encode("UTF-8") + "\n")
            file.write("\n")
        file.flush()
        file.close()

    bootstrap.logger.info("Finished, Written output to " + file.name)

if __name__ == '__main__':
    main()
<commit_msg>changing -f to -q<commit_after>#!/usr/bin/python
# -*- coding: utf-8 -*-

'''sabd-cli

Because I need gurbani on the cli.

Usage:
  sabd-cli [options] -q <query>
  sabd-cli [options] -s <sabd_id>
  sabd-cli (-h | --help)

Examples:
  sabd-cli -q hjkk <to find "har jio kirpa karo">
  sabd-cli -s 2289 <to bring all lines of "har jio kirpa karo" back>

Options:
  -q <query>, --first-letter-search <query>     To do a first letter search
  -s <sabad_id>, --sabd <sabad_id>              To bring a specific sabd back
  -o, --output-as html|txt|md                   Output sabd as html|txt|md
  --html-template line-by-line|presentation     If html then output as line-by-line[default] mode or presentation
  -h, --help                                    Show this screen
'''

import docopt
from clint.textui import colored
import tempfile
# for html output
from jinja2 import Environment, PackageLoader
import webbrowser

import application.bootstrap as bootstrap
from application.models.IsgGurbaniDb import IsgGurbaniDb

def main():
    '''Main entry point for sabd-cli.'''
    args = docopt.docopt(__doc__, help=True, version=None, options_first=False)

    if args['--output-as'] and args['--output-as'] not in ['md', 'txt', 'html']:
        bootstrap.logger.info("we don't support that file format yet")
        quit()

    model_gurbani = IsgGurbaniDb(bootstrap.config)
    if args['--first-letter-search']:
        query = args['--first-letter-search'].strip()
        bootstrap.logger.info("Doing Gurmukhi first letter search with " + query)

        try:
            data = model_gurbani.first_letter_search(query)
        except Exception as e:
            bootstrap.logger.critical("There was an issue searching for " + query, e.message)

        if data:
            print (getattr(colored, 'red')("sabd_id") + "\t")
            print ("Gurmukhi" + "\t")
            for line in data:
                print (getattr(colored, 'red')(line[1]) + " " + line[2] + "\t")

        else:
            bootstrap.logger.warn("couldn't find any sabds with first letter search: " + query)

    # get a sabad
    if args['--sabd']:
        sabd_id = int(args['--sabd'])
        bootstrap.logger.info("getting sabad by id " + str(sabd_id))
        bootstrap.logger.info("Searching for sabd by " + str(sabd_id))

        try:
            data = model_gurbani.get_sabad_by_sabad_id(sabd_id)
        except Exception as e:
            bootstrap.logger.critical("There was an issue searching for " + str(sabd_id), e.message)

        if data:
            if args['--output-as'] == 'txt':
                for line in data:
                    # print (getattr(colored, 'cyan')(line[0])) + "\t", #ID
                    #print (getattr(colored, 'cyan')(line[1])) + "\t", #ANG
                    #print (getattr(colored, 'cyan')(line[2])) + "\t", #source
                    #print (getattr(colored, 'cyan')(line[3])) + "\t", # line
                    ##print (getattr(colored, 'red')(line[4])) + "\t", # sabd_id
                    print (line[5])  # gurmukhi
                    print (line[6])  #transliteration
                    print (getattr(colored, 'cyan')(line[7]))  #english translation
                    print ("\n")
                    #print " " + (getattr(colored, 'magenta')line[1].ljust(40)
            elif args['--output-as'] == 'md':
                bootstrap.logger.warn("md output no yet implemented")
            else:  # default is presentation mode html
                html_mode = 'line-by-line'
                if args['--html-template']:
                    if args['--html-template'].strip() == 'presentation':
                        html_mode = 'presentation'

                output_file = tempfile.NamedTemporaryFile(mode='w+b', buffering=-1, encoding=None, newline=None, suffix='.html', prefix='sabd-cli-', dir=None, delete=False)
                bootstrap.logger.info('output filename has not been specified, will write to ' + str(output_file.name))
                outputSabadAsFile(output_file, 'html', data, html_mode)

        else:
            bootstrap.logger.warn("couldn't find a sabad by " + str(sabd_id))

def outputSabadAsFile(file, type, data, html_template_type='line-by-line'):
    """
    Output sabd as certain filetypes
    @param file the file object
    @param string md|html|txt
    @param dictionary data
    @param string html template type
    """
    if type == 'md':
        for line in data:
            file.write("**" + line[5].encode("UTF-8") + "**" + "\n")
            file.write("*" + line[6].encode("UTF-8") + "*" + "\n")
            file.write(line[7].encode("UTF-8") + "\n")
            file.write("\n")
        file.flush()
        file.close()

    if type == 'html':
        # use jinja templates
        env = Environment(loader=PackageLoader('application', 'views'))
        template = env.get_template(html_template_type + '.html')
        template_data = {'title': 'Sabad CLI output [' + html_template_type + ']', 'gurbani': data}
        output = template.render(template_data)
        file.write(output.encode("UTF-8"))
        file.flush()
        file.close()
        webbrowser.open(file.name, 2)  #open file in new browser 2=new tab
        quit()

    if type == 'txt':
        for line in data:
            file.write(line[5].encode("UTF-8") + "\n")
            file.write(line[6].encode("UTF-8") + "\n")
            file.write(line[7].encode("UTF-8") + "\n")
            file.write("\n")
        file.flush()
        file.close()

    bootstrap.logger.info("Finished, Written output to " + file.name)

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


<commit_before># Copyright 2021 DeepMind Technologies Limited
#
# 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
#
#      http://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.
"""Resources specification for use in the API.

Various classes defined to support resources specification for jobs.
"""

import enum
import functools
import itertools
import operator
import re
from typing import Any, Iterable, Iterator, Dict, Mapping, MutableMapping, Optional, Tuple, Union

import immutabledict
from xmanager.xm import pattern_matching as pm

class _CaseInsensetiveResourceTypeMeta(enum.EnumMeta):
  """Metaclass which allows case-insensetive enum lookup.

  Enum keys are upper case, but we allow other cases for the input. For
  example existing flags and JobRequirements use lower case for resource names.
  """

  def __getitem__(cls, resource_name: str) -> 'ResourceType':
    try:
      return super().__getitem__(resource_name.upper())
    except KeyError:
      raise KeyError(f'Unknown {cls.__name__} {resource_name!r}')  # pylint: disable=raise-missing-from

class ResourceType(enum.Enum, metaclass=_CaseInsensetiveResourceTypeMeta):
  """Type of a countable resource (e.g., CPU, memory, accelerators etc).

  We use a schema in which every particular accelerator has its own type. This
  way all countable resources required for a job could be represented by a
  simple dictionary.
  """

  # Amount of required CPU resources in vCPUs.
  CPU = 100002
  # Amount of required memory resources in bytes.
  MEMORY = 39
  RAM = 39
  # Amount of required disk resources in bytes.
  EPHEMERAL_STORAGE = 100003
  DISK = 100003

  # GPUs
  LOCAL_GPU = 100006
  P4 = 21
  T4 = 22
  P100 = 14
  V100 = 17
  A100 = 46

  # TPUs
  TPU_V2 = 3
  TPU_V3 = 16

  # TODO: do we need V2_DONUT and V3_DONUT?

  def __str__(self):
    return self._name_

class _CaseInsensetiveServiceTierMeta(enum.EnumMeta):
  """Metaclass which allows case-insensetive enum lookup.

  Enum keys are upper case, but we allow other cases for the input. For
  example existing flags and JobRequirements use lower case for resource names.
  """

  def __getitem__(cls, resource_name: str) -> 'ServiceTier':
    try:
      return super().__getitem__(resource_name.upper())
    except KeyError:
      raise KeyError(f'Unknown {cls.__name__} {resource_name!r}')  # pylint: disable=raise-missing-from

class ServiceTier(enum.Enum, metaclass=_CaseInsensetiveServiceTierMeta):
  """The job availability guarantees which underlying platfrom should provide.

  Most cloud platforms offer a selection of availability/price tradeoff options.
  Usually there are at least "Take my money, this workload is important" and
  "Buy excess compute for cheap" offerings. This enum provides a classification
  of such offerings and allows matching comparable (but not necessary identical)
  options from different runtimes.
  """

  # Highly available tier. The job is expected to be scheduled fast once sent to
  # the cloud. Recommended tier for multi-job work units as lower tiers may lead
  # to partially-scheduled work units.
  PROD = 200
  # A cheaper tier with guaranteed average throughput. Jobs may spend hours
  # awaiting scheduling by the cloud and can be preempted.
  BATCH = 100
  # The cheapest tier. No guarantees, but it often works.
  FREEBIE = 25

def _enum_subset(class_name: str, values: Iterable[ResourceType]) -> type:  # pylint: disable=g-bare-generic
  """Returns an enum subset class.

  The class is syntactically equivalent to an enum with the given resource
  types. But the concrete constants are the same as in the ResourceType enum,
  making all equivalence comparisons work correctly. Additionally operator `in`
  is supported for checking if a resource belongs to the subset.

  Args:
    class_name: Class name of the subset enum.
    values: A list of resources that belong to the subset.
  """
  values = set(values)

  class EnumSubsetMetaclass(type):  # pylint: disable=g-bare-generic
    """Metaclass which implements enum subset operations."""

    def __new__(
        cls,
        name: str,
        bases: Tuple[type],  # pylint: disable=g-bare-generic
        dct: Dict[str, Any],
    ) -> type:  # pylint: disable=g-bare-generic
      # Add constants to the class dict.
      for name, member in ResourceType.__members__.items():
        if member in values:
          dct[name] = member

      return super().__new__(cls, class_name, bases, dct)

    def __getitem__(cls, item: str) -> ResourceType:
      result = ResourceType[item]
      if result not in cls:  # pylint: disable=unsupported-membership-test
        raise AttributeError(
            f"type object '{cls.__name__}' has no attribute '{item}'")
      return result

    def __iter__(cls) -> Iterator[ResourceType]:
      return iter(values)

    def contains(cls, value: ResourceType) -> bool:
      return value in values

  class EnumSubset(metaclass=EnumSubsetMetaclass):

    def __new__(cls, value: int) -> ResourceType:
      resource = ResourceType(value)
      if resource not in cls:
        raise ValueError(f'{value} is not a valid {cls.__name__}')
      return resource

  return EnumSubset

# TODO: Use centralized resource metadata.
TpuType = _enum_subset(
    'TpuType',
    [
        ResourceType.TPU_V2,
        ResourceType.TPU_V3,
    ],
)

GpuType = _enum_subset(
    'GpuType',
    [
        # LOCAL_GPU is missing as only specific GPU types should be added.
        ResourceType.P4,
        ResourceType.T4,
        ResourceType.P100,
        ResourceType.V100,
        ResourceType.A100,
    ],
)

_AcceleratorType = _enum_subset(
    '_AcceleratorType',
    [
        ResourceType.LOCAL_GPU,
        *list(TpuType),
        *list(GpuType),
    ],
)

class ResourceDict(MutableMapping):
  """Internal class to represent amount of countable resources.

  A mapping from ResourceType to amount of the resource combined with
  convenience methods. This class only tracks amounts of the resources, but not
  their topologies, locations or constraints.

  This class is rather generic and is designed be used internally as job
  requirements as well as in the executors. API users should not use it
  explicitly.

  Usage:
    # Construct (implicitly) from user code using JobRequirements:
    requirements = JobRequirements(cpu=0.5 * xm.vCPU, memory=2 * xm.GiB, v100=8)
    resources = requirements.task_requirements
    # Resources are available by their canonical names.
    assert(resources[ResourceType.V100], 8)
    # Print user-friendly representation:
    print(f'The task needs {resources}')
  """

  def __init__(self) -> None:
    self.__dict: Dict[ResourceType, float] = {}

  def __setitem__(self, key: ResourceType, value: float) -> None:
    self.__dict.__setitem__(key, value)

  def __getitem__(self, key: ResourceType) -> float:
    return self.__dict.__getitem__(key)

  def __delitem__(self, key: ResourceType) -> None:
    self.__dict.__delitem__(key)

  def __iter__(self):
    return self.__dict.__iter__()

  def __len__(self) -> int:
    return self.__dict.__len__()

  def __str__(self) -> str:
    """Returns user-readable text representation.

    Such as "V100: 8, CPU: 1.2, MEMORY: 5.4GiB".
    """
    # TODO: We do not aggregate memory yet, update this method to be more
    # user-friendly.
    return ', '.join(
        sorted([f'{key}: {value}' for (key, value) in self.items()]))

  def __add__(self: 'ResourceDict', rhs: 'ResourceDict') -> 'ResourceDict':
    """Returns a sum of two ResourceDicts."""
    result = ResourceDict()
    for key in [*self.keys(), *rhs.keys()]:
      result[key] = self.get(key, 0) + rhs.get(key, 0)
    return result

class InvalidTpuTopologyError(Exception):
  """An unrecognized TPU topology has been provided."""

class Topology:
  """Accelerator topology configuration.

  Describes accelerator interconnection. For example could be a TPU topology or
  several GPUs connected with NVLink. Topologies have a form of 'NxM_suffix'
  where N & M are the number of accelerators across the dimension and suffix
  corresponds to a specific interconnect type. Number of dimensions may vary.

  Examples of valid topologies:
    '1' - a single device.
    '4' - 4 GPUs on one host.
    '4x4' - a 4x4 TPU grid.
  """

  def __init__(self, name: str) -> None:
    if not re.fullmatch('([\\d]+x?)+(_(un)?twisted)?', name):
      raise InvalidTpuTopologyError(f'Invalid TPU topology: {name}.')

    self._name = name

    dimensions_str = name.split('_')[0]
    self.dimensions = list(map(int, dimensions_str.split('x')))

  @property
  def chip_count(self) -> int:
    """Returns the number of chips of the TPU topology."""
    return functools.reduce(operator.mul, self.dimensions)

  @property
  def name(self) -> str:
    """Returns the topology as a string."""
    return self._name

  def __repr__(self) -> str:
    return f'xm.Topology({self.name!r})'

  def __eq__(self, other: 'Topology') -> bool:
    return self.name == other.name

  def __hash__(self) -> int:
    return hash(self.name)

ResourceQuantity = Union[int, float, str, Topology]

def _parse_resource_quantity(
    resource_name: str,
    value: ResourceQuantity) -> Tuple[float, Optional[Topology]]:
  """Parses a string representation of a resource quantity."""

  def parse_string(value: str):
    if 'x' in value:
      topology = Topology(value)
      return topology.chip_count, topology
    else:
      # TODO: Parse SI suffixes, like GiB.
      return float(value), None

  def parse_topology(topology: Topology):
    return topology.chip_count, topology

  def parse_number(value: Any):
    return float(value), None

  try:
    return pm.match(parse_string, parse_topology, parse_number)(value)
  except Exception as e:
    raise ValueError(f"Couldn't parse resource quantity for {resource_name}. "
                     f'{value!r} was given.') from e

class JobRequirements:
  """Describes the resource requirements of a Job.

  Attributes:
    task_requirements: Amount of resources needed for a single task within a
      job.
    accelerator: The accelearator the jobs uses, if there is one. Jobs using
      multiple accelerators are not supported because different kinds of
      accelerators are usually not installed on the same host.
    topology: Accelerator topology, if an accelerator is used.
    location: Place where the job should run. For example a cluster name or
      a Borg cell.
    service_tier: A service tier at which the job should run.
    replicas: Number of identical tasks to run within a job
  """

  task_requirements: ResourceDict
  accelerator: Optional[ResourceType]
  topology: Optional[Topology]

  location: Optional[str]
  _service_tier: ServiceTier

  def __init__(
      self,
      resources: Mapping[Union[ResourceType, str],
                         ResourceQuantity] = immutabledict.immutabledict(),
      *,
      location: Optional[str] = None,
      replicas: Optional[int] = None,
      service_tier: Optional[ServiceTier] = None,
      **kw_resources: ResourceQuantity) -> None:
    """Define a set of resources.

    Args:
      resources: resource amounts as a dictionary,
        for example {xm.ResourceType.V100: 2}.
      location: Place where the job should run. For example a cluster name or a
        Borg cell.
      replicas: Number of identical tasks to run within a job. 1 by default.
      service_tier: A service tier at which the job should run.
      **kw_resources: resource amounts as a kwargs, for example `v100=2` or
        `ram=1 * xm.GiB`. See xm.ResourceType enum for the list of supported
        types and aliases.

    Raises:
      ValueError:
        If several accelerator resources are supplied (i.e. GPU and TPU).
        If the same resource is passed in a `resources` dictionary and as
          a command line argument.
        If topology is supplied for a non acceelerator resource.
    """
    self.location = location
    self._service_tier = service_tier or ServiceTier.PROD

    self.task_requirements = ResourceDict()
    self.accelerator = None
    self.topology = None

    for resource_name, value in itertools.chain(resources.items(),
                                                kw_resources.items()):
      scalar, topology = _parse_resource_quantity(resource_name, value)
      resource = pm.match(
          pm.Case([str], lambda r: ResourceType[r]),
          pm.Case([ResourceType], lambda r: r))(
              resource_name)

      if resource in _AcceleratorType:
        if self.accelerator is not None:
          raise ValueError('Accelerator already set.')
        self.accelerator = resource
        self.topology = topology or Topology(f'{scalar:g}')
      elif topology is not None:
        raise ValueError(
            f'A topology specified for non accelerator resource {resource}.')

      if resource in self.task_requirements:
        raise ValueError(f'{resource} has been specified twice.')
      self.task_requirements[resource] = scalar

    if (self.accelerator in GpuType and self.topology and
        len(self.topology.dimensions) == 2):
      if replicas is not None and replicas != self.topology.dimensions[1]:
        raise ValueError(
            f'For multihost GPUs with topology {self.topology}, replicas should'
            f'be either None or {self.topology.dimensions[1]}. Found: '
            f'{replicas}')
      replicas = self.topology.dimensions[1]

    self.replicas = replicas or 1
    self._validate_replicas()

  @property
  def service_tier(self):
    return self._service_tier

  @service_tier.setter
  def service_tier(self, new_service_tier):
    self._service_tier = new_service_tier

  def _validate_replicas(self) -> None:
    """Raises ValueError if replication is not supported."""
    if self.replicas > 1 and self.accelerator in TpuType:
      raise ValueError(
          f'Replicated jobs are not supported for {self.accelerator}.')

  def __repr__(self) -> str:
    """Returns string representation of the requirements."""
    args = []

    for resource, value in self.task_requirements.items():
      if resource in TpuType:
        args.append(f'{resource.name.lower()}={self.topology!r}')
      else:
        args.append(f'{resource.name.lower()}={value!r}')

    if self.location:
      args.append(f'location={self.location!r}')
    if self.service_tier != ServiceTier.PROD:
      args.append(f'service_tier=xm.{self.service_tier}')
    if self.replicas != 1:
      args.append(f'replicas={self.replicas}')

    return f'xm.JobRequirements({", ".join(args)})'
<commit_msg>Fix a typo in the docstring.<commit_after># Copyright 2021 DeepMind Technologies Limited
#
# 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
#
#      http://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.
"""Resources specification for use in the API.

Various classes defined to support resources specification for jobs.
"""

import enum
import functools
import itertools
import operator
import re
from typing import Any, Iterable, Iterator, Dict, Mapping, MutableMapping, Optional, Tuple, Union

import immutabledict
from xmanager.xm import pattern_matching as pm

class _CaseInsensetiveResourceTypeMeta(enum.EnumMeta):
  """Metaclass which allows case-insensetive enum lookup.

  Enum keys are upper case, but we allow other cases for the input. For
  example existing flags and JobRequirements use lower case for resource names.
  """

  def __getitem__(cls, resource_name: str) -> 'ResourceType':
    try:
      return super().__getitem__(resource_name.upper())
    except KeyError:
      raise KeyError(f'Unknown {cls.__name__} {resource_name!r}')  # pylint: disable=raise-missing-from

class ResourceType(enum.Enum, metaclass=_CaseInsensetiveResourceTypeMeta):
  """Type of a countable resource (e.g., CPU, memory, accelerators etc).

  We use a schema in which every particular accelerator has its own type. This
  way all countable resources required for a job could be represented by a
  simple dictionary.
  """

  # Amount of required CPU resources in vCPUs.
  CPU = 100002
  # Amount of required memory resources in bytes.
  MEMORY = 39
  RAM = 39
  # Amount of required disk resources in bytes.
  EPHEMERAL_STORAGE = 100003
  DISK = 100003

  # GPUs
  LOCAL_GPU = 100006
  P4 = 21
  T4 = 22
  P100 = 14
  V100 = 17
  A100 = 46

  # TPUs
  TPU_V2 = 3
  TPU_V3 = 16

  # TODO: do we need V2_DONUT and V3_DONUT?

  def __str__(self):
    return self._name_

class _CaseInsensetiveServiceTierMeta(enum.EnumMeta):
  """Metaclass which allows case-insensetive enum lookup.

  Enum keys are upper case, but we allow other cases for the input. For
  example existing flags and JobRequirements use lower case for resource names.
  """

  def __getitem__(cls, resource_name: str) -> 'ServiceTier':
    try:
      return super().__getitem__(resource_name.upper())
    except KeyError:
      raise KeyError(f'Unknown {cls.__name__} {resource_name!r}')  # pylint: disable=raise-missing-from

class ServiceTier(enum.Enum, metaclass=_CaseInsensetiveServiceTierMeta):
  """The job availability guarantees which underlying platfrom should provide.

  Most cloud platforms offer a selection of availability/price tradeoff options.
  Usually there are at least "Take my money, this workload is important" and
  "Buy excess compute for cheap" offerings. This enum provides a classification
  of such offerings and allows matching comparable (but not necessary identical)
  options from different runtimes.
  """

  # Highly available tier. The job is expected to be scheduled fast once sent to
  # the cloud. Recommended tier for multi-job work units as lower tiers may lead
  # to partially-scheduled work units.
  PROD = 200
  # A cheaper tier with guaranteed average throughput. Jobs may spend hours
  # awaiting scheduling by the cloud and can be preempted.
  BATCH = 100
  # The cheapest tier. No guarantees, but it often works.
  FREEBIE = 25

def _enum_subset(class_name: str, values: Iterable[ResourceType]) -> type:  # pylint: disable=g-bare-generic
  """Returns an enum subset class.

  The class is syntactically equivalent to an enum with the given resource
  types. But the concrete constants are the same as in the ResourceType enum,
  making all equivalence comparisons work correctly. Additionally operator `in`
  is supported for checking if a resource belongs to the subset.

  Args:
    class_name: Class name of the subset enum.
    values: A list of resources that belong to the subset.
  """
  values = set(values)

  class EnumSubsetMetaclass(type):  # pylint: disable=g-bare-generic
    """Metaclass which implements enum subset operations."""

    def __new__(
        cls,
        name: str,
        bases: Tuple[type],  # pylint: disable=g-bare-generic
        dct: Dict[str, Any],
    ) -> type:  # pylint: disable=g-bare-generic
      # Add constants to the class dict.
      for name, member in ResourceType.__members__.items():
        if member in values:
          dct[name] = member

      return super().__new__(cls, class_name, bases, dct)

    def __getitem__(cls, item: str) -> ResourceType:
      result = ResourceType[item]
      if result not in cls:  # pylint: disable=unsupported-membership-test
        raise AttributeError(
            f"type object '{cls.__name__}' has no attribute '{item}'")
      return result

    def __iter__(cls) -> Iterator[ResourceType]:
      return iter(values)

    def contains(cls, value: ResourceType) -> bool:
      return value in values

  class EnumSubset(metaclass=EnumSubsetMetaclass):

    def __new__(cls, value: int) -> ResourceType:
      resource = ResourceType(value)
      if resource not in cls:
        raise ValueError(f'{value} is not a valid {cls.__name__}')
      return resource

  return EnumSubset

# TODO: Use centralized resource metadata.
TpuType = _enum_subset(
    'TpuType',
    [
        ResourceType.TPU_V2,
        ResourceType.TPU_V3,
    ],
)

GpuType = _enum_subset(
    'GpuType',
    [
        # LOCAL_GPU is missing as only specific GPU types should be added.
        ResourceType.P4,
        ResourceType.T4,
        ResourceType.P100,
        ResourceType.V100,
        ResourceType.A100,
    ],
)

_AcceleratorType = _enum_subset(
    '_AcceleratorType',
    [
        ResourceType.LOCAL_GPU,
        *list(TpuType),
        *list(GpuType),
    ],
)

class ResourceDict(MutableMapping):
  """Internal class to represent amount of countable resources.

  A mapping from ResourceType to amount of the resource combined with
  convenience methods. This class only tracks amounts of the resources, but not
  their topologies, locations or constraints.

  This class is rather generic and is designed be used internally as job
  requirements as well as in the executors. API users should not use it
  explicitly.

  Usage:
    # Construct (implicitly) from user code using JobRequirements:
    requirements = JobRequirements(cpu=0.5 * xm.vCPU, memory=2 * xm.GiB, v100=8)
    resources = requirements.task_requirements
    # Resources are available by their canonical names.
    assert(resources[ResourceType.V100], 8)
    # Print user-friendly representation:
    print(f'The task needs {resources}')
  """

  def __init__(self) -> None:
    self.__dict: Dict[ResourceType, float] = {}

  def __setitem__(self, key: ResourceType, value: float) -> None:
    self.__dict.__setitem__(key, value)

  def __getitem__(self, key: ResourceType) -> float:
    return self.__dict.__getitem__(key)

  def __delitem__(self, key: ResourceType) -> None:
    self.__dict.__delitem__(key)

  def __iter__(self):
    return self.__dict.__iter__()

  def __len__(self) -> int:
    return self.__dict.__len__()

  def __str__(self) -> str:
    """Returns user-readable text representation.

    Such as "V100: 8, CPU: 1.2, MEMORY: 5.4GiB".
    """
    # TODO: We do not aggregate memory yet, update this method to be more
    # user-friendly.
    return ', '.join(
        sorted([f'{key}: {value}' for (key, value) in self.items()]))

  def __add__(self: 'ResourceDict', rhs: 'ResourceDict') -> 'ResourceDict':
    """Returns a sum of two ResourceDicts."""
    result = ResourceDict()
    for key in [*self.keys(), *rhs.keys()]:
      result[key] = self.get(key, 0) + rhs.get(key, 0)
    return result

class InvalidTpuTopologyError(Exception):
  """An unrecognized TPU topology has been provided."""

class Topology:
  """Accelerator topology configuration.

  Describes accelerator interconnection. For example could be a TPU topology or
  several GPUs connected with NVLink. Topologies have a form of 'NxM_suffix'
  where N & M are the number of accelerators across the dimension and suffix
  corresponds to a specific interconnect type. Number of dimensions may vary.

  Examples of valid topologies:
    '1' - a single device.
    '4' - 4 GPUs on one host.
    '4x4' - a 4x4 TPU grid.
  """

  def __init__(self, name: str) -> None:
    if not re.fullmatch('([\\d]+x?)+(_(un)?twisted)?', name):
      raise InvalidTpuTopologyError(f'Invalid TPU topology: {name}.')

    self._name = name

    dimensions_str = name.split('_')[0]
    self.dimensions = list(map(int, dimensions_str.split('x')))

  @property
  def chip_count(self) -> int:
    """Returns the number of chips of the TPU topology."""
    return functools.reduce(operator.mul, self.dimensions)

  @property
  def name(self) -> str:
    """Returns the topology as a string."""
    return self._name

  def __repr__(self) -> str:
    return f'xm.Topology({self.name!r})'

  def __eq__(self, other: 'Topology') -> bool:
    return self.name == other.name

  def __hash__(self) -> int:
    return hash(self.name)

ResourceQuantity = Union[int, float, str, Topology]

def _parse_resource_quantity(
    resource_name: str,
    value: ResourceQuantity) -> Tuple[float, Optional[Topology]]:
  """Parses a string representation of a resource quantity."""

  def parse_string(value: str):
    if 'x' in value:
      topology = Topology(value)
      return topology.chip_count, topology
    else:
      # TODO: Parse SI suffixes, like GiB.
      return float(value), None

  def parse_topology(topology: Topology):
    return topology.chip_count, topology

  def parse_number(value: Any):
    return float(value), None

  try:
    return pm.match(parse_string, parse_topology, parse_number)(value)
  except Exception as e:
    raise ValueError(f"Couldn't parse resource quantity for {resource_name}. "
                     f'{value!r} was given.') from e

class JobRequirements:
  """Describes the resource requirements of a Job.

  Attributes:
    task_requirements: Amount of resources needed for a single task within a
      job.
    accelerator: The accelerator the jobs uses, if there is one. Jobs using
      multiple accelerators are not supported because different kinds of
      accelerators are usually not installed on the same host.
    topology: Accelerator topology, if an accelerator is used.
    location: Place where the job should run. For example a cluster name or
      a Borg cell.
    service_tier: A service tier at which the job should run.
    replicas: Number of identical tasks to run within a job
  """

  task_requirements: ResourceDict
  accelerator: Optional[ResourceType]
  topology: Optional[Topology]

  location: Optional[str]
  _service_tier: ServiceTier

  def __init__(
      self,
      resources: Mapping[Union[ResourceType, str],
                         ResourceQuantity] = immutabledict.immutabledict(),
      *,
      location: Optional[str] = None,
      replicas: Optional[int] = None,
      service_tier: Optional[ServiceTier] = None,
      **kw_resources: ResourceQuantity) -> None:
    """Define a set of resources.

    Args:
      resources: resource amounts as a dictionary,
        for example {xm.ResourceType.V100: 2}.
      location: Place where the job should run. For example a cluster name or a
        Borg cell.
      replicas: Number of identical tasks to run within a job. 1 by default.
      service_tier: A service tier at which the job should run.
      **kw_resources: resource amounts as a kwargs, for example `v100=2` or
        `ram=1 * xm.GiB`. See xm.ResourceType enum for the list of supported
        types and aliases.

    Raises:
      ValueError:
        If several accelerator resources are supplied (i.e. GPU and TPU).
        If the same resource is passed in a `resources` dictionary and as
          a command line argument.
        If topology is supplied for a non acceelerator resource.
    """
    self.location = location
    self._service_tier = service_tier or ServiceTier.PROD

    self.task_requirements = ResourceDict()
    self.accelerator = None
    self.topology = None

    for resource_name, value in itertools.chain(resources.items(),
                                                kw_resources.items()):
      scalar, topology = _parse_resource_quantity(resource_name, value)
      resource = pm.match(
          pm.Case([str], lambda r: ResourceType[r]),
          pm.Case([ResourceType], lambda r: r))(
              resource_name)

      if resource in _AcceleratorType:
        if self.accelerator is not None:
          raise ValueError('Accelerator already set.')
        self.accelerator = resource
        self.topology = topology or Topology(f'{scalar:g}')
      elif topology is not None:
        raise ValueError(
            f'A topology specified for non accelerator resource {resource}.')

      if resource in self.task_requirements:
        raise ValueError(f'{resource} has been specified twice.')
      self.task_requirements[resource] = scalar

    if (self.accelerator in GpuType and self.topology and
        len(self.topology.dimensions) == 2):
      if replicas is not None and replicas != self.topology.dimensions[1]:
        raise ValueError(
            f'For multihost GPUs with topology {self.topology}, replicas should'
            f'be either None or {self.topology.dimensions[1]}. Found: '
            f'{replicas}')
      replicas = self.topology.dimensions[1]

    self.replicas = replicas or 1
    self._validate_replicas()

  @property
  def service_tier(self):
    return self._service_tier

  @service_tier.setter
  def service_tier(self, new_service_tier):
    self._service_tier = new_service_tier

  def _validate_replicas(self) -> None:
    """Raises ValueError if replication is not supported."""
    if self.replicas > 1 and self.accelerator in TpuType:
      raise ValueError(
          f'Replicated jobs are not supported for {self.accelerator}.')

  def __repr__(self) -> str:
    """Returns string representation of the requirements."""
    args = []

    for resource, value in self.task_requirements.items():
      if resource in TpuType:
        args.append(f'{resource.name.lower()}={self.topology!r}')
      else:
        args.append(f'{resource.name.lower()}={value!r}')

    if self.location:
      args.append(f'location={self.location!r}')
    if self.service_tier != ServiceTier.PROD:
      args.append(f'service_tier=xm.{self.service_tier}')
    if self.replicas != 1:
      args.append(f'replicas={self.replicas}')

    return f'xm.JobRequirements({", ".join(args)})'
<|endoftext|>


<commit_before># vim:ts=4:sw=4:softtabstop=4:smarttab:expandtab

# Copyright (C) 2014  Sangoma Technologies Corp.
# All Rights Reserved.
# Author(s)
# Simon Gomizelj <sgomizelj@sangoma.com>

import re
import json
import keyword
import requests
import logging
from .url import url_builder, unpack_rest_response
from .library import CommitIncomplete, parse_messages
from .parser import parse
from .utils import deprecated

__all__ = ['api']

logger = logging.getLogger('safepy2')

def make_typename(name):
    '''Sanitize a name to remove spaces and replace all instances of
    symbols which are not valid for python types with underscore. Should
    the name be None, return RestObject for compatibility. Eventually
    a strict more will raise an exception instead'''

    if name is None:
        return str('RestObject')

    name = name.replace(' ', '')
    name = re.sub('[^a-zA-Z0-9_]', '_', name)

    if not all(c.isalnum() or c == '_' for c in name):
        raise ValueError('Type names and field names can only contain '
                         'alphanumeric characters and underscores: '
                         '{!r}'.format(name))
    if name[0].isdigit():
        raise ValueError('Type names and field names cannot start with '
                         'a number: {!r}'.format(name))
    if keyword.iskeyword(name):
        raise ValueError('Type names and field names cannot be a '
                         'keyword: {!r}'.format(name))

    return str(name)

def make_docstring(description):
    '''Make a valid docstring. Needs to be handled carefully since
    sometimes the json specification returns an array instead of a list.'''

    if description is None:
        return None

    if isinstance(description, list):
        description = '\n'.join(description)

    return str(description)

class APIWrapper(object):
    def __init__(self, node, version, session, builder):
        self.node = node
        self.version = version
        self.session = session
        self.builder = builder

    @property
    def interface(self):
        if self.node:
            return [node.tag for node in self.node.cls]

    @property
    def methods(self):
        if self.node:
            return [node.tag for node in self.node.methods]

    def __contains__(self, key):
        return key in self.interface

    def get_child(self, key):
        new_builder = self.builder.join(key)
        new_api = APIWrapper(self.node, self.version,
                             self.session, new_builder)
        return compile_object(self.node, new_api)(key)

    def get_config(self):
        safe_url = self.builder.url(None, section='config')
        return unpack_rest_response(self.session.get(safe_url))

    def upload(self, filename, payload=None):
        if not payload:
            with open(filename) as fp:
                payload = fp.read()

        files = {'archive': (filename, payload)}
        safe_url = self.builder.url('upload')
        return unpack_rest_response(self.session.post(safe_url, files=files))

    def get(self, method, path=None, params=None):
        safe_url = self.builder.url(method, path=path)
        return unpack_rest_response(self.session.get(safe_url, params=params))

    def post(self, method, path=None, data=None, params=None):
        postdata = json.dumps(data) if data else None
        safe_url = self.builder.url(method, path=path)
        data = self.session.post(safe_url, data=postdata,
                                 params=params,
                                 headers={'Content-Type': 'application/json'})
        return unpack_rest_response(data)

def api_wrapper(session, builder):
    version_url = builder.url('retrieve', path=['nsc', 'version'])
    version_data = unpack_rest_response(session.get(version_url)).data
    version = (int(version_data['major_version']),
               int(version_data['minor_version']),
               int(version_data['patch_version']))

    return APIWrapper(None, version, session, builder)

class API(object):
    def config(self):
        return self.api.get_config().content

    def changelog(self):
        return parse_messages(self.nsc.configuration.status())

    def commit(self):
        if 'smartapply' in self.nsc.configuration.api.methods:
            logger.info('Applying configuration')
            self.nsc.configuration.smartapply()
            state = self.nsc.configuration.status()
        else:
            logger.info('Attempting to apply configuration...')
            state = self.nsc.configuration.status()

            # Attempt to just reload the NSC configuration
            if state['modified'] and state['can_reload']:
                logger.info('Trying reload...')
                self.nsc.configuration.reload()
                state = self.nsc.configuration.status()

            # Changes may still now require us to restart the nsc service
            # to apply
            if state['modified']:
                status = self.nsc.service.status()['status_text']
                if status == 'RUNNING':
                    logger.info('Suspending NSC')
                    self.nsc.service.stop()
                logger.info('Trying apply...')
                self.nsc.configuration.apply()
                if status == 'RUNNING':
                    self.nsc.service.start()
                state = self.nsc.configuration.status()

        if state['modified']:
            raise CommitIncomplete(parse_messages(state))

class APICollection(object):
    def create(self, key, data):
        if 'display-name' in self.api.interface and 'display-name' not in data:
            data['display-name'] = key

        self.api.post('create', path=[key], data=data)
        return self[key]

    def delete(self, key):
        self.api.post('delete', path=[key])

    def update(self, key, data):
        self.api.post('update', path=[key], data=data)

    def retrieve(self, key):
        return self.api.get('retrieve', path=[key]).data

    def keys(self):
        return self.api.get('list').data

    def find(self, filter_expr):
        if not filter_expr:
            keys = self.api.get('list').data
        elif self.api.version >= (2, 1, 13):
            expression = {'filter': filter_expr}
            keys = self.api.post('list', data=expression).data
        else:
            raise NotImplementedError("No REST support for searching on 2.1")

        return iter(self[key] for key in keys)

    search = deprecated('Method renamed to find')(find)

    def __getitem__(self, key):
        return self.api.get_child(key)

    def __contains__(self, key):
        return key in self.keys()

    def __iter__(self):
        return iter(self[key] for key in self.keys())

    def __len__(self):
        return len(self.keys())

    def __bool__(self):
        return bool(self.keys())

    def __repr__(self):
        return '{}({!r})'.format(self.__class__.__name__, self.keys())

class APIObject(object):
    def __init__(self, name=None):
        if name:
            self.ident = name

    def retrieve(self):
        return self.api.get('retrieve').data

    def update(self, data):
        return self.api.post('update', data=data).data

    def __contains__(self, key):
        return key in self.api.interface

    def __getitem__(self, key):
        return self.retrieve()[key]

    def __setitem__(self, key, value):
        self.update({key: value})

    def __repr__(self):
        return '{}({!r})'.format(self.__class__.__name__, self.retrieve())

class APIModule(object):
    def __getitem__(self, key):
        return self.retrieve()[key]

    def __setitem__(self, key, value):
        try:
            self.update({key: value})
        except AttributeError:
            raise TypeError('Module does not support updating')

    def __contains__(self, key):
        return key in self.api.interface

    def __repr__(self):
        try:
            data = self.retrieve()
            return '{}({!r})'.format(self.__class__.__name__, data)
        except AttributeError:
            return '{}()'.format(self.__class__.__name__)

def add_methods(ast, reserved=None):
    '''Compile all the methods specified in the json 'methods' section.
    Prefer specialized implementations of common and important rest
    functions, falling back to a generic implementation for others.'''

    def make_upload_method(nodeid):
        def upload(self, filename, payload=None):
            return self.api.upload(filename, payload=payload).data
        return upload

    def make_download_method(nodeid):
        def download(self, *args):
            return self.api.get(nodeid, path=args).content
        return download

    def make_get_method(nodeid):
        def get(self, *args, **kwargs):
            r = self.api.get(nodeid, path=args, params=kwargs)
            assert r.mimetype == 'application/json'
            return r.data
        return get

    def make_post_method(nodeid):
        def post(self, *args, **kwargs):
            if args and isinstance(args[-1], dict):
                r = self.api.post(nodeid, path=args[:-1], data=args[-1],
                                  params=kwargs)
            else:
                r = self.api.post(nodeid, path=args)
            assert r.mimetype == 'application/json'
            return r.data
        return post

    for node in ast:
        if node.tag == 'list' or (reserved and node.tag in reserved):
            continue

        if node.tag == 'upload':
            method = make_upload_method(node.tag)
        elif node.tag == 'download':
            method = make_download_method(node.tag)
        elif node['request'] == 'GET':
            method = make_get_method(node.tag)
        elif node['request'] == 'POST':
            method = make_post_method(node.tag)

        method.__name__ = make_typename(node.tag)
        method.__doc__ = make_docstring(node.get('description', None))
        yield method.__name__, method

def build_type(node, api, base):
    typename = make_typename(node.get('name', None))
    docstring = make_docstring(node.get('description'))

    namespace = {'__doc__': docstring, 'api': api}
    namespace.update(add_children(node.objs, api))
    namespace.update(add_methods(node.methods, base.__dict__))
    return type(typename, (base,), namespace)

def compile_object(node, api):
    return build_type(node, api, APIObject)

def compile_module(node, api):
    return build_type(node, api, APIModule)

def compile_collection(node, api):
    return build_type(node, api, APICollection)

def add_children(ast, api):
    for node in ast:
        typename = make_typename(node.tag)
        new_builder = api.builder.join(node.tag)
        new_api = APIWrapper(node, api.version, api.session, new_builder)

        node_type = node.node_type
        if node_type == 'collection':
            yield typename, compile_collection(node, new_api)()
        elif node_type == 'object':
            yield typename, compile_object(node, new_api)()
        elif node_type == 'module':
            yield typename, compile_module(node, new_api)()

def api(host, port=80, scheme='http', token=None, specfile=None, timeout=None):
    '''Connects to a remote device, download the json specification
    describing the supported rest calls and dynamically compile a new
    object to wrap the rest.

    :param name: The hostname of the device to connect to.
    :type name: str
    :param port: The port the SAFe framework is listening on.
    :type port: int
    :param scheme: Specify the scheme of the request url.
    :type scheme: str
    :returns: the dynamically generated code.
    '''
    builder = url_builder(host, port, scheme)
    session = requests.session()
    if timeout:
        session.timeout = timeout
    if token:
        session.headers['X-API-KEY'] = token

    api = api_wrapper(session, builder)
    if not specfile:
        logger.info('Retrieving specification from NSC')
        r = session.get(builder.url(None, section='doc'))
        spec = unpack_rest_response(r).content
    else:
        with open(specfile) as fp:
            spec = json.load(fp)

    namespace = dict(add_children(parse(spec), api))
    product_cls = type('API', (API,), namespace)
    return product_cls()
<commit_msg>Simplify APIModule since it'll never have an update method<commit_after># vim:ts=4:sw=4:softtabstop=4:smarttab:expandtab

# Copyright (C) 2014  Sangoma Technologies Corp.
# All Rights Reserved.
# Author(s)
# Simon Gomizelj <sgomizelj@sangoma.com>

import re
import json
import keyword
import requests
import logging
from .url import url_builder, unpack_rest_response
from .library import CommitIncomplete, parse_messages
from .parser import parse
from .utils import deprecated

__all__ = ['api']

logger = logging.getLogger('safepy2')

def make_typename(name):
    '''Sanitize a name to remove spaces and replace all instances of
    symbols which are not valid for python types with underscore. Should
    the name be None, return RestObject for compatibility. Eventually
    a strict more will raise an exception instead'''

    if name is None:
        return str('RestObject')

    name = name.replace(' ', '')
    name = re.sub('[^a-zA-Z0-9_]', '_', name)

    if not all(c.isalnum() or c == '_' for c in name):
        raise ValueError('Type names and field names can only contain '
                         'alphanumeric characters and underscores: '
                         '{!r}'.format(name))
    if name[0].isdigit():
        raise ValueError('Type names and field names cannot start with '
                         'a number: {!r}'.format(name))
    if keyword.iskeyword(name):
        raise ValueError('Type names and field names cannot be a '
                         'keyword: {!r}'.format(name))

    return str(name)

def make_docstring(description):
    '''Make a valid docstring. Needs to be handled carefully since
    sometimes the json specification returns an array instead of a list.'''

    if description is None:
        return None

    if isinstance(description, list):
        description = '\n'.join(description)

    return str(description)

class APIWrapper(object):
    def __init__(self, node, version, session, builder):
        self.node = node
        self.version = version
        self.session = session
        self.builder = builder

    @property
    def interface(self):
        if self.node:
            return [node.tag for node in self.node.cls]

    @property
    def methods(self):
        if self.node:
            return [node.tag for node in self.node.methods]

    def __contains__(self, key):
        return key in self.interface

    def get_child(self, key):
        new_builder = self.builder.join(key)
        new_api = APIWrapper(self.node, self.version,
                             self.session, new_builder)
        return compile_object(self.node, new_api)(key)

    def get_config(self):
        safe_url = self.builder.url(None, section='config')
        return unpack_rest_response(self.session.get(safe_url))

    def upload(self, filename, payload=None):
        if not payload:
            with open(filename) as fp:
                payload = fp.read()

        files = {'archive': (filename, payload)}
        safe_url = self.builder.url('upload')
        return unpack_rest_response(self.session.post(safe_url, files=files))

    def get(self, method, path=None, params=None):
        safe_url = self.builder.url(method, path=path)
        return unpack_rest_response(self.session.get(safe_url, params=params))

    def post(self, method, path=None, data=None, params=None):
        postdata = json.dumps(data) if data else None
        safe_url = self.builder.url(method, path=path)
        data = self.session.post(safe_url, data=postdata,
                                 params=params,
                                 headers={'Content-Type': 'application/json'})
        return unpack_rest_response(data)

def api_wrapper(session, builder):
    version_url = builder.url('retrieve', path=['nsc', 'version'])
    version_data = unpack_rest_response(session.get(version_url)).data
    version = (int(version_data['major_version']),
               int(version_data['minor_version']),
               int(version_data['patch_version']))

    return APIWrapper(None, version, session, builder)

class API(object):
    def config(self):
        return self.api.get_config().content

    def changelog(self):
        return parse_messages(self.nsc.configuration.status())

    def commit(self):
        if 'smartapply' in self.nsc.configuration.api.methods:
            logger.info('Applying configuration')
            self.nsc.configuration.smartapply()
            state = self.nsc.configuration.status()
        else:
            logger.info('Attempting to apply configuration...')
            state = self.nsc.configuration.status()

            # Attempt to just reload the NSC configuration
            if state['modified'] and state['can_reload']:
                logger.info('Trying reload...')
                self.nsc.configuration.reload()
                state = self.nsc.configuration.status()

            # Changes may still now require us to restart the nsc service
            # to apply
            if state['modified']:
                status = self.nsc.service.status()['status_text']
                if status == 'RUNNING':
                    logger.info('Suspending NSC')
                    self.nsc.service.stop()
                logger.info('Trying apply...')
                self.nsc.configuration.apply()
                if status == 'RUNNING':
                    self.nsc.service.start()
                state = self.nsc.configuration.status()

        if state['modified']:
            raise CommitIncomplete(parse_messages(state))

class APICollection(object):
    def create(self, key, data):
        if 'display-name' in self.api.interface and 'display-name' not in data:
            data['display-name'] = key

        self.api.post('create', path=[key], data=data)
        return self[key]

    def delete(self, key):
        self.api.post('delete', path=[key])

    def update(self, key, data):
        self.api.post('update', path=[key], data=data)

    def retrieve(self, key):
        return self.api.get('retrieve', path=[key]).data

    def keys(self):
        return self.api.get('list').data

    def find(self, filter_expr):
        if not filter_expr:
            keys = self.api.get('list').data
        elif self.api.version >= (2, 1, 13):
            expression = {'filter': filter_expr}
            keys = self.api.post('list', data=expression).data
        else:
            raise NotImplementedError("No REST support for searching on 2.1")

        return iter(self[key] for key in keys)

    search = deprecated('Method renamed to find')(find)

    def __getitem__(self, key):
        return self.api.get_child(key)

    def __contains__(self, key):
        return key in self.keys()

    def __iter__(self):
        return iter(self[key] for key in self.keys())

    def __len__(self):
        return len(self.keys())

    def __bool__(self):
        return bool(self.keys())

    def __repr__(self):
        return '{}({!r})'.format(self.__class__.__name__, self.keys())

class APIObject(object):
    def __init__(self, name=None):
        if name:
            self.ident = name

    def retrieve(self):
        return self.api.get('retrieve').data

    def update(self, data):
        return self.api.post('update', data=data).data

    def __contains__(self, key):
        return key in self.api.interface

    def __getitem__(self, key):
        return self.retrieve()[key]

    def __setitem__(self, key, value):
        self.update({key: value})

    def __repr__(self):
        return '{}({!r})'.format(self.__class__.__name__, self.retrieve())

class APIModule(object):
    def __getitem__(self, key):
        return self.retrieve()[key]

    def __contains__(self, key):
        return key in self.api.interface

    def __repr__(self):
        return '{}()'.format(self.__class__.__name__)

def add_methods(ast, reserved=None):
    '''Compile all the methods specified in the json 'methods' section.
    Prefer specialized implementations of common and important rest
    functions, falling back to a generic implementation for others.'''

    def make_upload_method(nodeid):
        def upload(self, filename, payload=None):
            return self.api.upload(filename, payload=payload).data
        return upload

    def make_download_method(nodeid):
        def download(self, *args):
            return self.api.get(nodeid, path=args).content
        return download

    def make_get_method(nodeid):
        def get(self, *args, **kwargs):
            r = self.api.get(nodeid, path=args, params=kwargs)
            assert r.mimetype == 'application/json'
            return r.data
        return get

    def make_post_method(nodeid):
        def post(self, *args, **kwargs):
            if args and isinstance(args[-1], dict):
                r = self.api.post(nodeid, path=args[:-1], data=args[-1],
                                  params=kwargs)
            else:
                r = self.api.post(nodeid, path=args)
            assert r.mimetype == 'application/json'
            return r.data
        return post

    for node in ast:
        if node.tag == 'list' or (reserved and node.tag in reserved):
            continue

        if node.tag == 'upload':
            method = make_upload_method(node.tag)
        elif node.tag == 'download':
            method = make_download_method(node.tag)
        elif node['request'] == 'GET':
            method = make_get_method(node.tag)
        elif node['request'] == 'POST':
            method = make_post_method(node.tag)

        method.__name__ = make_typename(node.tag)
        method.__doc__ = make_docstring(node.get('description', None))
        yield method.__name__, method

def build_type(node, api, base):
    typename = make_typename(node.get('name', None))
    docstring = make_docstring(node.get('description'))

    namespace = {'__doc__': docstring, 'api': api}
    namespace.update(add_children(node.objs, api))
    namespace.update(add_methods(node.methods, base.__dict__))
    return type(typename, (base,), namespace)

def compile_object(node, api):
    return build_type(node, api, APIObject)

def compile_module(node, api):
    return build_type(node, api, APIModule)

def compile_collection(node, api):
    return build_type(node, api, APICollection)

def add_children(ast, api):
    for node in ast:
        typename = make_typename(node.tag)
        new_builder = api.builder.join(node.tag)
        new_api = APIWrapper(node, api.version, api.session, new_builder)

        node_type = node.node_type
        if node_type == 'collection':
            yield typename, compile_collection(node, new_api)()
        elif node_type == 'object':
            yield typename, compile_object(node, new_api)()
        elif node_type == 'module':
            yield typename, compile_module(node, new_api)()

def api(host, port=80, scheme='http', token=None, specfile=None, timeout=None):
    '''Connects to a remote device, download the json specification
    describing the supported rest calls and dynamically compile a new
    object to wrap the rest.

    :param name: The hostname of the device to connect to.
    :type name: str
    :param port: The port the SAFe framework is listening on.
    :type port: int
    :param scheme: Specify the scheme of the request url.
    :type scheme: str
    :returns: the dynamically generated code.
    '''
    builder = url_builder(host, port, scheme)
    session = requests.session()
    if timeout:
        session.timeout = timeout
    if token:
        session.headers['X-API-KEY'] = token

    api = api_wrapper(session, builder)
    if not specfile:
        logger.info('Retrieving specification from NSC')
        r = session.get(builder.url(None, section='doc'))
        spec = unpack_rest_response(r).content
    else:
        with open(specfile) as fp:
            spec = json.load(fp)

    namespace = dict(add_children(parse(spec), api))
    product_cls = type('API', (API,), namespace)
    return product_cls()
<|endoftext|>


<commit_before>"""Device, context and memory management on CuPy.

Chainer uses CuPy (with very thin wrapper) to exploit the speed of GPU
computation. Following modules and classes are imported to :mod:`cuda`
module for convenience (refer to this table when reading chainer's source
codes).

============================ =================================
 imported name                original name
============================ =================================
 ``chainer.cuda.cupy``        :mod:`cupy`
 ``chainer.cuda.ndarray``     :mod:`cupy.ndarray`
 ``chainer.cuda.cupy.cuda``   :mod:`cupy.cuda`
 ``chainer.cuda.Device``      :mod:`cupy.cuda.Device`
 ``chainer.cuda.Event``       :mod:`cupy.cuda.Event`
 ``chainer.cuda.Stream``      :mod:`cupy.cuda.Stream`
============================ =================================

Chainer replaces the default allocator of CuPy by its memory pool
implementation. It enables us to reuse the device memory over multiple
forward/backward computations, and temporary arrays for consecutive elementwise
operations.
"""

import functools
import os
import warnings

import numpy

available = False
cudnn_enabled = False

try:
    import cupy
    import cupy.cuda
    import cupy.cuda.cublas

    cuda = cupy.cuda
    cublas = cuda.cublas

    ndarray = cupy.ndarray
    Device = cuda.Device
    Event = cuda.Event
    Stream = cuda.Stream

    available = True
except Exception as e:
    _resolution_error = e

    class ndarray(object):
        pass  # for type testing

if available:
    try:
        import cupy.cudnn
        cudnn = cupy.cudnn
        cudnn_enabled = int(os.environ.get('CHAINER_CUDNN', '1')) != 0
    except Exception as e:
        _resolution_error = e

def init(arg=None):
    warnings.warn(
        'chainer.cuda.init is deprecated. You need to call nothing to '
        'initialize your environment. Call chainer.cuda.check_cuda_available '
        'to check availability of CUDA.',
        DeprecationWarning)
    check_cuda_available()

def check_cuda_available():
    """Checks if CUDA is available.

    When CUDA is correctly set up, nothing happens.
    Otherwise it raises ``RuntimeError``.
    """
    if not available:
        global _resolution_error
        msg = 'CUDA environment is not correctly set up.\n'
        msg += str(_resolution_error)
        raise RuntimeError(msg)
    if not cudnn_enabled:
        warnings.warn('cuDNN is not enabled')

class DummyDeviceType(object):

    """Dummy device class that does nothing with cupy.cuda.Device interface.

    This class is used to represent CPU device.

    """
    def __int__(self):
        return -1

    def __enter__(self):
        return self

    def __exit__(self, *args):
        pass

    def use(self):
        pass

    def synchronize(self):
        pass

    def __eq__(self, other):
        return isinstance(other, DummyDevice)

    def __ne__(self, other):
        return not (self == other)

DummyDevice = DummyDeviceType()

# ------------------------------------------------------------------------------
# Global states
# ------------------------------------------------------------------------------
if available:
    cuda.set_allocator(cuda.MemoryPool().malloc)

# ------------------------------------------------------------------------------
# Global states
# ------------------------------------------------------------------------------
def get_device(*args):
    """Gets the device from an ID integer or an array object.

    This is a convenient utility to select a correct device if the type of
    ``arg`` is unknown (i.e., one can use this function on arrays that may be
    on CPU or GPU). The returned device object supports the context management
    protocol of Python for the *with* statement.

    Args:
        args: Values to specify a GPU device. :class:`numpy.ndarray` objects
            are skipped. If all arguments are numpy.ndarray objects, it returns
            a dummy device object. Otherwise, the first non-numpy object is
            used to select a device. If it is a :class:`cupy.ndarray` object,
            its device is returned. Otherwise, the argument is passed to the
            initializer of :class:`~cupy.cuda.Device` and it is returned.

    Returns:
        Device object specified by given ``args``.

    .. seealso::
       See :class:`cupy.cuda.Device` for the device selection not by arrays.

    """
    for arg in args:
        if arg is None:
            continue
        if not isinstance(arg, numpy.ndarray):
            check_cuda_available()
            if isinstance(arg, cupy.ndarray):
                if arg.device is None:
                    continue
                return arg.device
            else:
                return Device(arg)

    return DummyDevice

# ------------------------------------------------------------------------------
# cupy.ndarray allocation and copy
# ------------------------------------------------------------------------------

def to_gpu(array, device=None, stream=None):
    """Copies the given CPU array to specified device.

    Args:
        array: Array to be sent to GPU.
        device: Device specifier.
        stream (cupy.cuda.Stream): CUDA stream.

    Returns:
        cupy.ndarray: Array on GPU.

        If ``array`` is already on GPU, then this function just returns
        ``array`` without performing any copy. Note that this function does not
        copy cupy.ndarray into specified device.

    """
    check_cuda_available()
    assert stream is None  # TODO(beam2d): FIX IT
    with get_device(device):
        return cupy.asarray(array)

def to_cpu(array, stream=None):
    """Copies the given GPU array to host CPU.

    Args:
        array: Array to be sent to CPU.
        stream (cupy.cuda.Stream): CUDA stream.

    Returns:
        numpy.ndarray: Array on CPU.

        If given ``array`` is already on CPU, then this function just returns
        ``array`` without performing any copy.

    """
    if isinstance(array, ndarray):
        return array.get(stream)
    elif isinstance(array, numpy.ndarray):
        return array
    else:
        raise TypeError(
            'The array sent to cpu must be numpy.ndarray or cupy.ndarray')

def empty(shape, dtype=numpy.float32):
    """Creates an uninitialized cupy.ndarray object.

    Args:
        shape (tuple of ints): The shape of array.
        dtype (numpy.dtype): Element type.

    Returns:
        cupy.ndarray: Uninitialized GPU array allocated by the memory pool.

    """
    warnings.warn(
        'chainer.cuda.empty is deprecated. Use cupy.empty instead.',
        DeprecationWarning)
    check_cuda_available()
    return cupy.empty(shape, dtype)

def full(shape, fill_value, dtype=numpy.float32, stream=None):
    """Creates a constant-filled cupy.ndarray object.

    Args:
        shape (tuple of ints): The shape of array.
        fill_value: Constant to fill the array by.
        dtype (numpy.dtype): Element type.
        stream (cupy.cuda.Stream): CUDA stream.

    Returns:
        cupy.ndarray: Constant-filled GPU array allocated by the memory pool.

    """
    warnings.warn(
        'chainer.cuda.full is deprecated. Use cupy.full instead.',
        DeprecationWarning)
    check_cuda_available()
    assert stream is None
    return cupy.full(shape, fill_value, dtype=dtype)

def zeros(shape, dtype=numpy.float32, stream=None):
    """Creates a zero-filled cupy.ndarray object.

    This function is equivalent to ``full(shape, 0, dtype, stream)``.

    """
    warnings.warn(
        'chainer.cuda.zeros is deprecated. Use cupy.zeros instead.',
        DeprecationWarning)
    check_cuda_available()
    assert stream is None
    return cupy.zeros(shape, dtype=dtype)

def ones(shape, dtype=numpy.float32, stream=None):
    """Creates a zero-filled cupy.ndarray object.

    This function is equivalent to ``full(shape, 1, dtype, stream)``.

    """
    warnings.warn(
        'chainer.cuda.ones is deprecated. Use cupy.ones instead.',
        DeprecationWarning)
    check_cuda_available()
    assert stream is None
    return cupy.ones(shape, dtype=dtype)

def empty_like(array):
    """Creates an uninitialized GPU array like the given one.

    Args:
        array (cupy.ndarray or numpy.ndarray): Base array.

    Returns:
        cupy.ndarray: GPU array of the same shape and dtype as `array`.

    """
    warnings.warn(
        'chainer.cuda.empty_like is deprecated. Use cupy.empty_like instead.',
        DeprecationWarning)
    check_cuda_available()
    if isinstance(array, cupy.ndarray):
        return cupy.empty_like(array)
    return cupy.empty(array.shape, dtype=array.dtype)

def full_like(array, fill_value, stream=None):
    """Creates a constant-filled cupy.ndarray object like the given array.

    Args:
        array (cupy.ndarray or numpy.ndarray): Base array.
        fill_value: Constant value to fill the array by.
        stream (cupy.cuda.Stream): CUDA stream.

    Returns:
        cupy.ndarray: Constant-filled array.

    """
    warnings.warn(
        'chainer.cuda.full_like is deprecated. Use cupy.full_like instead.',
        DeprecationWarning)
    check_cuda_available()
    assert stream is None
    if isinstance(array, cupy.ndarray):
        return cupy.full_like(array, fill_value)
    return cupy.full(array.shape, fill_value, dtype=array.dtype)

def zeros_like(array, stream=None):
    """Creates a zero-filled cupy.ndarray object like the given array.

    Args:
        array (cupy.ndarray or numpy.ndarray): Base array.
        stream (cupy.cuda.Stream): CUDA stream.

    Returns:
        cupy.ndarray: Zero-filled array.

    """
    warnings.warn(
        'chainer.cuda.zeros_like is deprecated. Use cupy.zeros_like instead.',
        DeprecationWarning)
    check_cuda_available()
    assert stream is None
    if isinstance(array, cupy.ndarray):
        return cupy.zeros_like(array)
    return cupy.zeros(array.shape, dtype=array.dtype)

def ones_like(array, stream=None):
    """Creates a one-filled cupy.ndarray object like the given array.

    Args:
        array (cupy.ndarray or numpy.ndarray): Base array.
        stream (cupy.cuda.Stream): CUDA stream.

    Returns:
        cupy.ndarray: One-filled array.

    """
    warnings.warn(
        'chainer.cuda.ones_like is deprecated. Use cupy.ones_like instead.',
        DeprecationWarning)
    check_cuda_available()
    assert stream is None
    if isinstance(array, cupy.ndarray):
        return cupy.ones_like(array)
    return cupy.ones(array.shape, dtype=array.dtype)

def copy(array, out=None, out_device=None, stream=None):
    """Copies a cupy.ndarray object using the default stream.

    This function can copy the device array to the destination array on another
    device.

    Args:
        array (cupy.ndarray): Array to be copied.
        out (cupy.ndarray): Destination array.
            If it is not ``None``, then ``out_device`` argument is ignored.
        out_device: Destination device specifier. Actual device object is
            obtained by passing this value to :func:`get_device`.
        stream (cupy.cuda.Stream): CUDA stream.

    Returns:
        cupy.ndarray: Copied array.

        If ``out`` is not specified, then the array is allocated on the device
        specified by ``out_device`` argument.

    """
    check_cuda_available()
    assert stream is None  # TODO(beam2d): FIX IT

    if out is None:
        if out_device is None:
            out_device = array
        with get_device(out_device):
            out = cupy.empty_like(array)

    with get_device(array):
        cupy.copyto(out, array)

    return out

# ------------------------------------------------------------------------------
# Function result memoization
# ------------------------------------------------------------------------------
def memoize(for_each_device=False):
    """Makes a function memoizing the result for each argument and device.

    This is a similar version of :func:`cupy.memoize`. The difference is that
    this function can be used in the global scope even if CUDA is not
    available. In such case, this function does nothing.

    .. note::
       This decorator acts as a dummy if CUDA is not available. It cannot be
       used for general purpose memoization even if ``for_each_device`` is set
       to False.

    """
    if available:
        return cupy.memoize(for_each_device)

    def dummy_decorator(f):
        @functools.wraps(f)
        def ret(*args, **kwargs):
            return f(*args, **kwargs)
        return ret
    return dummy_decorator

# ------------------------------------------------------------------------------
# Kernel definition utility
# ------------------------------------------------------------------------------
@memoize(for_each_device=True)
def elementwise(in_params, out_params, operation, name, **kwargs):
    """Creates an elementwise kernel function.

    This function uses :func:`~chainer.cuda.memoize` to cache the
    kernel object, i.e. the resulting kernel object is cached for each argument
    combination and CUDA device.

    The arguments are the same as those for
    :class:`cupy.ElementwiseKernel`, except that the ``name`` argument is
    mandatory.

    """
    check_cuda_available()
    return cupy.ElementwiseKernel(
        in_params, out_params, operation, name, **kwargs)

@memoize(for_each_device=True)
def reduce(in_params, out_params, map_expr, reduce_expr, post_map_expr,
           identity, name,  **kwargs):
    """Creates a global reduction kernel function.

    This function uses :func:`~chainer.cuda.memoize` to cache the resulting
    kernel object, i.e. the resulting kernel object is cached for each argument
    combination and CUDA device.

    The arguments are the same as those for
    :class:`cupy.ReductionKernel`, except that the ``name`` argument is
    mandatory.

    """
    check_cuda_available()
    return cupy.ReductionKernel(
        in_params, out_params, map_expr, reduce_expr, post_map_expr,
        identity, name, **kwargs)

# ------------------------------------------------------------------------------
# numpy/cupy compatible coding
# ------------------------------------------------------------------------------
def get_array_module(*args):
    """Gets an appropriate one from :mod:`numpy` or :mod:`cupy`.

    This is almost equivalent to :func:`cupy.get_array_module`. The only
    difference is that this function can be used even if CUDA is not available.

    Args:
        args: Values to determine whether NumPy or CuPy should be used.

    Returns:
        module: :mod:`cupy` or :mod:`numpy` is returned based on the types of
        the arguments.

    """
    if available:
        return cupy.get_array_module(*args)
    else:
        return numpy
<commit_msg>Remove unnecessary line<commit_after>"""Device, context and memory management on CuPy.

Chainer uses CuPy (with very thin wrapper) to exploit the speed of GPU
computation. Following modules and classes are imported to :mod:`cuda`
module for convenience (refer to this table when reading chainer's source
codes).

============================ =================================
 imported name                original name
============================ =================================
 ``chainer.cuda.cupy``        :mod:`cupy`
 ``chainer.cuda.ndarray``     :mod:`cupy.ndarray`
 ``chainer.cuda.cupy.cuda``   :mod:`cupy.cuda`
 ``chainer.cuda.Device``      :mod:`cupy.cuda.Device`
 ``chainer.cuda.Event``       :mod:`cupy.cuda.Event`
 ``chainer.cuda.Stream``      :mod:`cupy.cuda.Stream`
============================ =================================

Chainer replaces the default allocator of CuPy by its memory pool
implementation. It enables us to reuse the device memory over multiple
forward/backward computations, and temporary arrays for consecutive elementwise
operations.
"""

import functools
import os
import warnings

import numpy

available = False
cudnn_enabled = False

try:
    import cupy
    import cupy.cuda
    import cupy.cuda.cublas

    cuda = cupy.cuda
    cublas = cuda.cublas

    ndarray = cupy.ndarray
    Device = cuda.Device
    Event = cuda.Event
    Stream = cuda.Stream

    available = True
except Exception as e:
    _resolution_error = e

    class ndarray(object):
        pass  # for type testing

if available:
    try:
        import cupy.cudnn
        cudnn = cupy.cudnn
        cudnn_enabled = int(os.environ.get('CHAINER_CUDNN', '1')) != 0
    except Exception as e:
        _resolution_error = e

def init(arg=None):
    warnings.warn(
        'chainer.cuda.init is deprecated. You need to call nothing to '
        'initialize your environment. Call chainer.cuda.check_cuda_available '
        'to check availability of CUDA.',
        DeprecationWarning)
    check_cuda_available()

def check_cuda_available():
    """Checks if CUDA is available.

    When CUDA is correctly set up, nothing happens.
    Otherwise it raises ``RuntimeError``.
    """
    if not available:
        msg = 'CUDA environment is not correctly set up.\n'
        msg += str(_resolution_error)
        raise RuntimeError(msg)
    if not cudnn_enabled:
        warnings.warn('cuDNN is not enabled')

class DummyDeviceType(object):

    """Dummy device class that does nothing with cupy.cuda.Device interface.

    This class is used to represent CPU device.

    """
    def __int__(self):
        return -1

    def __enter__(self):
        return self

    def __exit__(self, *args):
        pass

    def use(self):
        pass

    def synchronize(self):
        pass

    def __eq__(self, other):
        return isinstance(other, DummyDevice)

    def __ne__(self, other):
        return not (self == other)

DummyDevice = DummyDeviceType()

# ------------------------------------------------------------------------------
# Global states
# ------------------------------------------------------------------------------
if available:
    cuda.set_allocator(cuda.MemoryPool().malloc)

# ------------------------------------------------------------------------------
# Global states
# ------------------------------------------------------------------------------
def get_device(*args):
    """Gets the device from an ID integer or an array object.

    This is a convenient utility to select a correct device if the type of
    ``arg`` is unknown (i.e., one can use this function on arrays that may be
    on CPU or GPU). The returned device object supports the context management
    protocol of Python for the *with* statement.

    Args:
        args: Values to specify a GPU device. :class:`numpy.ndarray` objects
            are skipped. If all arguments are numpy.ndarray objects, it returns
            a dummy device object. Otherwise, the first non-numpy object is
            used to select a device. If it is a :class:`cupy.ndarray` object,
            its device is returned. Otherwise, the argument is passed to the
            initializer of :class:`~cupy.cuda.Device` and it is returned.

    Returns:
        Device object specified by given ``args``.

    .. seealso::
       See :class:`cupy.cuda.Device` for the device selection not by arrays.

    """
    for arg in args:
        if arg is None:
            continue
        if not isinstance(arg, numpy.ndarray):
            check_cuda_available()
            if isinstance(arg, cupy.ndarray):
                if arg.device is None:
                    continue
                return arg.device
            else:
                return Device(arg)

    return DummyDevice

# ------------------------------------------------------------------------------
# cupy.ndarray allocation and copy
# ------------------------------------------------------------------------------

def to_gpu(array, device=None, stream=None):
    """Copies the given CPU array to specified device.

    Args:
        array: Array to be sent to GPU.
        device: Device specifier.
        stream (cupy.cuda.Stream): CUDA stream.

    Returns:
        cupy.ndarray: Array on GPU.

        If ``array`` is already on GPU, then this function just returns
        ``array`` without performing any copy. Note that this function does not
        copy cupy.ndarray into specified device.

    """
    check_cuda_available()
    assert stream is None  # TODO(beam2d): FIX IT
    with get_device(device):
        return cupy.asarray(array)

def to_cpu(array, stream=None):
    """Copies the given GPU array to host CPU.

    Args:
        array: Array to be sent to CPU.
        stream (cupy.cuda.Stream): CUDA stream.

    Returns:
        numpy.ndarray: Array on CPU.

        If given ``array`` is already on CPU, then this function just returns
        ``array`` without performing any copy.

    """
    if isinstance(array, ndarray):
        return array.get(stream)
    elif isinstance(array, numpy.ndarray):
        return array
    else:
        raise TypeError(
            'The array sent to cpu must be numpy.ndarray or cupy.ndarray')

def empty(shape, dtype=numpy.float32):
    """Creates an uninitialized cupy.ndarray object.

    Args:
        shape (tuple of ints): The shape of array.
        dtype (numpy.dtype): Element type.

    Returns:
        cupy.ndarray: Uninitialized GPU array allocated by the memory pool.

    """
    warnings.warn(
        'chainer.cuda.empty is deprecated. Use cupy.empty instead.',
        DeprecationWarning)
    check_cuda_available()
    return cupy.empty(shape, dtype)

def full(shape, fill_value, dtype=numpy.float32, stream=None):
    """Creates a constant-filled cupy.ndarray object.

    Args:
        shape (tuple of ints): The shape of array.
        fill_value: Constant to fill the array by.
        dtype (numpy.dtype): Element type.
        stream (cupy.cuda.Stream): CUDA stream.

    Returns:
        cupy.ndarray: Constant-filled GPU array allocated by the memory pool.

    """
    warnings.warn(
        'chainer.cuda.full is deprecated. Use cupy.full instead.',
        DeprecationWarning)
    check_cuda_available()
    assert stream is None
    return cupy.full(shape, fill_value, dtype=dtype)

def zeros(shape, dtype=numpy.float32, stream=None):
    """Creates a zero-filled cupy.ndarray object.

    This function is equivalent to ``full(shape, 0, dtype, stream)``.

    """
    warnings.warn(
        'chainer.cuda.zeros is deprecated. Use cupy.zeros instead.',
        DeprecationWarning)
    check_cuda_available()
    assert stream is None
    return cupy.zeros(shape, dtype=dtype)

def ones(shape, dtype=numpy.float32, stream=None):
    """Creates a zero-filled cupy.ndarray object.

    This function is equivalent to ``full(shape, 1, dtype, stream)``.

    """
    warnings.warn(
        'chainer.cuda.ones is deprecated. Use cupy.ones instead.',
        DeprecationWarning)
    check_cuda_available()
    assert stream is None
    return cupy.ones(shape, dtype=dtype)

def empty_like(array):
    """Creates an uninitialized GPU array like the given one.

    Args:
        array (cupy.ndarray or numpy.ndarray): Base array.

    Returns:
        cupy.ndarray: GPU array of the same shape and dtype as `array`.

    """
    warnings.warn(
        'chainer.cuda.empty_like is deprecated. Use cupy.empty_like instead.',
        DeprecationWarning)
    check_cuda_available()
    if isinstance(array, cupy.ndarray):
        return cupy.empty_like(array)
    return cupy.empty(array.shape, dtype=array.dtype)

def full_like(array, fill_value, stream=None):
    """Creates a constant-filled cupy.ndarray object like the given array.

    Args:
        array (cupy.ndarray or numpy.ndarray): Base array.
        fill_value: Constant value to fill the array by.
        stream (cupy.cuda.Stream): CUDA stream.

    Returns:
        cupy.ndarray: Constant-filled array.

    """
    warnings.warn(
        'chainer.cuda.full_like is deprecated. Use cupy.full_like instead.',
        DeprecationWarning)
    check_cuda_available()
    assert stream is None
    if isinstance(array, cupy.ndarray):
        return cupy.full_like(array, fill_value)
    return cupy.full(array.shape, fill_value, dtype=array.dtype)

def zeros_like(array, stream=None):
    """Creates a zero-filled cupy.ndarray object like the given array.

    Args:
        array (cupy.ndarray or numpy.ndarray): Base array.
        stream (cupy.cuda.Stream): CUDA stream.

    Returns:
        cupy.ndarray: Zero-filled array.

    """
    warnings.warn(
        'chainer.cuda.zeros_like is deprecated. Use cupy.zeros_like instead.',
        DeprecationWarning)
    check_cuda_available()
    assert stream is None
    if isinstance(array, cupy.ndarray):
        return cupy.zeros_like(array)
    return cupy.zeros(array.shape, dtype=array.dtype)

def ones_like(array, stream=None):
    """Creates a one-filled cupy.ndarray object like the given array.

    Args:
        array (cupy.ndarray or numpy.ndarray): Base array.
        stream (cupy.cuda.Stream): CUDA stream.

    Returns:
        cupy.ndarray: One-filled array.

    """
    warnings.warn(
        'chainer.cuda.ones_like is deprecated. Use cupy.ones_like instead.',
        DeprecationWarning)
    check_cuda_available()
    assert stream is None
    if isinstance(array, cupy.ndarray):
        return cupy.ones_like(array)
    return cupy.ones(array.shape, dtype=array.dtype)

def copy(array, out=None, out_device=None, stream=None):
    """Copies a cupy.ndarray object using the default stream.

    This function can copy the device array to the destination array on another
    device.

    Args:
        array (cupy.ndarray): Array to be copied.
        out (cupy.ndarray): Destination array.
            If it is not ``None``, then ``out_device`` argument is ignored.
        out_device: Destination device specifier. Actual device object is
            obtained by passing this value to :func:`get_device`.
        stream (cupy.cuda.Stream): CUDA stream.

    Returns:
        cupy.ndarray: Copied array.

        If ``out`` is not specified, then the array is allocated on the device
        specified by ``out_device`` argument.

    """
    check_cuda_available()
    assert stream is None  # TODO(beam2d): FIX IT

    if out is None:
        if out_device is None:
            out_device = array
        with get_device(out_device):
            out = cupy.empty_like(array)

    with get_device(array):
        cupy.copyto(out, array)

    return out

# ------------------------------------------------------------------------------
# Function result memoization
# ------------------------------------------------------------------------------
def memoize(for_each_device=False):
    """Makes a function memoizing the result for each argument and device.

    This is a similar version of :func:`cupy.memoize`. The difference is that
    this function can be used in the global scope even if CUDA is not
    available. In such case, this function does nothing.

    .. note::
       This decorator acts as a dummy if CUDA is not available. It cannot be
       used for general purpose memoization even if ``for_each_device`` is set
       to False.

    """
    if available:
        return cupy.memoize(for_each_device)

    def dummy_decorator(f):
        @functools.wraps(f)
        def ret(*args, **kwargs):
            return f(*args, **kwargs)
        return ret
    return dummy_decorator

# ------------------------------------------------------------------------------
# Kernel definition utility
# ------------------------------------------------------------------------------
@memoize(for_each_device=True)
def elementwise(in_params, out_params, operation, name, **kwargs):
    """Creates an elementwise kernel function.

    This function uses :func:`~chainer.cuda.memoize` to cache the
    kernel object, i.e. the resulting kernel object is cached for each argument
    combination and CUDA device.

    The arguments are the same as those for
    :class:`cupy.ElementwiseKernel`, except that the ``name`` argument is
    mandatory.

    """
    check_cuda_available()
    return cupy.ElementwiseKernel(
        in_params, out_params, operation, name, **kwargs)

@memoize(for_each_device=True)
def reduce(in_params, out_params, map_expr, reduce_expr, post_map_expr,
           identity, name,  **kwargs):
    """Creates a global reduction kernel function.

    This function uses :func:`~chainer.cuda.memoize` to cache the resulting
    kernel object, i.e. the resulting kernel object is cached for each argument
    combination and CUDA device.

    The arguments are the same as those for
    :class:`cupy.ReductionKernel`, except that the ``name`` argument is
    mandatory.

    """
    check_cuda_available()
    return cupy.ReductionKernel(
        in_params, out_params, map_expr, reduce_expr, post_map_expr,
        identity, name, **kwargs)

# ------------------------------------------------------------------------------
# numpy/cupy compatible coding
# ------------------------------------------------------------------------------
def get_array_module(*args):
    """Gets an appropriate one from :mod:`numpy` or :mod:`cupy`.

    This is almost equivalent to :func:`cupy.get_array_module`. The only
    difference is that this function can be used even if CUDA is not available.

    Args:
        args: Values to determine whether NumPy or CuPy should be used.

    Returns:
        module: :mod:`cupy` or :mod:`numpy` is returned based on the types of
        the arguments.

    """
    if available:
        return cupy.get_array_module(*args)
    else:
        return numpy
<|endoftext|>


<commit_before><commit_msg>disable forking in devserver to enable the in-browser debugging to work<commit_after><|endoftext|>