Repository: jcrist/skein (python)

Issue:
Human-readable memory strings in skein.Resources
What do you think about changing the constructor of `skein.Resources` to accept memory as either a number of megabytes, or a string, e.g. "5G" or "1024M"? This could be a bit more readable for bigger containers. We routinely allocate 20-30G, and 20 * 2**10 does not look as nice.

Pull request description:
Add support for human-readable memory specs
Previously we requested memory as an integer in MiB, we now allow (and
prefer) that memory is specified as a string with units.

Note that the default unit is MiB (mebibytes) not MB (megabytes). This
is a carryover from YARN, and is unlikely to change.

Fixes #86.

Patch:
diff --git a/docs/source/key-value-store.rst b/docs/source/key-value-store.rst
index bd3a7bb..2eab797 100644
--- a/docs/source/key-value-store.rst
+++ b/docs/source/key-value-store.rst
@@ -49,7 +49,7 @@ forever.
     ... services:
     ...   sleeper:
     ...     resources:
-    ...       memory: 128
+    ...       memory: 128 MiB
     ...       vcores: 1
     ...     commands:
     ...       - sleep infinity
diff --git a/docs/source/quickstart.rst b/docs/source/quickstart.rst
index 21fabcc..5c495b7 100644
--- a/docs/source/quickstart.rst
+++ b/docs/source/quickstart.rst
@@ -68,7 +68,7 @@ Here we create a simple "Hello World" application as a YAML file:
       my_service:
         resources:
           vcores: 1
-          memory: 128
+          memory: 128 MiB
         files:
           payload.txt: payload.txt
         commands:
@@ -94,7 +94,7 @@ Walking through this specification:
      queue: default
 
 - The application starts a single service ``my_service`` on a container using 1
-  virtual core and 128 MB of memory.
+  virtual core and 128 MiB of memory.
 
   .. code-block:: yaml
 
@@ -102,7 +102,7 @@ Walking through this specification:
        my_service:
          resources:
            vcores: 1
-           memory: 128
+           memory: 128 MiB
 
 
 - The file ``payload.txt`` is distributed with the application, and is named
diff --git a/docs/source/recipes-ipython-kernel.rst b/docs/source/recipes-ipython-kernel.rst
index 6fbbdec..f694db6 100644
--- a/docs/source/recipes-ipython-kernel.rst
+++ b/docs/source/recipes-ipython-kernel.rst
@@ -92,7 +92,7 @@ the `specification docs <specification.html>`__.
     services:
       ipython:
         resources:
-          memory: 1024
+          memory: 1 GiB
           vcores: 1
         files:
           # Distribute the bundled environment as part of the application.
diff --git a/docs/source/recipes-jupyter-notebook.rst b/docs/source/recipes-jupyter-notebook.rst
index 51bb4d3..556d0b2 100644
--- a/docs/source/recipes-jupyter-notebook.rst
+++ b/docs/source/recipes-jupyter-notebook.rst
@@ -140,7 +140,7 @@ the `specification docs <specification.html>`__.
     services:
       jupyter:
         resources:
-          memory: 1024
+          memory: 1 GiB
           vcores: 1
         files:
           # Distribute the bundled environment as part of the application.
diff --git a/docs/source/specification.rst b/docs/source/specification.rst
index 5ea0b50..dd1d8f0 100644
--- a/docs/source/specification.rst
+++ b/docs/source/specification.rst
@@ -254,9 +254,16 @@ following fields:
 
 - ``memory``
 
-  The amount of memory to request, in MB. Requests smaller than the minimum
-  allocation will receive the minimum allocation (usually 1024). Requests
-  larger than the maximum allocation will error on application submission.
+  The amount of memory to request. Can be either a string with units (e.g. ``"5
+  GiB"``), or numeric. If numeric, specifies the amount of memory in *MiB*.
+  Note that the units are in mebibytes (MiB) *not* megabytes (MB) - the former
+  being binary based (1024 MiB in a GiB), the latter being decimal based (1000
+  MB in a GB). See `here <https://en.wikipedia.org/wiki/Mebibyte>`__ for more
+  information on this distinction.
+
+  Requests smaller than the minimum allocation will receive the minimum
+  allocation (1024 MiB by default). Requests larger than the maximum allocation
+  will error on application submission.
 
 - ``vcores``
 
@@ -272,7 +279,7 @@ following fields:
 
     services:
       my_service:
-        memory: 2048  # 2 GB
+        memory: 2 GiB
         vcores: 2
 
 
@@ -414,7 +421,7 @@ applications are free to package files any way they see fit.
     services:
       jupyter:
         resources:
-          memory: 1024
+          memory: 1 GiB
           vcores: 1
         files:
           conda_env: env.zip
@@ -425,7 +432,7 @@ applications are free to package files any way they see fit.
 
       dask.scheduler:
         resources:
-          memory: 2048
+          memory: 2 GiB
           vcores: 1
         files:
           conda_env: env.zip
@@ -436,7 +443,7 @@ applications are free to package files any way they see fit.
       dask.worker:
         instances: 4
         resources:
-          memory: 2048
+          memory: 2 GiB
           vcores: 4
         max_restarts: 8  # Restart workers a maximum of 8 times
         files:
diff --git a/examples/echo_server/spec.yaml b/examples/echo_server/spec.yaml
index 2eaf6e8..ccd47f4 100644
--- a/examples/echo_server/spec.yaml
+++ b/examples/echo_server/spec.yaml
@@ -4,7 +4,7 @@ services:
     server:
         resources:
             vcores: 1
-            memory: 256
+            memory: 256 MiB
         files:
             # A packaged conda environment to be distributed with the
             # application. During YARN resource localization this will be
diff --git a/examples/pricing_dashboard/spec.yaml b/examples/pricing_dashboard/spec.yaml
index c23d807..a8a2d86 100644
--- a/examples/pricing_dashboard/spec.yaml
+++ b/examples/pricing_dashboard/spec.yaml
@@ -4,7 +4,7 @@ services:
     dashboard:
         resources:
             vcores: 1
-            memory: 256
+            memory: 256 MiB
         files:
             # A packaged conda or virtual environment to be distributed with
             # the application. During YARN resource localization this will be
diff --git a/skein/compatibility.py b/skein/compatibility.py
index d323472..50388d0 100644
--- a/skein/compatibility.py
+++ b/skein/compatibility.py
@@ -10,6 +10,7 @@ if PY2:
     import os
     import re
     import types
+    from math import ceil as _ceil
     from urlparse import urlparse, urlsplit  # noqa
     from Queue import Queue  # noqa
     unicode = unicode  # noqa
@@ -45,7 +46,10 @@ if PY2:
     def isidentifier(s):
         return bool(_name_re.match(s))
 
+    def math_ceil(x):
+        return int(_ceil(x))
 else:
+    from math import ceil as math_ceil  # noqa
     from urllib.parse import urlparse, urlsplit  # noqa
     from os import makedirs  # noqa
     from queue import Queue  # noqa
diff --git a/skein/model.py b/skein/model.py
index be63112..bd9c6d6 100644
--- a/skein/model.py
+++ b/skein/model.py
@@ -8,7 +8,7 @@ from getpass import getuser
 import yaml
 
 from . import proto as _proto
-from .compatibility import urlparse, string, integer
+from .compatibility import urlparse, string, integer, math_ceil
 from .objects import Enum, ProtobufMessage, Specification, required
 from .exceptions import context
 from .utils import implements, format_list, datetime_from_millis, runtime
@@ -42,6 +42,70 @@ def container_instance_to_string(id):
     return '%s_%d' % (id.service_name, id.instance)
 
 
+def parse_memory(s):
+    """Converts bytes expression to number of mebibytes.
+
+    If no unit is specified, ``MiB`` is used."""
+    if isinstance(s, integer):
+        out = s
+    elif isinstance(s, float):
+        out = math_ceil(s)
+    elif isinstance(s, string):
+        s = s.replace(' ', '')
+
+        if not s:
+            raise context.ValueError("Could not interpret %r as a byte unit" % s)
+
+        if s[0].isdigit():
+            for i, c in enumerate(reversed(s)):
+                if not c.isalpha():
+                    break
+
+            index = len(s) - i
+            prefix = s[:index]
+            suffix = s[index:]
+
+            try:
+                n = float(prefix)
+            except ValueError:
+                raise context.ValueError("Could not interpret %r as a number" % prefix)
+        else:
+            n = 1
+            suffix = s
+
+        try:
+            multiplier = _byte_sizes[suffix.lower()]
+        except KeyError:
+            raise context.ValueError("Could not interpret %r as a byte unit" % suffix)
+
+        out = math_ceil(n * multiplier / (2 ** 20))
+    else:
+        raise context.TypeError("memory must be an integer, got %r"
+                                % type(s).__name__)
+
+    if out < 0:
+        raise context.ValueError("memory must be positive")
+
+    return out
+
+
+_byte_sizes = {
+    'kb': 10**3,
+    'mb': 10**6,
+    'gb': 10**9,
+    'tb': 10**12,
+    'pb': 10**15,
+    'kib': 2**10,
+    'mib': 2**20,
+    'gib': 2**30,
+    'tib': 2**40,
+    'pib': 2**50,
+    'b': 1,
+    '': 2 ** 20
+}
+_byte_sizes.update({k[:-1]: v for k, v in _byte_sizes.items() if len(k) >= 2})
+
+
 def check_no_cycles(dependencies):
     completed = set()
     seen = set()
@@ -157,18 +221,24 @@ class Resources(Specification):
 
     Parameters
     ----------
-    memory : int
-        The amount of memory to request, in MB. Requests smaller than the
-        minimum allocation will receive the minimum allocation (usually 1024).
-        Requests larger than the maximum allocation will error on application
-        submission.
+    memory : string or int
+        The amount of memory to request. Can be either a string with units
+        (e.g. ``"5 GiB"``), or numeric. If numeric, specifies the amount of
+        memory in *MiB*. Note that the units are in mebibytes (MiB) NOT
+        megabytes (MB) - the former being binary based (1024 MiB in a GiB), the
+        latter being decimal based (1000 MB in a GB).
+
+        Requests smaller than the minimum allocation will receive the minimum
+        allocation (1024 MiB by default). Requests larger than the maximum
+        allocation will error on application submission.
     vcores : int
         The number of virtual cores to request. Depending on your system
         configuration one virtual core may map to a single actual core, or a
         fraction of a core. Requests larger than the maximum allocation will
         error on application submission.
     """
-    __slots__ = ('memory', 'vcores')
+    __slots__ = ('_memory', 'vcores')
+    _params = ('memory', 'vcores')
     _protobuf_cls = _proto.Resources
 
     def __init__(self, memory=required, vcores=required):
@@ -176,6 +246,14 @@ class Resources(Specification):
         self._assign_required('vcores', vcores)
         self._validate()
 
+    @property
+    def memory(self):
+        return self._memory
+
+    @memory.setter
+    def memory(self, value):
+        self._memory = parse_memory(value)
+
     def __repr__(self):
         return 'Resources<memory=%d, vcores=%d>' % (self.memory, self.vcores)
 


Test patch:
diff --git a/skein/test/test_model.py b/skein/test/test_model.py
index 0462282..2049746 100644
--- a/skein/test/test_model.py
+++ b/skein/test/test_model.py
@@ -7,11 +7,11 @@ import pickle
 
 import pytest
 
-from skein.compatibility import UTC
+from skein.compatibility import UTC, math_ceil
 from skein.model import (ApplicationSpec, Service, Resources, File,
                          ApplicationState, FinalStatus, FileType, ACLs, Master,
                          Container, ApplicationReport, ResourceUsageReport,
-                         LogLevel)
+                         LogLevel, parse_memory)
 
 
 def indent(s, n):
@@ -23,7 +23,7 @@ service_spec = """\
 node_label: gpu
 resources:
     vcores: 1
-    memory: 1024
+    memory: 1 GiB
 files:
     testfile: /path/to/testfile
     testarchive: /path/to/testarchive.zip
@@ -88,6 +88,39 @@ def check_specification_methods(obj, obj2):
             assert obj == obj2
 
 
+def test_parse_memory():
+    mib = 2 ** 20
+    assert parse_memory('100') == 100
+    assert parse_memory('100 MiB') == 100
+    assert parse_memory('100 MB') == math_ceil(100 * 1e6 / mib)
+    assert parse_memory('100M') == math_ceil(100 * 1e6 / mib)
+    assert parse_memory('100Mi') == 100
+    assert parse_memory('5kB') == 1
+    assert parse_memory('5.4 MiB') == 6
+    assert parse_memory('0.9 MiB') == 1
+    assert parse_memory('1e3') == 1000
+    assert parse_memory('1e6 kB') == math_ceil(1e6 * 1e3 / mib)
+    assert parse_memory('MiB') == 1
+
+    with pytest.raises(ValueError):
+        parse_memory('5 foos')
+
+    with pytest.raises(ValueError):
+        parse_memory('')
+
+    with pytest.raises(ValueError):
+        parse_memory('1.1.1GB')
+
+    with pytest.raises(ValueError):
+        parse_memory(-1)
+
+    with pytest.raises(ValueError):
+        parse_memory('-1.5 MiB')
+
+    with pytest.raises(TypeError):
+        parse_memory([])
+
+
 def test_resources():
     r = Resources(memory=1024, vcores=1)
     r2 = Resources(memory=1024, vcores=2)
@@ -95,14 +128,36 @@ def test_resources():
 
 
 def test_resources_invariants():
+    assert Resources(memory='1024 MiB', vcores=1).memory == 1024
+    assert Resources(memory='5 GiB', vcores=1).memory == 5 * 1024
+    assert Resources(memory='512', vcores=1).memory == 512
+    assert Resources(memory=1e3, vcores=1).memory == 1000
+    assert Resources(memory='0.9 MiB', vcores=1).memory == 1
+
+    r = Resources(memory='1 GiB', vcores=1)
+    assert r.memory == 1024
+    r.memory = '2 GiB'
+    assert r.memory == 2048
+    r.memory = 1e3
+    assert r.memory == 1000
+
+    with pytest.raises(ValueError):
+        r.memory = -1
+
+    with pytest.raises(ValueError):
+        r.memory = -1
+
     with pytest.raises(TypeError):
         Resources()
 
     with pytest.raises(TypeError):
         Resources(None, None)
 
+    with pytest.raises(ValueError):
+        Resources(memory='foo', vcores=1)
+
     with pytest.raises(TypeError):
-        Resources(memory='foo', vcores='bar')
+        Resources(memory=1, vcores='foo')
 
     with pytest.raises(ValueError):
         Resources(memory=-1, vcores=-1)