|
| 1 | +import re |
| 2 | + |
| 3 | +import click |
| 4 | + |
| 5 | +from runpod.api.rest import run_rest_request |
| 6 | +from runpod.rp_cli import compute, console as ui |
| 7 | +from runpod.rp_cli.api import api_errors, list_all |
| 8 | +from runpod.rp_cli.output import ( |
| 9 | + completer, |
| 10 | + confirm, |
| 11 | + json_option, |
| 12 | + pick, |
| 13 | + print_fields, |
| 14 | + print_json, |
| 15 | + print_table, |
| 16 | +) |
| 17 | + |
| 18 | +NETWORK_FIELDS = [ |
| 19 | + ("id", lambda row: row["id"]), |
| 20 | + ("name", lambda row: row["name"]), |
| 21 | + ("size", lambda row: f"{row['size']} GB"), |
| 22 | + ("data center", lambda row: row["dataCenter"]), |
| 23 | + ("type", lambda row: row["type"].lower()), |
| 24 | +] |
| 25 | +GLOBAL_FIELDS = [ |
| 26 | + ("id", lambda row: row["id"]), |
| 27 | + ("name", lambda row: row["name"]), |
| 28 | + ("created", lambda row: row["createdAt"]), |
| 29 | + ("created by", lambda row: row["createdByUserId"]), |
| 30 | +] |
| 31 | +global_option = click.option( |
| 32 | + "--global", |
| 33 | + "kind", |
| 34 | + flag_value="global", |
| 35 | + default="network", |
| 36 | + help="Manage global volumes, which mount in any data center.", |
| 37 | +) |
| 38 | + |
| 39 | + |
| 40 | +def _label(record): |
| 41 | + if "size" not in record: |
| 42 | + return record["name"] |
| 43 | + return ( |
| 44 | + f"{record['name']} {record['size']} GB " |
| 45 | + f"{record['dataCenter']} {record['type'].lower()}" |
| 46 | + ) |
| 47 | + |
| 48 | + |
| 49 | +def _choices(kind): |
| 50 | + return [ |
| 51 | + (row["id"], _label(row)) |
| 52 | + for row in list_all(f"/v2/{kind}-volumes", f"{kind}Volumes") |
| 53 | + ] |
| 54 | + |
| 55 | + |
| 56 | +def _complete_ids(ctx, param, incomplete): |
| 57 | + return completer(lambda: _choices(ctx.params.get("kind", "network")))( |
| 58 | + ctx, param, incomplete |
| 59 | + ) |
| 60 | + |
| 61 | + |
| 62 | +def _show(record, kind, as_json, action=None): |
| 63 | + if as_json: |
| 64 | + print_json(record) |
| 65 | + return |
| 66 | + if action: |
| 67 | + ui.success(f"{action} {kind} volume [white]{record['id']}[/white]") |
| 68 | + print_fields(record, GLOBAL_FIELDS if kind == "global" else NETWORK_FIELDS) |
| 69 | + |
| 70 | + |
| 71 | +def _data_center(regions, volume_type): |
| 72 | + datacenters = compute.list_datacenters() |
| 73 | + if regions: |
| 74 | + allowed = compute.resolve_regions(regions, datacenters) |
| 75 | + datacenters = [dc for dc in datacenters if dc["id"] in allowed] |
| 76 | + candidates = [ |
| 77 | + dc |
| 78 | + for dc in datacenters |
| 79 | + if dc["networkVolumeTypes"] |
| 80 | + and (not volume_type or volume_type in dc["networkVolumeTypes"]) |
| 81 | + ] |
| 82 | + scope = f"region {', '.join(regions)}" if regions else "the catalog" |
| 83 | + offered = ( |
| 84 | + f"{volume_type.lower()} network volumes" if volume_type else "network volumes" |
| 85 | + ) |
| 86 | + if not candidates: |
| 87 | + raise compute.ResolutionError(f"no data center in {scope} offers {offered}") |
| 88 | + if len(candidates) == 1: |
| 89 | + return candidates[0]["id"] |
| 90 | + return pick( |
| 91 | + "data center", |
| 92 | + [ |
| 93 | + ( |
| 94 | + dc["id"], |
| 95 | + f"{dc['id']} {dc['region'].lower().replace('_', ' ')} " |
| 96 | + + ", ".join(t.lower() for t in dc["networkVolumeTypes"]), |
| 97 | + ) |
| 98 | + for dc in candidates |
| 99 | + ], |
| 100 | + f"several data centers in {scope} offer {offered} " |
| 101 | + f"({', '.join(dc['id'] for dc in candidates)}). pass one with --data-center", |
| 102 | + ) |
| 103 | + |
| 104 | + |
| 105 | +def _growth(volume_id, size, expand=False): |
| 106 | + text = str(size).strip() |
| 107 | + if not re.fullmatch(r"\+?\d+", text): |
| 108 | + raise click.BadParameter( |
| 109 | + f"'{size}' is not a size. use N for a total or +N to add N GB", |
| 110 | + param_hint="--size", |
| 111 | + ) |
| 112 | + current = run_rest_request("GET", f"/v2/network-volumes/{volume_id}")["size"] |
| 113 | + target = current + int(text) if text.startswith("+") else int(text) |
| 114 | + if target < current or (expand and target == current): |
| 115 | + raise click.UsageError( |
| 116 | + f"network volumes can only grow: {volume_id} is {current} GB, " |
| 117 | + f"requested {target} GB" |
| 118 | + ) |
| 119 | + if target > 4096: |
| 120 | + raise click.UsageError( |
| 121 | + f"network volumes hold at most 4096 GB, requested {target} GB" |
| 122 | + ) |
| 123 | + return current, target |
| 124 | + |
| 125 | + |
| 126 | +@click.group() |
| 127 | +def volume(): |
| 128 | + """Create, list, resize, and delete network and global volumes.""" |
| 129 | + |
| 130 | + |
| 131 | +@volume.command(name="list") |
| 132 | +@global_option |
| 133 | +@json_option |
| 134 | +def volume_list(kind, as_json): |
| 135 | + """List volumes with their size, data center, and type.""" |
| 136 | + with api_errors(): |
| 137 | + rows = list_all(f"/v2/{kind}-volumes", f"{kind}Volumes") |
| 138 | + if as_json: |
| 139 | + print_json(rows) |
| 140 | + return |
| 141 | + fields = NETWORK_FIELDS |
| 142 | + if kind == "global": |
| 143 | + fields = GLOBAL_FIELDS[:2] + [("created", lambda row: row["createdAt"][:10])] |
| 144 | + print_table( |
| 145 | + rows, |
| 146 | + [(label.upper(), value) for label, value in fields], |
| 147 | + f"no {kind} volumes. create one with rp volume create", |
| 148 | + ) |
| 149 | + |
| 150 | + |
| 151 | +@volume.command(name="get") |
| 152 | +@click.argument("volume_id", required=False, shell_complete=_complete_ids) |
| 153 | +@global_option |
| 154 | +@json_option |
| 155 | +def volume_get(volume_id, kind, as_json): |
| 156 | + """Show one volume (pick one when VOLUME_ID is omitted).""" |
| 157 | + with api_errors(): |
| 158 | + if volume_id is None: |
| 159 | + volume_id = pick(f"{kind} volume", _choices(kind), "pass VOLUME_ID") |
| 160 | + record = run_rest_request("GET", f"/v2/{kind}-volumes/{volume_id}") |
| 161 | + _show(record, kind, as_json) |
| 162 | + |
| 163 | + |
| 164 | +@volume.command(name="create") |
| 165 | +@click.option("--name", required=True, help="Human-readable name.") |
| 166 | +@click.option( |
| 167 | + "--size", |
| 168 | + type=click.IntRange(10, 4096), |
| 169 | + help="Storage in GB (10-4096). Required for network volumes.", |
| 170 | +) |
| 171 | +@click.option( |
| 172 | + "--region", |
| 173 | + "regions", |
| 174 | + multiple=True, |
| 175 | + help="Region (eu, na), id prefix (us, eu-ro), or data center id. Repeatable.", |
| 176 | +) |
| 177 | +@click.option("--data-center", help="Exact data center id, like EU-RO-1.") |
| 178 | +@click.option( |
| 179 | + "--type", |
| 180 | + "volume_type", |
| 181 | + type=click.Choice(["standard", "high_performance"], case_sensitive=False), |
| 182 | + help="Storage tier, immutable after creation. Defaults to the data center's " |
| 183 | + "primary tier.", |
| 184 | +) |
| 185 | +@global_option |
| 186 | +@json_option |
| 187 | +def volume_create(name, size, regions, data_center, volume_type, kind, as_json): |
| 188 | + """Create a network volume in one data center, or a global volume. |
| 189 | +
|
| 190 | + A region matching several data centers narrows to those offering |
| 191 | + the requested --type, then prompts for one. |
| 192 | +
|
| 193 | + \b |
| 194 | + rp volume create --name models --size 100 --region eu |
| 195 | + rp volume create --name scratch --size 50 --data-center EU-RO-1 |
| 196 | + rp volume create --name shared --global |
| 197 | + """ |
| 198 | + body = {"name": name} |
| 199 | + with api_errors(): |
| 200 | + if kind == "global": |
| 201 | + given = [ |
| 202 | + flag |
| 203 | + for flag, value in ( |
| 204 | + ("--size", size), |
| 205 | + ("--region", regions), |
| 206 | + ("--data-center", data_center), |
| 207 | + ("--type", volume_type), |
| 208 | + ) |
| 209 | + if value |
| 210 | + ] |
| 211 | + if given: |
| 212 | + raise click.UsageError( |
| 213 | + f"{', '.join(given)} do not apply to global volumes" |
| 214 | + ) |
| 215 | + else: |
| 216 | + if size is None: |
| 217 | + raise click.UsageError("pass --size in GB for a network volume") |
| 218 | + if regions and data_center: |
| 219 | + raise click.UsageError("use either --region or --data-center") |
| 220 | + volume_type = volume_type.upper() if volume_type else None |
| 221 | + body.update( |
| 222 | + size=size, |
| 223 | + dataCenter=data_center or _data_center(regions, volume_type), |
| 224 | + ) |
| 225 | + if volume_type: |
| 226 | + body["type"] = volume_type |
| 227 | + record = run_rest_request("POST", f"/v2/{kind}-volumes", json=body) |
| 228 | + _show(record, kind, as_json, "created") |
| 229 | + |
| 230 | + |
| 231 | +@volume.command(name="update") |
| 232 | +@click.argument("volume_id", shell_complete=_complete_ids) |
| 233 | +@click.option("--name", help="New name.") |
| 234 | +@click.option( |
| 235 | + "--size", |
| 236 | + type=click.IntRange(10, 4096), |
| 237 | + help="New size in GB; network volumes can only grow.", |
| 238 | +) |
| 239 | +@global_option |
| 240 | +@json_option |
| 241 | +def volume_update(volume_id, name, size, kind, as_json): |
| 242 | + """Rename and/or resize a volume.""" |
| 243 | + if kind == "global" and size is not None: |
| 244 | + raise click.UsageError("global volumes have no size to change") |
| 245 | + body = { |
| 246 | + key: value |
| 247 | + for key, value in (("name", name), ("size", size)) |
| 248 | + if value is not None |
| 249 | + } |
| 250 | + if not body: |
| 251 | + raise click.UsageError("pass --name and/or --size") |
| 252 | + with api_errors(): |
| 253 | + if size is not None: |
| 254 | + _growth(volume_id, size) |
| 255 | + record = run_rest_request("PATCH", f"/v2/{kind}-volumes/{volume_id}", json=body) |
| 256 | + _show(record, kind, as_json, "updated") |
| 257 | + |
| 258 | + |
| 259 | +@volume.command(name="expand") |
| 260 | +@click.argument("volume_id", shell_complete=_complete_ids) |
| 261 | +@click.option( |
| 262 | + "--size", |
| 263 | + required=True, |
| 264 | + metavar="N|+N", |
| 265 | + help="New total size in GB, or +N to add N GB.", |
| 266 | +) |
| 267 | +@json_option |
| 268 | +def volume_expand(volume_id, size, as_json): |
| 269 | + """Grow a network volume; storage can never shrink.""" |
| 270 | + with api_errors(): |
| 271 | + current, target = _growth(volume_id, size, expand=True) |
| 272 | + record = run_rest_request( |
| 273 | + "PATCH", f"/v2/network-volumes/{volume_id}", json={"size": target} |
| 274 | + ) |
| 275 | + if as_json: |
| 276 | + print_json(record) |
| 277 | + else: |
| 278 | + ui.success( |
| 279 | + f"expanded [white]{volume_id}[/white] from {current} GB to {record['size']} GB" |
| 280 | + ) |
| 281 | + |
| 282 | + |
| 283 | +@volume.command(name="delete") |
| 284 | +@click.argument("volume_id", shell_complete=_complete_ids) |
| 285 | +@click.option("--yes", "-y", is_flag=True, help="Skip the confirmation prompt.") |
| 286 | +@global_option |
| 287 | +def volume_delete(volume_id, yes, kind): |
| 288 | + """Delete a volume and everything stored on it.""" |
| 289 | + path = f"/v2/{kind}-volumes/{volume_id}" |
| 290 | + with api_errors(): |
| 291 | + if not yes: |
| 292 | + record = run_rest_request("GET", path) |
| 293 | + confirm( |
| 294 | + f"delete {kind} volume {volume_id} ({_label(record)}) and all its data?", |
| 295 | + yes, |
| 296 | + ) |
| 297 | + run_rest_request("DELETE", path) |
| 298 | + ui.success(f"deleted {kind} volume [white]{volume_id}[/white]") |
0 commit comments