diff --git a/pyproject.toml b/pyproject.toml
index 6c66b0696..d664c4c1e 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -50,7 +50,6 @@ dev = [
     "ruff",
 ]
 tests = [
-    "flaky",
     "pytest-asyncio>=0.18.2,!=0.22.0,<0.23.0",
     "pytest-env>=0.8,<2.0",
     "pytest-rerunfailures>=10,<14",
diff --git a/tests/ai/test_classify.py b/tests/ai/test_classify.py
index d98dbf947..a21b2c7cd 100644
--- a/tests/ai/test_classify.py
+++ b/tests/ai/test_classify.py
@@ -123,14 +123,14 @@ def test_return_index(self):
             assert result == 1
 
     class TestExamples:
-        @pytest.mark.flaky(max_runs=3)
+        @pytest.mark.flaky(reruns=3)
         async def test_hogwarts_sorting_hat(self):
-            description = "Brave, daring, chivalrous, and sometimes a bit reckless -- just like Harry Potter."
+            description = "Brave, daring, chivalrous -- it's Harry Potter!"
 
             house = marvin.classify(
                 description,
                 labels=["Gryffindor", "Hufflepuff", "Ravenclaw", "Slytherin"],
-                instructions="Select the best house for the student",
+                instructions="You're the sorting hat. Select the house for the student",
             )
 
             assert house == "Gryffindor"
diff --git a/tests/ai/test_extract.py b/tests/ai/test_extract.py
index c68b4898b..8656790e9 100644
--- a/tests/ai/test_extract.py
+++ b/tests/ai/test_extract.py
@@ -42,7 +42,7 @@ def test_extract_names(self):
             )
             assert result == ["John", "Mary", "Bob"]
 
-        @pytest.mark.flaky(max_runs=3)
+        @pytest.mark.flaky(reruns=3)
         def test_float_to_int(self):
             # gpt 3.5 sometimes struggles with this test, marked as flaky
             # pydantic no longer casts floats to ints, but gpt-3.5 assumes it's
diff --git a/tests/ai/test_models.py b/tests/ai/test_models.py
index 6d457712a..ea7135b10 100644
--- a/tests/ai/test_models.py
+++ b/tests/ai/test_models.py
@@ -59,7 +59,7 @@ class RentalHistory(BaseModel):
         """
         )
 
-    @pytest.mark.flaky(max_runs=3)
+    @pytest.mark.flaky(reruns=3)
     def test_resume(self):
         class Experience(BaseModel):
             technology: str
diff --git a/tests/ai/vision/test_caption.py b/tests/ai/vision/test_caption.py
index a8b2e7bbc..25b6a3174 100644
--- a/tests/ai/vision/test_caption.py
+++ b/tests/ai/vision/test_caption.py
@@ -8,7 +8,7 @@ def use_gpt4o_for_all_tests(gpt_4):
     pass
 
 
-@pytest.mark.flaky(max_runs=2)
+@pytest.mark.flaky(reruns=2)
 class TestVisionCaption:
     def test_ny(self):
         img = marvin.Image(
diff --git a/tests/ai/vision/test_cast.py b/tests/ai/vision/test_cast.py
index 1d466b30d..4fe70fb7c 100644
--- a/tests/ai/vision/test_cast.py
+++ b/tests/ai/vision/test_cast.py
@@ -14,7 +14,7 @@ def use_gpt4o_for_all_tests(gpt_4):
     pass
 
 
-@pytest.mark.flaky(max_runs=3)
+@pytest.mark.flaky(reruns=3)
 class TestVisionCast:
     def test_cast_ny(self):
         img = marvin.Image(
@@ -78,7 +78,7 @@ def test_map(self):
         assert_locations_equal(result[0], Location(city="New York", state="NY"))
         assert_locations_equal(result[1], Location(city="Washington", state="DC"))
 
-    @pytest.mark.flaky(max_runs=2)
+    @pytest.mark.flaky(reruns=2)
     async def test_async_map(self):
         ny = marvin.Image(
             "https://images.unsplash.com/photo-1568515387631-8b650bbcdb90"
diff --git a/tests/ai/vision/test_classify.py b/tests/ai/vision/test_classify.py
index d4d8fece5..d5552bbbb 100644
--- a/tests/ai/vision/test_classify.py
+++ b/tests/ai/vision/test_classify.py
@@ -7,7 +7,7 @@ def use_gpt4o_for_all_tests(gpt_4):
     pass
 
 
-@pytest.mark.flaky(max_runs=2)
+@pytest.mark.flaky(reruns=2)
 class TestVisionClassify:
     def test_ny_image_input(self):
         img = marvin.Image(
diff --git a/tests/ai/vision/test_extract.py b/tests/ai/vision/test_extract.py
index 01202340c..382a98469 100644
--- a/tests/ai/vision/test_extract.py
+++ b/tests/ai/vision/test_extract.py
@@ -16,7 +16,7 @@ def use_gpt4o_for_all_tests(gpt_4):
     pass
 
 
-@pytest.mark.flaky(max_runs=2)
+@pytest.mark.flaky(reruns=2)
 class TestVisionExtract:
     def test_ny_image_input(self):
         img = marvin.Image(
@@ -50,7 +50,7 @@ def test_two_cities(self):
         assert_locations_equal(locations[0], Location(city="New York", state="NY"))
         assert_locations_equal(locations[1], Location(city="Washington", state="DC"))
 
-    @pytest.mark.flaky(max_runs=3)
+    @pytest.mark.flaky(reruns=3)
     def test_dog(self):
         class Animal(BaseModel, frozen=True):
             type: Literal["cat", "dog", "bird", "frog", "horse", "pig"]
diff --git a/tests/beta/assistants/test_assistants.py b/tests/beta/assistants/test_assistants.py
index 70fc60007..cab5c3805 100644
--- a/tests/beta/assistants/test_assistants.py
+++ b/tests/beta/assistants/test_assistants.py
@@ -37,7 +37,7 @@ def test_code_interpreter(self):
         assert 85 <= output <= 95
 
 
-@pytest.mark.flaky(max_runs=2)
+@pytest.mark.flaky(reruns=2)
 class TestLifeCycle:
     @patch.object(client.beta.assistants, "delete", wraps=client.beta.assistants.delete)
     @patch.object(client.beta.assistants, "create", wraps=client.beta.assistants.create)