I am done

rl/Lib/site-packages/torch/_logging/__init__.py

@@ -0,0 +1,17 @@
+# Top level logging module for torch logging
+# Design doc: https://docs.google.com/document/d/1ZRfTWKa8eaPq1AxaiHrq4ASTPouzzlPiuquSBEJYwS8/edit#
+# Simple setup for onboarding (see above doc for more detail):
+# 1. register any top-level log qualified name for your module in torch._logging._registrations (see there for examples)
+# 2. register any artifacts (<artifact_name> below) in torch._logging._registrations
+#   a. call getArtifactLogger(__name__, <artifact_name>) at your logging site instead of the standard logger to log your artifact
+import torch._logging._registrations
+
+from ._internal import (
+    _init_logs,
+    DEFAULT_LOGGING,
+    getArtifactLogger,
+    LazyString,
+    set_logs,
+    trace_structured,
+    warning_once,
+)

rl/Lib/site-packages/torch/_logging/_registrations.py

@@ -0,0 +1,192 @@
+# flake8: noqa: B950
+from ._internal import register_artifact, register_log
+
+
+DYNAMIC = [
+    "torch.fx.experimental.symbolic_shapes",
+    "torch.fx.experimental.sym_node",
+    "torch.fx.experimental.recording",
+]
+DISTRIBUTED = [
+    "torch.distributed",
+    "torch._dynamo.backends.distributed",
+    "torch.nn.parallel.distributed",
+]
+
+register_log("dynamo", ["torch._dynamo", *DYNAMIC])
+register_log("fake_tensor", ["torch._subclasses.fake_tensor"])
+register_log("aot", ["torch._functorch.aot_autograd", "torch._functorch._aot_autograd"])
+register_log("autograd", "torch.autograd")
+register_log("inductor", ["torch._inductor", "torch._inductor.cudagraph_trees"])
+
+register_artifact(
+    "cudagraphs",
+    "Logs information from wrapping inductor generated code with cudagraphs.",
+)
+
+register_log("dynamic", DYNAMIC)
+register_log("torch", "torch")
+register_log("distributed", DISTRIBUTED)
+register_log(
+    "c10d", ["torch.distributed.distributed_c10d", "torch.distributed.rendezvous"]
+)
+register_log(
+    "ddp", ["torch.nn.parallel.distributed", "torch._dynamo.backends.distributed"]
+)
+register_log("pp", ["torch.distributed.pipelining"])
+register_log("fsdp", ["torch.distributed.fsdp", "torch.distributed._composable.fsdp"])
+register_log("dtensor", ["torch.distributed._tensor", "torch.distributed.tensor"])
+register_log("onnx", "torch.onnx")
+register_log(
+    "export",
+    [
+        "torch._dynamo",
+        "torch.export",
+        "torch.export.dynamic_shapes",
+        *DYNAMIC,
+        "torch._export.converter",
+        "torch._export.non_strict_utils",
+    ],
+)
+
+register_artifact(
+    "guards",
+    "This prints the guards for every compiled Dynamo frame. It does not tell you where the guards come from.",
+    visible=True,
+)
+register_artifact("verbose_guards", "", off_by_default=True)
+register_artifact(
+    "bytecode",
+    "Prints the original and modified bytecode from Dynamo. Mostly useful if you're debugging our bytecode generation in Dynamo.",
+    off_by_default=True,
+)
+register_artifact(
+    "graph",
+    "Prints the dynamo traced graph (prior to AOTDispatch) in a table. If you prefer python code use `graph_code` instead. ",
+)
+register_artifact("graph_code", "Like `graph`, but gives you the Python code instead.")
+register_artifact(
+    "graph_sizes", "Prints the sizes of all FX nodes in the dynamo graph."
+)
+register_artifact(
+    "trace_source",
+    "As we execute bytecode, prints the file name / line number we are processing and the actual source code. Useful with `bytecode`",
+)
+register_artifact(
+    "trace_call",
+    "Like trace_source, but it will give you the per-expression blow-by-blow if your Python is recent enough.",
+)
+register_artifact(
+    "trace_bytecode",
+    "As we trace bytecode, prints the instruction and the current stack.",
+)
+register_artifact(
+    "aot_graphs",
+    "Prints the FX forward and backward graph generated by AOTDispatch, after partitioning. Useful to understand what's being given to Inductor",
+    visible=True,
+)
+register_artifact(
+    "aot_joint_graph",
+    "Print FX joint graph from AOTAutograd, prior to partitioning. Useful for debugging partitioning",
+)
+register_artifact(
+    "aot_graphs_effects",
+    "Prints the FX forward and backward graph generated by AOTDispatch, useful for debugging effects processing.",
+    visible=True,
+)
+register_artifact(
+    "post_grad_graphs",
+    "Prints the FX graph generated by post grad passes. Useful to understand what's being given to Inductor after post grad passes",
+)
+register_artifact(
+    "compiled_autograd",
+    "Prints various logs in compiled_autograd, including but not limited to the graphs. Useful for debugging compiled_autograd.",
+    visible=True,
+)
+register_artifact(
+    "compiled_autograd_verbose",
+    "Will affect performance. Prints compiled_autograd logs with C++ info e.g. autograd node -> fx node mapping",
+    off_by_default=True,
+)
+register_artifact(
+    "ddp_graphs",
+    "Only relevant for compiling DDP. DDP splits into multiple graphs to trigger comms early. This will print each individual graph here.",
+)
+register_artifact(
+    "recompiles",
+    "Prints the reason why we recompiled a graph. Very, very useful.",
+    visible=True,
+)
+register_artifact(
+    "recompiles_verbose",
+    "Prints all guard checks that fail during a recompilation. "
+    "At runtime, Dynamo will stop at the first failed check for each failing guard. "
+    "So not all logged failing checks are actually ran by Dynamo.",
+    visible=True,
+    off_by_default=True,
+)
+register_artifact(
+    "graph_breaks",
+    "Prints whenever Dynamo decides that it needs to graph break (i.e. create a new graph). Useful for debugging why torch.compile has poor performance",
+    visible=True,
+)
+register_artifact(
+    "not_implemented",
+    "Prints log messages whenever we return NotImplemented in a multi-dispatch, letting you trace through each object we attempted to dispatch to",
+)
+register_artifact(
+    "output_code",
+    "Prints the code that Inductor generates (either Triton or C++)",
+    off_by_default=True,
+    visible=True,
+)
+register_artifact(
+    "kernel_code",
+    "Prints the code that Inductor generates (on a per-kernel basis)",
+    off_by_default=True,
+    visible=True,
+)
+register_artifact(
+    "schedule",
+    "Inductor scheduler information. Useful if working on Inductor fusion algo",
+    off_by_default=True,
+)
+register_artifact("perf_hints", "", off_by_default=True)
+register_artifact("onnx_diagnostics", "", off_by_default=True)
+register_artifact(
+    "fusion",
+    "Detailed Inductor fusion decisions. More detailed than 'schedule'",
+    off_by_default=True,
+)
+register_artifact(
+    "loop_ordering",
+    "Logs related to loop ordering",
+    off_by_default=True,
+)
+register_artifact(
+    "overlap",
+    "Detailed Inductor compute/comm overlap decisions",
+    off_by_default=True,
+)
+register_artifact(
+    "sym_node",
+    "Logs extra info for various SymNode operations",
+    off_by_default=True,
+)
+register_artifact(
+    "trace_shape_events",
+    "Logs traces for every ShapeEnv operation that we record for replay",
+    off_by_default=True,
+)
+register_artifact(
+    "cudagraph_static_inputs",
+    "Logs static inputs handling in dynamo, AOT, and cudagraphs",
+    off_by_default=True,
+)
+register_artifact(
+    "benchmarking",
+    "Detailed Inductor benchmarking information.",
+    off_by_default=True,
+)
+
+register_artifact("custom_format_test_artifact", "Testing only", log_format="")

rl/Lib/site-packages/torch/_logging/scribe.py

@@ -0,0 +1,61 @@
+from typing import Callable, List, Union
+from typing_extensions import TypeAlias
+
+
+try:
+    from fbscribelogger import make_scribe_logger  # type: ignore[import-untyped]
+except ImportError:
+    TAtom: TypeAlias = Union[int, float, bool, str]
+    TField: TypeAlias = Union[TAtom, List[TAtom]]
+    TLazyField: TypeAlias = Union[TField, Callable[[], TField]]
+
+    def make_scribe_logger(name: str, thrift_src: str) -> Callable[..., None]:
+        def inner(**kwargs: TLazyField) -> None:
+            pass
+
+        return inner
+
+
+open_source_signpost = make_scribe_logger(
+    "TorchOpenSourceSignpost",
+    """
+struct TorchOpenSourceSignpostLogEntry {
+
+  # The commit SHA that triggered the workflow, e.g., 02a6b1d30f338206a71d0b75bfa09d85fac0028a. Derived from GITHUB_SHA.
+  4: optional string commit_sha;
+
+  # Commit date (not author date) of the commit in commit_sha as timestamp, e.g., 1724208105.  Increasing if merge bot is used, though not monotonic; duplicates occur when stack is landed.
+  5: optional i64 commit_date;
+
+  # The fully-formed ref of the branch or tag that triggered the workflow run, e.g., refs/pull/133891/merge or refs/heads/main. Derived from GITHUB_REF.
+  6: optional string github_ref;
+
+  # Indicates if branch protections or rulesets are configured for the ref that triggered the workflow run. Derived from GITHUB_REF_PROTECTED.
+  7: optional bool github_ref_protected;
+
+  # A unique number for each attempt of a particular workflow run in a repository, e.g., 1. Derived from GITHUB_RUN_ATTEMPT.
+  8: optional string github_run_attempt;
+
+  # A unique number for each workflow run within a repository, e.g., 19471190684. Derived from GITHUB_RUN_ID.
+  9: optional string github_run_id;
+
+  # A unique number for each run of a particular workflow in a repository, e.g., 238742. Derived from GITHUB_RUN_NUMBER.
+  10: optional string github_run_number_str;
+
+  # The name of the current job. Derived from JOB_NAME, e.g., linux-jammy-py3.8-gcc11 / test (default, 3, 4, amz2023.linux.2xlarge).
+  11: optional string job_name;
+
+  # The GitHub user who triggered the job.  Derived from GITHUB_TRIGGERING_ACTOR.
+  12: optional string github_triggering_actor;
+  13: optional string name; # Event name
+  14: optional string parameters; # Parameters (JSON data)
+  16: optional string subsystem; # Subsystem the event is associated with
+
+  # The unit timestamp in second for the Scuba Time Column override
+  17: optional i64 time;
+
+  # The weight of the record according to current sampling rate
+  18: optional i64 weight;
+}
+""",  # noqa: B950
+)

rl/Lib/site-packages/torch/_logging/structured.py

@@ -0,0 +1,57 @@
+"""
+Utilities for converting data types into structured JSON for dumping.
+"""
+
+import traceback
+from typing import Any, Dict, List, Sequence, Set
+
+import torch._logging._internal
+
+
+INTERN_TABLE: Dict[str, int] = {}
+
+
+DUMPED_FILES: Set[str] = set()
+
+
+def intern_string(s: str) -> int:
+    r = INTERN_TABLE.get(s, None)
+    if r is None:
+        r = len(INTERN_TABLE)
+        INTERN_TABLE[s] = r
+        torch._logging._internal.trace_structured(
+            "str", lambda: (s, r), suppress_context=True
+        )
+    return r
+
+
+def dump_file(filename: str) -> None:
+    if "eval_with_key" not in filename:
+        return
+    if filename in DUMPED_FILES:
+        return
+    DUMPED_FILES.add(filename)
+    from torch.fx.graph_module import _loader
+
+    torch._logging._internal.trace_structured(
+        "dump_file",
+        metadata_fn=lambda: {
+            "name": filename,
+        },
+        payload_fn=lambda: _loader.get_source(filename),
+    )
+
+
+def from_traceback(tb: Sequence[traceback.FrameSummary]) -> List[Dict[str, Any]]:
+    r = []
+    for frame in tb:
+        # dict naming convention here coincides with
+        # python/combined_traceback.cpp
+        r.append(
+            {
+                "line": frame.lineno,
+                "name": frame.name,
+                "filename": intern_string(frame.filename),
+            }
+        )
+    return r