from __future__ import annotations
from collections.abc import Sequence
from .. import p2p_pb2
def tensor_source_metadata(
tensors: Sequence[p2p_pb2.TensorDescriptor],
) -> p2p_pb2.TensorSourceMetadata:
return p2p_pb2.TensorSourceMetadata(tensors=list(tensors))
def worker_tensor_descriptors(worker: p2p_pb2.WorkerMetadata):
payload = worker.WhichOneof("source_payload")
if payload == "tensor_source":
return worker.tensor_source.tensors
if payload == "artifact_source":
return []
return worker.tensors
def worker_tensor_count(worker: p2p_pb2.WorkerMetadata) -> int:
return len(worker_tensor_descriptors(worker))
def accelerators_compatible(target: str, source: str) -> bool:
if not target or not source:
return True
return target == source