|
4 | 4 | import ipaddress
|
5 | 5 | import time
|
6 | 6 |
|
7 |
| -from traitlets import Integer, TraitError |
| 7 | +from traitlets import Unicode, Integer, TraitError |
8 | 8 |
|
9 | 9 |
|
10 | 10 | # default _request_timeout for kubernetes api requests
|
@@ -42,6 +42,86 @@ def rendezvous_rank(buckets, key):
|
42 | 42 | return [b for (s, b) in sorted(ranking, reverse=True)]
|
43 | 43 |
|
44 | 44 |
|
| 45 | +class CPUSpecification(Unicode): |
| 46 | + """ |
| 47 | + Allows specifying CPU limits |
| 48 | +
|
| 49 | + Suffixes allowed are: |
| 50 | + - m -> millicore |
| 51 | +
|
| 52 | + """ |
| 53 | + |
| 54 | + # Default to allowing None as a value |
| 55 | + allow_none = True |
| 56 | + |
| 57 | + def validate(self, obj, value): |
| 58 | + """ |
| 59 | + Validate that the passed in value is a valid cpu specification |
| 60 | + in the K8s CPU meaning. |
| 61 | + |
| 62 | + See https://kubernetes.io/docs/concepts/configuration/manage-resources-containers/#meaning-of-cpu |
| 63 | +
|
| 64 | + It could either be a pure int or float, when it is taken as a value. |
| 65 | + In case of integer it can optionally have 'm' suffix to designate millicores. |
| 66 | + """ |
| 67 | + |
| 68 | + def raise_error(value): |
| 69 | + raise TraitError( |
| 70 | + "{val} is not a valid cpu specification".format( |
| 71 | + val=value |
| 72 | + ) |
| 73 | + ) |
| 74 | + |
| 75 | + # Positive filter for numberic values |
| 76 | + only_positive = lambda v : v if v >= 0 else raise_error(v) |
| 77 | + |
| 78 | + if value is None: |
| 79 | + return 0 |
| 80 | + |
| 81 | + if isinstance(value, bool): |
| 82 | + raise_error(value) |
| 83 | + |
| 84 | + if isinstance(value, int): |
| 85 | + return only_positive(int(value)) |
| 86 | + |
| 87 | + if isinstance(value, float): |
| 88 | + return only_positive(float(value)) |
| 89 | + |
| 90 | + # Must be string |
| 91 | + if not isinstance(value, str): |
| 92 | + raise_error(value) |
| 93 | + |
| 94 | + # Try treat it as integer |
| 95 | + _int_value = None |
| 96 | + try: |
| 97 | + _int_value = int(value) |
| 98 | + except ValueError: |
| 99 | + pass |
| 100 | + |
| 101 | + if isinstance(_int_value, int): |
| 102 | + return only_positive(_int_value) |
| 103 | + |
| 104 | + # Try treat it as float |
| 105 | + _float_value = None |
| 106 | + try: |
| 107 | + _float_value = float(value) |
| 108 | + except ValueError: |
| 109 | + pass |
| 110 | + |
| 111 | + if isinstance(_float_value, float): |
| 112 | + return only_positive(_float_value) |
| 113 | + |
| 114 | + # Try treat it as millicore spec |
| 115 | + try: |
| 116 | + _unused = only_positive(int(value[:-1])) |
| 117 | + except ValueError: |
| 118 | + raise_error(value) |
| 119 | + |
| 120 | + if value[-1] not in ['m']: |
| 121 | + raise_error(value) |
| 122 | + |
| 123 | + return value |
| 124 | + |
45 | 125 | class ByteSpecification(Integer):
|
46 | 126 | """
|
47 | 127 | Allow easily specifying bytes in units of 1024 with suffixes
|
|
0 commit comments