main.tf

terraform {
  required_version = ">= 1.6"
  required_providers {
    aws = {
      source  = "hashicorp/aws"
      version = "~> 5.60"
    }
  }
}

provider "aws" {
  region = var.region
  default_tags {
    tags = {
      Project = var.project_name
    }
  }
}

# --- Network: use the region's default VPC and all its subnets (one per AZ) ---
data "aws_vpc" "default" {
  default = true
}

data "aws_subnets" "default" {
  filter {
    name   = "vpc-id"
    values = [data.aws_vpc.default.id]
  }
}

# --- SSH key: public half of the operator key generated by `make key` ---
resource "aws_key_pair" "operator" {
  key_name   = "${var.project_name}-key"
  public_key = file("${path.module}/${local.public_key_path}")
}

# --- Security group: inbound SSH from your IP only; EGRESS DENIED ---
# No egress rule => all outbound denied. The grind host has no channel to
# exfiltrate the secret. We pull via inbound SSH; SG is stateful so the reply
# traffic on our established connection is allowed regardless of egress rules.
resource "aws_security_group" "grind" {
  name        = var.project_name
  description = "Vanity grind: inbound SSH from operator only, no egress."
  vpc_id      = data.aws_vpc.default.id

  ingress {
    description = "SSH from operator"
    from_port   = 22
    to_port     = 22
    protocol    = "tcp"
    cidr_blocks = [var.ssh_cidr]
  }

  # Intentionally NO egress block => zero outbound.
}

locals {
  # Default key path derives from project_name (variable defaults cannot reference other variables).
  public_key_path = coalesce(var.public_key_path, "${var.project_name}-key.pub")

  # vanity flags: --prefix/--suffix as set, plus --count so each node keeps grinding
  # (no node goes idle after its first find) until the fleet total reaches the target.
  vanity_args = join(" ", concat(
    var.prefix != "" ? ["--prefix", var.prefix] : [],
    var.suffix != "" ? ["--suffix", var.suffix] : [],
    ["--count", tostring(var.num_keys)],
    var.case_insensitive ? ["--case-insensitive"] : [],
  ))

  user_data = base64encode(templatefile("${path.module}/userdata.sh.tftpl", {
    vanity_args = local.vanity_args
  }))

  # One override per (instance_type x subnet); weighted_capacity = GPUs of that type.
  # sort() the subnet ids: the data source does NOT guarantee a stable order, and an
  # unstable override order makes aws_ec2_fleet see a diff and REPLACE the fleet
  # (killing a running grind). Sorting makes re-applies idempotent.
  overrides = [
    for pair in setproduct(var.instance_types, sort(data.aws_subnets.default.ids)) : {
      instance_type     = pair[0]
      subnet_id         = pair[1]
      weighted_capacity = lookup(var.gpus_per_type, pair[0], 1)
    }
  ]
}

resource "aws_launch_template" "grind" {
  name_prefix   = "${var.project_name}-"
  image_id      = var.ami_id
  key_name      = aws_key_pair.operator.key_name
  user_data     = local.user_data
  ebs_optimized = true
  # Top-level SG (not a network_interfaces block) so per-override subnet_id is
  # allowed. Default-VPC subnets auto-assign a public IP, which the poller needs.
  vpc_security_group_ids = [aws_security_group.grind.id]

  tag_specifications {
    resource_type = "instance"
    tags          = { Project = var.project_name, Name = var.project_name }
  }

  lifecycle {
    precondition {
      condition     = var.prefix != "" || var.suffix != ""
      error_message = "Set at least one of `prefix` or `suffix`."
    }
  }
}

# --- EC2 Fleet (spot, maintain): fills gpu_target from whichever pools have stock ---
resource "aws_ec2_fleet" "grind" {
  type                                = "maintain"
  terminate_instances                 = true
  terminate_instances_with_expiration = true
  replace_unhealthy_instances         = true

  # Credential-free kill switch: at valid_until AWS terminates the whole fleet even
  # if your SSO session expired and `tofu destroy` can no longer run. Empty = none.
  valid_until = var.valid_until != "" ? var.valid_until : null

  launch_template_config {
    launch_template_specification {
      launch_template_id = aws_launch_template.grind.id
      version            = aws_launch_template.grind.latest_version
    }

    dynamic "override" {
      for_each = local.overrides
      content {
        instance_type     = override.value.instance_type
        subnet_id         = override.value.subnet_id
        weighted_capacity = override.value.weighted_capacity
      }
    }
  }

  spot_options {
    allocation_strategy            = var.allocation_strategy
    instance_interruption_behavior = "terminate"
  }

  target_capacity_specification {
    default_target_capacity_type = "spot"
    total_target_capacity        = var.gpu_target
    spot_target_capacity         = var.gpu_target
  }
}