def _AnnotateInSubprocess(
analysis: str,
graph: Union[programl_pb2.ProgramGraph, bytes],
n: int = 0,
timeout: int = 120,
binary_graph: bool = False,
) -> programl_pb2.ProgramGraphs:
"""Run this script in a subprocess.
This is the most robust method for enforcing the timeout, but has a huge
overhead in starting up a new python interpreter for every invocation.
DISCLAIMER: Because a target cannot depend on itself, all calling code must
add //deeplearning/ml4pl/graphs/labelled/dataflow:annotate to its list of
data dependencies.
Args:
analysis: The name of the analysis to run.
graph: The unlabelled ProgramGraph protocol buffer to to annotate, either
as a proto instance or as binary-encoded byte array.
n: The maximum number of labelled graphs to produce.
timeout: The maximum number of seconds to run the analysis for.
binary_graph: If true, treat the graph argument as a binary byte array.
Returns:
A ProgramGraphs protocol buffer.
Raises:
IOError: If serializing the input or output protos fails.
ValueError: If an invalid analysis is requested.
data_flow_graphs.AnalysisFailed: If the analysis raised an error.
data_flow_graphs.AnalysisTimeout: If the analysis did not complete within
the requested timeout.
"""
process = subprocess.Popen(
[
"timeout",
"-s9",
str(timeout),
str(SELF),
"--analysis",
analysis,
"--n",
str(n),
"--stdin_fmt",
"pb",
"--stdout_fmt",
"pb",
],
stdin=subprocess.PIPE,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
)
# Encode the input if required.
if binary_graph:
stdin = graph
else:
stdin = programl.ToBytes(graph, fmt=programl.StdoutGraphFormat.PB)
# Run this analysis script.
stdout, stderr = process.communicate(stdin)
if process.returncode == 9 or process.returncode == -9:
# Process was killed. We assume this is because of timeout, though it could
# be the user.
raise data_flow_graphs.AnalysisTimeout(timeout)
elif process.returncode == E_INVALID_INPUT:
raise IOError("Failed to serialize input graph")
elif process.returncode == E_INVALID_STDOUT:
raise IOError("Analysis failed to write stdout")
elif process.returncode:
raise data_flow_graphs.AnalysisFailed(
f"Analysis failed with returncode {process.returncode}: "
f"{stderr.decode('utf-8')}")
# Construct the protocol buffer from stdout.
output = programl.FromBytes(
stdout,
programl.StdinGraphFormat.PB,
proto=programl_pb2.ProgramGraphs(),
empty_okay=True,
)
return output
def test_fuzz_proto_bytes_equivalence(fmt: programl.InputOutputFormat):
"""Test that conversion to and from bytes does not change the proto."""
input = random_programl_generator.CreateRandomProto()
output = programl.FromBytes(programl.ToBytes(input, fmt), fmt)
assert input == output