Add SSH private key support for connecting to Lambda instances

The Docker container running the action needs the SSH private key to connect
to Lambda instances. Added:
- `ssh_private_key` input to action.yml
- `LAMBDA_SSH_PRIVATE_KEY` secret to runner.yml workflow
- Code to write key to temp file and use `-i` flag in SSH commands
- Better error logging for Lambda API errors (log response body)
- Reverted e2e-test.yml to use `gpu_1x_a10` in `us-east-1`

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

Author: ryan-williams

.github/workflows/e2e-test.yml

@@ -8,8 +8,8 @@ jobs:
     uses: ./.github/workflows/runner.yml
     secrets: inherit
     with:
-      instance_type: gpu_1x_a100_sxm4
-      region: us-west-2
+      instance_type: gpu_1x_a10
+      region: us-east-1
       runner_grace_period: "30"
       runner_initial_grace_period: "120"
       debug: "true"

.github/workflows/runner.yml

@@ -16,6 +16,9 @@ on:
       LAMBDA_API_KEY:
         description: "Lambda Labs API key"
         required: true
+      LAMBDA_SSH_PRIVATE_KEY:
+        description: "SSH private key for connecting to Lambda instances"
+        required: true
     inputs:
       action_ref:
         description: "lambda-gha Git ref (branch/tag/SHA) to checkout"
@@ -121,6 +124,7 @@ jobs:
           runner_poll_interval: ${{ inputs.runner_poll_interval || vars.RUNNER_POLL_INTERVAL }}
           runner_registration_timeout: ${{ inputs.runner_registration_timeout || vars.RUNNER_REGISTRATION_TIMEOUT }}
           ssh_key_names: ${{ inputs.ssh_key_names || vars.LAMBDA_SSH_KEY_NAMES }}
+          ssh_private_key: ${{ secrets.LAMBDA_SSH_PRIVATE_KEY }}
           userdata: ${{ inputs.userdata }}
         env:
           GH_PAT: ${{ secrets.GH_SA_TOKEN }}

action.yml

@@ -48,6 +48,9 @@ inputs:
   ssh_key_names:
     description: "SSH key names registered in Lambda Labs (comma-separated)"
     required: false
+  ssh_private_key:
+    description: "SSH private key for connecting to Lambda instances"
+    required: false
   userdata:
     description: "Additional script to run before runner setup"
     required: false

src/lambda_gha/__main__.py

@@ -56,6 +56,7 @@ def main():
         .update_state("INPUT_RUNNER_INITIAL_GRACE_PERIOD", "runner_initial_grace_period")
         .update_state("INPUT_RUNNER_POLL_INTERVAL", "runner_poll_interval")
         .update_state("INPUT_SSH_KEY_NAMES", "ssh_key_names")
+        .update_state("INPUT_SSH_PRIVATE_KEY", "ssh_private_key")
         .update_state("INPUT_USERDATA", "userdata")
         .update_state("GITHUB_REPOSITORY", "repo")
         .update_state("INPUT_REPO", "repo")

src/lambda_gha/start.py

@@ -93,6 +93,7 @@ class StartLambdaLabs:
     runner_initial_grace_period: str = "180"
     runner_poll_interval: str = "10"
     runner_release: str = ""
+    ssh_private_key: str = ""
     userdata: str = ""
 
     def _api_request(
@@ -106,7 +107,14 @@ def _api_request(
         headers = {"Authorization": f"Bearer {self.api_key}"}
 
         resp = requests.request(method, url, headers=headers, json=json_data)
-        resp.raise_for_status()
+        if not resp.ok:
+            # Log the actual error body before raising
+            try:
+                error_body = resp.json()
+                print(f"Lambda API error: {error_body}")
+            except Exception:
+                print(f"Lambda API error (raw): {resp.text}")
+            resp.raise_for_status()
         return resp.json()
 
     def _get_template_vars(self, idx: int = None) -> dict:
@@ -341,15 +349,32 @@ def execute_setup_via_ssh(
         retry_delay : int
             Seconds between retry attempts.
         """
+        import os
+        import stat
+        import tempfile
+
         print(f"Connecting to {ssh_user}@{ip} to execute setup...")
 
+        # Write SSH private key to temporary file if provided
+        key_file = None
+        if self.ssh_private_key:
+            key_file = tempfile.NamedTemporaryFile(mode='w', suffix='_key', delete=False)
+            key_file.write(self.ssh_private_key)
+            if not self.ssh_private_key.endswith('\n'):
+                key_file.write('\n')
+            key_file.close()
+            os.chmod(key_file.name, stat.S_IRUSR)  # 0400
+            print(f"Using SSH key from secret")
+
         # SSH options for non-interactive, key-based auth
         ssh_opts = [
             "-o", "StrictHostKeyChecking=no",
             "-o", "UserKnownHostsFile=/dev/null",
             "-o", "ConnectTimeout=10",
             "-o", "BatchMode=yes",
         ]
+        if key_file:
+            ssh_opts.extend(["-i", key_file.name])
 
         # Wait for SSH to be available
         for attempt in range(1, max_retries + 1):