Repository: LMCache/LMCache
File: lmcache/v1/storage_backend/__init__.py (Python)

Code before:
    if dst_device == "cuda":
        dst_device = f"cuda:{torch.cuda.current_device()}"

    storage_backends: OrderedDict[str, StorageBackendInterface] = OrderedDict()

    if config.enable_nixl:
        if config.enable_xpyd:
            # First Party
            from lmcache.v1.storage_backend.nixl_backend_v3 import NixlBackend

            storage_backends["NixlBackend"] = NixlBackend.CreateNixlBackend(
                config, metadata, memory_allocator
            )
        else:
            # First Party
            from lmcache.v1.storage_backend.nixl_backend import NixlBackend

            storage_backends["NixlBackend"] = NixlBackend.CreateNixlBackend(
                config, metadata
            )
            assert config.nixl_buffer_device is not None

    # TODO(Jiayi): The hierarchy is fixed for now
    # NOTE(Jiayi): The local_cpu backend is always created because
    # other backends might need it as a buffer.
    if config.local_cpu:
        local_cpu_backend = LocalCPUBackend(
            config,
            memory_allocator,
            lookup_server,
            lmcache_worker,
        )
        backend_name = str(local_cpu_backend)
        storage_backends[backend_name] = local_cpu_backend

    if config.local_disk and config.max_local_disk_size > 0:
        local_disk_backend = LocalDiskBackend(
            config,
            loop,
            local_cpu_backend,
            dst_device,
            lmcache_worker,
            lookup_server,
        )
        backend_name = str(local_disk_backend)
        storage_backends[backend_name] = local_disk_backend

    if config.weka_path is not None:
        weka_backend = WekaGdsBackend(config, loop, memory_allocator, dst_device)
        # TODO(Serapheim): there's a chance we don't want the local
        # CPU cache in front of ours. Let's experiment and potentially

Review comment (suggestion):
this is a bit unclear: previously, I could have omitted to specify the inclusion of local_cpu and the rest of the engines would have received it anyway as a buffer, now it requires explicit indication.

Code after:
    if dst_device == "cuda":
        dst_device = f"cuda:{torch.cuda.current_device()}"

    storage_backends: OrderedDict[str, StorageBackendInterface] = OrderedDict()

    if config.enable_nixl:
        if config.enable_xpyd:
            # First Party
            from lmcache.v1.storage_backend.nixl_backend_v3 import NixlBackend

            storage_backends["NixlBackend"] = NixlBackend.CreateNixlBackend(
                config, metadata, memory_allocator
            )
        else:
            # First Party
            from lmcache.v1.storage_backend.nixl_backend import NixlBackend

            storage_backends["NixlBackend"] = NixlBackend.CreateNixlBackend(
                config, metadata
            )
            assert config.nixl_buffer_device is not None

    # TODO(Jiayi): The hierarchy is fixed for now
    # NOTE(Jiayi): The local_cpu backend is always created because
    # other backends might need it as a buffer.
    if config.enable_nixl and not config.local_cpu:
        pass
    else:
        local_cpu_backend = LocalCPUBackend(
            config,
            memory_allocator,
            lookup_server,
            lmcache_worker,
        )
        backend_name = str(local_cpu_backend)
        storage_backends[backend_name] = local_cpu_backend

    if config.local_disk and config.max_local_disk_size > 0:
        local_disk_backend = LocalDiskBackend(
            config,
            loop,
            local_cpu_backend,
            dst_device,
            lmcache_worker,
            lookup_server,
        )
        backend_name = str(local_disk_backend)
        storage_backends[backend_name] = local_disk_backend

    if config.weka_path is not None:
        weka_backend = WekaGdsBackend(config, loop, memory_allocator, dst_device)