fix multiple tests. with latest monarch and rdma, ts.put and ts.get only works inside an actor (#75)

Summary: Pull Request resolved: #75

Reviewed By: LucasLLC, amirafzali

Differential Revision: D86156904

Author: casteryh

tests/test_keys.py

@@ -10,7 +10,6 @@
 
 import pytest
 import torch
-
 import torchstore as ts
 from monarch.actor import Actor, current_rank, endpoint
 from pytest_unordered import unordered
@@ -21,24 +20,36 @@
 @pytest.mark.asyncio
 async def test_keys_basic():
     """Test basic put/get functionality"""
+
+    class TestActor(Actor):
+        @endpoint
+        async def test(self) -> Exception or None:
+            try:
+                await ts.put("", torch.tensor([1, 2, 3]))
+                await ts.put(".x", torch.tensor([1, 2, 3]))
+                await ts.put("v0.x", torch.tensor([1, 2, 3]))
+                await ts.put("v0.y", torch.tensor([4, 5, 6]))
+                await ts.put("v0.x.z", torch.tensor([7, 8, 9]))
+                await ts.put("v1.x", torch.tensor([7, 8, 9]))
+                await ts.put("v1.y", torch.tensor([10, 11, 12]))
+
+                assert await ts.keys() == unordered(
+                    ["", ".x", "v0.x", "v0.y", "v0.x.z", "v1.x", "v1.y"]
+                )
+                assert await ts.keys("v0") == unordered(["v0.x", "v0.y", "v0.x.z"])
+                assert await ts.keys("v0.x") == unordered(["v0.x", "v0.x.z"])
+                assert await ts.keys("v0.x.z") == unordered(["v0.x.z"])
+                assert await ts.keys("") == unordered(["", ".x"])
+                assert await ts.keys("v1") == unordered(["v1.x", "v1.y"])
+            except Exception as e:
+                return e
+
     await ts.initialize()
 
-    await ts.put("", torch.tensor([1, 2, 3]))
-    await ts.put(".x", torch.tensor([1, 2, 3]))
-    await ts.put("v0.x", torch.tensor([1, 2, 3]))
-    await ts.put("v0.y", torch.tensor([4, 5, 6]))
-    await ts.put("v0.x.z", torch.tensor([7, 8, 9]))
-    await ts.put("v1.x", torch.tensor([7, 8, 9]))
-    await ts.put("v1.y", torch.tensor([10, 11, 12]))
+    actor = await spawn_actors(1, TestActor, "actor_0")
+    err = await actor.test.call_one()
 
-    assert await ts.keys() == unordered(
-        ["", ".x", "v0.x", "v0.y", "v0.x.z", "v1.x", "v1.y"]
-    )
-    assert await ts.keys("v0") == unordered(["v0.x", "v0.y", "v0.x.z"])
-    assert await ts.keys("v0.x") == unordered(["v0.x", "v0.x.z"])
-    assert await ts.keys("v0.x.z") == unordered(["v0.x.z"])
-    assert await ts.keys("") == unordered(["", ".x"])
-    assert await ts.keys("v1") == unordered(["v1.x", "v1.y"])
+    assert err is None
 
     await ts.shutdown()
 

tests/test_large_tensors.py

@@ -10,7 +10,6 @@
 
 import pytest
 import torch
-
 import torchstore as ts
 from monarch.actor import Actor, endpoint
 from torchstore.logging import init_logging

tests/test_models.py

@@ -12,7 +12,6 @@
 
 import pytest
 import torch
-
 import torchstore as ts
 from monarch.actor import Actor, current_rank, endpoint
 from torch.distributed.device_mesh import init_device_mesh

tests/test_resharding_basic.py

@@ -11,11 +11,8 @@
 from typing import List, Tuple, Union
 
 import pytest
-
 import torch
-
 import torchstore as ts
-
 from torch.distributed._tensor import Replicate, Shard
 from torch.distributed.tensor._utils import _compute_local_shape_and_global_offset
 from torchstore.utils import get_local_tensor, spawn_actors

tests/test_resharding_ext.py

@@ -7,11 +7,9 @@
 from logging import getLogger
 
 import pytest
-
 from torch.distributed._tensor import Shard
 
 from .test_resharding_basic import _test_resharding
-
 from .utils import main, transport_plus_strategy_params
 
 logger = getLogger(__name__)

tests/test_state_dict.py

@@ -12,13 +12,10 @@
 from typing import Union
 
 import pytest
-
 import torch
 import torch.distributed.checkpoint as dcp
 import torch.nn as nn
-
 import torchstore as ts
-
 from monarch.actor import Actor, current_rank, endpoint
 from torch.distributed.checkpoint._nested_dict import flatten_state_dict
 from torch.distributed.checkpoint.state_dict import (
@@ -279,7 +276,6 @@ def _assert_equal_state_dict(state_dict1, state_dict2):
         flattened_state_dict_2
     ), f"{flattened_state_dict_1.keys()=}\n{flattened_state_dict_2.keys()=}"
     for key in flattened_state_dict_1:
-
         assert key in flattened_state_dict_2
         if isinstance(flattened_state_dict_1[key], torch.Tensor):
             t1, t2 = flattened_state_dict_1[key], flattened_state_dict_2[key]

tests/test_store.py

@@ -8,13 +8,9 @@
 from logging import getLogger
 
 import pytest
-
 import torch
-
 import torchstore as ts
-
 from monarch.actor import Actor, current_rank, endpoint
-
 from torchstore.logging import init_logging
 from torchstore.utils import spawn_actors
 
@@ -303,19 +299,31 @@ async def get(self, key):
 @pytest.mark.asyncio
 async def test_key_miss():
     """Test the behavior of get() when the key is missing."""
-    await ts.initialize()
 
-    key = "foo"
-    value = torch.tensor([1, 2, 3])
-    await ts.put(key, value)
+    class TestActor(Actor):
+        @endpoint
+        async def test(self) -> Exception or None:
+            try:
+                key = "foo"
+                value = torch.tensor([1, 2, 3])
+                await ts.put(key, value)
+
+                # Get the value back
+                retrieved_value = await ts.get(key)
+                assert torch.equal(value, retrieved_value)
+
+                # Get a missing key
+                with pytest.raises(KeyError):
+                    await ts.get("bar")
+            except Exception as e:
+                return e
+
+    await ts.initialize()
 
-    # Get the value back
-    retrieved_value = await ts.get(key)
-    assert torch.equal(value, retrieved_value)
+    actor = await spawn_actors(1, TestActor, "actor_0")
+    err = await actor.test.call_one()
 
-    # Get a missing key
-    with pytest.raises(KeyError):
-        await ts.get("bar")
+    assert err is None
 
     await ts.shutdown()