def test_annotate(analysis: str, real_proto: programl_pb2.ProgramGraph, n: int): """Test all annotators over all real protos.""" try: # Use a lower timeout for testing. annotated = annotate.Annotate(analysis, real_proto, n, timeout=30) # Check that up to 'n' annotated graphs were generated. assert 0 <= len(annotated.protos) <= n # Check that output graphs have the same shape as the input graphs. for graph in annotated.protos: assert len(graph.node) == len(real_proto.node) assert len(graph.edge) == len(real_proto.edge) except data_flow_graphs.AnalysisTimeout: # A timeout error is acceptable. pass
def ProcessWorker(packed_args) -> AnnotationResult: """The process pool worker function. Accepts a batch of unlabelled graphs as inputs, labels them, and returns a list of graph tuples. """ start_time = time.time() # Unpack the args generated by ProcessWorkerArgsGenerator(). # Index into the tuple rather than arg unpacking so that we can assign # type annotations. worker_id: str = f"{packed_args[0]:06d}" max_mem_size: int = packed_args[1] analysis: str = packed_args[2] program_graphs: List[ProgramGraphProto] = packed_args[3] ctx: progress.ProgressBarContext = packed_args[4] # Set the hard limit on the memory size. Exceeding this limit will raise # a MemoryError. if FLAGS.limit_worker_mem: resource.setrlimit(resource.RLIMIT_DATA, (max_mem_size, max_mem_size)) resource.setrlimit(resource.RLIMIT_AS, (max_mem_size, max_mem_size)) graph_tuples = [] ctx.Log( 2, "[worker %s] received %s unlabelled graphs to process", worker_id, len(program_graphs), ) with ctx.Profile( 2, lambda t: (f"[worker {worker_id}] processed {len(program_graphs)} protos " f"({len(graph_tuples)} graphs, {humanize.Duration(t / len(program_graphs))} /proto)" ), ): for i, program_graph in enumerate(program_graphs): try: annotated_graphs = annotate.Annotate( analysis, programl.FromBytes(program_graph.serialized_proto, programl.StdinGraphFormat.PB), n=FLAGS.n, timeout=FLAGS.annotator_timeout, ) if annotated_graphs.graphs: # Record the annotated analysis results. for annotated_graph in annotated_graphs.graphs: graph_tuples.append( graph_tuple_database.GraphTuple.CreateFromNetworkX( annotated_graph, ir_id=program_graph.ir_id)) else: # Analysis produced no outputs, so just record an empty graph. graph_tuples.append( graph_tuple_database.GraphTuple.CreateEmpty( ir_id=program_graph.ir_id)) except Exception as e: _, _, tb = sys.exc_info() tb = traceback.extract_tb(tb, 2) filename, line_number, function_name, *_ = tb[-1] filename = pathlib.Path(filename).name ctx.Error( "Failed to annotate graph for ProgramGraph.ir_id=%d: %s " "(%s:%s:%s() -> %s)", program_graph.ir_id, e, filename, line_number, function_name, type(e).__name__, ) graph_tuples.append( graph_tuple_database.GraphTuple.CreateEmpty( ir_id=program_graph.ir_id)) return AnnotationResult( runtime=time.time() - start_time, proto_count=len(program_graphs), graph_tuples=graph_tuples, )
def test_timeout(one_proto: programl_pb2.ProgramGraph): """Test that error is raised if the analysis times out.""" with test.Raises(data_flow_graphs.AnalysisTimeout): annotate.Annotate("test_timeout", one_proto, timeout=1)
def test_invalid_analysis(one_proto: programl_pb2.ProgramGraph, n: int): """Test that error is raised if the input is invalid.""" with test.Raises(ValueError) as e_ctx: annotate.Annotate("invalid_analysis", one_proto, n) assert str(e_ctx.value).startswith("Unknown analysis: invalid_analysis. ")