From d4f7fbc67ef8245192c1ef340ab85a36385389d6 Mon Sep 17 00:00:00 2001 From: Edward Yang Date: Mon, 19 Jan 2026 09:58:35 -0500 Subject: [PATCH 1/2] Add --gist flag to upload tlparse output to GitHub Gist This allows open source contributors to easily share their tlparse artifacts with PyTorch developers by uploading to a secret GitHub Gist. The feature: - Checks for gh CLI availability and authentication - Shows an interactive privacy warning explaining that the output contains model architecture details and that secret gists are accessible to anyone with the link - Requires explicit y/N confirmation before uploading - Uploads all output files using `gh gist create` - Displays the gist URL for sharing Fixes https://github.com/meta-pytorch/tlparse/issues/161 --- src/cli.rs | 144 ++++++++++++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 143 insertions(+), 1 deletion(-) diff --git a/src/cli.rs b/src/cli.rs index 587de96..5f2e9a4 100644 --- a/src/cli.rs +++ b/src/cli.rs @@ -2,8 +2,9 @@ use clap::Parser; use anyhow::{bail, Context}; use std::fs; -use std::io::Read; +use std::io::{self, Read, Write}; use std::path::PathBuf; +use std::process::Command; use tlparse::{ // New reusable library API for multi-rank landing generation @@ -64,6 +65,9 @@ pub struct Cli { /// Port for the HTTP server (used with --serve). If not specified, finds an available port. #[arg(long)] port: Option, + /// Upload output to GitHub Gist using gh CLI. Requires gh to be installed and authenticated. + #[arg(long)] + gist: bool, } fn main() -> anyhow::Result<()> { @@ -125,6 +129,10 @@ fn main() -> anyhow::Result<()> { )?; } + if cli.gist { + handle_gist_upload(&cli.out)?; + } + if cli.serve { serve_directory(&cli.out, cli.port)?; } @@ -403,3 +411,137 @@ fn guess_content_type(path: &PathBuf) -> String { } .to_string() } + +/// Check if gh CLI is installed and authenticated +fn check_gh_cli() -> anyhow::Result<()> { + // Check if gh is available + let output = Command::new("gh") + .arg("--version") + .output() + .context("GitHub CLI (gh) is not installed. Install it from https://cli.github.com/")?; + + if !output.status.success() { + bail!("GitHub CLI (gh) is not working properly"); + } + + // Check if gh is authenticated + let output = Command::new("gh") + .args(["auth", "status"]) + .output() + .context("Failed to check gh authentication status")?; + + if !output.status.success() { + bail!( + "GitHub CLI is not authenticated. Run 'gh auth login' to authenticate.\n{}", + String::from_utf8_lossy(&output.stderr) + ); + } + + Ok(()) +} + +/// Prompt user with privacy warning and get confirmation +fn confirm_gist_upload() -> anyhow::Result { + println!(); + println!("╔════════════════════════════════════════════════════════════════════╗"); + println!("║ ⚠️ PRIVACY WARNING ⚠️ ║"); + println!("╠════════════════════════════════════════════════════════════════════╣"); + println!("║ The tlparse output contains detailed information about your model, ║"); + println!("║ including: ║"); + println!("║ • Model architecture and structure ║"); + println!("║ • Graph operations and transformations ║"); + println!("║ • Compilation traces and debug information ║"); + println!("║ ║"); + println!("║ This gist will be SECRET but ANYONE WITH THE LINK can view it. ║"); + println!("║ GitHub secret gists are not truly private - they are unlisted ║"); + println!("║ but accessible to anyone who has the URL. ║"); + println!("╚════════════════════════════════════════════════════════════════════╝"); + println!(); + print!("Do you want to proceed with uploading to GitHub Gist? [y/N]: "); + io::stdout().flush()?; + + let mut input = String::new(); + io::stdin().read_line(&mut input)?; + + let input = input.trim().to_lowercase(); + Ok(input == "y" || input == "yes") +} + +/// Collect all files from a directory recursively +fn collect_files_recursive(dir: &PathBuf, base: &PathBuf) -> anyhow::Result> { + let mut files = Vec::new(); + + for entry in fs::read_dir(dir)? { + let entry = entry?; + let path = entry.path(); + + if path.is_dir() { + files.extend(collect_files_recursive(&path, base)?); + } else if path.is_file() { + files.push(path); + } + } + + Ok(files) +} + +/// Upload the output directory to GitHub Gist +fn upload_to_gist(output_dir: &PathBuf) -> anyhow::Result { + println!("Collecting files from {}...", output_dir.display()); + + let files = collect_files_recursive(output_dir, output_dir)?; + + if files.is_empty() { + bail!("No files found in output directory"); + } + + println!("Uploading {} files to GitHub Gist...", files.len()); + + // Build the gh gist create command with all files + let mut cmd = Command::new("gh"); + cmd.arg("gist") + .arg("create") + .arg("--desc") + .arg("tlparse output - PyTorch compilation trace"); + + for file in &files { + cmd.arg(file); + } + + let output = cmd.output().context("Failed to execute gh gist create")?; + + if !output.status.success() { + bail!( + "Failed to create gist: {}", + String::from_utf8_lossy(&output.stderr) + ); + } + + let gist_url = String::from_utf8_lossy(&output.stdout).trim().to_string(); + + Ok(gist_url) +} + +/// Handle the --gist flag: check prerequisites, confirm, and upload +fn handle_gist_upload(output_dir: &PathBuf) -> anyhow::Result<()> { + // Check gh CLI is available and authenticated + check_gh_cli()?; + + // Prompt for confirmation + if !confirm_gist_upload()? { + println!("Gist upload cancelled."); + return Ok(()); + } + + // Upload to gist + let gist_url = upload_to_gist(output_dir)?; + + println!(); + println!("✅ Successfully uploaded to GitHub Gist!"); + println!(); + println!(" {}", gist_url); + println!(); + println!("Share this URL with PyTorch developers to help debug your issue."); + + Ok(()) +} From 62ecde4044797aea95957b294ddf3042a8b344d1 Mon Sep 17 00:00:00 2001 From: "Edward Z. Yang" Date: Mon, 19 Jan 2026 10:04:16 -0500 Subject: [PATCH 2/2] Fix formatting of privacy warning message --- src/cli.rs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/cli.rs b/src/cli.rs index 5f2e9a4..cbabc2c 100644 --- a/src/cli.rs +++ b/src/cli.rs @@ -444,7 +444,7 @@ fn check_gh_cli() -> anyhow::Result<()> { fn confirm_gist_upload() -> anyhow::Result { println!(); println!("╔════════════════════════════════════════════════════════════════════╗"); - println!("║ ⚠️ PRIVACY WARNING ⚠️ ║"); + println!("║ ⚠️ PRIVACY WARNING ⚠️ ║"); println!("╠════════════════════════════════════════════════════════════════════╣"); println!("║ The tlparse output contains detailed information about your model, ║"); println!("║ including: ║"); @@ -452,9 +452,9 @@ fn confirm_gist_upload() -> anyhow::Result { println!("║ • Graph operations and transformations ║"); println!("║ • Compilation traces and debug information ║"); println!("║ ║"); - println!("║ This gist will be SECRET but ANYONE WITH THE LINK can view it. ║"); - println!("║ GitHub secret gists are not truly private - they are unlisted ║"); - println!("║ but accessible to anyone who has the URL. ║"); + println!("║ This gist will be SECRET but ANYONE WITH THE LINK can view it. ║"); + println!("║ GitHub secret gists are not truly private - they are unlisted ║"); + println!("║ but accessible to anyone who has the URL. ║"); println!("╚════════════════════════════════════════════════════════════════════╝"); println!(); print!("Do you want to proceed with uploading to GitHub Gist? [y/N]: ");