Skip to content

Commit c0fd47d

Browse files
authored
feat: add single-operator subgraph dataset generation script (#635)
* feat: add single-operator subgraph dataset generation script This commit introduces `generate_single_op_dataset.sh` to automate the workflow for generating single-operator subgraph datasets. * Refactor generate_single_op_dataset.sh for dynamic paths Refactor script for dynamic path detection and improved error handling. Added logging and workspace setup enhancements. * Refactor dataset generation script to strict serial execution mode * extract raw_list.txt generation as separate step and copy to final output * remove redundant clean_list.txt and rename raw_list.txt to generated_subgraphs_list.txt * move generate_single_op_dataset.sh to tools directory
1 parent be4075c commit c0fd47d

1 file changed

Lines changed: 135 additions & 0 deletions

File tree

Lines changed: 135 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,135 @@
1+
#!/bin/bash
2+
set -e
3+
4+
# ==============================================================================
5+
# Configuration Area
6+
# ==============================================================================
7+
8+
# Dynamic Path Retrieval
9+
PYTHON_EXEC=$(which python3)
10+
if [ -z "$PYTHON_EXEC" ]; then
11+
echo "Error: 'python3' not found in PATH. Please activate your virtualenv."
12+
exit 1
13+
fi
14+
15+
GRAPH_NET_ROOT=$($PYTHON_EXEC -c "import graph_net; import os; print(os.path.dirname(os.path.dirname(graph_net.__file__)))")
16+
if [ -z "$GRAPH_NET_ROOT" ]; then
17+
echo "Error: Could not determine GRAPH_NET_ROOT. Ensure 'graph_net' is installed or in PYTHONPATH."
18+
exit 1
19+
fi
20+
21+
RESUME="false"
22+
23+
# Workspace Setup
24+
TIMESTAMP=$(date +%Y%m%d_%H%M)
25+
WORKSPACE="/tmp/single_op_workspace_${TIMESTAMP}"
26+
MODEL_LIST="${MODEL_LIST:-${GRAPH_NET_ROOT}/graph_net/config/small100_torch_samples_list.txt}"
27+
28+
# Output Directories
29+
OP_NAMES_DIR="${WORKSPACE}/01_op_names"
30+
RANGES_DIR="${WORKSPACE}/02_ranges"
31+
RAW_SUBGRAPH_DIR="${WORKSPACE}/03_raw_subgraphs"
32+
RENAMED_DIR="${WORKSPACE}/04_renamed"
33+
DEDUPLICATED_DIR="${WORKSPACE}/05_deduplicated"
34+
35+
mkdir -p "$WORKSPACE"
36+
37+
# ==============================================================================
38+
# Main Pipeline
39+
# ==============================================================================
40+
41+
echo ">>> Starting Pipeline..."
42+
echo " Python: $PYTHON_EXEC"
43+
echo " Root: $GRAPH_NET_ROOT"
44+
45+
# 1. Prepare Data
46+
if [ ! -f "$MODEL_LIST" ]; then
47+
echo "Error: Model list not found at $MODEL_LIST"
48+
exit 1
49+
fi
50+
51+
# 2. Stage 1: Op Names
52+
echo ">>> Running Stage 1: Op Names..."
53+
python3 -m graph_net.model_path_handler \
54+
--model-path-list "${MODEL_LIST}" \
55+
--handler-config=$(base64 -w 0 <<EOF
56+
{
57+
"handler_path": "$GRAPH_NET_ROOT/graph_net/torch/sample_pass/op_names_extractor.py",
58+
"handler_class_name": "OpNamesExtractor",
59+
"handler_config": { "resume": $RESUME, "model_path_prefix": "$GRAPH_NET_ROOT", "output_dir": "$OP_NAMES_DIR" }
60+
}
61+
EOF
62+
)
63+
64+
# 3. Stage 2: Ranges
65+
echo ">>> Running Stage 2: Ranges..."
66+
python3 -m graph_net.apply_sample_pass \
67+
--model-path-list "${MODEL_LIST}" \
68+
--sample-pass-file-path "$GRAPH_NET_ROOT/graph_net/sample_pass/op_extract_points_generator.py" \
69+
--sample-pass-class-name "OpExtractPointsGenerator" \
70+
--sample-pass-config=$(base64 -w 0 <<EOF
71+
{
72+
"model_path_prefix": "$GRAPH_NET_ROOT", "op_names_path_prefix": "$OP_NAMES_DIR",
73+
"output_dir": "$RANGES_DIR", "subgraph_ranges_file_name": "subgraph_ranges.json"
74+
}
75+
EOF
76+
)
77+
78+
# 4. Stage 3: Decompose
79+
echo ">>> Running Stage 3: Decompose..."
80+
python3 -m graph_net.model_path_handler \
81+
--model-path-list "${MODEL_LIST}" \
82+
--handler-config=$(base64 -w 0 <<EOF
83+
{
84+
"handler_path": "$GRAPH_NET_ROOT/graph_net/torch/sample_pass/subgraph_generator.py",
85+
"handler_class_name": "SubgraphGenerator",
86+
"handler_config": {
87+
"resume": $RESUME, "model_path_prefix": "$GRAPH_NET_ROOT", "output_dir": "$RAW_SUBGRAPH_DIR",
88+
"subgraph_ranges_json_root": "$RANGES_DIR", "subgraph_ranges_json_file_name": "subgraph_ranges.json",
89+
"group_head_and_tail": false, "chain_style": false, "device": "cuda"
90+
}
91+
}
92+
EOF
93+
)
94+
95+
# 5. Generate generated_subgraphs_list.txt
96+
echo ">>> Generating generated_subgraphs_list.txt..."
97+
find ${RAW_SUBGRAPH_DIR} -name "model.py" \
98+
| xargs dirname \
99+
| xargs realpath --relative-to=${RAW_SUBGRAPH_DIR} \
100+
> "${WORKSPACE}/generated_subgraphs_list.txt"
101+
102+
# 6. Post-processing: Rename
103+
echo ">>> Running Post-processing: Rename..."
104+
python3 -m graph_net.model_path_handler \
105+
--model-path-list "${WORKSPACE}/generated_subgraphs_list.txt" \
106+
--handler-config=$(base64 -w 0 <<EOF
107+
{
108+
"handler_path": "$GRAPH_NET_ROOT/graph_net/sample_pass/ast_graph_variable_renamer.py",
109+
"handler_class_name": "AstGraphVariableRenamer",
110+
"handler_config": {
111+
"device": "cuda", "try_run": false, "resume": $RESUME,
112+
"model_path_prefix": "${RAW_SUBGRAPH_DIR}",
113+
"data_input_predicator_filepath": "$GRAPH_NET_ROOT/graph_net/torch/constraint_util.py",
114+
"data_input_predicator_class_name": "NaiveDataInputPredicator",
115+
"model_runnable_predicator_filepath": "$GRAPH_NET_ROOT/graph_net/torch/constraint_util.py",
116+
"model_runnable_predicator_class_name": "ModelRunnablePredicator",
117+
"output_dir": "${RENAMED_DIR}"
118+
}
119+
}
120+
EOF
121+
)
122+
123+
# 7. Post-processing: Deduplicate
124+
echo ">>> Running Post-processing: Deduplicate..."
125+
if [ -d "${DEDUPLICATED_DIR}" ]; then rm -rf "${DEDUPLICATED_DIR}"; fi
126+
127+
python3 -m graph_net.tools.deduplicated \
128+
--samples-dir ${RENAMED_DIR} \
129+
--target-dir ${DEDUPLICATED_DIR}
130+
131+
# Copy generated_subgraphs_list.txt to final output
132+
cp "${WORKSPACE}/generated_subgraphs_list.txt" "${DEDUPLICATED_DIR}/"
133+
134+
echo ">>> ALL DONE. Final dataset located at: ${DEDUPLICATED_DIR}"
135+
echo ">>> generated_subgraphs_list.txt also saved to: ${DEDUPLICATED_DIR}/generated_subgraphs_list.txt"

0 commit comments

Comments
 (0)