feat(clp): Add CLP support to the OSS trace parsing workflow; Add yscope_clp_core dependency.

[Bill-hbrhbr]

Changes

• Write CLP archives directly to the requested .clp file path instead of creating an archive directory.
• Add support for CLP compressed trace output from structured logging.
• Add support for reading CLP compressed trace output in unified_parse.
  • The workflow now cleans up the logging handler before running unified_parse.
• Add support for writing parsed output as CLP archives.

Utilities

• Add CLP single-file archive detection to the compression utilities.
• Add enumerate_json() to skip line parsing when JSON objects are expected.

Misc

• Add yscope_clp_core to pyproject.toml.
• Add a CLP parsed output public asset for the UI.

Tests

• The workflow can be tested with:

TMPDIR=<tmp_dir> TRITON_TRACE_COMPRESSION=clp TEST_KEEP_OUTPUT=1 \
  uv run python3 -m unittest tests.gpu.test_complex_kernels

• Add CLP coverage for compression helpers.
• Add enumerate_json() coverage for gzip, zstd, and CLP.

[FindHao]

The overall design direction is sound: CLP is repositioned from a "parsed-output format" to "a compression format for raw traces" (alongside gzip/zstd/none); enumerate_json() is a clean abstraction; the large parse_single_rank refactor is logically equivalent — verified line by line, nothing dropped; and the test coverage is reasonable.

However, there are 2 issues I'd consider merge blockers (① a mandatory dependency that breaks cross-platform installs; ② in structured_logging.py, the handler's close() is wrongly replaced with the private _ensure_stream_closed()), plus a few semantic / readability concerns. More importantly: this PR changes too many independent concerns at once and should be split into several smaller, independently reviewable PRs — in particular, the extremely core structured_logging.py must be carved out on its own (see the next section).

---

⭐ Top recommendation: split into smaller PRs

This PR bundles several independent concerns into a single commit: dependency management, compression utilities, CLP read, CLP write + log-handler lifecycle, and a parsing-workflow refactor. That makes review hard and rollback risky — the .close() downgrade bug above slipped in precisely because it was buried in a big PR.

Splitting principle: split by concern, not by file. In particular, carve out the "handler lifecycle" change — the most sensitive one, and one that is fundamentally unrelated to CLP.

PR-1: Pure refactor — introduce enumerate_json and adopt it in trace_processor (no CLP)

• tools/compression.py: add enumerate_json() (gzip/zstd/plain branches only)
• parse/trace_processor.py: switch _prescan_for_fake_compilations and parse_single_rank to enumerate_json
• Corresponding enumerate_json tests (gzip/zstd/plain/malformed-line)
• ⚠️ Do not include the unrelated bidirectional-mapping change (see nit 5)
• Value: a clean de-duplication refactor that does not depend on CLP, and can land first.

PR-2 (standalone, sensitive ⚠️): finalize the trace stream before parse (handler lifecycle)

• structured_logging.py: extract clean_up_log_handler (using .close(), see High 3)
• context_manager.py: call it before unified_parse
• Nature: this is actually a CLP-agnostic, general correctness fix — "flush+close the trace stream before reading the files, so on-disk files are complete." It's harmless for gzip/zstd/none and merely required for CLP. Because it is sensitive, self-contained, and unrelated to the contentious dependency, it must be its own PR and reviewed line by line.
• ⚠️ Do not include the _compress() removal from close() here — that's a different concern, see PR-3.

PR-3: CLP support (dependency + read + write)

• pyproject.toml: add yscope-clp-core as an optional extra (see High 1)
• clp.py: ClpTextStream, clp_open
• tools/compression.py: the clp branch in detect_compression, is_clp_file, the clp branches in open_compressed_file / enumerate_json
• structured_logging.py: remove _compress() from close() — this is tightly coupled to the dependency version bump: with >=0.12.1b0, open_archive writes directly to .clp and _compress() is no longer needed; the two must ship in the same PR (merging either alone breaks things), so it belongs here, not in PR-2
• parse/common.py: the clp write branch in compress_single_file
• CLP read/write tests (with a skipUnless guard, see High 2)

PR-4: Parsing-workflow / source-flow cleanup

• parse/utils.py: oss_run flow rewrite, temp-dir cleanup
• parse/common.py: drop the LOCAL_CLP directory copy in copy_local_to_tmpdir, temp-dir prefixes
• parse/source_type.py: decide what to do with LOCAL_CLP (see Medium 4)

Rough dependency order: PR-1 / PR-2 can go in parallel; PR-3 depends on PR-1 (it uses the clp branch of enumerate_json) and ideally lands after PR-2 (so CLP write works end to end through the context manager); PR-4 last.

---

🔴 High priority (suggested blocker)

1. Making yscope-clp-core a mandatory dependency breaks cross-platform installs

pyproject.toml:11 adds it to the required dependencies. On PyPI, yscope-clp-core 0.12.1b0 ships exactly one wheel (verified via the PyPI JSON API + local curl):

yscope_clp_core-0.12.1b0-py3-none-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl

• ❌ No macOS / Windows wheel
• ❌ No Linux aarch64 wheel (Grace-Hopper / GH200 and other ARM GPU machines can't install it)
• ❌ No sdist (no fallback to building from source)
• ⚠️ Every published version is a pre-release (beta)

Consequence: pip install tritonparse will fail outright on Mac, Windows, and Linux-ARM. Yet tritonparse's parsing + Web UI are cross-platform usable (Mac users often just parse / view traces), so this is a real regression.

Also, CLP is one of several optional compression formats, and the runtime code is already written to treat it as optional:

• tools/compression.py detect_compression uses try: from yscope_clp_core import ...; except ImportError: pass
• The clp branches in open_compressed_file / enumerate_json are only reached when detection returns "clp"
• structured_logging.py:84 is also if TRITON_TRACE_COMPRESSION == "clp": from .clp import clp_open

So pyproject contradicts the runtime code. Recommend making it an optional extra (the repo already uses the optional-dependencies pattern):

[project.optional-dependencies]
clp = ["yscope-clp-core>=0.12.1b0"]

Change CI to pip install -e ".[test,clp]" (CI is Linux x86_64, so it installs fine). Default installs stay cross-platform; users who need CLP run pip install tritonparse[clp]. (Belongs to PR-3.)

2. The test module hard-imports CLP at top level, coupling all compression tests

tests/cpu/test_compression.py:16 does from tritonparse.clp import clp_open at module top level, and clp.py:9 does from yscope_clp_core import ... at top level. Verified: when yscope_clp_core is missing, the entire test_compression module fails to import — even the gzip/zstd/plain-text tests all error out:

ModuleNotFoundError: No module named 'yscope_clp_core'
  File ".../tests/cpu/test_compression.py", line 16, in <module>
    from tritonparse.clp import clp_open

If 1 is adopted (optional dependency), this should be guarded with unittest.skipUnless(...), and the clp import moved into the CLP-specific test cases so non-CLP tests aren't dragged down. (Belongs to PR-3.)

3. clean_up_log_handler replaces close() with the private _ensure_stream_closed() (a wrong substitution)

The new clean_up_log_handler at structured_logging.py:2013-2023 (called from both context_manager.__exit__ and clear_logging_config) calls the private _ensure_stream_closed(), whereas pre-PR clear_logging_config called the public TRITON_TRACE_HANDLER.close() (see the diff: - TRITON_TRACE_HANDLER.close()).

Why it's wrong — this is an extract-method refactor, so it must be behavior-preserving. Since the original implementation called .close(), the extracted function should still call .close(); changing it to ._ensure_stream_closed() is an undeclared behavior regression smuggled into a refactor. It has two concrete harms:

1. It drops the lock → races with emit(). close() runs inside the self.acquire()/release() lock (:1302-1317); _ensure_stream_closed() does not (:1335-1342). emit() is invoked by logging while it holds the handler lock (Handler.handle()), so calling _ensure_stream_closed() inside emit is safe — but clean_up_log_handler is called outside that lock. If another thread is still in emit() (background threads still compiling/launching at exit, or multi-threaded tracing), teardown's stream.close(); stream=None collides with emit's stream.write(...) → writing to a closed stream / dereferencing None / a CLP archive that never gets finalized (corrupt).

2. It bypasses the deliberately-overridden close(). _ensure_stream_closed() is a private helper for the internal "close-then-reopen" path inside emit() / init_logs() (where the lock is already held). Using it for teardown skips TritonTraceHandler.close(), whose finally: logging.StreamHandler.close(self) guarantees the parent Handler.close() always runs (the comment attributes this to PyTorch PR #120289).

Fix (approach A, keep the extraction): inside clean_up_log_handler, call TRITON_TRACE_HANDLER.close(), and keep clear_logging_config's behavior byte-for-byte identical to pre-PR. Note the distinction: the newly added "call before parse" in context_manager is intentional and correct (CLP must be finalized before reading) — what should be reverted is only the close mechanism (.close()), not the call timing.

This is exactly why structured_logging.py must be its own PR (see PR-2): it's the core of the tracing write path, and a seemingly harmless close() → _ensure_stream_closed() swap introduced both a race and a bypass of an existing fix.

---

🟡 Medium

4. SourceType.LOCAL_CLP is now dead/misleading

parse/source_type.py:50-51 still returns LOCAL_CLP for "a directory containing a file with .clp in its name", but oss_run has removed every branch that handled it (parse/utils.py). In OSS it is now fully equivalent to LOCAL (from_cli_args also falls into the else branch, see common.py:366-374). The old "CLP dir = already-parsed output, skip parsing" semantics have been removed (correctly, per the new model), but leaving the enum value + detection in place is confusing.

Suggestion: either remove LOCAL_CLP and its detection, or add a comment explaining why it's retained (e.g., fbcode's from_cli_args still branches on it). Also, the detection uses substring ".clp" in p.name, which would falsely match things like foo.clp.bak. (Belongs to PR-4.)

---

🟢 Nits

5. Unrelated readability regression: the bidirectional-mapping loop

parse/trace_processor.py:575-576 changes the loop from enumerate(stage_names) + stage_names[i+1:] to enumerate(stage_names, start=1) + stage_names[i:]. The two are functionally identical (verified), but afterwards i no longer equals the index of src_stage, which is harder to read, and it's unrelated to CLP. Recommend reverting to keep the diff focused.

6. Redundant detect_compression

enumerate_json calls detect_compression once, and in the non-CLP branch open_compressed_file calls it again — two opens per file. Combined with the _prescan_for_fake_compilations + main pass (two scans), that's ~4 detections per trace file. The detection result could be threaded through, or detection done once. Minor impact.

7. clp.py minor blemish

__init__ has a superfluous explicit return None (harmless but non-idiomatic); the rest of the style is fine.

8. oss_run cleanup is not in a try/finally

The new shutil.rmtree(logs, ...) at parse/utils.py:189-192 won't run if parsing raises midway, leaking the staging temp dir. The old behavior leaked unconditionally, so this is not a regression, but a try/finally would be cleaner.