{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e957187d-9efd-41ac-a96f-5c80582bfba5",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-01-26T13:23:50.184464Z",
     "iopub.status.busy": "2026-01-26T13:23:50.183876Z",
     "iopub.status.idle": "2026-01-26T13:23:50.192578Z",
     "shell.execute_reply": "2026-01-26T13:23:50.191034Z"
    }
   },
   "outputs": [],
   "source": [
    "import subprocess\n",
    "from concurrent.futures import ThreadPoolExecutor, as_completed\n",
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "b4ad0906-3be6-4361-a2eb-975bf10e3b39",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-01-26T13:23:50.198097Z",
     "iopub.status.busy": "2026-01-26T13:23:50.197572Z",
     "iopub.status.idle": "2026-01-26T13:23:50.204772Z",
     "shell.execute_reply": "2026-01-26T13:23:50.203171Z"
    }
   },
   "outputs": [],
   "source": [
    "# task_name = \"var-copy\"\n",
    "# task_name = \"assoc-recall-mk\"\n",
    "# task_name = \"decode-recall\"\n",
    "# task_name = \"decode-recall-last\"\n",
    "task_name = \"needle\"\n",
    "\n",
    "mixed = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "9857d742-ffcd-4ad8-a16c-f3cf7ed3ad35",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-01-26T13:23:50.211018Z",
     "iopub.status.busy": "2026-01-26T13:23:50.210463Z",
     "iopub.status.idle": "2026-01-26T13:23:50.229232Z",
     "shell.execute_reply": "2026-01-26T13:23:50.228666Z"
    }
   },
   "outputs": [],
   "source": [
    "# For variable copy\n",
    "if task_name == \"var-copy\":\n",
    "    experiment1_params = {'layer1': ['TF', 'SSM'], 'layer2': ['TF', 'SSM'], 'd': [12, 24], 'window': [5, 10, 15, 20, 30, 50, 100], 'nh': [1], 'sd': [1]}\n",
    "    experiment2_params = {'layer1': ['TF', 'SSM'], 'layer2': ['TF', 'SSM'], 'd': [4, 8, 12, 16, 20, 24, 32, 48], 'window': [20, 100], 'nh': [1], 'sd': [1]}\n",
    "    experiment3_params = {'layer1': ['TF', 'SSM'], 'layer2': ['TF', 'SSM'], 'd': [8, 24], 'window': [20, 100], 'nh': [2,4], 'sd': [1]}\n",
    "    experiment4_params = {'layer1': ['TF', 'SSM'], 'layer2': ['TF', 'SSM'], 'd': [8, 24], 'window': [20, 100], 'nh': [1], 'sd': [2,4,6,12]}\n",
    "    # experiment1_params = {'layer1': ['TF', 'SSM'], 'layer2': ['TF', 'SSM'], 'd': [8, 24], 'window': [5, 10, 15, 20, 30, 50], 'nh': [1], 'sd': [1]}\n",
    "    # experiment2_params = {'layer1': ['TF', 'SSM'], 'layer2': ['TF', 'SSM'], 'd': [4, 8, 12, 16, 20, 24], 'window': [20], 'nh': [1], 'sd': [1]}\n",
    "    # experiment3_params = {'layer1': ['TF', 'SSM'], 'layer2': ['TF', 'SSM'], 'd': [8, 24], 'window': [20], 'nh': [2,4], 'sd': [1]}\n",
    "    # experiment4_params = {'layer1': ['TF', 'SSM'], 'layer2': ['TF', 'SSM'], 'd': [8, 24], 'window': [20], 'nh': [1], 'sd': [2,4,6,12]}\n",
    "    # experiments = [experiment1_params, experiment2_params, experiment3_params, experiment4_params]\n",
    "    experiments = [experiment2_params]\n",
    "\n",
    "if task_name == \"needle\":\n",
    "    # experiment1_params = {'layer1': ['TF', 'SSM'], 'layer2': ['TF', 'SSM'], 'd': [12, 24], 'window': [5, 10, 15, 20, 30, 50, 100], 'nh': [1], 'sd': [1]}\n",
    "    experiment2_params = {'layer1': ['TF', 'SSM'], 'layer2': ['TF', 'SSM'], 'd': [4, 8, 12, 16, 20, 24, 32, 48], 'window': [20, 100], 'nh': [1], 'sd': [1]}\n",
    "    # experiment3_params = {'layer1': ['TF', 'SSM'], 'layer2': ['TF', 'SSM'], 'd': [8, 24], 'window': [20, 100], 'nh': [2,4], 'sd': [1]}\n",
    "    # experiment4_params = {'layer1': ['TF', 'SSM'], 'layer2': ['TF', 'SSM'], 'd': [8, 24], 'window': [20, 100], 'nh': [1], 'sd': [2,4,6,12]}\n",
    "    experiments = [experiment2_params]\n",
    "\n",
    "# For multi-key associative recall\n",
    "if task_name == \"assoc-recall-mk\":\n",
    "    experiment1_params = {'layer1': ['TF', 'SSM'], 'layer2': ['TF', 'SSM'], 'd': [16, 24], 'window': [5, 10, 15, 20, 30, 50, 100], 'nh': [1], 'sd': [1]}\n",
    "    experiment2_params = {'layer1': ['TF', 'SSM'], 'layer2': ['TF', 'SSM'], 'd': [4, 8, 12, 16, 20, 24], 'window': [20, 100], 'nh': [1], 'sd': [1]}\n",
    "    experiment3_params = {'layer1': ['TF', 'SSM'], 'layer2': ['TF', 'SSM'], 'd': [16, 24], 'window': [20, 100], 'nh': [2,4], 'sd': [1]}\n",
    "    experiment4_params = {'layer1': ['TF', 'SSM'], 'layer2': ['TF', 'SSM'], 'd': [16, 24], 'window': [20, 100], 'nh': [1], 'sd': [2,4,6,12]}\n",
    "    # experiment5_params = {'layer1': ['TF-nC'], 'layer2': ['TF-nC'], 'd': [2, 4, 6, 8, 10, 12, 15, 20, 24], 'window': [20, 100], 'nh': [1], 'sd': [1]}\n",
    "    experiment5_params = {'model': ['TF', 'SSM', 'hybrid'], 'num_layers': [1], 'd': [8, 24], 'window': [20, 100], 'nh': [1], 'sd': [1]}\n",
    "    # experiment5_params = {'model': ['TF', 'SSM', 'hybrid'], 'num_layers': [1, 2, 3, 4, 5], 'd': [8, 24], 'window': [20, 100], 'nh': [1], 'sd': [1]}\n",
    "    # experiments = [experiment1_params, experiment2_params, experiment3_params, experiment4_params]\n",
    "    experiments = [experiment5_params]\n",
    "\n",
    "# For multi-key binary decoding + recall\n",
    "if task_name == \"decode-recall\" or task_name == \"decode-recall-last\":\n",
    "    experiment1_params = {'layer1': ['TF', 'SSM'], 'layer2': ['TF', 'SSM'], 'd': [16, 24], 'window': [5, 10, 15, 20, 30, 50, 100], 'nh': [1], 'sd': [1]}\n",
    "    experiment2_params = {'layer1': ['TF', 'SSM'], 'layer2': ['TF', 'SSM'], 'd': [4, 8, 12, 16, 20, 24], 'window': [20, 100], 'nh': [1], 'sd': [1]}\n",
    "    experiment3_params = {'layer1': ['TF', 'SSM'], 'layer2': ['TF', 'SSM'], 'd': [16, 24], 'window': [20, 100], 'nh': [2,4], 'sd': [1]}\n",
    "    experiment4_params = {'layer1': ['TF', 'SSM'], 'layer2': ['TF', 'SSM'], 'd': [16, 24], 'window': [20, 100], 'nh': [1], 'sd': [2,4,6,12]}\n",
    "    experiment5_params = {'layer1': ['TF', 'SSM'], 'layer2': ['TF', 'SSM'], 'd': [100], 'window': [5, 10, 15, 20, 30, 50, 100], 'nh': [1], 'sd': [1]}\n",
    "    experiment6_params = {'layer1': ['TF'], 'layer2': ['TF'], 'layer3': ['TF'], 'd': [24], 'window': [5, 10, 15, 20, 30, 50, 100], 'nh': [1], 'sd': [1]}\n",
    "    experiment7_params = {'layer1': ['SSM'], 'layer2': ['SSM'], 'layer3': ['TF'], 'd': [24], 'window': [5, 10, 15, 20, 30, 50, 100], 'nh': [1], 'sd': [1]}\n",
    "    experiment8_params = {'layer1': ['SSM'], 'layer2': ['SSM'], 'layer3': ['SSM'], 'd': [24], 'window': [5, 10, 15, 20, 30, 50, 100], 'nh': [1], 'sd': [1]}\n",
    "    experiment9_params = {'layer1': ['TF'], 'layer2': ['TF'], 'layer3': ['TF'], 'd': [24, 48, 96, 192, 384, 768], 'window': [20, 100], 'nh': [1], 'sd': [1]}\n",
    "    experiment10_params = {'layer1': ['SSM'], 'layer2': ['SSM'], 'layer3': ['TF'], 'd': [24, 48, 96, 192, 384, 768], 'window': [20, 100], 'nh': [1], 'sd': [1]}\n",
    "    experiment11_params = {'layer1': ['SSM'], 'layer2': ['SSM'], 'layer3': ['SSM'], 'd': [24, 48, 96, 192, 384, 768], 'window': [20, 100], 'nh': [1], 'sd': [1]}\n",
    "    # experiment9_params = {'layer1': ['TF'], 'layer2': ['TF'], 'layer3': ['TF'], 'd': [384, 768], 'window': [20, 100], 'nh': [1], 'sd': [1]}\n",
    "    # experiment10_params = {'layer1': ['SSM'], 'layer2': ['SSM'], 'layer3': ['TF'], 'd': [384, 768], 'window': [20, 100], 'nh': [1], 'sd': [1]}\n",
    "    # experiment11_params = {'layer1': ['SSM'], 'layer2': ['SSM'], 'layer3': ['SSM'], 'd': [384, 768], 'window': [20, 100], 'nh': [1], 'sd': [1]}\n",
    "    \n",
    "    # experiments = [experiment1_params, experiment2_params, experiment3_params, experiment4_params]\n",
    "    # experiments = [experiment5_params]\n",
    "    # experiments = [experiment6_params, experiment7_params, experiment8_params]\n",
    "    experiments = [experiment9_params, experiment10_params, experiment11_params]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "501791fd-5cd8-4086-86bf-655500db978c",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-01-26T13:23:50.231984Z",
     "iopub.status.busy": "2026-01-26T13:23:50.231745Z",
     "iopub.status.idle": "2026-01-26T13:23:50.259037Z",
     "shell.execute_reply": "2026-01-26T13:23:50.258529Z"
    }
   },
   "outputs": [],
   "source": [
    "# List of commands to run\n",
    "commands = []\n",
    "\n",
    "def make_command_model(model, num_layers, embed_dim, window, num_heads, state_dim, run_number=0):\n",
    "    if task_name == \"var-copy\":\n",
    "        # return f\"python3 main.py --run_number {run_number} --train_task {task_name} --eval_task {task_name} --model {model} --num_layers {num_layers} --hidden_size {embed_dim} --window {window} --heads {num_heads} --state_dim {state_dim} --save_results True --lr 1e-2 --num_vocab 26\"\n",
    "        return f\"python3 main.py --run_number {run_number} --train_task {task_name} --eval_task {task_name} --model {model} --num_layers {num_layers} --hidden_size {embed_dim} --window {window} --heads {num_heads} --state_dim {state_dim} --save_results True --auto_lr True --num_vocab 26\"\n",
    "    if task_name == \"decode-recall\" or task_name == \"decode-recall-last\":\n",
    "        # return f\"python3 main.py --run_number {run_number} --train_task {task_name} --eval_task {task_name} --model {model} --num_layers {num_layers} --hidden_size {embed_dim} --window {window} --heads {num_heads} --state_dim {state_dim} --save_results True --lr 1e-2 --num_vocab 33\" # Above 32 for any issues with rounding\n",
    "        return f\"python3 main.py --run_number {run_number} --train_task {task_name} --eval_task {task_name} --model {model} --num_layers {num_layers} --hidden_size {embed_dim} --window {window} --heads {num_heads} --state_dim {state_dim} --save_results True --auto_lr True --num_vocab 33\" # Above 32 for any issues with rounding\n",
    "    if task_name == \"assoc-recall-mk\":\n",
    "        # return f\"python3 main.py --run_number {run_number} --train_task {task_name} --eval_task {task_name} --model {model} --num_layers {num_layers} --hidden_size {embed_dim} --window {window} --heads {num_heads} --state_dim {state_dim} --save_results True --lr 1e-2 --num_vocab 8\"\n",
    "        return f\"python3 main.py --run_number {run_number} --train_task {task_name} --eval_task {task_name} --model {model} --num_layers {num_layers} --hidden_size {embed_dim} --window {window} --heads {num_heads} --state_dim {state_dim} --save_results True --auto_lr True --num_vocab 8\"\n",
    "\n",
    "def make_command(layer1, layer2, embed_dim, window, num_heads, state_dim, run_number=0):\n",
    "    if task_name == \"var-copy\":\n",
    "        # return f\"python3 main.py --run_number {run_number} --train_task {task_name} --eval_task {task_name} --layer1 {layer1} --layer2 {layer2} --hidden_size {embed_dim} --window {window} --heads {num_heads} --state_dim {state_dim} --save_results True --lr 1e-2 --num_vocab 26\"\n",
    "        # return f\"python3 main.py --run_number {run_number} --train_task {task_name} --eval_task {task_name} --layer1 {layer1} --layer2 {layer2} --hidden_size {embed_dim} --window {window} --heads {num_heads} --state_dim {state_dim} --save_results True --auto_lr True --num_vocab 26\"\n",
    "        # return f\"python3 main.py --run_number {run_number} --train_task {task_name} --eval_task {task_name} --layer1 {layer1} --layer2 {layer2} --hidden_size {embed_dim} --window {window} --heads {num_heads} --state_dim {state_dim} --save_results True --auto_lr True --num_vocab 200\"\n",
    "        return f\"python3 main.py --run_number {run_number} --train_task {task_name} --eval_task {task_name} --layer1 {layer1} --layer2 {layer2} --hidden_size {embed_dim} --window {window} --heads {num_heads} --state_dim {state_dim} --save_results True --auto_lr True --num_vocab 1000\"\n",
    "    if task_name == \"decode-recall\" or task_name == \"decode-recall-last\":\n",
    "        # return f\"python3 main.py --run_number {run_number} --train_task {task_name} --eval_task {task_name} --layer1 {layer1} --layer2 {layer2} --hidden_size {embed_dim} --window {window} --heads {num_heads} --state_dim {state_dim} --save_results True --lr 1e-2 --num_vocab 33\" # Above 32 for any issues with rounding\n",
    "        return f\"python3 main.py --run_number {run_number} --train_task {task_name} --eval_task {task_name} --layer1 {layer1} --layer2 {layer2} --hidden_size {embed_dim} --window {window} --heads {num_heads} --state_dim {state_dim} --save_results True --auto_lr True --num_vocab 33\" # Above 32 for any issues with rounding\n",
    "    if task_name == \"assoc-recall-mk\":\n",
    "        # return f\"python3 main.py --run_number {run_number} --train_task {task_name} --eval_task {task_name} --layer1 {layer1} --layer2 {layer2} --hidden_size {embed_dim} --window {window} --heads {num_heads} --state_dim {state_dim} --save_results True --lr 1e-2 --num_vocab 8\"\n",
    "        return f\"python3 main.py --run_number {run_number} --train_task {task_name} --eval_task {task_name} --layer1 {layer1} --layer2 {layer2} --hidden_size {embed_dim} --window {window} --heads {num_heads} --state_dim {state_dim} --save_results True --auto_lr True --num_vocab 8\"\n",
    "    if task_name == \"needle\":\n",
    "        # return f\"python3 main.py --run_number {run_number} --train_task {task_name} --eval_task {task_name} --layer1 {layer1} --layer2 {layer2} --hidden_size {embed_dim} --window {window} --heads {num_heads} --state_dim {state_dim} --save_results True --auto_lr True --num_vocab 100 --sequence_length 500 --eval_sequence_length 500 --run_anyways True\"\n",
    "        return f\"python3 main.py --run_number {run_number} --train_task {task_name} --eval_task {task_name} --layer1 {layer1} --layer2 {layer2} --hidden_size {embed_dim} --window {window} --heads {num_heads} --state_dim {state_dim} --save_results True --auto_lr True --force_do_lr True --num_vocab 100 --sequence_length 100 --eval_sequence_length 100 --run_anyways True --epochs 10\"\n",
    "\n",
    "\n",
    "def make_command_3(layer1, layer2, layer3, embed_dim, window, num_heads, state_dim, run_number=0):\n",
    "    if task_name == \"var-copy\":\n",
    "        # return f\"python3 main.py --run_number {run_number} --train_task {task_name} --eval_task {task_name} --layer1 {layer1} --layer2 {layer2} --layer3 {layer3} --hidden_size {embed_dim} --window {window} --heads {num_heads} --state_dim {state_dim} --save_results True --lr 1e-2 --num_vocab 26\"\n",
    "        return f\"python3 main.py --run_number {run_number} --train_task {task_name} --eval_task {task_name} --layer1 {layer1} --layer2 {layer2} --layer3 {layer3} --hidden_size {embed_dim} --window {window} --heads {num_heads} --state_dim {state_dim} --save_results True --auto_lr True --num_vocab 26\"\n",
    "    if task_name == \"decode-recall\" or task_name == \"decode-recall-last\":\n",
    "        # return f\"python3 main.py --run_number {run_number} --train_task {task_name} --eval_task {task_name} --layer1 {layer1} --layer2 {layer2} --layer3 {layer3} --hidden_size {embed_dim} --window {window} --heads {num_heads} --state_dim {state_dim} --save_results True --lr 1e-2 --num_vocab 33\" # Above 32 for any issues with rounding\n",
    "        return f\"python3 main.py --run_number {run_number} --train_task {task_name} --eval_task {task_name} --layer1 {layer1} --layer2 {layer2} --layer3 {layer3} --hidden_size {embed_dim} --window {window} --heads {num_heads} --state_dim {state_dim} --save_results True --auto_lr True --num_vocab 33\" # Above 32 for any issues with rounding\n",
    "    if task_name == \"assoc-recall-mk\":\n",
    "        # return f\"python3 main.py --run_number {run_number} --train_task {task_name} --eval_task {task_name} --layer1 {layer1} --layer2 {layer2} --layer3 {layer3} --hidden_size {embed_dim} --window {window} --heads {num_heads} --state_dim {state_dim} --save_results True --lr 1e-2 --num_vocab 8\"\n",
    "        return f\"python3 main.py --run_number {run_number} --train_task {task_name} --eval_task {task_name} --layer1 {layer1} --layer2 {layer2} --layer3 {layer3} --hidden_size {embed_dim} --window {window} --heads {num_heads} --state_dim {state_dim} --save_results True --auto_lr True --num_vocab 8\"\n",
    "\n",
    "def make_command_mixed(layer1, layer2, embed_dim, window, num_heads, state_dim, run_number=0):\n",
    "    if task_name == \"var-copy\":\n",
    "        # return f\"python3 main.py --run_number {run_number} --train_task {task_name} --eval_task {task_name} --layer1 {layer1} --layer2 {layer2} --hidden_size {embed_dim} --window {window} --heads {num_heads} --state_dim {state_dim} --save_results True --lr 1e-2 --num_vocab 26\"\n",
    "        return f\"python3 main.py --run_number {run_number} --train_task {task_name} --eval_task {task_name} --layer1 {layer1} --layer2 {layer2} --hidden_size {embed_dim} --window {window} --heads {num_heads} --state_dim {state_dim} --save_results True --auto_lr True --num_vocab 26 --mixed True\"\n",
    "    if task_name == \"decode-recall\" or task_name == \"decode-recall-last\":\n",
    "        # return f\"python3 main.py --run_number {run_number} --train_task {task_name} --eval_task {task_name} --layer1 {layer1} --layer2 {layer2} --hidden_size {embed_dim} --window {window} --heads {num_heads} --state_dim {state_dim} --save_results True --lr 1e-2 --num_vocab 33\" # Above 32 for any issues with rounding\n",
    "        return f\"python3 main.py --run_number {run_number} --train_task {task_name} --eval_task {task_name} --layer1 {layer1} --layer2 {layer2} --hidden_size {embed_dim} --window {window} --heads {num_heads} --state_dim {state_dim} --save_results True --auto_lr True --num_vocab 33 --mixed True\" # Above 32 for any issues with rounding\n",
    "    if task_name == \"assoc-recall-mk\":\n",
    "        # return f\"python3 main.py --run_number {run_number} --train_task {task_name} --eval_task {task_name} --layer1 {layer1} --layer2 {layer2} --hidden_size {embed_dim} --window {window} --heads {num_heads} --state_dim {state_dim} --save_results True --lr 1e-2 --num_vocab 8\"\n",
    "        return f\"python3 main.py --run_number {run_number} --train_task {task_name} --eval_task {task_name} --layer1 {layer1} --layer2 {layer2} --hidden_size {embed_dim} --window {window} --heads {num_heads} --state_dim {state_dim} --save_results True --auto_lr True --num_vocab 8 --mixed True\"\n",
    "\n",
    "\n",
    "for run_number in range(5): # 11\n",
    "    for experiment in experiments:\n",
    "        if \"model\" in experiment.keys():\n",
    "            for model in experiment[\"model\"]:\n",
    "                for num_layers in experiment[\"num_layers\"]:\n",
    "                    for embed_dim in experiment[\"d\"]:\n",
    "                        for window in experiment[\"window\"]:\n",
    "                            for num_heads in experiment[\"nh\"]:\n",
    "                                for state_dim in experiment[\"sd\"]:\n",
    "                                    commands.append(make_command_model(model, num_layers, embed_dim, window, num_heads, state_dim, run_number))\n",
    "        \n",
    "        else:\n",
    "            for layer1 in experiment[\"layer1\"]:\n",
    "                for layer2 in experiment[\"layer2\"]:\n",
    "                    for embed_dim in experiment[\"d\"]:\n",
    "                        for window in experiment[\"window\"]:\n",
    "                            for num_heads in experiment[\"nh\"]:\n",
    "                                for state_dim in experiment[\"sd\"]:\n",
    "                                    if \"layer3\" in experiment.keys():\n",
    "                                        for layer3 in experiment[\"layer3\"]:\n",
    "                                            commands.append(make_command_3(layer1, layer2, layer3, embed_dim, window, num_heads, state_dim, run_number))\n",
    "                                    elif mixed:\n",
    "                                        commands.append(make_command_mixed(layer1, layer2, embed_dim, window, num_heads, state_dim, run_number))\n",
    "                                    else:\n",
    "                                        commands.append(make_command(layer1, layer2, embed_dim, window, num_heads, state_dim, run_number))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "b9e6a6e5-bfbd-44b8-9c76-591e8cd47ec4",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-01-26T13:23:50.260673Z",
     "iopub.status.busy": "2026-01-26T13:23:50.260525Z",
     "iopub.status.idle": "2026-01-26T13:23:50.265985Z",
     "shell.execute_reply": "2026-01-26T13:23:50.265427Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['python3 main.py --run_number 0 --train_task needle --eval_task needle --layer1 TF --layer2 TF --hidden_size 4 --window 20 --heads 1 --state_dim 1 --save_results True --auto_lr True --force_do_lr True --num_vocab 100 --sequence_length 100 --eval_sequence_length 100 --run_anyways True --epochs 10',\n",
       " 'python3 main.py --run_number 0 --train_task needle --eval_task needle --layer1 TF --layer2 TF --hidden_size 4 --window 100 --heads 1 --state_dim 1 --save_results True --auto_lr True --force_do_lr True --num_vocab 100 --sequence_length 100 --eval_sequence_length 100 --run_anyways True --epochs 10',\n",
       " 'python3 main.py --run_number 0 --train_task needle --eval_task needle --layer1 TF --layer2 TF --hidden_size 8 --window 20 --heads 1 --state_dim 1 --save_results True --auto_lr True --force_do_lr True --num_vocab 100 --sequence_length 100 --eval_sequence_length 100 --run_anyways True --epochs 10',\n",
       " 'python3 main.py --run_number 0 --train_task needle --eval_task needle --layer1 TF --layer2 TF --hidden_size 8 --window 100 --heads 1 --state_dim 1 --save_results True --auto_lr True --force_do_lr True --num_vocab 100 --sequence_length 100 --eval_sequence_length 100 --run_anyways True --epochs 10',\n",
       " 'python3 main.py --run_number 0 --train_task needle --eval_task needle --layer1 TF --layer2 TF --hidden_size 12 --window 20 --heads 1 --state_dim 1 --save_results True --auto_lr True --force_do_lr True --num_vocab 100 --sequence_length 100 --eval_sequence_length 100 --run_anyways True --epochs 10',\n",
       " 'python3 main.py --run_number 0 --train_task needle --eval_task needle --layer1 TF --layer2 TF --hidden_size 12 --window 100 --heads 1 --state_dim 1 --save_results True --auto_lr True --force_do_lr True --num_vocab 100 --sequence_length 100 --eval_sequence_length 100 --run_anyways True --epochs 10',\n",
       " 'python3 main.py --run_number 0 --train_task needle --eval_task needle --layer1 TF --layer2 TF --hidden_size 16 --window 20 --heads 1 --state_dim 1 --save_results True --auto_lr True --force_do_lr True --num_vocab 100 --sequence_length 100 --eval_sequence_length 100 --run_anyways True --epochs 10',\n",
       " 'python3 main.py --run_number 0 --train_task needle --eval_task needle --layer1 TF --layer2 TF --hidden_size 16 --window 100 --heads 1 --state_dim 1 --save_results True --auto_lr True --force_do_lr True --num_vocab 100 --sequence_length 100 --eval_sequence_length 100 --run_anyways True --epochs 10',\n",
       " 'python3 main.py --run_number 0 --train_task needle --eval_task needle --layer1 TF --layer2 TF --hidden_size 20 --window 20 --heads 1 --state_dim 1 --save_results True --auto_lr True --force_do_lr True --num_vocab 100 --sequence_length 100 --eval_sequence_length 100 --run_anyways True --epochs 10',\n",
       " 'python3 main.py --run_number 0 --train_task needle --eval_task needle --layer1 TF --layer2 TF --hidden_size 20 --window 100 --heads 1 --state_dim 1 --save_results True --auto_lr True --force_do_lr True --num_vocab 100 --sequence_length 100 --eval_sequence_length 100 --run_anyways True --epochs 10']"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\n",
    "# commands = ['python3 train.py --run_number 2 --task_name var_copy --layer1 SSM --layer2 SSM --embed_dim 12 --window 20 --num_heads 1 --state_dim 2 --save True']\n",
    "commands[:10]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "30f4f897-7954-4843-a4a3-c1abe3a66ea3",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-01-26T13:23:50.267929Z",
     "iopub.status.busy": "2026-01-26T13:23:50.267799Z",
     "iopub.status.idle": "2026-01-26T15:20:38.824170Z",
     "shell.execute_reply": "2026-01-26T15:20:38.823840Z"
    }
   },
   "outputs": [],
   "source": [
    "def run_command(cmd):\n",
    "    \"\"\"Run a single shell command and return (cmd, returncode, stdout, stderr).\"\"\"\n",
    "    result = subprocess.run(cmd, shell=True, capture_output=True, text=True)\n",
    "    return cmd, result.returncode, result.stdout, result.stderr\n",
    "\n",
    "max_workers = 5 # 5, 10\n",
    "\n",
    "with ThreadPoolExecutor(max_workers=max_workers) as executor:\n",
    "    futures = {executor.submit(run_command, cmd): cmd for cmd in commands}\n",
    "\n",
    "    for future in as_completed(futures):\n",
    "        cmd, returncode, stdout, stderr = future.result()\n",
    "        if returncode != 0:\n",
    "            print(f\"[{cmd}] exited with {returncode}\")\n",
    "        \n",
    "        if returncode == 1:\n",
    "            executor.shutdown()\n",
    "            print(stdout)\n",
    "            print(stderr)\n",
    "            assert False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4bc2226-523f-42de-9fc1-d46f96103d4d",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "hybrid",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}