Merge branch 'main' into binyli/cupy

python/mscclpp/comm.py

@@ -220,3 +220,24 @@ def make_proxy_channels_with_scratch(
                 proxy_service.proxy_channel(semaphore_ids[rank]), memory_ids[rank], memory_ids[self.my_rank]
             )
         return channels
+
+    def register_semaphore_with_proxy(
+        self, proxy_service: ProxyService, connections: dict[int, Connection]
+    ) -> dict[int, SmChannel]:
+        semaphores = self.make_semaphore(connections, Host2DeviceSemaphore)
+        semaphore_ids = {}
+        for rank in semaphores:
+            semaphore_ids[rank] = proxy_service.add_semaphore(semaphores[rank])
+        channels = {}
+        for rank in semaphores:
+            channels[rank] = proxy_service.proxy_channel(semaphore_ids[rank])
+        return channels
+
+    def register_memory_with_proxy(
+        self, proxy_service: ProxyService, tensor: cp.ndarray, connections: dict[int, Connection]
+    ) -> dict[int, int]:
+        registered_memories = self.register_tensor_with_connections(tensor, connections)
+        memory_ids = {}
+        for rank in registered_memories:
+            memory_ids[rank] = proxy_service.add_memory(registered_memories[rank])
+        return memory_ids