aws-bootstrap-g4dn 0.6.0__py3-none-any.whl → 0.8.0__py3-none-any.whl

aws_bootstrap/ec2.py

@@ -7,6 +7,7 @@ import botocore.exceptions
 import click
 
 from .config import LaunchConfig
+from .output import echo, is_text, secho
 
 
 EBS_DEVICE_NAME = "/dev/sdf"
@@ -81,7 +82,7 @@ def ensure_security_group(ec2_client, name: str, tag_value: str, ssh_port: int =
     if existing["SecurityGroups"]:
         sg_id = existing["SecurityGroups"][0]["GroupId"]
         msg = "  Security group " + click.style(f"'{name}'", fg="bright_white")
-        click.echo(msg + f" already exists ({sg_id}), reusing.")
+        echo(msg + f" already exists ({sg_id}), reusing.")
         return sg_id
 
     # Create new SG
@@ -113,7 +114,7 @@ def ensure_security_group(ec2_client, name: str, tag_value: str, ssh_port: int =
             }
         ],
     )
-    click.secho(f"  Created security group '{name}' ({sg_id}) with SSH ingress.", fg="green")
+    secho(f"  Created security group '{name}' ({sg_id}) with SSH ingress.", fg="green")
     return sg_id
 
 
@@ -163,8 +164,8 @@ def launch_instance(ec2_client, config: LaunchConfig, ami_id: str, sg_id: str) -
         if code in ("MaxSpotInstanceCountExceeded", "VcpuLimitExceeded"):
             _raise_quota_error(code, config)
         elif code in ("InsufficientInstanceCapacity", "SpotMaxPriceTooLow") and config.spot:
-            click.secho(f"\n  Spot request failed: {e.response['Error']['Message']}", fg="yellow")
-            if click.confirm("  Retry as on-demand instance?"):
+            secho(f"\n  Spot request failed: {e.response['Error']['Message']}", fg="yellow")
+            if not is_text() or click.confirm("  Retry as on-demand instance?"):
                 launch_params.pop("InstanceMarketOptions", None)
                 try:
                     response = ec2_client.run_instances(**launch_params)
@@ -329,15 +330,15 @@ def terminate_tagged_instances(ec2_client, instance_ids: list[str]) -> list[dict
 
 def wait_instance_ready(ec2_client, instance_id: str) -> dict:
     """Wait for the instance to be running and pass status checks."""
-    click.echo("  Waiting for instance " + click.style(instance_id, fg="bright_white") + " to enter 'running' state...")
+    echo("  Waiting for instance " + click.style(instance_id, fg="bright_white") + " to enter 'running' state...")
     waiter = ec2_client.get_waiter("instance_running")
     waiter.wait(InstanceIds=[instance_id], WaiterConfig={"Delay": 10, "MaxAttempts": 60})
-    click.secho("  Instance running.", fg="green")
+    secho("  Instance running.", fg="green")
 
-    click.echo("  Waiting for instance status checks to pass...")
+    echo("  Waiting for instance status checks to pass...")
     waiter = ec2_client.get_waiter("instance_status_ok")
     waiter.wait(InstanceIds=[instance_id], WaiterConfig={"Delay": 15, "MaxAttempts": 60})
-    click.secho("  Status checks passed.", fg="green")
+    secho("  Status checks passed.", fg="green")
 
     # Refresh instance info to get public IP
     desc = ec2_client.describe_instances(InstanceIds=[instance_id])
@@ -467,3 +468,39 @@ def find_ebs_volumes_for_instance(ec2_client, instance_id: str, tag_value: str)
             }
         )
     return volumes
+
+
+def find_orphan_ebs_volumes(ec2_client, tag_value: str, live_instance_ids: set[str]) -> list[dict]:
+    """Find aws-bootstrap EBS volumes whose linked instance no longer exists.
+
+    Only returns volumes in ``available`` state (not attached to any instance).
+    Volumes that are ``in-use`` are never considered orphans, even if their
+    tagged instance ID is not in *live_instance_ids*.
+
+    Returns a list of dicts with VolumeId, Size, State, and InstanceId
+    (the instance ID from the ``aws-bootstrap-instance`` tag).
+    """
+    try:
+        response = ec2_client.describe_volumes(
+            Filters=[
+                {"Name": "tag:created-by", "Values": [tag_value]},
+                {"Name": "status", "Values": ["available"]},
+            ]
+        )
+    except botocore.exceptions.ClientError:
+        return []
+
+    orphans = []
+    for vol in response.get("Volumes", []):
+        tags = {t["Key"]: t["Value"] for t in vol.get("Tags", [])}
+        linked_instance = tags.get("aws-bootstrap-instance", "")
+        if linked_instance and linked_instance not in live_instance_ids:
+            orphans.append(
+                {
+                    "VolumeId": vol["VolumeId"],
+                    "Size": vol["Size"],
+                    "State": vol["State"],
+                    "InstanceId": linked_instance,
+                }
+            )
+    return orphans

aws_bootstrap/output.py

@@ -0,0 +1,106 @@
+"""Output formatting for structured CLI output (JSON, YAML, table, text)."""
+
+from __future__ import annotations
+import json
+from datetime import datetime
+from enum import StrEnum
+from pathlib import Path
+from typing import Any
+
+import click
+
+
+class OutputFormat(StrEnum):
+    TEXT = "text"
+    JSON = "json"
+    YAML = "yaml"
+    TABLE = "table"
+
+
+def get_format(ctx: click.Context | None = None) -> OutputFormat:
+    """Return the current output format from the click context."""
+    if ctx is None:
+        ctx = click.get_current_context(silent=True)
+    if ctx is None or ctx.obj is None:
+        return OutputFormat.TEXT
+    return ctx.obj.get("output_format", OutputFormat.TEXT)
+
+
+def is_text(ctx: click.Context | None = None) -> bool:
+    """Return True if the current output format is text (default)."""
+    return get_format(ctx) == OutputFormat.TEXT
+
+
+def _default_serializer(obj: Any) -> Any:
+    """JSON serializer for objects not serializable by default."""
+    if isinstance(obj, datetime):
+        return obj.isoformat()
+    if isinstance(obj, Path):
+        return str(obj)
+    raise TypeError(f"Object of type {type(obj).__name__} is not JSON serializable")
+
+
+def emit(data: dict | list, *, headers: dict[str, str] | None = None, ctx: click.Context | None = None) -> None:
+    """Emit structured data in the configured output format.
+
+    For JSON/YAML: serializes the data directly.
+    For TABLE: renders using tabulate. If *data* is a list of dicts, uses
+    *headers* mapping ``{dict_key: column_label}`` for column selection/ordering.
+    If *data* is a single dict, renders as key-value pairs.
+    """
+    fmt = get_format(ctx)
+
+    if fmt == OutputFormat.JSON:
+        click.echo(json.dumps(data, indent=2, default=_default_serializer))
+        return
+
+    if fmt == OutputFormat.YAML:
+        import yaml  # noqa: PLC0415
+
+        # Convert datetime/Path objects before YAML dump
+        prepared = json.loads(json.dumps(data, default=_default_serializer))
+        click.echo(yaml.dump(prepared, default_flow_style=False, sort_keys=False).rstrip())
+        return
+
+    if fmt == OutputFormat.TABLE:
+        from tabulate import tabulate  # noqa: PLC0415
+
+        table_data = data
+        # Unwrap dict-wrapped lists (e.g. {"instances": [...]}) for table rendering
+        if isinstance(data, dict) and headers:
+            for v in data.values():
+                if isinstance(v, list):
+                    table_data = v
+                    break
+
+        if isinstance(table_data, list) and table_data and isinstance(table_data[0], dict):
+            if headers:
+                keys = list(headers.keys())
+                col_labels = list(headers.values())
+                rows = [[row.get(k, "") for k in keys] for row in table_data]
+            else:
+                col_labels = list(table_data[0].keys())
+                keys = col_labels
+                rows = [[row.get(k, "") for k in keys] for row in table_data]
+            click.echo(tabulate(rows, headers=col_labels, tablefmt="simple"))
+        elif isinstance(table_data, dict):
+            rows = [[k, v] for k, v in table_data.items()]
+            click.echo(tabulate(rows, headers=["Key", "Value"], tablefmt="simple"))
+        elif isinstance(table_data, list):
+            # Empty list
+            click.echo("(no data)")
+        return
+
+    # TEXT format: emit() is a no-op in text mode (text output is handled inline)
+
+
+def echo(msg: str = "", **kwargs: Any) -> None:
+    """Wrap ``click.echo``; silent in non-text output modes."""
+    if is_text():
+        click.echo(msg, **kwargs)
+
+
+def secho(msg: str = "", **kwargs: Any) -> None:
+    """Wrap ``click.secho``; silent in non-text output modes."""
+    if is_text():
+        click.secho(msg, **kwargs)

aws_bootstrap/ssh.py

@@ -13,6 +13,7 @@ from pathlib import Path
 import click
 
 from .gpu import _GPU_ARCHITECTURES, GpuInfo
+from .output import echo, secho
 
 
 # ---------------------------------------------------------------------------
@@ -54,7 +55,7 @@ def import_key_pair(ec2_client, key_name: str, key_path: Path) -> str:
     # Check if key pair already exists
     try:
         existing = ec2_client.describe_key_pairs(KeyNames=[key_name])
-        click.echo("  Key pair " + click.style(f"'{key_name}'", fg="bright_white") + " already exists, reusing.")
+        echo("  Key pair " + click.style(f"'{key_name}'", fg="bright_white") + " already exists, reusing.")
         return existing["KeyPairs"][0]["KeyName"]
     except ec2_client.exceptions.ClientError as e:
         if "InvalidKeyPair.NotFound" not in str(e):
@@ -70,7 +71,7 @@ def import_key_pair(ec2_client, key_name: str, key_path: Path) -> str:
             }
         ],
     )
-    click.secho(f"  Imported key pair '{key_name}' from {key_path}", fg="green")
+    secho(f"  Imported key pair '{key_name}' from {key_path}", fg="green")
     return key_name
 
 
@@ -88,7 +89,7 @@ def wait_for_ssh(host: str, user: str, key_path: Path, retries: int = 30, delay:
             sock = socket.create_connection((host, port), timeout=5)
             sock.close()
         except (TimeoutError, ConnectionRefusedError, OSError):
-            click.echo("  SSH not ready " + click.style(f"(attempt {attempt}/{retries})", dim=True) + ", waiting...")
+            echo("  SSH not ready " + click.style(f"(attempt {attempt}/{retries})", dim=True) + ", waiting...")
             time.sleep(delay)
             continue
 
@@ -106,10 +107,10 @@ def wait_for_ssh(host: str, user: str, key_path: Path, retries: int = 30, delay:
         ]
         result = subprocess.run(cmd, capture_output=True, text=True)
         if result.returncode == 0:
-            click.secho("  SSH connection established.", fg="green")
+            secho("  SSH connection established.", fg="green")
             return True
 
-        click.echo("  SSH not ready " + click.style(f"(attempt {attempt}/{retries})", dim=True) + ", waiting...")
+        echo("  SSH not ready " + click.style(f"(attempt {attempt}/{retries})", dim=True) + ", waiting...")
         time.sleep(delay)
 
     return False
@@ -125,89 +126,89 @@ def run_remote_setup(
     requirements_path = script_path.parent / "requirements.txt"
 
     # SCP the requirements file
-    click.echo("  Uploading requirements.txt...")
+    echo("  Uploading requirements.txt...")
     req_result = subprocess.run(
         ["scp", *ssh_opts, *scp_port_opts, str(requirements_path), f"{user}@{host}:/tmp/requirements.txt"],
         capture_output=True,
         text=True,
     )
     if req_result.returncode != 0:
-        click.secho(f"  SCP failed: {req_result.stderr}", fg="red", err=True)
+        secho(f"  SCP failed: {req_result.stderr}", fg="red", err=True)
         return False
 
     # SCP the GPU benchmark script
     benchmark_path = script_path.parent / "gpu_benchmark.py"
-    click.echo("  Uploading gpu_benchmark.py...")
+    echo("  Uploading gpu_benchmark.py...")
     bench_result = subprocess.run(
         ["scp", *ssh_opts, *scp_port_opts, str(benchmark_path), f"{user}@{host}:/tmp/gpu_benchmark.py"],
         capture_output=True,
         text=True,
     )
     if bench_result.returncode != 0:
-        click.secho(f"  SCP failed: {bench_result.stderr}", fg="red", err=True)
+        secho(f"  SCP failed: {bench_result.stderr}", fg="red", err=True)
         return False
 
     # SCP the GPU smoke test notebook
     notebook_path = script_path.parent / "gpu_smoke_test.ipynb"
-    click.echo("  Uploading gpu_smoke_test.ipynb...")
+    echo("  Uploading gpu_smoke_test.ipynb...")
     nb_result = subprocess.run(
         ["scp", *ssh_opts, *scp_port_opts, str(notebook_path), f"{user}@{host}:/tmp/gpu_smoke_test.ipynb"],
         capture_output=True,
         text=True,
     )
     if nb_result.returncode != 0:
-        click.secho(f"  SCP failed: {nb_result.stderr}", fg="red", err=True)
+        secho(f"  SCP failed: {nb_result.stderr}", fg="red", err=True)
         return False
 
     # SCP the CUDA example source
     saxpy_path = script_path.parent / "saxpy.cu"
-    click.echo("  Uploading saxpy.cu...")
+    echo("  Uploading saxpy.cu...")
     saxpy_result = subprocess.run(
         ["scp", *ssh_opts, *scp_port_opts, str(saxpy_path), f"{user}@{host}:/tmp/saxpy.cu"],
         capture_output=True,
         text=True,
     )
     if saxpy_result.returncode != 0:
-        click.secho(f"  SCP failed: {saxpy_result.stderr}", fg="red", err=True)
+        secho(f"  SCP failed: {saxpy_result.stderr}", fg="red", err=True)
         return False
 
     # SCP the VSCode launch.json
     launch_json_path = script_path.parent / "launch.json"
-    click.echo("  Uploading launch.json...")
+    echo("  Uploading launch.json...")
     launch_result = subprocess.run(
         ["scp", *ssh_opts, *scp_port_opts, str(launch_json_path), f"{user}@{host}:/tmp/launch.json"],
         capture_output=True,
         text=True,
     )
     if launch_result.returncode != 0:
-        click.secho(f"  SCP failed: {launch_result.stderr}", fg="red", err=True)
+        secho(f"  SCP failed: {launch_result.stderr}", fg="red", err=True)
         return False
 
     # SCP the VSCode tasks.json
     tasks_json_path = script_path.parent / "tasks.json"
-    click.echo("  Uploading tasks.json...")
+    echo("  Uploading tasks.json...")
     tasks_result = subprocess.run(
         ["scp", *ssh_opts, *scp_port_opts, str(tasks_json_path), f"{user}@{host}:/tmp/tasks.json"],
         capture_output=True,
         text=True,
     )
     if tasks_result.returncode != 0:
-        click.secho(f"  SCP failed: {tasks_result.stderr}", fg="red", err=True)
+        secho(f"  SCP failed: {tasks_result.stderr}", fg="red", err=True)
         return False
 
     # SCP the script
-    click.echo("  Uploading remote_setup.sh...")
+    echo("  Uploading remote_setup.sh...")
     scp_result = subprocess.run(
         ["scp", *ssh_opts, *scp_port_opts, str(script_path), f"{user}@{host}:/tmp/remote_setup.sh"],
         capture_output=True,
         text=True,
     )
     if scp_result.returncode != 0:
-        click.secho(f"  SCP failed: {scp_result.stderr}", fg="red", err=True)
+        secho(f"  SCP failed: {scp_result.stderr}", fg="red", err=True)
         return False
 
     # Execute the script, passing PYTHON_VERSION as an inline env var if specified
-    click.echo("  Running remote_setup.sh on instance...")
+    echo("  Running remote_setup.sh on instance...")
     remote_cmd = "chmod +x /tmp/remote_setup.sh && "
     if python_version:
         remote_cmd += f"PYTHON_VERSION={python_version} "