Repository: ray-project/ray
File: python/ray/tests/gpu_objects/test_gpu_objects_nixl.py (Python)

Code before:
    # this tensor may overlap with the first tensor (de-registered).
    ref2 = src_actor.echo.remote(torch.tensor([1, 2, 3]).to("cuda"), "cuda")

    # Create the third tensor at the sender. The memory address of
    # this tensor may overlap with the first tensor (de-registered).
    ref3 = src_actor.echo.remote(torch.tensor([4, 5, 6]).to("cuda"), "cuda")
    # Trigger the transfer. The receiver successfully gets
    # the third tensor (nixl agent is reset internally).
    assert ray.get(dst_actor.sum.remote(ref3, "cuda")) == 15

    del ref2, ref3


@ray.remote(num_gpus=1, num_cpus=0, enable_tensor_transport=True)
class StorageTestActor:
    def produce_views(self):
        tensor = torch.tensor([[1, 1], [2, 2]], dtype=torch.float32).to("cuda")
        views = [tensor[0], tensor[1]]
        return ray.put(views, _tensor_transport="nixl")

    def consume(self, ref):
        return ray.get(ref)


@pytest.mark.parametrize("ray_start_regular", [{"num_gpus": 2}], indirect=True)
def test_storage_level_registration(ray_start_regular):
    """Test that views sharing the same storage result in a single NIXL registration."""

    actors = [StorageTestActor.remote() for _ in range(2)]
    src_actor, dst_actor = actors[0], actors[1]

    refs = src_actor.produce_views.remote()

    tensors = ray.get(dst_actor.consume.remote(refs))
    assert [tensor.tolist() for tensor in tensors] == [[1, 1], [2, 2]]


if __name__ == "__main__":
    sys.exit(pytest.main(["-sv", __file__]))


Review comment (suggestion):
Hmm again we would need to check internal state here for the test to really test anything. It might be better to write a unit test here that directly calls the methods on the nixl tensor transport, instead of using integration tests through high-level ray calls like `ray.put`.

Code after:
    # this tensor may overlap with the first tensor (de-registered).
    ref2 = src_actor.echo.remote(torch.tensor([1, 2, 3]).to("cuda"), "cuda")

    # Create the third tensor at the sender. The memory address of
    # this tensor may overlap with the first tensor (de-registered).
    ref3 = src_actor.echo.remote(torch.tensor([4, 5, 6]).to("cuda"), "cuda")
    # Trigger the transfer. The receiver successfully gets
    # the third tensor (nixl agent is reset internally).
    assert ray.get(dst_actor.sum.remote(ref3, "cuda")) == 15

    del ref2, ref3


@pytest.mark.parametrize("ray_start_regular", [{"num_gpus": 1}], indirect=True)
def test_storage_level_overlapping_views_reference_count(ray_start_regular):
    """Test that two overlapping tensors sharing the same underlying storage produce a
    single NIXL registration. When each tensor's ref goes out of scope via
    garbage_collect, the metadata_count decrements. After both are freed,
    the registration is removed."""
    from ray.experimental.gpu_object_manager.nixl_tensor_transport import (
        NixlTensorTransport,
    )

    transport = NixlTensorTransport()

    tensor = torch.tensor([[1, 1], [2, 2], [3, 3]], dtype=torch.float32).to("cuda")
    view0 = tensor[0:2]
    view1 = tensor[1:3]
    storage_key = tensor.untyped_storage().data_ptr()

    assert view0.untyped_storage().data_ptr() == storage_key
    assert view1.untyped_storage().data_ptr() == storage_key
    assert view0.data_ptr() != view1.data_ptr()

    # Simulate ray.put(view0)
    obj_id1 = "test_obj_id_1"
    meta1 = transport.extract_tensor_transport_metadata(obj_id1, [view0])
    assert len(transport._tensor_desc_cache) == 1
    assert transport._tensor_desc_cache[storage_key].metadata_count == 1

    # Simulate ray.put(view1) and check that the a new entry is not created in the tensor desc cache
    # since they share the same storage key and the metadata_count is incremented by 1
    obj_id2 = "test_obj_id_2"
    meta2 = transport.extract_tensor_transport_metadata(obj_id2, [view1])
    assert len(transport._tensor_desc_cache) == 1
    assert transport._tensor_desc_cache[storage_key].metadata_count == 2

    # Simulate the obj ref for view0 going out of scope and check that the nixl memory registration is
    # not cleared since the object ref for view1 is still in scope
    transport.garbage_collect(obj_id1, meta1, [view0])
    assert storage_key in transport._tensor_desc_cache