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))