def __init__(self, staging_dir=None): super(LocalJobServicer, self).__init__() self._cleanup_staging_dir = staging_dir is None self._staging_dir = staging_dir or tempfile.mkdtemp() self._artifact_service = artifact_service.ArtifactStagingService( artifact_service.BeamFilesystemHandler(self._staging_dir).file_writer) self._artifact_staging_endpoint = None # type: Optional[endpoints_pb2.ApiServiceDescriptor]
def _start_artifact_service(self, jar, requested_port): self._artifact_manager = JarArtifactManager(self._jar, self.ARTIFACT_FOLDER) self._artifact_staging_service = artifact_service.ArtifactStagingService( self._artifact_manager.file_writer) self._artifact_staging_service.register_job( self._job_id, { env_id: env.dependencies for (env_id, env ) in self._pipeline_proto.components.environments.items() }) self._artifact_staging_server = grpc.server( futures.ThreadPoolExecutor()) port = self._artifact_staging_server.add_insecure_port('[::]:%s' % requested_port) beam_artifact_api_pb2_grpc.add_ArtifactStagingServiceServicer_to_server( self._artifact_staging_service, self._artifact_staging_server) self._artifact_staging_endpoint = endpoints_pb2.ApiServiceDescriptor( url='localhost:%d' % port) self._artifact_staging_server.start() _LOGGER.info('Artifact server started on port %s', port) return port
def test_push_artifacts(self): unresolved = beam_runner_api_pb2.ArtifactInformation( type_urn='unresolved') resolved_a = self.embedded_artifact(data=b'a', name='a.txt') resolved_b = self.embedded_artifact(data=b'bb', name='b.txt') dep_big = self.embedded_artifact(data=b'big ' * 100, name='big.txt') class TestArtifacts(object): def ResolveArtifacts(self, request): replacements = [] for artifact in request.artifacts: if artifact.type_urn == 'unresolved': replacements += [resolved_a, resolved_b] else: replacements.append(artifact) return beam_artifact_api_pb2.ResolveArtifactsResponse( replacements=replacements) def GetArtifact(self, request): if request.artifact.type_urn == common_urns.artifact_types.EMBEDDED.urn: content = proto_utils.parse_Bytes( request.artifact.type_payload, beam_runner_api_pb2.EmbeddedFilePayload).data for k in range(0, len(content), 13): yield beam_artifact_api_pb2.GetArtifactResponse( data=content[k:k + 13]) else: raise NotImplementedError file_manager = InMemoryFileManager() server = artifact_service.ArtifactStagingService( file_manager.file_writer) server.register_job('staging_token', {'env': [unresolved, dep_big]}) # "Push" artifacts as if from a client. t = threading.Thread(target=lambda: artifact_service.offer_artifacts( server, TestArtifacts(), 'staging_token')) t.daemon = True t.start() resolved_deps = server.resolved_deps('staging_token', timeout=5)['env'] expected = { 'a.txt': b'a', 'b.txt': b'bb', 'big.txt': b'big ' * 100, } for dep in resolved_deps: self.assertEqual(dep.type_urn, common_urns.artifact_types.FILE.urn) self.assertEqual(dep.role_urn, common_urns.artifact_roles.STAGING_TO.urn) type_payload = proto_utils.parse_Bytes( dep.type_payload, beam_runner_api_pb2.ArtifactFilePayload) role_payload = proto_utils.parse_Bytes( dep.role_payload, beam_runner_api_pb2.ArtifactStagingToRolePayload) self.assertTrue( type_payload.path.endswith(role_payload.staged_name), type_payload.path) self.assertEqual(file_manager.get(type_payload.path), expected.pop(role_payload.staged_name)) self.assertEqual(expected, {})