uvms: add --mem

Yeah, yeah, I know, it's stupid, adding add_argument() statements
manually, et c. I'll throw the whole Python thing away whenever I might
have the time

pkgs/uvms/uvms.py

@@ -10,6 +10,7 @@ import os
 import subprocess
 import socket
 import json
+import re
 from argparse import ArgumentParser
 from contextlib import contextmanager, closing, ExitStack
 
@@ -20,6 +21,7 @@ parser.add_argument("--prefix", default="$HOME/uvms/$VM")
 parser.add_argument("--vm-config", default="@BASE_CONFIG@")  # noqa: E501
 parser.add_argument("--persist-home", action="store_true")
 parser.add_argument("--run", action="append")
+parser.add_argument("--mem", default=None)
 parser.add_argument("app", nargs="*", default=())
 
 TOOLS_DIR = "@TOOLS@"  # noqa: E501
@@ -469,6 +471,21 @@ def connect_ch_vsock(
         yield s
 
 
+BYTES_PATTERN = re.compile(r"^([0-9]+)([MmGgKk]?)$")
+BYTES_UNITS = {
+    "k": 1024,
+    "m": 1048576,
+    "g": 1024 * 1048576,
+}
+
+
+def parse_bytes(s):
+    m = BYTES_PATTERN.match(s)
+    assert m, s
+    size, unit = m.groups()
+    return int(size) * BYTES_UNITS.get(unit.lower(), 1)
+
+
 @contextmanager
 def listen_ch_vsock(
     vsock_sock_path,
@@ -564,6 +581,9 @@ def main(args, args_next, cleanup, ps):
         )
         virtiofs_socks.append(("home", sock_path))
         config["payload"]["cmdline"] += " uvms.persist-home=1"
+    if args.mem is not None:
+        config["memory"]["size"] = parse_bytes(args.mem)
+        config["memory"]["hotplug_size"] = parse_bytes(args.mem)
 
     gpud, gpud_path = cleanup.enter_context(ps.start_gpu())