lablink_allocator_service.providers.aws

lablink_allocator_service.providers.aws

AWSProvider — behavior-preserving wrapper over existing AWS utilities.

Classes

AWSProvider(*, region=None, tofu_dir=None, **_)

Source code in packages/allocator/src/lablink_allocator_service/providers/aws.py

def __init__(self, *, region=None, tofu_dir=None, **_):
    self._region = region
    self._tofu_dir = tofu_dir
    self.client_connectivity = AllocatorProxiedClientConnectivity()

Attributes

can_destroy_hosts = True class-attribute instance-attribute

can_provision_hosts = True class-attribute instance-attribute

can_recover_hosts = True class-attribute instance-attribute

client_connectivity = AllocatorProxiedClientConnectivity() instance-attribute

name = 'aws' class-attribute instance-attribute

Methods:

destroy_hosts(handles, progress_callback=None)

Plan, then apply, a full-workspace destroy. handles is accepted for protocol consistency but not used for per-handle filtering — OpenTofu destroy operates against the whole state file. Mirrors provision_hosts's plan+show+apply sequence (added so resources_total can be computed exactly the same way provision_hosts computes it for creates) instead of the single tofu destroy call this method used before.

Raises FileNotFoundError if no runtime tfvars exists (signals "no client VMs were ever launched" — route handler maps this to 404).

Source code in packages/allocator/src/lablink_allocator_service/providers/aws.py

def destroy_hosts(
    self,
    handles: list[ClientHandle],
    progress_callback: Optional[Callable[[int, int], None]] = None,
) -> DestroyResult:
    """Plan, then apply, a full-workspace destroy. `handles` is
    accepted for protocol consistency but not used for per-handle
    filtering — OpenTofu destroy operates against the whole state
    file. Mirrors provision_hosts's plan+show+apply sequence (added
    so resources_total can be computed exactly the same way
    provision_hosts computes it for creates) instead of the single
    `tofu destroy` call this method used before.

    Raises FileNotFoundError if no runtime tfvars exists (signals
    "no client VMs were ever launched" — route handler maps this to 404).
    """
    if self._tofu_dir is None:
        raise RuntimeError(
            "AWSProvider not configured with tofu_dir — cannot destroy."
        )
    tofu_dir = Path(self._tofu_dir)
    runtime_file = tofu_dir / "terraform.runtime.tfvars"
    if not runtime_file.exists():
        raise FileNotFoundError(
            "tfvars does not exist — no client VMs were launched"
        )

    var_args = ["-var-file=terraform.runtime.tfvars"]
    try:
        sg_id = current_instance_security_group(region=self._region)
        var_args.append(f"-var=allocator_sg_id={sg_id}")
    except NotOnEC2Error:
        # Caller may log; we just skip the SG var.
        pass

    plan_file = "tfplan-destroy.binary"
    plan_file_path = tofu_dir / plan_file
    try:
        subprocess.run(
            ["tofu", "plan", "-destroy", "-no-color",
             "-out", plan_file, *var_args],
            cwd=tofu_dir, check=True, capture_output=True, text=True,
        )
        show = subprocess.run(
            ["tofu", "show", "-json", plan_file],
            cwd=tofu_dir, check=True, capture_output=True, text=True,
        )
        plan_json = json.loads(show.stdout)
        resources_total = sum(
            1
            for rc in plan_json.get("resource_changes", [])
            if "delete" in (rc.get("change") or {}).get("actions", [])
        )
        if progress_callback:
            progress_callback(0, resources_total)

        completed = 0

        def _on_resource_complete() -> None:
            nonlocal completed
            completed += 1
            if progress_callback:
                progress_callback(completed, resources_total)

        result = _run_streamed(
            ["tofu", "apply", "-auto-approve", plan_file],
            cwd=tofu_dir,
            resource_complete_re=_DESTROY_COMPLETE_RE,
            on_resource_complete=_on_resource_complete,
        )
    finally:
        plan_file_path.unlink(missing_ok=True)

    return DestroyResult(stdout=strip_ansi(result.stdout))

get_host_access(hostname)

Return (instance_id, public_ip, ssh_key_path) for hostname.

Looks up the EC2 instance by Name tag, fetches its public IP, and reads the SSH private key from the OpenTofu state directory. Any component may be None if unavailable (instance not found, no public IP, or tofu_dir not set / key file absent).

Source code in packages/allocator/src/lablink_allocator_service/providers/aws.py

def get_host_access(
    self, hostname: str
) -> tuple[str | None, str | None, str | None]:
    """Return (instance_id, public_ip, ssh_key_path) for *hostname*.

    Looks up the EC2 instance by Name tag, fetches its public IP, and
    reads the SSH private key from the OpenTofu state directory.  Any
    component may be None if unavailable (instance not found, no public
    IP, or tofu_dir not set / key file absent).
    """
    instance_id = get_instance_id_by_name(hostname, region=self._region)
    if not instance_id:
        return (None, None, None)

    ip = get_instance_public_ip(instance_id, region=self._region)

    key_path: str | None = None
    if self._tofu_dir:
        try:
            key_path = get_ssh_private_key(self._tofu_dir)
        except Exception:
            pass

    return (instance_id, ip, key_path)

list_hosts()

Source code in packages/allocator/src/lablink_allocator_service/providers/aws.py

def list_hosts(self) -> list[ClientHandle]:
    ids = get_instance_ids(tofu_dir=self._tofu_dir)
    names = get_instance_names(tofu_dir=self._tofu_dir)
    return [
        ClientHandle(id=i, hostname=n, provider_metadata={"region": self._region})
        for i, n in zip(ids, names)
    ]

provision_hosts(count, spec, progress_callback=None)

Run tofu plan + audit + apply for count new client hosts.

Moves the inline logic that used to live in main.py's /api/launch handler behind the provider seam (SR-F1). spec is a dict of runtime values that used to be assembled inline in the route.

Raises:

Type	Description
RuntimeError	if tofu_dir is None
SGAuditFailure	propagated from audit_tofu_plan when the plan would expose :6080 / :7070 to the internet
CalledProcessError	propagated from any tofu invocation failure

Source code in packages/allocator/src/lablink_allocator_service/providers/aws.py

def provision_hosts(
    self,
    count: int,
    spec: dict,
    progress_callback: Optional[Callable[[int, int], None]] = None,
) -> ProvisionResult:
    """Run `tofu plan + audit + apply` for `count` new client hosts.

    Moves the inline logic that used to live in main.py's /api/launch
    handler behind the provider seam (SR-F1). `spec` is a dict of
    runtime values that used to be assembled inline in the route.

    Raises:
        RuntimeError: if tofu_dir is None
        SGAuditFailure: propagated from audit_tofu_plan when
            the plan would expose :6080 / :7070 to the internet
        subprocess.CalledProcessError: propagated from any
            tofu invocation failure
    """
    if self._tofu_dir is None:
        raise RuntimeError(
            "AWSProvider not configured with tofu_dir — cannot provision."
        )
    tofu_dir = Path(self._tofu_dir)
    runtime_file = tofu_dir / "terraform.runtime.tfvars"

    # GPU detection (moved from main.py)
    gpu_support_bool = check_support_nvidia(
        machine_type=spec["machine_type"]
    )
    gpu_support = "true" if gpu_support_bool else "false"

    # Write runtime tfvars (moved verbatim from main.py)
    with runtime_file.open("w") as f:
        f.write(f'allocator_ip = "{spec["allocator_ip"]}"\n')
        f.write(f'allocator_url = "{spec["allocator_url"]}"\n')
        f.write(f'machine_type = "{spec["machine_type"]}"\n')
        f.write(f'image_name = "{spec["image_name"]}"\n')
        # cfg.machine.repository is Optional and defaults to None; an
        # f-string over Python None writes the literal string "None",
        # which reaches client start.sh as TUTORIAL_REPO_TO_CLONE="None"
        # and makes every VM boot attempt `git clone None` (and log its
        # failure). Empty string is the "no repo" value start.sh skips.
        f.write(f'repository = "{spec["repository"] or ""}"\n')
        f.write(f'client_ami_id = "{spec["client_ami_id"]}"\n')
        f.write(f'subject_software = "{spec["subject_software"]}"\n')
        f.write(f'resource_prefix = "{spec["resource_prefix"]}"\n')
        f.write(f'gpu_support = "{gpu_support}"\n')
        f.write(
            f'cloud_init_output_log_group = '
            f'"{spec["cloud_init_output_log_group"]}"\n'
        )
        f.write(f'region = "{self._region}"\n')
        f.write(f'startup_on_error = "{spec["startup_on_error"]}"\n')
        f.write(f'startup_max_attempts = {spec["startup_max_attempts"]}\n')
        f.write(
            f'startup_base_delay_seconds = '
            f'{spec["startup_base_delay_seconds"]}\n'
        )
        f.write(
            f'startup_success_check_b64 = '
            f'"{spec["startup_success_check_b64"]}"\n'
        )
        f.write(f'agent_token = "{spec["agent_token"]}"\n')
        f.write(f'register_token = "{spec["register_token"]}"\n')

    tf_vars = [
        "-var-file=terraform.runtime.tfvars",
        f"-var=instance_count={count}",
    ]
    try:
        sg_id = current_instance_security_group(region=self._region)
        tf_vars.append(f"-var=allocator_sg_id={sg_id}")
    except NotOnEC2Error:
        # Caller may log; we just skip the SG var.
        pass

    # Plan + audit + apply sequence (verbatim from main.py:542-602)
    plan_file = "tfplan.binary"
    plan_file_path = tofu_dir / plan_file
    try:
        subprocess.run(
            ["tofu", "plan", "-no-color", "-out", plan_file, *tf_vars],
            cwd=tofu_dir, check=True, capture_output=True, text=True,
        )
        show = subprocess.run(
            ["tofu", "show", "-json", plan_file],
            cwd=tofu_dir, check=True, capture_output=True, text=True,
        )
        plan_json = json.loads(show.stdout)
        audit_tofu_plan(plan_json)  # may raise SGAuditFailure

        resources_total = sum(
            1
            for rc in plan_json.get("resource_changes", [])
            if "create" in (rc.get("change") or {}).get("actions", [])
        )
        if progress_callback:
            progress_callback(0, resources_total)

        completed = 0

        def _on_resource_complete() -> None:
            nonlocal completed
            completed += 1
            if progress_callback:
                progress_callback(completed, resources_total)

        apply_result = _run_streamed(
            ["tofu", "apply", "-auto-approve", plan_file],
            cwd=tofu_dir,
            resource_complete_re=_CREATE_COMPLETE_RE,
            on_resource_complete=_on_resource_complete,
        )
    finally:
        plan_file_path.unlink(missing_ok=True)

    clean_stdout = strip_ansi(apply_result.stdout)

    # Upload runtime tfvars to S3 (moved from main.py:614-620)
    upload_to_s3(
        local_path=runtime_file,
        env=spec["environment"],
        bucket_name=spec["bucket_name"],
        region=self._region,
        deployment_name=spec.get("deployment_name", "lablink"),
    )

    # Read back the freshly-created instances + timings
    ids = get_instance_ids(tofu_dir=str(tofu_dir))
    names = get_instance_names(tofu_dir=str(tofu_dir))
    timings = get_instance_timings(tofu_dir=str(tofu_dir))

    handles = [
        ClientHandle(
            id=i, hostname=n,
            provider_metadata={"region": self._region},
        )
        for i, n in zip(ids, names)
    ]
    return ProvisionResult(
        handles=handles, timings=timings, apply_stdout=clean_stdout,
    )

recover_hosts(handles)

Source code in packages/allocator/src/lablink_allocator_service/providers/aws.py

def recover_hosts(self, handles: list[ClientHandle]) -> bool:
    # Verbatim of reboot.py's EC2 fallback: stop_start_ec2_instance.
    # Returns True iff EVERY recycle succeeded — the caller
    # (reboot.py) uses this to decide record_reboot vs error log,
    # so failure must NOT masquerade as success.
    all_ok = True
    for h in handles:
        region = h.provider_metadata.get("region", self._region)
        if not stop_start_ec2_instance(h.id, region=region):
            all_ok = False
    return all_ok

Functions: