Add parsing support for json and string artifacts (#40)

This lets us fulfill the following interface:

```
trace_structured(
  "artifact",
  lambda: {
    name: <name>,
    encoding: "json" | "string"
  },
  payload_fn= lambda: <output>
)
```

This gets logged as a json or txt artifact associated with a compile:

![image](https://github.com/ezyang/tlparse/assets/4811293/35d23136-bc79-4cdf-910d-abca19c4eb5c)

src/parsers.rs

@@ -5,6 +5,8 @@ use std::path::Path;
 use std::path::PathBuf;
 use tinytemplate::TinyTemplate;
 
+use serde_json::Value;
+
 use syntect::highlighting::ThemeSet;
 use syntect::parsing::SyntaxSet;
 
@@ -418,6 +420,45 @@ impl StructuredLogParser for AOTAutogradBackwardCompilationMetricsParser<'_> {
     }
 }
 
+pub struct ArtifactParser;
+impl StructuredLogParser for ArtifactParser {
+    fn name(&self) -> &'static str {
+        "artifact"
+    }
+    fn get_metadata<'e>(&self, e: &'e Envelope) -> Option<Metadata<'e>> {
+        e.artifact.as_ref().map(|m| Metadata::Artifact(m))
+    }
+    fn parse<'e>(
+        &self,
+        lineno: usize,
+        metadata: Metadata<'e>,
+        _rank: Option<u32>,
+        compile_id: &Option<CompileId>,
+        payload: &str,
+    ) -> anyhow::Result<ParserResults> {
+        if let Metadata::Artifact(metadata) = metadata {
+            match metadata.encoding.as_str() {
+                "string" => {
+                    let filename = format!("{}.txt", metadata.name);
+                    simple_file_output(&filename, lineno, compile_id, &payload)
+                }
+                "json" => {
+                    let filename = format!("{}.json", metadata.name);
+                    let value: Value = serde_json::from_str(&payload).unwrap();
+                    let pretty = serde_json::to_string_pretty(&value).unwrap();
+                    simple_file_output(&filename, lineno, compile_id, &pretty)
+                }
+                _ => Err(anyhow::anyhow!(
+                    "Unsupported encoding: {}",
+                    metadata.encoding
+                )),
+            }
+        } else {
+            Err(anyhow::anyhow!("Expected Artifact metadata"))
+        }
+    }
+}
+
 // Register your parser here
 pub fn default_parsers<'t>(tt: &'t TinyTemplate<'t>) -> Vec<Box<dyn StructuredLogParser + 't>> {
     // We need to use Box wrappers here because vecs in Rust need to have known size
@@ -447,6 +488,7 @@ pub fn default_parsers<'t>(tt: &'t TinyTemplate<'t>) -> Vec<Box<dyn StructuredLo
         Box::new(OptimizeDdpSplitChildParser),
         Box::new(AOTAutogradBackwardCompilationMetricsParser { tt }), // TODO: use own tt instances
         Box::new(LinkParser),
+        Box::new(ArtifactParser),
     ];
 
     result

src/types.rs

@@ -234,6 +234,12 @@ pub struct LinkMetadata {
     pub url: String,
 }
 
+#[derive(Debug, Deserialize)]
+pub struct ArtifactMetadata {
+    pub name: String,
+    pub encoding: String,
+}
+
 #[derive(Debug, Deserialize, Serialize)]
 pub struct CompilationMetricsMetadata {
     // Other information like frame_key, co_name, etc. are already in envelope
@@ -345,6 +351,7 @@ pub enum Metadata<'e> {
     OptimizeDdpSplitChild(&'e OptimizeDdpSplitChildMetadata),
     CompilationMetrics(&'e CompilationMetricsMetadata),
     AOTAutogradBackwardCompilationMetrics(&'e AOTAutogradBackwardCompilationMetricsMetadata),
+    Artifact(&'e ArtifactMetadata),
 }
 
 #[derive(Debug, Deserialize)]
@@ -374,6 +381,7 @@ pub struct Envelope {
     pub graph_dump: Option<GraphDumpMetadata>,
     pub link: Option<LinkMetadata>,
     pub symbolic_shape_specialization: Option<SymbolicShapeSpecializationMetadata>,
+    pub artifact: Option<ArtifactMetadata>,
 }
 
 #[derive(Debug, Deserialize, Serialize)]