services/nodemanager/arvnodeman/computenode/driver/ec2.py

   1 #!/usr/bin/env python
   2 # Copyright (C) The Arvados Authors. All rights reserved.
   3 #
   4 # SPDX-License-Identifier: AGPL-3.0
   5
   6 from __future__ import absolute_import, print_function
   7
   8 import time
   9
  10 import libcloud.compute.base as cloud_base
  11 import libcloud.compute.providers as cloud_provider
  12 import libcloud.compute.types as cloud_types
  13 from libcloud.compute.drivers import ec2 as cloud_ec2
  14
  15 from . import BaseComputeNodeDriver
  16 from .. import arvados_node_fqdn
  17
  18 ### Monkeypatch libcloud to support AWS' new SecurityGroup API.
  19 # These classes can be removed when libcloud support specifying
  20 # security groups with the SecurityGroupId parameter.
  21 class ANMEC2Connection(cloud_ec2.EC2Connection):
  22     def request(self, *args, **kwargs):
  23         params = kwargs.get('params')
  24         if (params is not None) and (params.get('Action') == 'RunInstances'):
  25             for key in params.keys():
  26                 if key.startswith('SecurityGroup.'):
  27                     new_key = key.replace('Group.', 'GroupId.', 1)
  28                     params[new_key] = params.pop(key).id
  29             kwargs['params'] = params
  30         return super(ANMEC2Connection, self).request(*args, **kwargs)
  31
  32
  33 class ANMEC2NodeDriver(cloud_ec2.EC2NodeDriver):
  34     connectionCls = ANMEC2Connection
  35
  36
  37 class ComputeNodeDriver(BaseComputeNodeDriver):
  38     """Compute node driver wrapper for EC2.
  39
  40     This translates cloud driver requests to EC2's specific parameters.
  41     """
  42     DEFAULT_DRIVER = ANMEC2NodeDriver
  43 ### End monkeypatch
  44     SEARCH_CACHE = {}
  45
  46     def __init__(self, auth_kwargs, list_kwargs, create_kwargs,
  47                  driver_class=DEFAULT_DRIVER):
  48         # We need full lists of keys up front because these loops modify
  49         # dictionaries in-place.
  50         for key in list_kwargs.keys():
  51             list_kwargs[key.replace('_', ':')] = list_kwargs.pop(key)
  52         self.tags = {key[4:]: value
  53                      for key, value in list_kwargs.iteritems()
  54                      if key.startswith('tag:')}
  55         # Tags are assigned at instance creation time
  56         create_kwargs.setdefault('ex_metadata', {})
  57         create_kwargs['ex_metadata'].update(self.tags)
  58         super(ComputeNodeDriver, self).__init__(
  59             auth_kwargs, {'ex_filters': list_kwargs}, create_kwargs,
  60             driver_class)
  61
  62     def _init_image_id(self, image_id):
  63         return 'image', self.search_for(image_id, 'list_images', ex_owner='self')
  64
  65     def _init_security_groups(self, group_names):
  66         return 'ex_security_groups', [
  67             self.search_for(gname.strip(), 'ex_get_security_groups')
  68             for gname in group_names.split(',')]
  69
  70     def _init_subnet_id(self, subnet_id):
  71         return 'ex_subnet', self.search_for(subnet_id, 'ex_list_subnets')
  72
  73     create_cloud_name = staticmethod(arvados_node_fqdn)
  74
  75     def arvados_create_kwargs(self, size, arvados_node):
  76         kw = {'name': self.create_cloud_name(arvados_node),
  77                 'ex_userdata': self._make_ping_url(arvados_node)}
  78         # libcloud/ec2 disk sizes are in GB, Arvados/SLURM "scratch" value is in MB
  79         scratch = int(size.scratch / 1000) + 1
  80         if scratch > size.disk:
  81             volsize = scratch - size.disk
  82             if volsize > 16384:
  83                 # Must be 1-16384 for General Purpose SSD (gp2) devices
  84                 # https://docs.aws.amazon.com/AWSEC2/latest/APIReference/API_EbsBlockDevice.html
  85                 self._logger.warning("Requested EBS volume size %d is too large, capping size request to 16384 GB", volsize)
  86                 volsize = 16384
  87             kw["ex_blockdevicemappings"] = [{
  88                 "DeviceName": "/dev/xvdt",
  89                 "Ebs": {
  90                     "DeleteOnTermination": True,
  91                     "VolumeSize": volsize,
  92                     "VolumeType": "gp2"
  93                 }}]
  94         if size.preemptible:
  95             # Request a Spot instance for this node
  96             kw['ex_spot_market'] = True
  97         return kw
  98
  99     def sync_node(self, cloud_node, arvados_node):
 100         self.real.ex_create_tags(cloud_node,
 101                                  {'Name': arvados_node_fqdn(arvados_node)})
 102
 103     def create_node(self, size, arvados_node):
 104         # Set up tag indicating the Arvados assigned Cloud Size id.
 105         self.create_kwargs['ex_metadata'].update({'arvados_node_size': size.id})
 106         return super(ComputeNodeDriver, self).create_node(size, arvados_node)
 107
 108     def list_nodes(self):
 109         # Need to populate Node.size
 110         nodes = super(ComputeNodeDriver, self).list_nodes()
 111         for n in nodes:
 112             if not n.size:
 113                 n.size = self.sizes[n.extra["instance_type"]]
 114             n.extra['arvados_node_size'] = n.extra.get('tags', {}).get('arvados_node_size')
 115         return nodes
 116
 117     @classmethod
 118     def node_fqdn(cls, node):
 119         return node.name
 120
 121     @classmethod
 122     def node_start_time(cls, node):
 123         time_str = node.extra['launch_time'].split('.', 2)[0] + 'UTC'
 124         return time.mktime(time.strptime(
 125                 time_str,'%Y-%m-%dT%H:%M:%S%Z')) - time.timezone
 126
 127     @classmethod
 128     def node_id(cls, node):
 129         return node.id