I've already posted twice about using AI/LLMs, and buckle up for another exciting installment. But! Before I get into the meat of the topic, I need to lay some groundwork.

The company I work for uses Amazon Web Services (AWS) for compute and storage. Within AWS you're very much encouraged to set up a Virtual Private Cloud (VPC) which is a way of isolating your cloud infrastructure away from the internet and other AWS users. Part of the VPC is a firewall that prevents outside access to the Virtual Private Network. Of course, as an employee, I need to be able to get inside the firewall, so we have a VPN Gateway that allows outside users on the public internet to access resources inside the VPC, but only if you can successfully authenticate with the gateway.

Below is a diagram showing this. My laptop connects to the VPN gateway, and from the gateway I can access a Linux machine. For example, if I want to ssh into a Linux box, that connection gets wrapped via the VPN connection to the gateway, which then unwraps it and passes the ssh connection to the Linux box.

                 ┌─────────────────────────────────────────┐
                 │                VPC                      │
                 │                                         │
┌─────────┐      │ ┌──────────────┐       ┌──────────────┐ │
│ Laptop  ├──────┼►│ VPN Gateway  ├──────►│  Linux Box   │ │
└─────────┘      │ └──────────────┘  ssh  └──────────────┘ │
         VPN(ssh)└─────────────────────────────────────────┘

For reasons that are not important here, we are overhauling how we do various things with respect to security. One part of this is we are turning off the ability to ssh into Linux machines in favor of AWS SSM. SSM gives you an encrypted shell/terminal connection to a host, just like ssh, but it doesn't use ssh at all. One motivation is that AWS SSM allows convenient logging of who connected to a particular machine, when they did, and what shell commands they ran. Another is that access to resources can be defined at the AWS permissions level, away from the individual Linux machines. With SSM we don't need a gateway to access things inside the VPN. It works like this: When the Linux box boots, an agent on the machine registers itself with SSM, which lives outside the VPN (it's run by AWS, not us). Next, when I wish to connect to the box, my laptop authenticates with SSM, which checks if I'm allowed to connect. If I am, SSM connects me to the machine inside the VPN by forwarding data between my laptop and the Linux box. At the same time, logs are kept of the connection and commands I run. See the diagram below.

                    ┌─────────────┐  ┌─────────────┐
      ┌────────────►│   AWS SSM   ├─►│  AWS Logs   │
      │             └─────────▲───┘  └─────────────┘
      │                       │                     
┌─────┴─────┐                 │                     
│  Laptop   │    ┌────────────┼───────┐             
└───────────┘    │   VPC      │       │             
                 │  ┌─────────┴────┐  │             
                 │  │  Linux Box   │  │             
                 │  └──────────────┘  │             
                 └────────────────────┘             

The official way to connect is via the AWS CLI tool, which is written in Python. Somewhat oddly, the command to connect, aws ssm start-session --target instance-id --profile somename --region us-east-1, actually hands off to the Go binary session-manager-plugin after authentication is completed. This all works, but it wasn't satisfactory. I found it clunky and inconvenient — just compare the long aws ssm command above to the much shorter command ssh user@host-ip!

Much of that clunky command can be scripted away, but the official client has some broken functionality. For example, I don't use the nano text editor often, but it's useful for quick edits (I also use vi, if you're wondering). It uses "Ctrl-key" actions, such as searching ("Ctrl-w"), saving ("Ctrl-o"), and quitting ("Ctrl-x"). There may be something wrong on my end I never diagnosed, but I couldn't get any of those "Ctrl-key" commands to work. Additionally, ssh has its companion program scp that transfers files between computers. You can imagine that this is a useful feature, and it's something that SSM doesn't include.

Because of this I decided to see what an LLM (Claude) could do. I asked it to write a program that worked with AWS SSM but acted more like ssh. I named it sshaws. I told it to write it in Rust because Rust is fast, secure, and I figured that all the required crates would already exist. To give away the ending, the LLM did a fine job! The program doesn't have any of the odd quirks (no "Ctrl-key" problems) that aws ssm start-session does. But more than that, I tailored it to work well with my workflow, and it's now more convenient than ssh ever was, amazingly. Here are a few examples of neat features it has:

• I had it add a Terminal User Interface (TUI) that auto discovers all the Linux boxes I have access to (see the screenshot below¹). It shows if the instance is running, or if I can access the instance (perhaps the SSM agent isn't set up on the Linux machine). Other metadata can also shown, such as instance type, IP address, or environment. I can filter the instances list with a simple text search. The instance listing is especially useful for me because much of my workflow uses temporary Linux machines that exist for minutes to hours, so being able to quickly narrow down to one is great.

• There is a "recent" hosts pane that is useful for reconnecting to a host I've worked with recently (bottom left). Instances I have connected to that have been deleted are shown with strike through text, while those simply powered off use greyed out text.

• The application has a config file that allows me to define quick shortcuts for machines I access frequently. These shortcuts are available in a TUI pane (top left), and also from the command line, e.g. sshaws webhost.

• Unlike VPN/ssh connections, SSM connections can time out after as little as six hours due to AWS authentication rules. Having a session die while using it is very annoying, so I added a feature where sshaws will auto-connect to a running tmux session, or start one if needed. This way, if you're cut off, what you were doing on the host is uninterrupted, and you can pick up where you left off by reconnecting.

• To replace scp I had the LLM add flags to transfer files: sshaws put file.txt webhost:/home/user/ or sshaws get webhost:/home/user/file.txt .. But even better, I had it add a second TUI (see below) that allows me to browse my local machine and remote, and mark files and directories for transfer either direction. It auto compresses and decompresses files and directories transparently to speed up transfers.

• I added a command each which will run a command on more than one host in parallel and return the result, for example sshaws each 'boxname*' 'uptime' will print out the result of running the uptime command on all hosts that match boxname*, e.g. boxname01, boxname02, and boxname-prod would all match.

Why was the LLM so successful at this? The answer is simple: LLMs are advanced prediction engines, like Google autocomplete, but fancier. If you tell it where to start and where to go, you're giving it a path to follow, and that's when they work best. In this case I asked it to put to things together (SSM & ssh) that it knows about because it's all open source, and what functionality I wanted. It didn't have to invent anything new, it just had to mix two things together by following my instructions.

Am I okay with using LLM like this? Yes, yes I am. My job is not writing tooling for accessing AWS SSM. As I said above, nothing new is being created here; existing things were just smooshed together. My job as a human is to create new things with my biological creativity, and offloading this task to an LLM was a win.

Finally, I don't know if I will try to make sshaws open source. There's nothing proprietary about it, so I could. I'm not super worried about major issues. I've passed it through Claude Fable², which has been reported as being excellent at finding code bugs, and it fixed a few minor ones. While the tool might be useful for people outside my company, I would feel more comfortable about releasing it if I and a few other humans went through it and audited it for hallucinations and confusing comments, Fable's findings notwithstanding.

1. All the screenshots use fake data ↩︎
2. I mainly used Claude Opus to write the code because it's cheaper than Fable ↩︎

One of my primary responsibilities at my current job is ownership of the ETL pipeline that brings in the data upon which we run our business. Every day it processes hundreds of gigabytes of data, cleaning and normalizing it, and outputting it in several different forms.

For a few years I have been using a Hadoop-based mrjob pipeline on top of AWS EMR. It replaced a (I kid you not) Redshift-based pipeline that was hugely expensive (that I didn't write). It's been very reliable. For the last year or two the only failures have been when AWS's systems have had issues, something I can't do anything about. Despite this reliability, it hasn't been perfect. The biggest problem with it is that it's expensive to run. There are few reasons why it's been so expensive:

• EMR adds a 25% upcharge on all resources used. It is reasonable for AWS to charge something because EMR is a useful service with added value. However, in my opinion, 25% is too high for what it does. AWS already gets paid plenty because you're using their other services underneath it (mainly EC2 virtual computers and EBS block storage), so the extra mark-up feels like a big cash-grab. It is, of course, entirely possible to run Hadoop on AWS without using EMR. But it is a hassle, and it's likely that AWS has figured out that 25% is the inflection point between too expensive and not worth the hassle

• Hadoop isn't the most efficient way of doing things. Newer tools, notably Spark, have surpassed Hadoop in both speed and features. I originally used Hadoop because I wasn't happy with how Spark needed quite a bit more memory than Hadoop for a similar operation, but over time I became less and less satisfied with my inability to speed up certain parts of the process

• AWS offers spot compute instances, which are virtual computers that are significantly cheaper than on-demand instances. The difference is that on-demand nodes are yours as long as you want them, while spot instances can be taken away at any time with only two minutes warning.

  One of Hadoop's killer features is that during a MapReduce cycle, if one or more worker nodes goes away (for whatever reason including spot node removal), in most cases it can recover and redo any lost results.

  However, the EMR pipeline used a multi-step MapReduce process. Unfortunately, Hadoop cannot recover lost results from earlier, fully completed MapReduce cycles. This means that in order to run the pipeline, it had to have enough on-demand instances that the data that needed to be preserved across cycles fit on them. This raised the cost considerably when compared to an all-spot pipeline run

Earlier this year I decided it was time to start looking at how to rewrite the pipeline to dramatically lower costs. I saw two possible ways forward:

• Choose a modern, high performance tool like Spark/PySpark or Dask that would still run on EMR but hopefully would be much faster

• Abandon EMR entirely (and its 25% surcharge) and write something that could run completely on spot instances and use S3 for storage, which is four to five times cheaper than EBS

After some thought, I decided that the second option was the better choice. If I could figure it out, it offered the best possible outcome. The pipeline runs once per day, meaning that it has 24 hours to finish before the next run needs to start. Ultimately, high performance was less important than lowering costs.

I have been using Polars quite a bit and have been (mostly¹) impressed with its speed and functionality. It is written in Rust, which is one of the fastest programming languages. Polars has a Python-facing API as a first-class member of the project. Rust has an ever-growing library of packages that I've found are high-quality and well documented (in contrast to my experience with cough R packages). The crucial difference between Polars and Hadoop/Spark/Dask is that Polars runs on only one node at a time (it can and does use all the CPU cores), while all of the others can run on multiple nodes. If I could figure out how to slice up the pipeline into chunks that would work on separate instances, I believed I could use Polars in place of Hadoop.

Jumping to the end of the story, I was able to convert the pipeline to use Polars, and to great success. I use a simple pattern for each step. An orchestration process builds a list of work which is submitted to SQS. An EC2 spot fleet is created which launches workers that consume the work. The input and output of the work is stored on S3. The workers send success or failure messages to a callback queue on SQS, which is monitored by the orchestration process. If a spot node goes away interrupting work, the work will be picked up by a different node once the message returns to availability. Once the work is done for a step (i.e. all work has generated a callback), the orchestration process kills the spot fleet and continues to the next step (or sends an error message for a human to figure out).

The bottom line is that the cost has dropped by roughly 85%, primarily due to the following reasons:

• Polars has the concept of LazyFrames which are data objects that are not realized in memory nor computation until Polars is told to do so. Operations and filtering can be applied to them and Polars can do the work in parallel with efficiency tricks that overall increases speed without loading the whole dataset into memory at once. The combination of sink_parquet with PartitionByKey is effectively a MapReduce operation that is much faster than Hadoop on similar hardware

• AWS has "regions" and "availability zones (AZs)" which are the physical locations where cloud compute happens. Each AZ is a distinct data center (or close by data centers) within a region. When running an EMR job, you are restricted to a single AZ largely because AWS charges for cross-AZ data transfer, and EMR jobs are very loquacious across the network. There's also increased network latency between AZs. Running and EMR job in multiple AZs would hugely impact performance and cost.

  Because the new pipeline reads from and saves data to S3, and there are no cross-AZ charges for accessing S3 within a region, it doesn't matter which AZ the workers run in. This means that the spot fleets can target all AZs within the region, unlike EMR

• When launching an EC2 fleet, you must specify one or more launch templates, which describe how to launch each instance in terms of OS and installed software. AWS EC2 offers instances using x86 processors from Intel and AMD, and ARM instances using AWS Graviton processors. Conveniently, the pipeline doesn't require any processor-specific features. Therefore, I created two AMIs, one for each of x86 and ARM, which allows the spot fleet to target any and all of Intel, AMD, and Graviton instances

• The pipeline requirements for each step basically comes down to number of CPU cores and amount of RAM, more or less of each depending on what the step is doing. The upshot of all of the above is that for a given step all the pipeline cares about is the resources of the node, not what kind of node it is. Of course, not all instances are the same speed, but the cost of an instance is roughly proportional with its speed, so it all works out. This means that for a given step, across all AZs and EC2 instance types, there can be over 100 distinct resource combinations to pick from. This basically guarantees spot availability at all times

• The pipeline uses a fair number of User Defined Functions. Polars supports UDFs written in Python, Numba, and Rust using PyO3. By using the latter two, basically all of the inner loops and heavy computation in the pipeline happens in compiled C or Rust. This, in my opinion, is a really nice way of doing things. Let Python handle moving data around and high-level stuff, and run all the heavy computation in compiled code.

Overall, I'm very pleased about the results of this work. The goal was to save money, and it has done that. I wasn't expecting 85% savings (I'm not sure what I was hoping for), but I feel quite good about that.

1. Polars definitely has some frustrating issues, like this one or this one ↩︎

Amazon Web Services (AWS) offers dozens (if not over one hundred) different services. I probably use about a dozen of them regularly, including Elastic Map Reduce (EMR) which is their platform for running big-data things like Hadoop and Spark.

EMR runs your work on EC2 instances, and you can pick which kind(s) you want when the job starts. You can also pick the "lifecycle" of these instances. This means you can pick some instances to run as "on-demand" where the instance is yours (barring a hardware failure), and other instances to run as "spot" which costs much less than on-demand but AWS can take away the instance at any time.

Luckily, Hadoop and Spark are designed to work on unreliable hardware, and if previously done work is unavailable (e.g. because the instance that did it is no longer running), Hadoop and Spark can re-run that work again. This means that you can use a mix of on-demand and spot instances that, as long as AWS doesn't take away too many spot instances, will run the job for lower cost than otherwise.

A big issue with running Hadoop on spot instances is that multi-stage Hadoop jobs save some data between stages that can't be redone. This data is stored in HDFS, which is where Hadoop stores (semi)permanent data. Because we don't want this data going away, we need to run HDFS on on-demand instances, and not run it on spot instances. Hadoop handles this by having two kind of "worker" instances: "CORE" instances that run HDFS and have the important data, and "TASK" types that do not run HDFS and store easily reproduced data. Both types share in the computational workload, the difference is what kind of data is allowed to be stored on them. It makes sense, then, to confine "TASK" instances to spot nodes.

The trick is to configure Hadoop such that the instances themselves know what kind of instance they are. Figuring this out was harder than it should have been because AWS EMR doesn't auto-configure nodes to work this way; the user needs to configure the job to do this including running scripts on the instances themselves.

I like to run my Hadoop jobs using mrjob which makes development and running Hadoop with Python easier. I assume that this can be done outside of mrjob, but its up to the reader to figure out how to do that.

There are three parts to this. The first two are two Python scripts that are run on the EC2 instances, and the third is modifying the mrjob configuration file. The Python scripts should be uploaded to a S3 bucket because they will be downloaded to each Hadoop node (see the bootstrap actions below).

With the changes below, you should be able to run a multi-step Hadoop job on AWS EMR using spot nodes and not lose any intermediate work. Good luck!

make_node_labels.py

This script tells yarn what kind of instance types are available. This only needs to run once.

#!/usr/bin/python3
import subprocess
import time

def run(cmd):
    proc = subprocess.Popen(cmd,
        stdout = subprocess.PIPE,
        stderr = subprocess.PIPE,
    )
    stdout, stderr = proc.communicate()

    return proc.returncode, stdout, stderr

if __name__ == '__main__':

    # Wait for the yarn stuff to be installed
    code, out, err = run(['which', 'yarn'])
    while code == 1:
        time.sleep(5)
        code, out, err = run(['which', 'yarn'])

    # Now we wait for things to be configured
    time.sleep(60)

    # Now set the node label types
    code, out, err = run(["yarn",
        "rmadmin",
        "-addToClusterNodeLabels",
        '"CORE(exclusive=false),TASK(exclusive=false)"'])

get_node_label.py

This script tells Hadoop what kind of instance this is. It is called by Hadoop and run as many times as needed.

#!/usr/bin/python3
import json
k='/mnt/var/lib/info/extraInstanceData.json'
with open(k) as f:
    response = json.load(f)
    print("NODE_PARTITION:", response['instanceRole'].upper())

mrjob.conf

This is not a complete mrjob configuration file. It shows the essential parts needed for setting up CORE/TASK nodes. You will need to fill in the rest for your specific situation.

runners:
  emr:

    instance_fleets:
    - InstanceFleetType: MASTER
      TargetOnDemandCapacity: 1
      InstanceTypeConfigs:
      - InstanceType: (smallish instance type)
        WeightedCapacity: 1
    - InstanceFleetType: CORE
      # Some nodes are launched on-demand which prevents the whole job from
      # dying if spot nodes are yanked
      TargetOnDemandCapacity: NNN (count of on-demand cores)
      InstanceTypeConfigs:
      - InstanceType: (bigger instance type)
        BidPriceAsPercentageOfOnDemandPrice: 100
        WeightedCapacity: (core count per instance)
    - InstanceFleetType: TASK
      # TASK means no HDFS is stored so loss of a node won't lose data
      # that can't be recovered relatively easily
      TargetOnDemandCapacity: 0
      TargetSpotCapacity: MMM (count of spot cores)
      LaunchSpecifications:
        SpotSpecification:
          TimeoutDurationMinutes: 60
          TimeoutAction: SWITCH_TO_ON_DEMAND
      InstanceTypeConfigs:
      - InstanceType: (bigger instance type)
        BidPriceAsPercentageOfOnDemandPrice: 100
        WeightedCapacity: (core count per instance)
      - InstanceType: (alternative instance type)
        BidPriceAsPercentageOfOnDemandPrice: 100
        WeightedCapacity: (core count per instance)

    bootstrap:
    # Download the Python scripts to the instance
    - /usr/bin/aws s3 cp s3://bucket/get_node_label.py /home/hadoop/
    - chmod a+x /home/hadoop/get_node_label.py
    - /usr/bin/aws s3 cp s3://bucket/make_node_labels.py /home/hadoop/
    - chmod a+x /home/hadoop/make_node_labels.py
    # nohup runs this until it quits on its own
    - nohup /home/hadoop/make_node_labels.py &

    emr_configurations:
      - Classification: yarn-site
        Properties:
          yarn.node-labels.enabled: true
          yarn.node-labels.am.default-node-label-expression: CORE
          yarn.node-labels.configuration-type: distributed
          yarn.nodemanager.node-labels.provider: script
          yarn.nodemanager.node-labels.provider.script.path: /home/hadoop/get_node_label.py