Class SharedMemoryTensorBridge

rl_engine/executors/bridge.py:1696–1990

Same-node shared-memory transport with zero-copy import semantics. Publishing creates a shared-memory snapshot of model state. Importing that manifest returns tensor aliases to the shared storage instead of cloning. This proves the bridge lifecycle can expose a complete version thr

class SharedMemoryTensorBridge(LocalTensorCopyBridge):
    """
    Same-node shared-memory transport with zero-copy import semantics.

    Publishing creates a shared-memory snapshot of model state. Importing that
    manifest returns tensor aliases to the shared storage instead of cloning.
    This proves the bridge lifecycle can expose a complete version through
    shared memory, while keeping the CUDA IPC transport as a separate follow-up.
    """

    transport = "shared-memory"

    def __init__(self, *, source_worker: str = "local-training", source_rank: int = 0):
        super().__init__(source_worker=source_worker, source_rank=source_rank)
        self._shared_memory_segments: dict[str, dict[str, shared_memory.SharedMemory]] = {}
        self._owned_shared_memory_update_ids: set[str] = set()

    def publish(
        self,
        model: torch.nn.Module,
        *,
        weight_version: int,
        metadata: Optional[Mapping[str, Any]] = None,
    ) -> WeightUpdateManifest:
        user_metadata = _validated_manifest_metadata(metadata)
        if _BRIDGE_METADATA_KEY in user_metadata:
            raise WeightManifestValidationError(
                f"metadata key {_BRIDGE_METADATA_KEY!r} is reserved for bridge internals"
            )

        version = int(weight_version)
        if version <= self._latest_published_weight_version:
            raise WeightManifestValidationError(
                "weight_version must increase monotonically "
                f"(got {version}, latest {self._latest_published_weight_version})"
            )

        update_id = str(uuid.uuid4())
        tensors, segments, shared_metadata = self._snapshot_to_shared_memory(model)
        if not tensors:
            for segment in segments.values():
                segment.close()
                segment.unlink()
            raise WeightManifestValidationError("model state_dict produced no tensors")

        descriptors = {
            name: TensorDescriptor.from_tensor(name, tensor) for name, tensor in tensors.items()
        }
        manifest = WeightUpdateManifest(
            update_id=update_id,
            source_worker=self.source_worker,
            source_rank=self.source_rank,
            weight_version=version,
            transport=self.transport,
            tensors=descriptors,
            created_at=time.perf_counter(),
            metadata={**user_metadata, _BRIDGE_METADATA_KEY: shared_metadata},
        )