services/nodemanager/arvnodeman/test/fake_driver.py

   1 # Copyright (C) The Arvados Authors. All rights reserved.
   2 #
   3 # SPDX-License-Identifier: AGPL-3.0
   4
   5 import re
   6 import urllib
   7 import ssl
   8 import time
   9
  10 from arvnodeman.computenode import ARVADOS_TIMEFMT
  11
  12 from libcloud.compute.base import NodeSize, Node, NodeDriver, NodeState, NodeImage
  13 from libcloud.compute.drivers.gce import GCEDiskType
  14 from libcloud.common.exceptions import BaseHTTPError, RateLimitReachedError
  15
  16 all_nodes = []
  17 create_calls = 0
  18 quota = 2
  19
  20 class FakeDriver(NodeDriver):
  21     def __init__(self, *args, **kwargs):
  22         self.name = "FakeDriver"
  23
  24     def list_sizes(self, **kwargs):
  25         return [NodeSize("Standard_D3", "Standard_D3", 3500, 200, 0, 0, self),
  26                 NodeSize("Standard_D4", "Standard_D4", 7000, 400, 0, 0, self)]
  27
  28     def list_nodes(self, **kwargs):
  29         return all_nodes
  30
  31     def create_node(self, name=None,
  32                     size=None,
  33                     image=None,
  34                     auth=None,
  35                     ex_storage_account=None,
  36                     ex_customdata=None,
  37                     ex_resource_group=None,
  38                     ex_user_name=None,
  39                     ex_tags=None,
  40                     ex_metadata=None,
  41                     ex_network=None,
  42                     ex_userdata=None):
  43         global all_nodes, create_calls
  44         create_calls += 1
  45         nodeid = "node%i" % create_calls
  46         if ex_tags is None:
  47             ex_tags = {}
  48         ex_tags.update({'arvados_node_size': size.id})
  49         n = Node(nodeid, nodeid, NodeState.RUNNING, [], [], self, size=size, extra={"tags": ex_tags})
  50         all_nodes.append(n)
  51         if ex_customdata:
  52             ping_url = re.search(r"echo '(.*)' > /var/tmp/arv-node-data/arv-ping-url", ex_customdata).groups(1)[0]
  53         if ex_userdata:
  54             ping_url = ex_userdata
  55         elif ex_metadata:
  56             ping_url = ex_metadata["arv-ping-url"]
  57         ping_url += "&instance_id=" + nodeid
  58         ctx = ssl.SSLContext(ssl.PROTOCOL_SSLv23)
  59         ctx.verify_mode = ssl.CERT_NONE
  60         f = urllib.urlopen(ping_url, "", context=ctx)
  61         f.close()
  62         return n
  63
  64     def destroy_node(self, cloud_node):
  65         global all_nodes
  66         all_nodes = [n for n in all_nodes if n.id != cloud_node.id]
  67         return True
  68
  69     def get_image(self, img):
  70         pass
  71
  72     def ex_create_tags(self, cloud_node, tags):
  73         pass
  74
  75 class QuotaDriver(FakeDriver):
  76     def create_node(self, name=None,
  77                     size=None,
  78                     image=None,
  79                     auth=None,
  80                     ex_storage_account=None,
  81                     ex_customdata=None,
  82                     ex_resource_group=None,
  83                     ex_user_name=None,
  84                     ex_tags=None,
  85                     ex_network=None):
  86         global all_nodes, create_calls, quota
  87         if len(all_nodes) >= quota:
  88             raise BaseHTTPError(503, "Quota exceeded")
  89         else:
  90             return super(QuotaDriver, self).create_node(name=name,
  91                     size=size,
  92                     image=image,
  93                     auth=auth,
  94                     ex_storage_account=ex_storage_account,
  95                     ex_customdata=ex_customdata,
  96                     ex_resource_group=ex_resource_group,
  97                     ex_user_name=ex_user_name,
  98                     ex_tags=ex_tags,
  99                     ex_network=ex_network)
 100
 101     def destroy_node(self, cloud_node):
 102         global all_nodes, quota
 103         all_nodes = [n for n in all_nodes if n.id != cloud_node.id]
 104         if len(all_nodes) == 0:
 105             quota = 4
 106         return True
 107
 108 class FailingDriver(FakeDriver):
 109     def create_node(self, name=None,
 110                     size=None,
 111                     image=None,
 112                     auth=None,
 113                     ex_storage_account=None,
 114                     ex_customdata=None,
 115                     ex_resource_group=None,
 116                     ex_user_name=None,
 117                     ex_tags=None,
 118                     ex_network=None):
 119         raise Exception("nope")
 120
 121 class RetryDriver(FakeDriver):
 122     def create_node(self, name=None,
 123                     size=None,
 124                     image=None,
 125                     auth=None,
 126                     ex_storage_account=None,
 127                     ex_customdata=None,
 128                     ex_resource_group=None,
 129                     ex_user_name=None,
 130                     ex_tags=None,
 131                     ex_network=None):
 132         global create_calls
 133         create_calls += 1
 134         if create_calls < 2:
 135             raise RateLimitReachedError(429, "Rate limit exceeded",
 136                                         headers={'retry-after': '2'})
 137         elif create_calls < 3:
 138             raise BaseHTTPError(429, "Rate limit exceeded",
 139                                 {'retry-after': '1'})
 140         else:
 141             return super(RetryDriver, self).create_node(name=name,
 142                     size=size,
 143                     image=image,
 144                     auth=auth,
 145                     ex_storage_account=ex_storage_account,
 146                     ex_customdata=ex_customdata,
 147                     ex_resource_group=ex_resource_group,
 148                     ex_user_name=ex_user_name,
 149                     ex_tags=ex_tags,
 150                     ex_network=ex_network)
 151
 152 class FakeAwsDriver(FakeDriver):
 153
 154     def create_node(self, name=None,
 155                     size=None,
 156                     image=None,
 157                     auth=None,
 158                     ex_userdata=None,
 159                     ex_metadata=None,
 160                     ex_blockdevicemappings=None):
 161         n = super(FakeAwsDriver, self).create_node(name=name,
 162                                                       size=size,
 163                                                       image=image,
 164                                                       auth=auth,
 165                                                       ex_metadata=ex_metadata,
 166                                                       ex_userdata=ex_userdata)
 167         n.extra = {
 168             "launch_time": time.strftime(ARVADOS_TIMEFMT, time.gmtime())[:-1],
 169             "tags" : {
 170                 "arvados_node_size": size.id
 171             }
 172         }
 173         return n
 174
 175     def list_sizes(self, **kwargs):
 176         return [NodeSize("m3.xlarge", "Extra Large Instance", 3500, 80, 0, 0, self),
 177                 NodeSize("m4.xlarge", "Extra Large Instance", 3500, 0, 0, 0, self),
 178                 NodeSize("m4.2xlarge", "Double Extra Large Instance", 7000, 0, 0, 0, self)]
 179
 180
 181 class FakeGceDriver(FakeDriver):
 182
 183     def create_node(self, name=None,
 184                     size=None,
 185                     image=None,
 186                     auth=None,
 187                     external_ip=None,
 188                     ex_metadata=None,
 189                     ex_tags=None,
 190                     ex_disks_gce_struct=None):
 191         n = super(FakeGceDriver, self).create_node(name=name,
 192                                                    size=size,
 193                                                    image=image,
 194                                                    auth=auth,
 195                                                    ex_metadata=ex_metadata)
 196         n.extra = {
 197             "metadata": {
 198                 "items": [{"key": k, "value": v} for k,v in ex_metadata.iteritems()],
 199                 "arvados_node_size": size.id
 200             },
 201             "zone": "fake"
 202         }
 203         return n
 204
 205     def list_images(self, ex_project=None):
 206         return [NodeImage("fake_image_id", "fake_image_id", self)]
 207
 208     def list_sizes(self, **kwargs):
 209         return [NodeSize("n1-standard-1", "Standard", 3750, None, 0, 0, self),
 210                 NodeSize("n1-standard-2", "Double standard", 7500, None, 0, 0, self)]
 211
 212     def ex_list_disktypes(self, zone=None):
 213         return [GCEDiskType("pd-standard", "pd-standard", zone, self,
 214                             extra={"selfLink": "pd-standard"}),
 215                 GCEDiskType("local-ssd", "local-ssd", zone, self,
 216                             extra={"selfLink": "local-ssd"})]
 217
 218     def ex_get_node(self, name, zone=None):
 219         global all_nodes
 220         for n in all_nodes:
 221             if n.id == name:
 222                 return n
 223         return None
 224
 225     def ex_set_node_metadata(self, n, items):
 226         n.extra["metadata"]["items"] = items