Repository: run-house/runhouse

Issue:
[KIT-38] Add "on_this_cluster()" method to cluster
Cc @caroline

<sub>From [SyncLinear.com](https://synclinear.com) | [KIT-38](https://linear.app/runhouse/issue/KIT-38/add-on-this-cluster-method-to-cluster)</sub>


Discussion:

Resolving patch:
diff --git a/runhouse/rns/folders/folder.py b/runhouse/rns/folders/folder.py
--- a/runhouse/rns/folders/folder.py
+++ b/runhouse/rns/folders/folder.py
@@ -189,8 +189,7 @@ def system(self, data_source):
     def data_config(self):
         if isinstance(self.system, Resource):  # if system is a cluster
             # handle case cluster is itself
-            rns_address = _current_cluster("name")
-            if rns_address and rns_address == self.system.rns_address:
+            if self.system.on_this_cluster():
                 return self._data_config
 
             if not self.system.address:
@@ -222,7 +221,7 @@ def data_config(self, data_config):
     @property
     def _fs_str(self):
         if isinstance(self.system, Resource):  # if system is a cluster
-            if self.system.rns_address == _current_cluster("name"):
+            if self.system.on_this_cluster():
                 return self.DEFAULT_FS
             return self.CLUSTER_FS
         else:
diff --git a/runhouse/rns/hardware/cluster.py b/runhouse/rns/hardware/cluster.py
--- a/runhouse/rns/hardware/cluster.py
+++ b/runhouse/rns/hardware/cluster.py
@@ -15,8 +15,8 @@
 from sshtunnel import HandlerSSHTunnelForwarderError, SSHTunnelForwarder
 
 from runhouse.rh_config import open_grpc_tunnels, rns_client
+from runhouse.rns.obj_store import _current_cluster
 from runhouse.rns.packages.package import Package
-
 from runhouse.rns.resource import Resource
 
 from runhouse.servers.grpc.unary_client import UnaryClient
@@ -233,11 +233,9 @@ def clear_pins(self, pins: Optional[List[str]] = None):
         self.client.clear_pins(pins)
         logger.info(f'Clearing pins on cluster {pins or ""}')
 
-    def on_same_cluster(self, resource: Resource):
-        """Whether the given resource is on the cluster."""
-        if hasattr(resource, "system") and isinstance(resource.system, Resource):
-            return resource.system.rns_address == self.rns_address
-        return False
+    def on_this_cluster(self):
+        """Whether this function is being called on the same cluster."""
+        return _current_cluster("name") == self.rns_address
 
     # ----------------- gRPC Methods ----------------- #
 


Test patch:
diff --git a/tests/test_cluster.py b/tests/test_cluster.py
--- a/tests/test_cluster.py
+++ b/tests/test_cluster.py
@@ -5,6 +5,10 @@
 from runhouse.rns.hardware import cluster, OnDemandCluster
 
 
+def is_on_cluster(cluster):
+    return cluster.on_this_cluster()
+
+
 def test_cluster_config():
     rh_cpu = cluster(name="^rh-cpu")
     if not rh_cpu.is_up():
@@ -58,31 +62,24 @@ def test_restart_grpc():
     assert codes
 
 
-def test_same_cluster():
-    hw = cluster(name="^rh-cpu")
+def test_on_same_cluster():
+    hw = cluster(name="^rh-cpu").up_if_not()
+    hw.restart_grpc_server()
     hw.up_if_not()
 
     hw_copy = cluster(name="^rh-cpu")
 
-    def dummy_func(a):
-        return a
-
-    func_hw = rh.function(dummy_func).to(hw)
-    assert hw.on_same_cluster(func_hw)
-    assert hw_copy.on_same_cluster(func_hw)
-
-
-def test_diff_cluster():
-    hw = cluster(name="^rh-cpu")
-    hw.up_if_not()
+    func_hw = rh.function(is_on_cluster).to(hw)
+    assert func_hw(hw)
+    assert func_hw(hw_copy)
 
-    def dummy_func(a):
-        return a
 
-    func_hw = rh.function(dummy_func).to(hw)
+def test_on_diff_cluster():
+    hw = cluster(name="^rh-cpu").up_if_not()
+    diff_hw = rh.cluster(name="test-byo-cluster").up_if_not()
 
-    new_hw = cluster(name="diff-cpu")
-    assert not new_hw.on_same_cluster(func_hw)
+    func_hw = rh.function(is_on_cluster).to(hw)
+    assert not func_hw(diff_hw)
 
 
 if __name__ == "__main__":