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