Restore from snapshot

vast.py

@@ -755,6 +755,15 @@ def deindent(message: str) -> str:
     ("args", "args", "{}", None, True),
 )
 
+snapshot_fields = (
+    ("id",                  "id",                  "{}",    None, True),
+    ("image",               "snapshot_image",      "{}",    None, True),
+    ("status",              "status",              "{}",    None, True),
+    ("execution_time",      "execution_time",      "{}",    None, True),
+    ("container_registry",  "container_registry",  "{}",    None, True),
+    ("repo",                "repo",                "{}",    None, True),
+)
+
 invoice_fields = (
     ("description", "Description", "{}", None, True),
     ("quantity", "Quantity", "{}", None, True),
@@ -1537,6 +1546,146 @@ def cloud__copy(args: argparse.Namespace):
         print(r.text);
         print("failed with error {r.status_code}".format(**locals()));
 
+@parser.command(
+    argument("instance_id",      help="instance_id of the container instance to snapshot",      type=str),
+    argument("--container_registry", help="Container registry to push the snapshot to. Default will be docker.io", type=str, default="docker.io"),
+    argument("--repo",    help="repo to push the snapshot to",     type=str),
+    argument("--docker_login_user",help="Username for container registry with repo",     type=str),
+    argument("--docker_login_pass",help="Password or token for container registry with repo",     type=str),
+    argument("--pause",            help="Pause container's processes being executed by the CPU to take snapshot (true/false). Default will be true", type=str, default="true"),
+    usage="vastai take snapshot INSTANCE_ID "
+          "--repo REPO --docker_login_user USER --docker_login_pass PASS"
+          "[--container_registry REGISTRY] [--pause true|false]",
+    help="Schedule a snapshot of a running container and push it to your repo in a container registry",
+    epilog=deindent("""
+        Takes a snapshot of a running container instance and pushes snapshot to the specified repository in container registry.
+
+        Use pause=true to pause the container during commit (safer but slower),
+        or pause=false to leave it running (faster but may produce a filesystem-
+// safer snapshot).
+    """),
+)
+def take__snapshot(args: argparse.Namespace):
+    """
+    Take a container snapshot and push.
+
+    @param instance_id: instance identifier.
+    @param repo: Docker repository for the snapshot.
+    @param container_registry: Container registry
+    @param docker_login_user: Docker registry username.
+    @param docker_login_pass: Docker registry password/token.
+    @param pause: "true" or "false" to pause the container during commit.
+    """
+    instance_id       = args.instance_id
+    repo              = args.repo
+    container_registry = args.container_registry
+    user              = args.docker_login_user
+    password          = args.docker_login_pass
+    pause_flag        = args.pause
+
+    print(f"Taking snapshot for instance {instance_id} and pushing to repo {repo} in container registry {container_registry}")
+    req_json = {
+        "id":               instance_id,
+        "container_registry": container_registry,
+        "personal_repo":    repo,
+        "docker_login_user":user,
+        "docker_login_pass":password,
+        "pause":            pause_flag
+    }
+
+    url = apiurl(args, f"/snapshot/take/")
+    if args.explain:
+        print("Request JSON:")
+        print(json.dumps(req_json, indent=2))
+
+    # POST to the snapshot endpoint
+    r = http_post(args, url, headers=headers, json=req_json)
+    r.raise_for_status()
+
+    if r.status_code == 200:
+        data = r.json()
+        if data.get("success"):
+            print(f"Snapshot request sent successfully. Please check your repo {repo} in container registry {container_registry} in 5-10 mins. It can take longer than 5-10 mins to push your snapshot image to your repo depending on the size of your image.")
+        else:
+            print(data.get("msg", "Unknown error with snapshot request"))
+    else:
+        print(r.text);
+        print("failed with error {r.status_code}".format(**locals()));
+
+def safe_int(value, default=-1):
+    try:
+        return int(value)
+    except (TypeError, ValueError):
+        return default
+
+@parser.command(
+    argument("snapshot_id",help="ID of the container snapshot to restore from",type=str),
+    argument("offer_id",help="Offer ID to use when creating the new instance",type=str),
+    usage="vastai restore from snapshot SNAPSHOT_ID OFFER_ID",
+    help="Restore from a container snapshot and create a new instance",
+    epilog="""
+Restores from a container snapshot and creates a new instance.  
+You must supply the snapshot ID and the offer ID you wish to use for the new instance.
+"""
+)
+def restore__snapshot(args: argparse.Namespace):
+    """
+    Restore a container instance from a snapshot.
+
+    @param snapshot_id: snapshot identifier.
+    @param offer_id: offer identifier for the new instance.
+    """
+    snapshot_id = safe_int(args.snapshot_id)
+    offer_id    = safe_int(args.offer_id)
+    if snapshot_id is None or offer_id is None:
+        print("Error: both snapshot_id and offer_id must be integers")
+        return 1
+
+    print(f"Restoring new instance from snapshot {snapshot_id} using offer {offer_id}")
+
+    # build request payload
+    req_json = {"snapshot_id": snapshot_id, "offer_id": offer_id}
+
+    url = apiurl(args, "/snapshot/restore/")
+
+    if args.explain:
+        print("Request URL:")
+        print(url)
+        print("Request JSON:")
+        print(json.dumps(req_json, indent=2))
+
+    # POST to the restore endpoint
+    r = http_put(args, url, headers=headers, json=req_json)
+    r.raise_for_status()
+
+    if r.status_code == 200:
+        if args.raw:
+            return r
+        else:
+            print("Restored from snapshot successfully. {}".format(r.json()))
+    else:
+        print(r.text);
+        print("failed with error {r.status_code}".format(**locals()));
+
+@parser.command(
+    usage="vastai show snapshots [--api-key API_KEY] [--raw]",
+    help="Show container snapshots"
+)
+def show__snapshots(args):
+    """
+    Shows the history of ip address accesses to console.vast.ai endpoints
+
+    :param argparse.Namespace args: should supply all the command-line options
+    :rtype:
+    """
+    req_url = apiurl(args, "/snapshot/")
+    r = http_get(args, req_url)
+    r.raise_for_status()
+    rows = r.json()
+    if args.raw:
+        return rows
+    else:
+        display_table(rows, snapshot_fields)
 
 @parser.command(
     argument("--name", help="name of the api-key", type=str),