feat: Add integration tests

Author: Kovbo

examples/art-e.ipynb

@@ -44,7 +44,7 @@
    "outputs": [],
    "source": [
     "%%capture\n",
-    "!uv pip install openpipe-art==0.3.11.post5 langchain-core tenacity datasets \"gql<4\" --prerelease allow --no-cache-dir"
+    "!uv pip install openpipe-art==0.4.7 vllm==0.9.2 langchain-core tenacity datasets \"gql<4\" --prerelease allow --no-cache-dir"
    ]
   },
   {
@@ -581,6 +581,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "import torch\n",
+    "\n",
     "import art\n",
     "from art.local import LocalBackend\n",
     "\n",
@@ -594,15 +596,16 @@
     ")\n",
     "\n",
     "# To run on a T4, we need to override some config defaults.\n",
-    "model._internal_config = art.dev.InternalModelConfig(\n",
-    "    init_args=art.dev.InitArgs(\n",
-    "        max_seq_length=8192,\n",
-    "    ),\n",
-    "    engine_args=art.dev.EngineArgs(\n",
-    "        enforce_eager=True,\n",
-    "        gpu_memory_utilization=0.8,\n",
-    "    ),\n",
-    ")\n",
+    "if torch.cuda.get_device_properties(0).major < 8:\n",
+    "    model._internal_config = art.dev.InternalModelConfig(\n",
+    "        init_args=art.dev.InitArgs(\n",
+    "            max_seq_length=8192,\n",
+    "        ),\n",
+    "        engine_args=art.dev.EngineArgs(\n",
+    "            enforce_eager=True,\n",
+    "            gpu_memory_utilization=0.8,\n",
+    "        ),\n",
+    "    )\n",
     "\n",
     "# Initialize the server\n",
     "backend = LocalBackend(\n",

examples/prisoners-dilemma.ipynb

@@ -18,6 +18,7 @@
     "\n",
     "BASE_MODEL = \"Qwen/Qwen2.5-7B-Instruct\"\n",
     "PRISONERS_DILEMMA_ROUNDS = 10\n",
+    "TRAINING_STEPS = 1_000\n",
     "\n",
     "backend = LocalBackend()\n",
     "model = art.TrainableModel(\n",
@@ -117,7 +118,7 @@
     "    return trajectories\n",
     "\n",
     "\n",
-    "for _ in range(await model.get_step(), 1_000):\n",
+    "for _ in range(await model.get_step(), TRAINING_STEPS):\n",
     "    # Simultaneously rollout self-play games, and games versus the base model.\n",
     "    self_play_trajectories, base_play_trajectories = await asyncio.gather(\n",
     "        art.gather_trajectories(\n",

examples/rock-paper-tool-use.ipynb

@@ -47,6 +47,7 @@
     "\n",
     "MODEL_NAME = \"001\"\n",
     "BASE_MODEL = \"Qwen/Qwen2.5-7B-Instruct\"\n",
+    "TRAINING_STEPS = 1_000\n",
     "\n",
     "model = art.TrainableModel(\n",
     "    name=MODEL_NAME, project=\"rock-paper-tool-use\", base_model=BASE_MODEL\n",
@@ -175,7 +176,7 @@
     "    return trajectories[0]\n",
     "\n",
     "\n",
-    "for i in range(await model.get_step(), 1_000):\n",
+    "for i in range(await model.get_step(), TRAINING_STEPS):\n",
     "    trajectories = await art.gather_trajectories(\n",
     "        (rollout() for _ in range(64)), max_exceptions=64\n",
     "    )\n",
@@ -202,7 +203,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.13"
+   "version": "3.10.16"
   }
  },
  "nbformat": 4,

examples/temporal_clue/temporal-clue.ipynb

@@ -40,7 +40,7 @@
    "outputs": [],
    "source": [
     "%%capture\n",
-    "!uv pip install openpipe-art==0.3.11.post3 \"gql<4\" --prerelease allow --no-cache-dir"
+    "!uv pip install openpipe-art==0.4.7 vllm==0.9.2 \"gql<4\" --prerelease allow --no-cache-dir"
    ]
   },
   {
@@ -221,8 +221,12 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "stride = 4\n",
-    "for i in range(await model.get_step(), 1_000):\n",
+    "STRIDE = 4\n",
+    "TRAINING_STEPS = 1_000\n",
+    "ROLLOUTS_PER_STEP = 50\n",
+    "LEARNING_RATE = 5e-5\n",
+    "\n",
+    "for i in range(await model.get_step(), TRAINING_STEPS):\n",
     "    val_groups, train_groups = await asyncio.gather(\n",
     "        art.gather_trajectory_groups(\n",
     "            (\n",
@@ -233,8 +237,10 @@
     "        ),\n",
     "        art.gather_trajectory_groups(\n",
     "            (\n",
-    "                art.TrajectoryGroup(rollout(model, puzzle) for _ in range(50))\n",
-    "                for puzzle in train_puzzles[i * stride : (i + 1) * stride]\n",
+    "                art.TrajectoryGroup(\n",
+    "                    rollout(model, puzzle) for _ in range(ROLLOUTS_PER_STEP)\n",
+    "                )\n",
+    "                for puzzle in train_puzzles[i * STRIDE : (i + 1) * STRIDE]\n",
     "            ),\n",
     "            pbar_desc=\"train\",\n",
     "        ),\n",
@@ -243,7 +249,7 @@
     "    await model.delete_checkpoints()\n",
     "    await model.train(\n",
     "        train_groups,\n",
-    "        config=art.TrainConfig(learning_rate=5e-5),\n",
+    "        config=art.TrainConfig(learning_rate=LEARNING_RATE),\n",
     "    )"
    ]
   },