misc.python.materialize.scratch

  1# Copyright Materialize, Inc. and contributors. All rights reserved.
  2#
  3# Use of this software is governed by the Business Source License
  4# included in the LICENSE file at the root of this repository.
  5#
  6# As of the Change Date specified in that file, in accordance with
  7# the Business Source License, use of this software will be governed
  8# by the Apache License, Version 2.0.
  9
 10"""Utilities for launching and interacting with scratch EC2 instances."""
 11
 12import asyncio
 13import csv
 14import datetime
 15import os
 16import re
 17import shlex
 18import subprocess
 19import sys
 20from subprocess import CalledProcessError
 21from typing import NamedTuple, cast
 22
 23import boto3
 24from botocore.exceptions import ClientError
 25from mypy_boto3_ec2.literals import InstanceTypeType
 26from mypy_boto3_ec2.service_resource import Instance
 27from mypy_boto3_ec2.type_defs import (
 28    FilterTypeDef,
 29    InstanceNetworkInterfaceSpecificationTypeDef,
 30    InstanceTypeDef,
 31    RunInstancesRequestServiceResourceCreateInstancesTypeDef,
 32)
 33from pydantic import BaseModel
 34
 35from materialize import MZ_ROOT, git, spawn, ui, util
 36
 37# Sane defaults for internal Materialize use in the scratch account
 38DEFAULT_SECURITY_GROUP_NAME = "scratch-security-group"
 39DEFAULT_INSTANCE_PROFILE_NAME = "admin-instance"
 40
 41# Ubuntu 26.04 LTS (Resolute Raccoon) base images, keyed by the CPU
 42# architecture that EC2 reports for an instance type.
 43#
 44# To update: open https://cloud-images.ubuntu.com/locator/ec2/, filter to the
 45# desired release, zone `us-east-1`, and instance type `hvm:ebs-ssd-gp3`. Copy
 46# the AMI id for each architecture (amd64 -> x86_64, arm64) below. Bumping the
 47# release also means updating the comment above and the `ami_user` default if
 48# Canonical changes it.
 49AMIS_BY_ARCH: dict[str, str] = {
 50    "x86_64": "ami-0b6d9d3d33ba97d99",
 51    "arm64": "ami-0bc7f2dbdcc6b5303",
 52}
 53
 54SSH_COMMAND = ["mssh", "-o", "StrictHostKeyChecking=off"]
 55SFTP_COMMAND = ["msftp", "-o", "StrictHostKeyChecking=off"]
 56
 57say = ui.speaker("scratch> ")
 58
 59
 60def instance_type_arch(instance_type: str) -> str:
 61    """Return the CPU architecture for an EC2 instance type as a key into
 62    `AMIS_BY_ARCH`. Queries EC2, which also validates that the type exists."""
 63    infos = boto3.client("ec2").describe_instance_types(
 64        InstanceTypes=[cast(InstanceTypeType, instance_type)]
 65    )["InstanceTypes"]
 66    if not infos:
 67        raise RuntimeError(f"unknown EC2 instance type: {instance_type}")
 68    archs = infos[0]["ProcessorInfo"]["SupportedArchitectures"]
 69    for arch in archs:
 70        if arch in AMIS_BY_ARCH:
 71            return arch
 72    raise RuntimeError(
 73        f"no scratch AMI for instance type {instance_type} "
 74        f"(architectures: {', '.join(archs)})"
 75    )
 76
 77
 78def resolve_ami(desc: "MachineDesc") -> str:
 79    """Resolve the AMI for a machine, deriving it from the instance type's
 80    architecture unless the config pins one explicitly."""
 81    if desc.ami:
 82        return desc.ami
 83    return AMIS_BY_ARCH[instance_type_arch(desc.instance_type)]
 84
 85
 86def tags(i: Instance) -> dict[str, str]:
 87    if not i.tags:
 88        return {}
 89    return {t["Key"]: t["Value"] for t in i.tags}
 90
 91
 92def instance_typedef_tags(i: InstanceTypeDef) -> dict[str, str]:
 93    return {t["Key"]: t["Value"] for t in i.get("Tags", [])}
 94
 95
 96def name(tags: dict[str, str]) -> str | None:
 97    return tags.get("Name")
 98
 99
100def launched_by(tags: dict[str, str]) -> str | None:
101    return tags.get("LaunchedBy")
102
103
104def ami_user(tags: dict[str, str]) -> str | None:
105    return tags.get("ami-user", "ubuntu")
106
107
108def delete_after(tags: dict[str, str]) -> datetime.datetime | None:
109    unix = tags.get("scratch-delete-after")
110    if not unix:
111        return None
112    unix = int(float(unix))
113    return datetime.datetime.fromtimestamp(unix)
114
115
116def instance_host(instance: Instance, user: str | None = None) -> str:
117    if user is None:
118        user = ami_user(tags(instance))
119    return f"{user}@{instance.id}"
120
121
122# Launch tags the instance Name as "{nonce}-{config name}"; the nonce is
123# redundant with the instance id in a flat listing, so strip it for display.
124_NONCE_PREFIX = re.compile(r"^[0-9a-f]{8}-")
125
126
127def _short_name(n: str | None) -> str:
128    if not n:
129        return "-"
130    return _NONCE_PREFIX.sub("", n, count=1)
131
132
133def render_table(headers: list[str], rows: list[list[str]]) -> None:
134    """Print a borderless, left-aligned table, like `docker ps` / `kubectl get`:
135    columns separated by two spaces, an uppercase header, no rules."""
136    widths = [max(len(r[c]) for r in [headers, *rows]) for c in range(len(headers))]
137    for r in [headers, *rows]:
138        print("  ".join(cell.ljust(w) for cell, w in zip(r, widths)).rstrip())
139
140
141def _format_expires(dt: datetime.datetime | None) -> str:
142    """Render a delete-after time as a compact relative duration."""
143    if dt is None:
144        return "-"
145    secs = (dt - datetime.datetime.now()).total_seconds()
146    if secs <= 0:
147        return "expired"
148    if secs >= 86400:
149        return f"{secs / 86400:.1f}d"
150    if secs >= 3600:
151        return f"{secs / 3600:.0f}h"
152    return f"{secs / 60:.0f}m"
153
154
155def print_instances(
156    ists: list[Instance],
157    format: str = "table",
158    numbered: bool = False,
159    show_launched_by: bool = False,
160) -> None:
161    if format == "csv":
162        # CSV is machine-readable: emit the full set of columns, unabridged.
163        field_names = [
164            "Name",
165            "Instance ID",
166            "Public IP Address",
167            "Private IP Address",
168            "Launched By",
169            "Delete After",
170            "State",
171        ]
172        if numbered:
173            field_names = ["#"] + field_names
174        w = csv.writer(sys.stdout)
175        w.writerow(field_names)
176        for idx, i in enumerate(ists, 1):
177            t = tags(i)
178            row = [
179                name(t),
180                i.instance_id,
181                i.public_ip_address,
182                i.private_ip_address,
183                launched_by(t),
184                delete_after(t),
185                i.state["Name"],
186            ]
187            if numbered:
188                row = [idx] + row
189            w.writerow(row)
190        return
191
192    if format != "table":
193        raise RuntimeError("Unknown format passed to print_instances")
194
195    # Table is for humans: a compact subset that fits a default terminal.
196    # Private IP is omitted (mssh connects by instance id, not IP) and the
197    # owner only shown when listing instances beyond your own.
198    headers = ["INSTANCE ID", "NAME", "STATE", "EXPIRES", "PUBLIC IP"]
199    if show_launched_by:
200        headers.append("LAUNCHED BY")
201    if numbered:
202        headers = ["#"] + headers
203    rows = []
204    for idx, i in enumerate(ists, 1):
205        t = tags(i)
206        state = i.state["Name"]
207        # A dead instance has no meaningful expiry; don't show a future time.
208        expires = (
209            "-"
210            if state in ("terminated", "shutting-down")
211            else _format_expires(delete_after(t))
212        )
213        row = [
214            i.instance_id,
215            _short_name(name(t)),
216            state,
217            expires,
218            i.public_ip_address or "-",
219        ]
220        if show_launched_by:
221            row.append(launched_by(t) or "-")
222        if numbered:
223            row = [str(idx)] + row
224        rows.append(row)
225
226    render_table(headers, rows)
227
228
229def mssh(
230    instance: Instance,
231    command: str,
232    *,
233    extra_ssh_args: list[str] | None = None,
234    input: bytes | None = None,
235    quiet: bool = False,
236) -> None:
237    """Runs a command over SSH via EC2 Instance Connect."""
238    extra_ssh_args = extra_ssh_args or []
239    host = instance_host(instance)
240    if command:
241        if not quiet:
242            print(f"{host}$ {command}", file=sys.stderr)
243        # Quote to work around:
244        # https://github.com/aws/aws-ec2-instance-connect-cli/pull/26
245        command = shlex.quote(command)
246    else:
247        print(f"$ mssh {host}")
248
249    result = subprocess.run(
250        [
251            *SSH_COMMAND,
252            *extra_ssh_args,
253            host,
254            command,
255        ],
256        input=input,
257        stdout=subprocess.DEVNULL if quiet else None,
258        stderr=subprocess.DEVNULL if quiet else None,
259    )
260    # Exit code 130 = SIGINT (Ctrl-C) — exit cleanly
261    if result.returncode == 130:
262        sys.exit(130)
263    if result.returncode != 0:
264        raise subprocess.CalledProcessError(result.returncode, result.args)
265
266
267def msftp(
268    instance: Instance,
269) -> None:
270    """Connects over SFTP via EC2 Instance Connect."""
271    host = instance_host(instance)
272    spawn.runv([*SFTP_COMMAND, host])
273
274
275def setup_ai_tools(instance: Instance) -> None:
276    """Transfer local Claude Code and Codex configs to a scratch instance
277    and install the materialize-docs skill."""
278    import io
279    import pathlib
280    import tarfile
281
282    home = pathlib.Path.home()
283
284    # Collect only essential config files — skip session logs, caches, etc.
285    items: list[tuple[str, pathlib.Path]] = []
286
287    # Claude Code: settings, credentials, and project memory files
288    claude_dir = home / ".claude"
289    if claude_dir.is_dir():
290        for cfg_name in ("settings.json", ".credentials.json"):
291            p = claude_dir / cfg_name
292            if p.is_file():
293                items.append((f".claude/{cfg_name}", p))
294        # Transfer memory directories from all projects
295        projects_dir = claude_dir / "projects"
296        if projects_dir.is_dir():
297            for proj in projects_dir.iterdir():
298                if not proj.is_dir():
299                    continue
300                memory_dir = proj / "memory"
301                if memory_dir.is_dir():
302                    items.append((f".claude/projects/{proj.name}/memory", memory_dir))
303
304    if (home / ".claude.json").is_file():
305        items.append((".claude.json", home / ".claude.json"))
306
307    # Codex: only config, not session logs
308    codex_dir = home / ".codex"
309    if codex_dir.is_dir():
310        for cfg_name in ("config.json", "instructions.md"):
311            p = codex_dir / cfg_name
312            if p.is_file():
313                items.append((f".codex/{cfg_name}", p))
314
315    if items:
316        buf = io.BytesIO()
317        with tarfile.open(fileobj=buf, mode="w:gz") as tar:
318            for arcname, path in items:
319                tar.add(str(path), arcname=arcname)
320        mssh(instance, "tar xzf - -C ~", input=buf.getvalue(), quiet=True)
321
322    mssh(
323        instance,
324        "cd materialize && "
325        "NPM_CONFIG_UPDATE_NOTIFIER=false "
326        "npx -q -y skills add MaterializeInc/agent-skills -g -a claude-code --copy -y --skill materialize-docs 2>/dev/null || true",
327        quiet=True,
328    )
329
330    mssh(
331        instance,
332        "cd materialize && "
333        "NPM_CONFIG_UPDATE_NOTIFIER=false "
334        "npx -q -y skills add MaterializeInc/agent-skills -g -a codex --copy -y --skill materialize-docs 2>/dev/null || true",
335        quiet=True,
336    )
337
338
339def mkrepo(
340    instance: Instance, rev: str, init: bool = True, force: bool = False
341) -> None:
342    if init:
343        mssh(instance, "git clone https://github.com/MaterializeInc/materialize.git")
344
345    rev = git.rev_parse(rev)
346
347    cmd: list[str] = [
348        "git",
349        "push",
350        "--no-verify",
351        f"{instance_host(instance)}:materialize/.git",
352        # Explicit refspec is required if the host repository is in detached
353        # HEAD mode.
354        f"{rev}:refs/heads/scratch",
355    ]
356    if force:
357        cmd.append("--force")
358
359    spawn.runv(
360        cmd,
361        cwd=MZ_ROOT,
362        env=dict(os.environ, GIT_SSH_COMMAND=" ".join(SSH_COMMAND)),
363    )
364    git_config_cmds = f"git checkout -f {rev}"
365
366    # Propagate local git user.name and user.email to the scratch instance.
367    if git_name := git.get_user_name():
368        git_config_cmds += f" && git config user.name {shlex.quote(git_name)}"
369    if git_email := git.get_user_email():
370        git_config_cmds += f" && git config user.email {shlex.quote(git_email)}"
371
372    mssh(
373        instance,
374        f"cd materialize && {git_config_cmds}",
375    )
376
377
378class MachineDesc(BaseModel):
379    name: str
380    launch_script: str | None = None
381    instance_type: str
382    # Optional: when unset, the AMI is derived from the instance type's
383    # architecture via `resolve_ami`.
384    ami: str | None = None
385    tags: dict[str, str] = {}
386    size_gb: int = 50
387    checkout: bool = True
388    ami_user: str = "ubuntu"
389
390
391def launch(
392    *,
393    key_name: str | None,
394    instance_type: str,
395    ami: str,
396    ami_user: str,
397    tags: dict[str, str],
398    display_name: str | None = None,
399    size_gb: int,
400    security_group_name: str,
401    instance_profile: str | None,
402    nonce: str,
403    delete_after: datetime.datetime,
404) -> Instance:
405    """Launch and configure an ec2 instance with the given properties."""
406
407    if display_name:
408        tags["Name"] = display_name
409    tags["scratch-delete-after"] = str(delete_after.timestamp())
410    tags["nonce"] = nonce
411    tags["git_ref"] = git.describe()
412    tags["ami-user"] = ami_user
413
414    ec2 = boto3.client("ec2")
415    groups = ec2.describe_security_groups()
416    security_group_id = None
417    for group in groups["SecurityGroups"]:
418        if group["GroupName"] == security_group_name:
419            security_group_id = group["GroupId"]
420            break
421
422    if security_group_id is None:
423        vpcs = ec2.describe_vpcs()
424        vpc_id = None
425        for vpc in vpcs["Vpcs"]:
426            if vpc["IsDefault"]:
427                vpc_id = vpc["VpcId"]
428                break
429        if vpc_id is None:
430            default_vpc = ec2.create_default_vpc()
431            vpc_id = default_vpc["Vpc"]["VpcId"]
432        securitygroup = ec2.create_security_group(
433            GroupName=security_group_name,
434            Description="Allows all.",
435            VpcId=vpc_id,
436        )
437        security_group_id = securitygroup["GroupId"]
438        ec2.authorize_security_group_ingress(
439            GroupId=security_group_id,
440            CidrIp="0.0.0.0/0",
441            IpProtocol="tcp",
442            FromPort=22,
443            ToPort=22,
444        )
445
446    network_interface: InstanceNetworkInterfaceSpecificationTypeDef = {
447        "AssociatePublicIpAddress": True,
448        "DeviceIndex": 0,
449        "Groups": [security_group_id],
450    }
451
452    say(f"launching instance {display_name or '(unnamed)'}")
453    with open(MZ_ROOT / "misc" / "scratch" / "provision.bash") as f:
454        provisioning_script = f.read()
455    kwargs: RunInstancesRequestServiceResourceCreateInstancesTypeDef = {
456        "MinCount": 1,
457        "MaxCount": 1,
458        "ImageId": ami,
459        "InstanceType": cast(InstanceTypeType, instance_type),
460        "UserData": provisioning_script,
461        "TagSpecifications": [
462            {
463                "ResourceType": "instance",
464                "Tags": [{"Key": k, "Value": v} for (k, v) in tags.items()],
465            }
466        ],
467        "NetworkInterfaces": [network_interface],
468        "BlockDeviceMappings": [
469            {
470                "DeviceName": "/dev/sda1",
471                "Ebs": {
472                    "VolumeSize": size_gb,
473                    "VolumeType": "gp3",
474                },
475            }
476        ],
477        "MetadataOptions": {
478            # Allow Docker containers to access IMDSv2.
479            "HttpPutResponseHopLimit": 2,
480        },
481    }
482    if key_name:
483        kwargs["KeyName"] = key_name
484    if instance_profile:
485        kwargs["IamInstanceProfile"] = {"Name": instance_profile}
486    i = boto3.resource("ec2").create_instances(**kwargs)[0]
487
488    return i
489
490
491class CommandResult(NamedTuple):
492    status: str
493    stdout: str
494    stderr: str
495
496
497async def setup(
498    i: Instance,
499    git_rev: str,
500) -> None:
501    def is_ready(i: Instance) -> bool:
502        return bool(
503            i.public_ip_address and i.state and i.state.get("Name") == "running"
504        )
505
506    done = False
507    async for remaining in ui.async_timeout_loop(60, 2):
508        print(
509            f"\rscratch> Waiting for instance to become ready: {remaining:0.0f}s remaining\033[K",
510            end="",
511            flush=True,
512            file=sys.stderr,
513        )
514        try:
515            i.reload()
516            if is_ready(i):
517                done = True
518                break
519        except ClientError:
520            pass
521    print(file=sys.stderr)
522    if not done:
523        raise RuntimeError(
524            f"Instance {i} did not become ready in a reasonable amount of time"
525        )
526
527    # Wait for SSH to be available
528    done = False
529    async for remaining in ui.async_timeout_loop(120, 2):
530        print(
531            f"\rscratch> Waiting for SSH access: {remaining:0.0f}s remaining\033[K",
532            end="",
533            flush=True,
534            file=sys.stderr,
535        )
536        try:
537            mssh(i, "true", quiet=True)
538            done = True
539            break
540        except CalledProcessError:
541            continue
542    print(file=sys.stderr)
543    if not done:
544        raise RuntimeError(
545            "SSH did not become available in a reasonable amount of time"
546        )
547
548    # Start git clone while provisioning continues
549    say("Cloning repo (provisioning continues in background)...")
550    mkrepo(i, git_rev)
551
552    # Wait for provisioning to finish
553    done = False
554    async for remaining in ui.async_timeout_loop(300, 2):
555        print(
556            f"\rscratch> Waiting for provisioning to complete: {remaining:0.0f}s remaining\033[K",
557            end="",
558            flush=True,
559            file=sys.stderr,
560        )
561        try:
562            mssh(i, "[[ -f /opt/provision/done ]]", quiet=True)
563            done = True
564            break
565        except CalledProcessError:
566            continue
567    print(file=sys.stderr)
568    if not done:
569        raise RuntimeError(
570            "Instance did not finish setup in a reasonable amount of time"
571        )
572
573
574def launch_cluster(
575    descs: list[MachineDesc],
576    *,
577    nonce: str | None = None,
578    key_name: str | None = None,
579    security_group_name: str = DEFAULT_SECURITY_GROUP_NAME,
580    instance_profile: str | None = DEFAULT_INSTANCE_PROFILE_NAME,
581    extra_tags: dict[str, str] | None = None,
582    delete_after: datetime.datetime,
583    git_rev: str = "HEAD",
584    extra_env: dict[str, str] | None = None,
585) -> list[Instance]:
586    """Launch a cluster of instances with a given nonce"""
587
588    extra_tags = extra_tags or {}
589    extra_env = extra_env or {}
590
591    if not nonce:
592        nonce = util.nonce(8)
593
594    instances = [
595        launch(
596            key_name=key_name,
597            instance_type=d.instance_type,
598            ami=resolve_ami(d),
599            ami_user=d.ami_user,
600            tags={**d.tags, **extra_tags},
601            display_name=f"{nonce}-{d.name}",
602            size_gb=d.size_gb,
603            security_group_name=security_group_name,
604            instance_profile=instance_profile,
605            nonce=nonce,
606            delete_after=delete_after,
607        )
608        for d in descs
609    ]
610
611    async def setup_all() -> None:
612        await asyncio.gather(
613            *(
614                setup(i, git_rev if d.checkout else "HEAD")
615                for (i, d) in zip(instances, descs)
616            )
617        )
618
619    asyncio.run(setup_all())
620
621    for i in instances:
622        i.reload()
623
624    hosts_str = "".join(
625        f"{i.private_ip_address}\t{d.name}\n" for (i, d) in zip(instances, descs)
626    )
627    for i in instances:
628        mssh(i, "sudo tee -a /etc/hosts", input=hosts_str.encode())
629
630    for i in instances:
631        setup_ai_tools(i)
632
633    env = " ".join(f"{k}={shlex.quote(v)}" for k, v in extra_env.items())
634    for i, d in zip(instances, descs):
635        if d.launch_script:
636            mssh(
637                i,
638                f"(cd materialize && {env} nohup bash -c {shlex.quote(d.launch_script)}) &> mzscratch.log &",
639            )
640
641    return instances
642
643
644def whoami() -> str:
645    return boto3.client("sts").get_caller_identity()["UserId"].split(":")[1]
646
647
648def get_instance(name: str) -> Instance:
649    """
650    Get an instance by instance id. The special name 'mine' resolves to a
651    unique running owned instance, if there is one; otherwise the name is
652    assumed to be an instance id.
653    """
654    if name == "mine":
655        filters: list[FilterTypeDef] = [
656            {"Name": "tag:LaunchedBy", "Values": [whoami()]},
657            {"Name": "instance-state-name", "Values": ["pending", "running"]},
658        ]
659        instances = [i for i in boto3.resource("ec2").instances.filter(Filters=filters)]
660        if not instances:
661            raise RuntimeError("can't understand 'mine': no owned instance?")
662        if len(instances) > 1:
663            raise RuntimeError(
664                f"can't understand 'mine': too many owned instances ({', '.join(i.id for i in instances)})"
665            )
666        instance = instances[0]
667        say(f"understanding 'mine' as unique owned instance {instance.id}")
668        return instance
669    return boto3.resource("ec2").Instance(name)
670
671
672def list_instances(
673    owners: list[str] | None = None, all: bool = False
674) -> list[Instance]:
675    """List AWS instances, optionally filtered by owner."""
676    filters: list[FilterTypeDef] = []
677    if not all:
678        if not owners:
679            owners = [whoami()]
680        filters.append({"Name": "tag:LaunchedBy", "Values": owners})
681    return list(boto3.resource("ec2").instances.filter(Filters=filters))
682
683
684def get_old_instances() -> list[InstanceTypeDef]:
685    def exists(i: InstanceTypeDef) -> bool:
686        return i["State"]["Name"] != "terminated"
687
688    def is_old(i: InstanceTypeDef) -> bool:
689        delete_after = instance_typedef_tags(i).get("scratch-delete-after")
690        if delete_after is None:
691            return False
692        delete_after = float(delete_after)
693        return datetime.datetime.now(datetime.timezone.utc).timestamp() > delete_after
694
695    return [
696        i
697        for r in boto3.client("ec2").describe_instances()["Reservations"]
698        for i in r["Instances"]
699        if exists(i) and is_old(i)
700    ]