diff --git a/.github/workflows/presubmit.yml b/.github/workflows/presubmit.yml index abf42ce13..419ae6fda 100644 --- a/.github/workflows/presubmit.yml +++ b/.github/workflows/presubmit.yml @@ -118,6 +118,12 @@ jobs: run: | python3 -m pip install --user -U pip python3 -m pip install --user -U ittapi + - name: Install PyTorch (CPU) + run: | + python3 -m pip install --user -U pip + python3 -m pip install --user -U torch + - name: Add $HOME/.local/bin to the GITHUB_PATH + run: | echo "$HOME/.local/bin" >> "$GITHUB_PATH" - name: Setup remote-ssh session uses: owenthereal/action-upterm@v2 diff --git a/backends/pytorch/Makefile.am b/backends/pytorch/Makefile.am index 44ccfceea..96192d5f2 100644 --- a/backends/pytorch/Makefile.am +++ b/backends/pytorch/Makefile.am @@ -4,7 +4,7 @@ else WERROR = endif -PYTORCH_STATIC_PROBES = pytorch +PYTORCH_STATIC_PROBES = pytorch_tracepoints PYTORCH_STATIC_PROBES_TP = $(PYTORCH_STATIC_PROBES:=.tp) @@ -13,7 +13,7 @@ PYTORCH_STATIC_PROBES_INCL = $(PYTORCH_STATIC_PROBES:=.h) PYTORCH_STATIC_PROBES_SRC = $(PYTORCH_STATIC_PROBES:=.c) $(PYTORCH_STATIC_PROBES_TP): %.tp: $(top_srcdir)/utils/gen_custom_probes.rb $(srcdir)/pytorch_events.yaml $(top_srcdir)/utils/gen_probe_base.rb - $(RUBY) $< $(srcdir)/pytorch_events.yaml lttng_ust_$* > $@ + $(RUBY) $< $(srcdir)/pytorch_events.yaml lttng_ust_pytorch > $@ %.h %.c: %.tp $(LTTNG_GEN_TP) $< -o $*.c -o $*.h @@ -41,8 +41,9 @@ CLEANFILES = \ BUILT_SOURCES = \ $(PYTORCH_STATIC_PROBES_INCL) -# Only -ltorch_cpu is needed: at::addGlobalCallback and at::RecordFunction::name -# are both defined in libtorch_cpu.so. +# Only -ltorch_cpu is needed: at::addGlobalCallback, at::RecordFunction::name +# and at::RecordFuntion::overload_name symbols must be resolved to register +# entry/exit callbacks at load time. DUMMY_TORCH_LIBS = \ dummy_libs/libtorch_cpu.so @@ -63,9 +64,6 @@ bin_SCRIPTS = \ lib_LTLIBRARIES = libTracerPytorch.la -# pytorch.c is only built once, into libpytorchtracepoints.la above; pull it -# in via LIBADD rather than compiling it again here (a second copy would -# clash with the first at link time: duplicate tracepoint symbols). nodist_libTracerPytorch_la_SOURCES = \ $(PYTORCH_STATIC_PROBES_INCL) @@ -73,14 +71,44 @@ libTracerPytorch_la_SOURCES = \ tracer_pytorch.cpp libTracerPytorch_la_CPPFLAGS = -I$(top_srcdir)/utils -I$(top_srcdir)/utils/include -I$(srcdir)/include -I./ -# -Wno-missing-field-initializers: LTTng's own tracepoint.h (not our code) -# aggregate-initializes structs without naming every field; clang is -# stricter than gcc about this under -Wextra. -# -Wno-unused-private-field: RecordFunctionCallback's fields exist only to -# match PyTorch's real class layout byte-for-byte (see the comment in -# include/ATen/record_function.h) -- the real libtorch reads them, we -# never do, by design. libTracerPytorch_la_CXXFLAGS = -std=c++17 -Wall -Wextra -Wno-unused-parameter -Wno-missing-field-initializers -Wno-unused-private-field $(WERROR) $(LTTNG_UST_CFLAGS) -libTracerPytorch_la_LDFLAGS = $(LTTNG_UST_LIBS) -Ldummy_libs -ltorch_cpu -avoid-version -module +libTracerPytorch_la_LDFLAGS = $(LTTNG_UST_LIBS) -Ldummy_libs -Wl,--no-as-needed,-ltorch_cpu,--as-needed -avoid-version -module libTracerPytorch_la_LIBADD = libpytorchtracepoints.la EXTRA_libTracerPytorch_la_DEPENDENCIES = $(DUMMY_TORCH_LIBS) + +BTX_PYTORCH_GENERATED = \ + btx_filter_pytorch/metababel/metababel.h \ + btx_filter_pytorch/metababel/btx_component.h \ + btx_filter_pytorch/metababel/btx_component.c \ + btx_filter_pytorch/metababel/btx_upstream.h \ + btx_filter_pytorch/metababel/btx_upstream.c \ + btx_filter_pytorch/metababel/btx_downstream.h \ + btx_filter_pytorch/metababel/btx_downstream.c \ + btx_filter_pytorch/btx_main.c + +EXTRA_DIST += \ + $(top_srcdir)/xprof/btx_interval_model.yaml \ + $(srcdir)/btx_pytorch_model.yaml + +$(BTX_PYTORCH_GENERATED) &: $(top_srcdir)/xprof/btx_interval_model.yaml $(srcdir)/btx_pytorch_model.yaml + $(METABABEL) -u $(srcdir)/btx_pytorch_model.yaml -d $(top_srcdir)/xprof/btx_interval_model.yaml -t FILTER -o btx_filter_pytorch -p pytorchinterval -c interval + +CLEANFILES += \ + $(BTX_PYTORCH_GENERATED) + +BUILT_SOURCES += \ + $(BTX_PYTORCH_GENERATED) + +nodist_libPyTorchInterval_la_SOURCES = \ + $(BTX_PYTORCH_GENERATED) + +libPyTorchInterval_la_SOURCES = \ + btx_pytorchinterval_callbacks.cpp + +libPyTorchInterval_la_CPPFLAGS = -I$(top_srcdir)/utils -I$(top_srcdir)/utils/include -I./ -I./btx_filter_pytorch +libPyTorchInterval_la_CFLAGS = -Wall -Wextra -Wno-unused-parameter $(WERROR) $(BABELTRACE2_CFLAGS) +libPyTorchInterval_la_CXXFLAGS = -std=c++17 -Wall -Wextra -Wno-unused-parameter $(WERROR) $(BABELTRACE2_CFLAGS) +libPyTorchInterval_la_LDFLAGS = $(BABELTRACE2_LIBS) -avoid-version -module + +bt2dir = $(pkglibdir)/bt2 +bt2_LTLIBRARIES = libPyTorchInterval.la diff --git a/backends/pytorch/btx_pytorch_model.yaml b/backends/pytorch/btx_pytorch_model.yaml new file mode 100644 index 000000000..3bcbe5346 --- /dev/null +++ b/backends/pytorch/btx_pytorch_model.yaml @@ -0,0 +1,53 @@ +:environment: + :entries: + - :name: hostname + :type: string +:stream_classes: +- :name: thapi_pytorch + :default_clock_class: {} + :packet_context_field_class: + :type: structure + :members: + - :name: cpu_id + :field_class: + :type: integer_unsigned + :cast_type: uint64_t + :field_value_range: 32 + :event_common_context_field_class: + :type: structure + :members: + - :name: vpid + :field_class: + :type: integer_signed + :cast_type: int64_t + :field_value_range: 64 + - :name: vtid + :field_class: + :type: integer_unsigned + :cast_type: uint64_t + :field_value_range: 64 + :event_classes: + - :name: lttng_ust_pytorch:op_entry + :payload_field_class: + :type: structure + :members: + - :name: name + :field_class: + :cast_type: char * + :type: string + - :name: overload_name + :field_class: + :cast_type: char * + :type: string + - :name: lttng_ust_pytorch:op_exit + :payload_field_class: + :type: structure + :members: + - :name: name + :field_class: + :cast_type: char * + :type: string + - :name: overload_name + :field_class: + :cast_type: char * + :type: string diff --git a/backends/pytorch/btx_pytorchinterval_callbacks.cpp b/backends/pytorch/btx_pytorchinterval_callbacks.cpp new file mode 100644 index 000000000..dd86900d9 --- /dev/null +++ b/backends/pytorch/btx_pytorchinterval_callbacks.cpp @@ -0,0 +1,67 @@ +#include "xprof_utils.hpp" +#include +#include +#include +#include + +// PyTorch requies a per-thread LIFO stack since there is reentrace and nesting calls. +struct data_s { + std::unordered_map> entry_stack; +}; +typedef struct data_s data_t; + +static void btx_initialize_component(void **usr_data) { *usr_data = new data_t; } + +static void btx_finalize_component(void *usr_data) { delete static_cast(usr_data); } + +// PyTorch identifies an operator by TWO strings: a schema name (e.g. +// "aten::abs") and an overload name (e.g. "" for the default overload, "out" +// Joining them here keeps e.g. "aten::abs" and "aten::abs.out" distinguishable +// in the trace instead of both showing up as plain "aten::abs". +static std::string qualified_name(const char *name, const char *overload_name) { + return (overload_name[0] == '\0') ? name : std::string(name) + "." + overload_name; +} + +static void lttng_ust_pytorch_op_entry_callback(void *btx_handle, + void *usr_data, + int64_t ts, + const char *hostname, + int64_t vpid, + uint64_t vtid, + char * /*name*/, + char * /*overload_name*/) { + static_cast(usr_data)->entry_stack[{hostname, vpid, vtid}].push_back(ts); +} + +static void lttng_ust_pytorch_op_exit_callback(void *btx_handle, + void *usr_data, + int64_t ts, + const char *hostname, + int64_t vpid, + uint64_t vtid, + char *name, + char *overload_name) { + auto *state = static_cast(usr_data); + auto &stack = state->entry_stack[{hostname, vpid, vtid}]; + + // Empty means an exit arrived with no matching entry + const bool err = stack.empty(); + int64_t entry_ts = ts; + if (!err) { + entry_ts = stack.back(); + stack.pop_back(); + } + + const std::string full_name = qualified_name(name, overload_name); + btx_push_message_lttng_host(btx_handle, hostname, vpid, vtid, entry_ts, BACKEND_PYTORCH, + full_name.c_str(), (ts - entry_ts), err); +} + +void btx_register_usr_callbacks(void *btx_handle) { + btx_register_callbacks_initialize_component(btx_handle, &btx_initialize_component); + btx_register_callbacks_finalize_component(btx_handle, &btx_finalize_component); + + btx_register_callbacks_lttng_ust_pytorch_op_entry(btx_handle, + <tng_ust_pytorch_op_entry_callback); + btx_register_callbacks_lttng_ust_pytorch_op_exit(btx_handle, <tng_ust_pytorch_op_exit_callback); +} diff --git a/backends/pytorch/include/ATen/record_function.h b/backends/pytorch/include/ATen/record_function.h index 6003887ed..d6d538057 100644 --- a/backends/pytorch/include/ATen/record_function.h +++ b/backends/pytorch/include/ATen/record_function.h @@ -1,28 +1,5 @@ -// Minimal hand-written stand-in for PyTorch's . -// -// Declares only the five names tracer_pytorch.cpp actually uses: -// at::RecordScope - FUNCTION / BACKWARD_FUNCTION values -// at::ObserverContext - empty base; we only ever return nullptr -// at::RecordFunction - only .name(); we never construct one -// ourselves, only receive a reference from -// the real library, so no field layout is -// needed here -- name() resolves against the -// real out-of-line symbol in libtorch_cpu. -// at::RecordFunctionCallback - constructed BY US and passed BY VALUE into -// addGlobalCallback. Its field layout below -// (order, types, and the scopes_ array sized -// by NUM_SCOPES) must match PyTorch's real -// class byte-for-byte, copied verbatim from -// the upstream header. If a future PyTorch -// release reorders/adds a field or changes -// RecordScope's member count, this header -// will still compile cleanly but -// addGlobalCallback will read the wrong -// bytes back -- silent corruption, not a -// build failure. Re-verify this layout -// against ATen/record_function.h whenever -// upstream PyTorch changes. -// at::addGlobalCallback - registers the callback pair +// Minimal hand-written PyTorch's . + #pragma once #include @@ -55,6 +32,7 @@ struct ObserverContext { struct RecordFunction { const char *name() const; + const char *overload_name() const; }; class RecordFunctionCallback { diff --git a/backends/pytorch/pytorch_events.yaml b/backends/pytorch/pytorch_events.yaml index 1d6807f42..0abce9e8c 100644 --- a/backends/pytorch/pytorch_events.yaml +++ b/backends/pytorch/pytorch_events.yaml @@ -3,10 +3,14 @@ lttng_ust_pytorch: - name: op_entry args: - ["const char *", name] + - ["const char *", overload_name] fields: - [ctf_string, name, name] + - [ctf_string, overload_name, overload_name] - name: op_exit args: - ["const char *", name] + - ["const char *", overload_name] fields: - [ctf_string, name, name] + - [ctf_string, overload_name, overload_name] diff --git a/backends/pytorch/tracer_pytorch.cpp b/backends/pytorch/tracer_pytorch.cpp index ac95849ae..ebbb659a0 100644 --- a/backends/pytorch/tracer_pytorch.cpp +++ b/backends/pytorch/tracer_pytorch.cpp @@ -1,19 +1,15 @@ +#include "pytorch_tracepoints.h" #include -#include "pytorch.h" - -// ENTRY: fires BEFORE the op runs. LTTng adds time + vpid/vtid via context. static std::unique_ptr on_entry(const at::RecordFunction &fn) { - tracepoint(lttng_ust_pytorch, op_entry, fn.name()); + tracepoint(lttng_ust_pytorch, op_entry, fn.name(), fn.overload_name()); return nullptr; } -// EXIT: fires AFTER the op returns. static void on_exit(const at::RecordFunction &fn, at::ObserverContext *) { - tracepoint(lttng_ust_pytorch, op_exit, fn.name()); + tracepoint(lttng_ust_pytorch, op_exit, fn.name(), fn.overload_name()); } -// Auto-register at library load (works under LD_PRELOAD, no python changes). __attribute__((constructor)) static void tracer_pytorch_init() { at::addGlobalCallback( at::RecordFunctionCallback(&on_entry, &on_exit) diff --git a/integration_tests/backend_pytorch.bats b/integration_tests/backend_pytorch.bats new file mode 100644 index 000000000..fe86d31aa --- /dev/null +++ b/integration_tests/backend_pytorch.bats @@ -0,0 +1,7 @@ +bats_require_minimum_version 1.5.0 + +@test "PyTorch: trace contains aten::empty.memory_format" { + iprof --backends pytorch --analysis-output ./pytorch_out.txt -- \ + python3 ./integration_tests/pytorch_example.py + grep "aten::empty.memory_format" ./pytorch_out.txt +} diff --git a/integration_tests/pytorch_example.py b/integration_tests/pytorch_example.py new file mode 100644 index 000000000..2cb34934b --- /dev/null +++ b/integration_tests/pytorch_example.py @@ -0,0 +1,2 @@ +import torch +torch.empty(3) diff --git a/utils/babeltrace_thapi.in b/utils/babeltrace_thapi.in index 3e3b187fb..0544b9a7c 100755 --- a/utils/babeltrace_thapi.in +++ b/utils/babeltrace_thapi.in @@ -401,7 +401,7 @@ class BabeltraceParserThapi < OptionParserWithDefaultAndValidation on('-h', '--help', 'Prints this help') { print_help_and_exit(self, exit_code: 0) } on('-b', '--backends BACKENDS', Array, "Select which and how backends' need to handled.", 'Format: backend_name[:backend_level],...', - default: ['mpi:3', 'omp:2', 'cl:1', 'ze:1', 'cuda:1', 'hip:1', 'cxi:4', 'itt:5']) + default: ['mpi:3', 'omp:2', 'cl:1', 'ze:1', 'cuda:1', 'hip:1', 'cxi:4', 'itt:5', 'pytorch:6']) on('--debug', default: false) on('--archive SESSION-NAME') on('--archive-session-found-file-path PATH') diff --git a/utils/xprof_utils.hpp b/utils/xprof_utils.hpp index 26c99aee7..daec319d1 100644 --- a/utils/xprof_utils.hpp +++ b/utils/xprof_utils.hpp @@ -23,6 +23,7 @@ enum backend_e { BACKEND_MPI = 7, BACKEND_CXI = 8, BACKEND_ITT = 9, + BACKEND_PYTORCH = 10, }; typedef enum backend_e backend_t; typedef unsigned backend_level_t; @@ -39,6 +40,7 @@ const std::unordered_map pretty_backend_name_g = { {"mpi", BACKEND_MPI}, {"cxi", BACKEND_CXI}, {"itt", BACKEND_ITT}, + {"pytorch", BACKEND_PYTORCH}, }; const std::unordered_map backend_levels_g = { @@ -52,6 +54,7 @@ const std::unordered_map backend_levels_g = { {BACKEND_MPI, 3}, {BACKEND_CXI, 4}, {BACKEND_ITT, 5}, + {BACKEND_PYTORCH, 6}, }; typedef std::string thapi_metadata_t; @@ -74,14 +77,18 @@ typedef std::tuple hpt_t; typedef std::tuple hpt_function_name_t; typedef std::tuple t_function_name_t; -typedef std::tuple hpt_device_function_name_t; typedef std::tuple hp_device_t; typedef std::tuple hp_dsd_t; -typedef std::tuple hi_t; // host + NIC interface -typedef std::tuple hic_t; // host + NIC interface + counter +typedef std::tuple hi_t; // host + NIC interface +typedef std::tuple hic_t; // host + NIC interface + counter typedef std::tuple sd_t; typedef std::tuple tfn_ts_t; typedef std::tuple fn_ts_t; diff --git a/xprof/xprof.rb.in b/xprof/xprof.rb.in index 35cf9b9cd..ebc16acde 100755 --- a/xprof/xprof.rb.in +++ b/xprof/xprof.rb.in @@ -76,12 +76,11 @@ end # \/ (_| | | (_) |_| _> |_ >< (/_ (_ # def exec(cmd, opts: {}, debug: true, ignore_exit_codes: []) - return Open3.capture3(opts, cmd).first unless debug + stdout_str, stderr_str, status = Open3.capture3(opts, cmd) + return [stdout_str, stderr_str, status] unless debug LOGGER.info { cmd } LOGGER.debug { opts } unless opts.empty? - - stdout_str, stderr_str, status = Open3.capture3(opts, cmd) LOGGER.debug { stdout_str.strip } unless stdout_str.empty? unless status.success? || ignore_exit_codes.include?(status.exitstatus) @@ -91,7 +90,8 @@ def exec(cmd, opts: {}, debug: true, ignore_exit_codes: []) end LOGGER.warn { stderr_str.strip } unless stderr_str.empty? - stdout_str + + [stdout_str, stderr_str, status] end def launch_usr_bin(env, cmd) @@ -153,9 +153,9 @@ def whichlib64(binary, *libs) whichlib64_bin = File.join(BINDIR, 'whichlib64') ld_path = [env_fetch_first('LD_LIBRARY_PATH'), env_fetch_first('CRAY_LD_LIBRARY_PATH')].compact.join(':') - stdout_str = exec("#{whichlib64_bin} #{binary} #{libs.join(' ')}", - opts: { 'LD_LIBRARY_PATH' => ld_path }, - ignore_exit_codes: [1, 2]) + stdout_str, = exec("#{whichlib64_bin} #{binary} #{libs.join(' ')}", + opts: { 'LD_LIBRARY_PATH' => ld_path }, + ignore_exit_codes: [1, 2]) libs.zip(stdout_str.lines.map(&:strip)).filter_map do |lib, path| [lib, path] unless path.end_with?('not found') @@ -886,6 +886,25 @@ end # _| # +def get_pytorch_library_path + torch_tested = Gem::Version.new('1.11.0')..Gem::Version.new('2.14.0') + torch_probe = 'python3 -c "import torch; print(torch.__file__); print(torch.__version__)"' + + lib = env_fetch_first('LTTNG_UST_PYTORCH_LIBRARY_PATH') || begin + # Raises for other error codes, but only warns for 127 (python3 command not found) and 1 (ModuleNotFoundError) + stdout_str, _, status = exec(torch_probe, ignore_exit_codes: [1, 127]) + if status.success? + file, version = stdout_str.split + version.slice(/\d+(\.\d+)*/).then do |v| + torch_tested.cover?(Gem::Version.new(v)) || LOGGER.warn("Untested PyTorch version #{version}") + end + File.join(File.dirname(file), 'lib') + end + end + + lib +end + def all_env_tracers(usr_binary) # Return the list of backends (used by local master to enable lttng events) # and the ENV used by any traced-ranks to preload THAPI tracers @@ -966,18 +985,9 @@ def all_env_tracers(usr_binary) end if OPTIONS[:'backend-names'].include?('pytorch') - # Supported calling: iprof -- python3 model.py (torch env is active). - # Unsupported calling: iprof -- venv/bin/python3 module.py (non-active torch env) - # TODO: We need env vars to support the second calling convention. - torch_file = exec('python3 -c "import torch; print(torch.__file__)"', debug: false).strip - torch_lib = torch_file.empty? ? '' : File.join(File.dirname(torch_file), 'lib') - if torch_lib.empty? - LOGGER.warn('No torch module found for python3, pytorch backend will not be enabled') - else + get_pytorch_library_path&.then do |lib| backends << 'pytorch' - # LibTracerPytorch.so calls at::addGlobalCallback() and RecordFunction::name() - # both implemented in libtorch_cpu.so. The lib is reachable via LD_LIBRARY_PATH. - h[%w[LD_LIBRARY_PATH prepend]] << torch_lib + h[%w[LD_LIBRARY_PATH prepend]] << lib h[%w[LD_PRELOAD prepend]] << File.join(LIBDIR, 'libTracerPytorch.so') end end @@ -1075,7 +1085,7 @@ if $thapi_launch || __FILE__ == $PROGRAM_NAME # General Options parser.on('-b', '--backends BACKENDS', Array, 'Select which backends to use and their grouping level.', 'Format: backend_name[:backend_level],...', - default: ['mpi:3', 'omp:2', 'cl:1', 'ze:1', 'cuda:1', 'hip:1', 'cxi:4', 'itt:5']) + default: ['mpi:3', 'omp:2', 'cl:1', 'ze:1', 'cuda:1', 'hip:1', 'cxi:4', 'itt:5', 'pytorch:6']) parser.on('--[no-]archive', 'Enable or disable archive support.', default: false) # Analysis