From 140821bc2a0df2ab35790064aa2b7d2b5987b3ae Mon Sep 17 00:00:00 2001 From: Aditya Singh Date: Sun, 24 May 2026 06:27:11 -0700 Subject: [PATCH] Use ast.literal_eval to parse checkpoint shape in colocated_python_benchmark 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. --- axlearn/cloud/gcp/examples/colocated_python_benchmark.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/axlearn/cloud/gcp/examples/colocated_python_benchmark.py b/axlearn/cloud/gcp/examples/colocated_python_benchmark.py index 67ae536aa..a66e91911 100644 --- a/axlearn/cloud/gcp/examples/colocated_python_benchmark.py +++ b/axlearn/cloud/gcp/examples/colocated_python_benchmark.py @@ -19,6 +19,7 @@ """ import argparse +import ast import os import time from contextlib import contextmanager @@ -82,8 +83,11 @@ def create_state_spec_from_checkpoint(ckpt_path: str): continue if isinstance(value, dict) and "shape" in value and "dtype" in value: - # pylint: disable=eval-used - shape = eval(value["shape"]) if isinstance(value["shape"], str) else value["shape"] + shape = ( + ast.literal_eval(value["shape"]) + if isinstance(value["shape"], str) + else value["shape"] + ) dtype_str = value["dtype"] # Convert dtype string to jax dtype