PyTorch's Execution Trace Observer Producing Invalid JSON

[wkaisertexas]

🐛 Describe the bug

I am using PyTorch's execution trace observer to collect traces about a basic MNIST program. While the program does produce valid data (everything is there), it does not produce valid json.

from torch.profiler import profile, schedule, tensorboard_trace_handler, ProfilerActivity, ExecutionTraceObserver

tracing_schedule = schedule(skip_first=5, wait=5, warmup=2, active=2, repeat=1)

et = ExecutionTraceObserver()
et.register_callback("pytorch_et.json")
et.start()
with profile(
  activities = [ProfilerActivity.CPU, ProfilerActivity.CUDA],
  schedule = tracing_schedule,
  on_trace_ready=lambda x: x.export_chrome_trace("./kineto_trace.json"),
  profile_memory = True,
  record_shapes = True,
  with_stack = True
) as prof:
  for epoch in range(1): # number of epochs
    for step, batch_data in list(enumerate(train_data))[:1]:
        train(batch_data)
        prof.step()

et.stop()
et.unregister_callback()

The JSON produced will be contain extra commas at the end of the file and omitted commas between objects. Despite these formatting errors, the data is still there and cleaning the trace with a basic script to add or remove commas yields correct data.

@briancoutinho

The full Google Collab of my issue is available for consideration.

Versions

PyTorch version: 2.2.1+cu121
Is debug build: False
CUDA used to build PyTorch: 12.1
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.3 LTS (x86_64)
GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0
Clang version: 14.0.0-1ubuntu1.1
CMake version: version 3.27.9
Libc version: glibc-2.35

Python version: 3.10.12 (main, Nov 20 2023, 15:14:05) [GCC 11.4.0] (64-bit runtime)
Python platform: Linux-6.1.85+-x86_64-with-glibc2.35
Is CUDA available: True
CUDA runtime version: 12.2.140
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: Tesla T4
Nvidia driver version: 535.104.05
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.9.6
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.9.6
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.9.6
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.9.6
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.9.6
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.9.6
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.9.6
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture: x86_64
CPU op-mode(s): 32-bit, 64-bit
Address sizes: 46 bits physical, 48 bits virtual
Byte Order: Little Endian
CPU(s): 2
On-line CPU(s) list: 0,1
Vendor ID: GenuineIntel
Model name: Intel(R) Xeon(R) CPU @ 2.30GHz
CPU family: 6
Model: 63
Thread(s) per core: 2
Core(s) per socket: 1
Socket(s): 1
Stepping: 0
BogoMIPS: 4599.99
Flags: fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ss ht syscall nx pdpe1gb rdtscp lm constant_tsc rep_good nopl xtopology nonstop_tsc cpuid tsc_known_freq pni pclmulqdq ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt aes xsave avx f16c rdrand hypervisor lahf_lm abm invpcid_single ssbd ibrs ibpb stibp fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid xsaveopt arat md_clear arch_capabilities
Hypervisor vendor: KVM
Virtualization type: full
L1d cache: 32 KiB (1 instance)
L1i cache: 32 KiB (1 instance)
L2 cache: 256 KiB (1 instance)
L3 cache: 45 MiB (1 instance)
NUMA node(s): 1
NUMA node0 CPU(s): 0,1
Vulnerability Gather data sampling: Not affected
Vulnerability Itlb multihit: Not affected
Vulnerability L1tf: Mitigation; PTE Inversion
Vulnerability Mds: Vulnerable; SMT Host state unknown
Vulnerability Meltdown: Vulnerable
Vulnerability Mmio stale data: Vulnerable
Vulnerability Reg file data sampling: Not affected
Vulnerability Retbleed: Vulnerable
Vulnerability Spec rstack overflow: Not affected
Vulnerability Spec store bypass: Vulnerable
Vulnerability Spectre v1: Vulnerable: __user pointer sanitization and usercopy barriers only; no swapgs barriers
Vulnerability Spectre v2: Vulnerable; IBPB: disabled; STIBP: disabled; PBRSB-eIBRS: Not affected; BHI: Vulnerable (Syscall hardening enabled)
Vulnerability Srbds: Not affected
Vulnerability Tsx async abort: Not affected

Versions of relevant libraries:
[pip3] numpy==1.25.2
[pip3] torch==2.2.1+cu121
[pip3] torchaudio==2.2.1+cu121
[pip3] torchdata==0.7.1
[pip3] torchsummary==1.5.1
[pip3] torchtext==0.17.1
[pip3] torchvision==0.17.1+cu121
[pip3] triton==2.2.0

cc @robieta @chaekit @aaronenyeshi @guotuofeng @guyang3532 @dzhulgakov @davidberard98 @briancoutinho @sraikund16 @sanrise

[wkaisertexas]

pytorch/torch/csrc/profiler/standalone/execution_trace_observer.cpp

Line 208 in 8777443

ExecutionTraceObserver() = default;

The JSON serialization code, particularly the escaping done for python looks to be suspect. Could be a simple fix

[wkaisertexas]

pytorch/torch/csrc/profiler/standalone/execution_trace_observer.cpp

Lines 573 to 640 in 8777443

	static void recordOperatorStart(
	ExecutionTraceObserver& ob,
	FunctionCallContext& fc,
	const RecordFunction& fn) {
	auto tid = fn.threadId();
	try {
	const std::lock_guard<std::recursive_mutex> lock(ob.g_mutex);
	// if current thread stack is empty, push the root node to the stack first
	if (ob.op_stack[tid].empty()) {
	auto thread_node_id = ob.getNewID();
	ob.op_stack[tid].push(thread_node_id);
	writeJsonNode(
	ob.out,
	"[pytorch|profiler|execution_trace|thread]",
	thread_node_id,
	0, // rf_id
	root_id,
	0, // fw_parent
	-1, // seq_id
	static_cast<std::underlying_type_t<RecordScope>>(
	RecordScope::USER_SCOPE),
	tid,
	0); // fw_tid
	ob.out << ",";
	}
	fc.name = fn.name();
	auto num_inputs = fn.num_inputs();
	const auto inputs = fn.inputs();
	VLOG(2) << "inputs: " << num_inputs << " " << inputs.size() << std::endl;
	// We have two cases: for unboxed kernel, we have num_inputs ==
	// inputs.size() for boxed kernel using stack, there could be more elements
	// on the stack from previous ops.
	// TORCH_INTERNAL_ASSERT(num_inputs <= inputs.size());
	if (num_inputs > inputs.size()) {
	LOG(WARNING) << "RecordFunction " << fc.name
	<< " expected num_inputs=" << num_inputs
	<< " > inputs.size()=" << inputs.size();
	return;
	}
	// need to account for Stack mode where the inputs are at the end.
	size_t input_start = inputs.size() - num_inputs;
	for (const auto i : c10::irange(input_start, inputs.size())) {
	appendValueInfo(
	ob, inputs[i], fc.input_values, fc.input_types, fc.input_shapes);
	}
	handleKernelBackendInfo(fc, fn);
	fc.parent_id = ob.op_stack[tid].top();
	// get parent id from the forward stack, this can be different for
	// autograd ops, which may execute on a different thread than the original
	// thread (which should have the parent op on the stack).
	auto fw_tid = fn.forwardThreadId();
	if (fw_tid != 0) {
	fc.fw_parent_id = ob.op_stack[fw_tid].top();
	}
	// all input nodes should have id > op_id
	fc.op_id = ob.getNewID();
	ob.op_stack[tid].push(fc.op_id);
	} catch (const std::exception& e) {
	LOG(WARNING) << "Exception in execution trace observer: " << e.what();
	}
	}

This segment in particular which naively adds commas in a list seems suspect upon initial glance

[wkaisertexas]

Additionally, there are no double ending curly braces to close off the json object

pytorch/torch/csrc/profiler/standalone/execution_trace_observer.cpp

Lines 280 to 291 in 8777443

	static inline std::string getAttrJson(
	const std::string& name,
	const std::string& type,
	const std::string& value) {
	// note name and type are not quoted but value should be if it is a string.
	return fmt::format(
	R"JSON(
	{{"name": "{}", "type": "{}", "value": {}}})JSON",
	name,
	type,
	value);
	}

Here we can see the tail produced by the execution trace observer where no double curly braces exist to close the JSON.

8833173  {
8833174   "id": 1, "name": "[pytorch|profiler|execution_trace|process]", "ctrl_deps": 1,
8833175   "inputs": {"values": [], "shapes": [], "types": []},
8833176   "outputs": {"values": [], "shapes": [], "types": []},
8833177   "attrs": [{"name": "rf_id", "type": "uint64", "value": 0}, {"name": "fw_parent", "type": "uint64", "value": 0}, {"name": "seq_id", "type": "int64", "value": -1}, {"name": "scope", "type": "uint64", "value": 7}, {"name": "tid", "type": "uint64", "value": 0}, {"name": "fw_tid", "type": "uint64", "value": 0}, {"name": "op_schema", "type": "string", "value": ""}]
8833178  }],
8833179  "finish_ts": 130558
8833180 }
8833181 
8833182 

[briancoutinho]

@wkaisertexas still not clear what the missing commas are? the addAttrs() will create a json blurb per attr and then it gets joined together. This has a ',' at start assuming there will be other attributed before it

pytorch/torch/csrc/profiler/standalone/execution_trace_observer.cpp

Lines 540 to 570 in 8777443

	auto addAttr =
	[&](const char* commsMetaName, const char* etMetaName, const char* type) {
	auto it = meta.find(commsMetaName);
	if (it != meta.end()) {
	attrs.push_back(getAttrJson(etMetaName, type, it->second));
	}
	};
	addAttr(kCommsName, kETCommsName, "string");
	addAttr(kDtype, kDtype, "string");
	addAttr(kInMsgNelems, kETInMsgNelems, "uint64");
	addAttr(kOutMsgNelems, kETOutMsgNelems, "uint64");
	// following two metadata are lists.
	addAttr(kInSplit, kETInSplit, "string");
	addAttr(kOutSplit, kETOutSplit, "string");
	addAttr(kGlobalRankStart, kETGlobalRankStart, "uint64");
	addAttr(kGlobalRankStride, kETGlobalRankStride, "uint64");
	// pg_name is a string.
	addAttr(kProcessGroupName, kETProcessGroupName, "string");
	addAttr(kProcessGroupDesc, kETProcessGroupDesc, "string");
	addAttr(kGroupSize, kETGroupSize, "uint64");
	#endif // USE_DISTRIBUTED
	// XXX consider using as string stream?
	return attrs.size() == 0 ? "" : fmt::format(", {}", fmt::join(attrs, ", "));

Maybe a bug here

pytorch/torch/csrc/profiler/standalone/execution_trace_observer.cpp

Line 320 in 8777443

"attrs": [{{"name": "rf_id", "type": "uint64", "value": {}}},{{"name": "fw_parent", "type": "uint64", "value": {}}},{{"name": "seq_id", "type": "int64", "value": {}}},{{"name": "scope", "type": "uint64", "value": {}}},{{"name": "tid", "type": "uint64", "value": {}}},{{"name": "fw_tid", "type": "uint64", "value": {}}},{{"name": "op_schema", "type": "string", "value": "{}"}},{{"name": "kernel_backend", "type": "string", "value": "{}"}},{{"name": "kernel_file", "type": "string", "value": "{}"}}{}]

[wkaisertexas]

That seems like the more likely explanation for missing commas between segments.

There is also a mismatch between opening and closing curly braces.

However, whole execution trace serialization seems error prone.

[alexseceks]

Also reproducing with torch 2.4.0.

[briancoutinho]

@alexseceks @wkaisertexas sorry about the delay, for your repro can you share the train() and data functions? or some smaller model for repro too

[wkaisertexas]

I am using a mnist model for training. Nothing fancy, I wanted to check if the execution trace observer was working. I think it is just a serialization issue however.

[briancoutinho]

@wkaisertexas just to be precise could you share the mnist model snippet, will be easier to have apple-apple repro

[wkaisertexas]

@briancoutinho

Sorry for the late response on this!

https://colab.research.google.com/drive/1t0IgZ8eZ5N7zwZci-jOMXycv78isUWS2

[wkaisertexas]

@briancoutinho I am wondering if there are any updates