From 782aec508057985f58e3e97b2da61a5a5bfb6265 Mon Sep 17 00:00:00 2001
From: zdevito <zdevito@meta.com>
Date: Fri, 11 Jul 2025 13:25:37 -0700
Subject: [PATCH 1/2] Use assert_type to check for type inference failures for
 actor endpoints

This requires running pyright on the file in the oss test because pyre doesn't seem to care about the assert_type statement.

Differential Revision: [D78185660](https://our.internmc.facebook.com/intern/diff/D78185660/)

[ghstack-poisoned]
---
 .github/workflows/test-cuda.yml    |  4 ++++
 python/tests/requirements.txt      |  1 +
 python/tests/test_python_actors.py | 18 +++++++++++++++---
 3 files changed, 20 insertions(+), 3 deletions(-)

diff --git a/.github/workflows/test-cuda.yml b/.github/workflows/test-cuda.yml
index d3ab71d8..65220590 100644
--- a/.github/workflows/test-cuda.yml
+++ b/.github/workflows/test-cuda.yml
@@ -47,6 +47,10 @@ jobs:
         # Install the built wheel from artifact
         install_wheel_from_artifact
 
+        # tests the type_assert statements in test_python_actor are correct
+        # pyre currently does not check these assertions
+        pyright python/tests/test_python_actors.py
+
         # Run CUDA tests
         LC_ALL=C pytest python/tests/ -s -v -m "not oss_skip"
         python python/tests/test_mock_cuda.py
diff --git a/python/tests/requirements.txt b/python/tests/requirements.txt
index e7c1a081..a560cd3d 100644
--- a/python/tests/requirements.txt
+++ b/python/tests/requirements.txt
@@ -2,3 +2,4 @@ pytest
 pytest-timeout
 pytest-asyncio
 pytest-xdist
+pyright
diff --git a/python/tests/test_python_actors.py b/python/tests/test_python_actors.py
index 821872e9..52b0f741 100644
--- a/python/tests/test_python_actors.py
+++ b/python/tests/test_python_actors.py
@@ -4,7 +4,6 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-# pyre-unsafe
 import asyncio
 import operator
 import threading
@@ -27,6 +26,8 @@
     proc_mesh,
 )
 from monarch.rdma import RDMABuffer
+from typing_extensions import assert_type
+
 
 needs_cuda = pytest.mark.skipif(
     not torch.cuda.is_available(),
@@ -46,6 +47,10 @@ async def incr(self):
     async def value(self) -> int:
         return self.v
 
+    @endpoint
+    def value_sync_endpoint(self) -> int:
+        return self.v
+
 
 class Indirect(Actor):
     @endpoint
@@ -79,10 +84,17 @@ async def test_choose():
     i = await proc.spawn("indirect", Indirect)
     v.incr.broadcast()
     result = await v.value.choose()
+
+    # Test that Pyre derives the correct type for result (int, not Any)
+    assert_type(result, int)
     result2 = await i.call_value.choose(v)
 
     assert result == result2
 
+    result3 = await v.value_sync_endpoint.choose()
+    assert_type(result, int)
+    assert result2 == result3
+
 
 async def test_stream():
     proc = await local_proc_mesh(gpus=2)
@@ -551,12 +563,12 @@ async def nope():
 
     assert v == 5
 
-    def nope():
+    def nope2():
         nonlocal v
         v += 1
         raise ValueError("nope")
 
-    f = Future(incr, nope)
+    f = Future(incr, nope2)
 
     with pytest.raises(ValueError):
         f.get()

From 67342774a6cfec4ec09e4d4b360ee538bf90d4c7 Mon Sep 17 00:00:00 2001
From: zdevito <zdevito@meta.com>
Date: Fri, 11 Jul 2025 15:54:48 -0700
Subject: [PATCH 2/2] Update on "Use assert_type to check for type inference
 failures for actor endpoints"

This requires running pyright on the file in the oss test because pyre doesn't seem to care about the assert_type statement.

Differential Revision: [D78185660](https://our.internmc.facebook.com/intern/diff/D78185660/)

[ghstack-poisoned]
---
 python/tests/test_python_actors.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/python/tests/test_python_actors.py b/python/tests/test_python_actors.py
index 52b0f741..cb9f357c 100644
--- a/python/tests/test_python_actors.py
+++ b/python/tests/test_python_actors.py
@@ -4,6 +4,7 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pyre-unsafe
 import asyncio
 import operator
 import threading