Skip to content

Use ast.literal_eval to parse checkpoint shape in colocated_python_benchmark#1367

Open
adityasingh2400 wants to merge 1 commit into
apple:mainfrom
adityasingh2400:literal-eval-checkpoint-shape
Open

Use ast.literal_eval to parse checkpoint shape in colocated_python_benchmark#1367
adityasingh2400 wants to merge 1 commit into
apple:mainfrom
adityasingh2400:literal-eval-checkpoint-shape

Conversation

@adityasingh2400
Copy link
Copy Markdown

Fixes #1364.

create_state_spec_from_checkpoint in colocated_python_benchmark.py reads the per-tensor shape from the checkpoint index returned by read_index_file, and when that value is a string it passes it to eval():

shape = eval(value["shape"]) if isinstance(value["shape"], str) else value["shape"]

The index file is loaded from the checkpoint path, which can come from an external or shared location, so this evaluates code from data that is not necessarily trusted.

The shape is only ever a literal tuple or list of integers, so ast.literal_eval parses it to the same value while rejecting anything that is not a Python literal. This removes the code-execution path and lets the # pylint: disable=eval-used suppression go away.

ast.literal_eval round-trips the shapes this code sees ("(2, 3)", "(1024,)", "[2, 3, 4]", "()"), so behavior is unchanged for valid checkpoints.

…nchmark

create_state_spec_from_checkpoint read the shape field from the checkpoint
index and ran it through eval() when it was a string. The shape is only ever
a literal tuple or list of ints, so ast.literal_eval parses it identically
while refusing to execute arbitrary code from an untrusted index file. This
also lets the eval-used pylint suppression go away.
@adityasingh2400 adityasingh2400 requested a review from a team as a code owner May 24, 2026 13:27
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Bug #8 — Potential Eval Injection in create_state_spec_from_checkpoint (axlearn)

1 participant