Add origin stack trace capture for DALI operators

dali/pipeline/operator/error_reporting.cc

@@ -0,0 +1,57 @@
+// Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <sstream>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "dali/pipeline/operator/error_reporting.h"
+#include "dali/pipeline/operator/op_spec.h"
+
+namespace dali {
+
+std::vector<PythonStackFrame> GetOperatorOriginInfo(const OpSpec &spec) {
+  auto origin_stack_filename = spec.GetRepeatedArgument<std::string>("_origin_stack_filename");
+  auto origin_stack_lineno = spec.GetRepeatedArgument<int>("_origin_stack_lineno");
+  auto origin_stack_name = spec.GetRepeatedArgument<std::string>("_origin_stack_name");
+  auto origin_stack_line = spec.GetRepeatedArgument<std::string>("_origin_stack_line");
+
+  std::vector<PythonStackFrame> origin_stack;
+  origin_stack.reserve(origin_stack_filename.size());
+  const char error[] = "Internal error, mismatch in origin stack trace data.";
+  DALI_ENFORCE(origin_stack_filename.size() == origin_stack_lineno.size(), error);
+  DALI_ENFORCE(origin_stack_filename.size() == origin_stack_name.size(), error);
+  DALI_ENFORCE(origin_stack_filename.size() == origin_stack_line.size(), error);
+  for (size_t i = 0; i < origin_stack_filename.size(); i++) {
+    origin_stack.emplace_back(std::move(origin_stack_filename[i]), origin_stack_lineno[i],
+                              std::move(origin_stack_name[i]), std::move(origin_stack_line[i]));
+  }
+  return origin_stack;
+}
+
+
+std::string FormatStack(const std::vector<PythonStackFrame> &stack_summary, bool include_context) {
+  std::stringstream s;
+  for (auto &frame_summary : stack_summary) {
+    s << "  File \"" << frame_summary.filename << "\", line " << frame_summary.lineno << ", in "
+      << frame_summary.name << "\n";
+    if (include_context) {
+      s << "    " << frame_summary.line << "\n";
+    }
+  }
+  return s.str();
+}
+
+}  // namespace dali

dali/pipeline/operator/error_reporting.h

@@ -0,0 +1,64 @@
+// Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef DALI_PIPELINE_OPERATOR_ERROR_REPORTING_H_
+#define DALI_PIPELINE_OPERATOR_ERROR_REPORTING_H_
+
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "dali/core/api_helper.h"
+#include "dali/pipeline/data/types.h"
+#include "dali/pipeline/operator/op_spec.h"
+namespace dali {
+
+// TODO(klecki): Throw this one into a namespace?
+
+/**
+ * @brief Direct equivalent of Python's traceback.FrameSummary:
+ * https://docs.python.org/3/library/traceback.html#traceback.FrameSummary
+ * Describes a stack frame in the Python stack trace.
+ */
+struct PythonStackFrame {
+  PythonStackFrame(std::string filename, int lineno, std::string name, std::string line)
+      : filename(std::move(filename)),
+        lineno(lineno),
+        name(std::move(name)),
+        line(std::move(line)) {}
+  /** @brief File name of the source code executed for this frame. */
+  std::string filename;
+  /** @brief The line number of the source code for this frame. */
+  int lineno;
+  /** @brief Name of the function being executed in this frame. */
+  std::string name;
+  /** @brief A string representing the source code for this frame, with leading and trailing
+   * whitespace stripped. */
+  std::string line;
+};
+
+/**
+ * @brief Get the origin stack trace for operator constructed with given spec.
+ * The stack trace defines frames between invocation of the pipeline definition and the operator
+ * call. The returned PythonStackFrame corresponds to Python traceback.FrameSummary, but the `line`
+ * context may be invalid in some autograph transformed code.
+ */
+DLL_PUBLIC std::vector<PythonStackFrame> GetOperatorOriginInfo(const OpSpec &spec);
+
+DLL_PUBLIC std::string FormatStack(const std::vector<PythonStackFrame> &stack_summary,
+                                   bool include_context);
+
+}  // namespace dali
+
+#endif  // DALI_PIPELINE_OPERATOR_ERROR_REPORTING_H_

dali/pipeline/operator/op_schema.cc

@@ -1,4 +1,4 @@
-// Copyright (c) 2017-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// Copyright (c) 2017-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -78,6 +78,37 @@ to accommodate a batch of samples of this size.)code",
   AddOptionalArg("preserve", R"code(Prevents the operator from being removed from the
 graph even if its outputs are not used.)code",
                  false);
+
+
+  // For simplicity we pass StackSummary as 4 separate arguments so we don't need to extend DALI
+  // with support for special FrameSummary type.
+  // List of FrameSummaries can be reconstructed using utility functions.
+  AddOptionalArg("_origin_stack_filename", R"code(Every operator defined in Python captures and
+processes the StackSummary (a List[FrameSummary], defined in Python traceback module) that describes
+the callstack between the start of pipeline definition tracing and the "call" to the operator
+(or full trace if the operator is defined outside the pipeline).
+This information is propagated to the backend, so it can be later used to produce meaningful error
+messages, pointing to the origin of the error in pipeline definition.
+
+The list of FrameSummaries is split into four parameters: each is the list containing corresponding
+parameters of FrameSummary. This parameter represents the `filename` member.)code",
+                 std::vector<std::string>{});
+
+  AddOptionalArg("_origin_stack_lineno", R"code(StackSummary - lineno member of FrameSummary, see
+_origin_stack_filename for more information.)code",
+                 std::vector<int>{});
+
+  AddOptionalArg("_origin_stack_name", R"code(StackSummary - name member of FrameSummary, see
+_origin_stack_filename for more information.)code",
+                 std::vector<std::string>{});
+
+  AddOptionalArg("_origin_stack_line", R"code(StackSummary - line member of FrameSummary, see
+_origin_stack_filename for more information.)code",
+                 std::vector<std::string>{});
+
+  AddOptionalArg("_pipeline_internal", R"code(Boolean specifying if this operator was defined within
+a pipeline scope. False if it was defined without pipeline being set as current.)code",
+                 true);
 }
 
 

dali/python/nvidia/dali/_autograph/__init__.py

@@ -43,6 +43,8 @@
 from nvidia.dali._autograph.impl.api import do_not_convert
 from nvidia.dali._autograph.impl.api import autograph_artifact
 from nvidia.dali._autograph.impl.api import is_autograph_artifact
+from nvidia.dali._autograph.impl.api import is_frame_ag_call_entrypoint
+from nvidia.dali._autograph.impl.api import is_frame_ag_call_unconverted
 
 # from nvidia.dali._autograph.impl.api import StackTraceMapper
 from nvidia.dali._autograph.impl.api import to_code
@@ -52,6 +54,7 @@
 from nvidia.dali._autograph.utils import ag_logging
 from nvidia.dali._autograph.utils.all_utils import _remove_undocumented
 from nvidia.dali._autograph.utils.hooks import OperatorBase
+from nvidia.dali._autograph.utils.tf_stack import CustomModuleFilter
 
 # TODO(mdan): Revisit this list once we finalize the generated code mechanism.
 _allowed_symbols = [

dali/python/nvidia/dali/_autograph/impl/api.py

@@ -58,8 +58,7 @@
 from nvidia.dali._autograph.utils import ag_logging as logging
 from nvidia.dali._autograph.utils import all_utils
 
-# TODO(klecki): replace missing functionality
-# from nvidia.dali._autograph.utils import tf_stack
+from nvidia.dali._autograph.utils import tf_stack
 from nvidia.dali._autograph.utils.all_utils import export_symbol
 
 
@@ -130,40 +129,47 @@
     e.ag_error_metadata = _ErrorMetadata(cause_tb, metadata, message, source_map, __file__)
 
 
-# class StackTraceMapper(tf_stack.StackTraceMapper):
-#   """Remaps generated code to code it originated from."""
+class StackTraceMapper(tf_stack.StackTraceMapper):
+    """Remaps generated code to code it originated from."""
 
-#   def __init__(self, converted_fn):
-#     super().__init__()
-#     self._source_map = converted_fn.ag_source_map
-#     # This may be called repeatedly: once on entry, by the superclass, then by
-#     # each child context manager.
-#     self._cached_map = None
+    def __init__(self, converted_fn):
+        super().__init__()
+        self._source_map = converted_fn.ag_source_map
+        # This may be called repeatedly: once on entry, by the superclass, then by
+        # each child context manager.
+        self._cached_map = None
 
-#   def get_effective_source_map(self):
-#     if self._cached_map is not None:
-#       return self._cached_map
+    def get_effective_source_map(self):
+        if self._cached_map is not None:
+            return self._cached_map
 
-#     parent_map = self.parent.get_effective_source_map()
+        parent_map = self.parent.get_effective_source_map()
 
-#     effective_source_map = {}
-#     for loc, origin in self._source_map.items():
-#       effective_source_map[(loc.filename, loc.lineno)] = (origin.loc.filename,
-#                                                           origin.loc.lineno,
-#                                                           origin.function_name)
+        effective_source_map = {}
+        for loc, origin in self._source_map.items():
+            effective_source_map[(loc.filename, loc.lineno)] = (
+                origin.loc.filename,
+                origin.loc.lineno,
+                origin.function_name,
+                origin.source_code_line,
+            )
 
-#     for key, value in parent_map.items():
-#       filename, lineno, _ = value
-#       value_loc = origin_info.LineLocation(filename=filename, lineno=lineno)
-#       if value_loc in self._source_map:
-#         origin = self._source_map[value_loc]
-#         effective_source_map[key] = (origin.loc.filename, origin.loc.lineno,
-#                                      origin.function_name)
-#       else:
-#         effective_source_map[key] = value
+        for key, value in parent_map.items():
+            filename, lineno, _, _ = value
+            value_loc = origin_info.LineLocation(filename=filename, lineno=lineno)
+            if value_loc in self._source_map:
+                origin = self._source_map[value_loc]
+                effective_source_map[key] = (
+                    origin.loc.filename,
+                    origin.loc.lineno,
+                    origin.function_name,
+                    origin.source_code_line,
+                )
+            else:
+                effective_source_map[key] = value
 
-#     self._cached_map = effective_source_map
-#     return effective_source_map
+        self._cached_map = effective_source_map
+        return effective_source_map
 
 
 #
@@ -283,6 +289,24 @@
     return hasattr(entity, "autograph_info__")
 
 
+def is_frame_ag_call_entrypoint(frame_info):
+    """
+    True if the given frame is start of a function call wrapped by AutoGraph (ag__.converted_call)
+    """
+    return (
+        frame_info.filename.endswith("nvidia/dali/_autograph/impl/api.py")
+        and frame_info.name == "converted_call"
+    )
+
+
+def is_frame_ag_call_unconverted(frame_info):
+    """True if the given frame exits autograph to call unconverted user code."""
+    return (
+        frame_info.filename.endswith("nvidia/dali/_autograph/impl/api.py")
+        and frame_info.name == "_call_unconverted"
+    )
+
+
 def converted_call(f, args, kwargs, caller_fn_scope=None, options=None):
     """Converts a function call inline.
 
@@ -418,16 +442,20 @@
             raise
         return _fall_back_unconverted(f, args, kwargs, options, e)
 
-    # TODO(klecki): Revert the stack trace mapping functionality
-    # with StackTraceMapper(converted_f), tf_stack.CurrentModuleFilter():
-    try:
-        if kwargs is not None:
-            result = converted_f(*effective_args, **kwargs)
-        else:
-            result = converted_f(*effective_args)
-    except Exception as e:
-        _attach_error_metadata(e, converted_f)
-        raise
+    # We no longer need CurrentModuleFilter here, as we filter whole autograph
+    # TODO(klecki): Filter them just once.
+    import nvidia.dali._conditionals as dc
+    import nvidia.dali._autograph as ag
+
+    with StackTraceMapper(converted_f), tf_stack.CustomModuleFilter([ag, dc]):
+        try:
+            if kwargs is not None:
+                result = converted_f(*effective_args, **kwargs)
+            else:
+                result = converted_f(*effective_args)
+        except Exception as e:
+            _attach_error_metadata(e, converted_f)
+            raise
 
     return result