[tlparse] MemoizerArtifacts metadata and parser

src/parsers.rs

@@ -997,6 +997,33 @@ impl StructuredLogParser for ArtifactParser {
     }
 }
 
+pub struct MemoizerArtifactsParser;
+impl StructuredLogParser for MemoizerArtifactsParser {
+    fn name(&self) -> &'static str {
+        "memoizer_artifacts"
+    }
+    fn get_metadata<'e>(&self, e: &'e Envelope) -> Option<Metadata<'e>> {
+        e.memoizer_artifacts
+            .as_ref()
+            .map(|m| Metadata::MemoizerArtifacts(m))
+    }
+    fn parse<'e>(
+        &self,
+        lineno: usize,
+        _metadata: Metadata<'e>,
+        _rank: Option<u32>,
+        compile_id: &Option<CompileId>,
+        _payload: &str,
+    ) -> anyhow::Result<ParserResults> {
+        payload_reformat_file_output(
+            "memoizer_artifacts.json",
+            lineno,
+            compile_id,
+            format_json_pretty,
+        )
+    }
+}
+
 fn render_sym_expr_trie(
     expr: u64,
     sym_expr_info_index: &SymExprInfoIndex,

src/types.rs

@@ -447,6 +447,12 @@ pub struct ArtifactMetadata {
     pub encoding: String,
 }
 
+#[derive(Debug, Deserialize)]
+pub struct MemoizerArtifactsMetadata {
+    pub aggregated: Option<bool>,
+    pub sub_key: Option<String>,
+}
+
 #[derive(Debug, Deserialize, Serialize, Clone)]
 pub struct CompilationMetricsMetadata {
     // Other information like frame_key are already in envelope
@@ -699,6 +705,7 @@ pub enum Metadata<'e> {
     AOTAutogradBackwardCompilationMetrics(&'e AOTAutogradBackwardCompilationMetricsMetadata),
     BwdCompilationMetrics(&'e BwdCompilationMetricsMetadata),
     Artifact(&'e ArtifactMetadata),
+    MemoizerArtifacts(&'e MemoizerArtifactsMetadata),
     DumpFile(&'e DumpFileMetadata),
     GuardAddedFast(&'e GuardAddedFastMetadata),
     SymbolicShapePropagateRealTensor(&'e SymbolicShapePropagateRealTensorMetadata),
@@ -754,6 +761,7 @@ pub struct Envelope {
     pub missing_fake_kernel: Option<FakeKernelMetadata>,
     pub mismatched_fake_kernel: Option<FakeKernelMetadata>,
     pub artifact: Option<ArtifactMetadata>,
+    pub memoizer_artifacts: Option<MemoizerArtifactsMetadata>,
     pub describe_storage: Option<StorageDesc>,
     pub describe_tensor: Option<TensorDesc>,
     pub describe_source: Option<SourceDesc>,