def GetArtifact(self, request, context=None):
with filesystems.FileSystems.open(
self._artifact_path(request.retrieval_token,
request.name)) as fin:
# This value is not emitted, but lets us yield a single empty
# chunk on an empty file.
chunk = True
while chunk:
chunk = fin.read(self._chunk_size)
yield beam_artifact_api_pb2.ArtifactChunk(data=chunk)
def GetArtifact(self, request, context=None):
for artifact in self._get_manifest_proxy(request.retrieval_token).location:
if artifact.name == request.name:
with self._open(artifact.uri, 'r') as fin:
# This value is not emitted, but lets us yield a single empty
# chunk on an empty file.
chunk = True
while chunk:
chunk = fin.read(self._chunk_size)
yield beam_artifact_api_pb2.ArtifactChunk(data=chunk)
break
else:
raise ValueError('Unknown artifact: %s' % request.name)
def artifact_request_generator():
artifact_metadata = beam_artifact_api_pb2.ArtifactMetadata(
name=artifact_name,
sha256=_get_file_hash(local_path_to_artifact))
metadata = beam_artifact_api_pb2.PutArtifactMetadata(
staging_session_token=self._staging_session_token,
metadata=artifact_metadata)
request = beam_artifact_api_pb2.PutArtifactRequest(metadata=metadata)
yield request
with open(local_path_to_artifact, 'rb') as f:
while True:
chunk = f.read(1 << 21) # 2MB
if not chunk:
break
request = beam_artifact_api_pb2.PutArtifactRequest(
data=beam_artifact_api_pb2.ArtifactChunk(data=chunk))
yield request
self._artifacts.append(artifact_metadata)
def put_data(chunk):
return beam_artifact_api_pb2.PutArtifactRequest(
data=beam_artifact_api_pb2.ArtifactChunk(data=chunk))
