Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you open a pull request for review, mark a draft as ready, or comment "@codex review". If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex fix this CI failure" or "@codex address that feedback".

source/api_cc/src/DeepPotPT.cc (3)

5-5: Use public profiler header and add required standard headers

Avoid internal headers; they change across PyTorch versions. Also add missing standard headers used below.

Apply this diff around includes:

-#include <torch/csrc/autograd/profiler.h>
+#include <torch/autograd/profiler.h>   // public C++ API
+#include <set>                         // std::set for activities
+#include <cstdlib>                     // std::getenv
+#include <filesystem>                  // path handling
+#ifdef _WIN32
+#  include <process.h>
+#  define getpid _getpid
+#else
+#  include <unistd.h>                  // getpid
+#endif

If you want to make profiler code optional at build time, gate the include and related code with a CMake-defined macro (e.g., DEEPMD_ENABLE_PYTORCH_PROFILER). This aligns with the earlier nit about conditional includes.

#ifdef DEEPMD_ENABLE_PYTORCH_PROFILER
#include <torch/autograd/profiler.h>
#endif

To confirm the correct public header for RecordProfile for your libtorch version, please check:

Which header exposes torch::autograd::profiler::RecordProfile in the PyTorch C++ API (v2.x)? Is <torch/autograd/profiler.h> the supported public include?

---

87-118: Switch to RecordProfile (RAII), append PID to filename, and harden path handling

This block uses impl:: types and manual enable/disable, which is fragile across libtorch releases and deviates from the issue’s ask to use RecordProfile. RAII also fixes leaks if init() throws after enabling. Additionally, append a PID suffix to avoid collisions across MPI ranks/processes and ensure parent dirs exist. Only add “.json” if missing.

Replace the whole block with:

-  // Configure PyTorch profiler
-  const char* env_profiler = std::getenv("DP_PROFILER");
-  if (env_profiler && *env_profiler) {
-    using torch::profiler::impl::ActivityType;
-    using torch::profiler::impl::ExperimentalConfig;
-    using torch::profiler::impl::ProfilerConfig;
-    using torch::profiler::impl::ProfilerState;
-    std::set<ActivityType> activities{ActivityType::CPU};
-    if (gpu_enabled) {
-      activities.insert(ActivityType::CUDA);
-    }
-    profiler_file = std::string(env_profiler);
-    if (gpu_enabled) {
-      profiler_file += "_gpu" + std::to_string(gpu_id);
-    }
-    profiler_file += ".json";
-    ExperimentalConfig exp_cfg;
-    ProfilerConfig cfg(ProfilerState::KINETO,
-                       false,  // report_input_shapes,
-                       false,  // profile_memory,
-                       true,   // with_stack,
-                       false,  // with_flops,
-                       true,   // with_modules,
-                       exp_cfg,
-                       std::string()  // trace_id
-    );
-    torch::autograd::profiler::prepareProfiler(cfg, activities);
-    torch::autograd::profiler::enableProfiler(cfg, activities);
-    std::cout << "PyTorch profiler enabled, output file: " << profiler_file
-              << std::endl;
-    profiler_enabled = true;
-  }
+  // Configure PyTorch profiler (RAII)
+  const char* env_profiler = std::getenv("DP_PROFILER");
+  if (env_profiler && *env_profiler) {
+    profiler_file = std::string(env_profiler);
+    // Avoid double ".json"
+    if (!std::filesystem::path(profiler_file).has_extension()) {
+      profiler_file += ".json";
+    }
+    // Add disambiguators: GPU id (if any) and PID
+    {
+      std::string stem = std::filesystem::path(profiler_file).stem().string();
+      std::string ext  = std::filesystem::path(profiler_file).extension().string();
+      if (gpu_enabled) stem += "_gpu" + std::to_string(gpu_id);
+      stem += "_pid" + std::to_string(static_cast<long long>(getpid()));
+      profiler_file = (std::filesystem::path(profiler_file).parent_path() /
+                       (stem + ext)).string();
+    }
+    // Best-effort create parent dirs
+    std::error_code ec;
+    auto parent = std::filesystem::path(profiler_file).parent_path();
+    if (!parent.empty()) std::filesystem::create_directories(parent, ec);
+    // Optional: very light path sanity to avoid surprises
+    // (keeps absolute paths, but normalizes .. segments)
+    profiler_file = std::filesystem::weakly_canonical(std::filesystem::path(profiler_file), ec).string();
+    // Engage RAII guard (saves to profiler_file on destruction)
+    profiler_guard = std::make_unique<torch::autograd::profiler::RecordProfile>(profiler_file);
+    std::cout << "PyTorch profiler enabled, output file: " << profiler_file << std::endl;
+  }

If you still need advanced knobs (shapes, modules), consider verifying the supported C++ API for your libtorch version and toggling them behind macros; avoid torch::profiler::impl where possible.

PyTorch C++ profiler: how to enable Kineto tracing with RecordProfile (RAII) and configure stack/modules in v2.x without touching torch::profiler::impl?

Note: Building on a prior comment about path safety, normalizing and restricting to intended dirs mitigates accidental directory traversals from env input.

---

153-162: Destructor must not throw; prefer relying on RAII or, at minimum, wrap in noexcept-safe try/catch

As previously noted, profiler operations can throw; throwing from a destructor during stack unwinding will call std::terminate. If you adopt RecordProfile above, you can delete this block entirely; otherwise, wrap it.

Option A (preferred, with RAII): delete the body.

-DeepPotPT::~DeepPotPT() {
-  if (profiler_enabled) {
-    auto result = torch::autograd::profiler::disableProfiler();
-    if (result) {
-      result->save(profiler_file);
-    }
-    std::cout << "PyTorch profiler result saved to " << profiler_file
-              << std::endl;
-  }
-}
+DeepPotPT::~DeepPotPT() = default;

Option B (keep manual control): add hard try/catch and silence all exceptions.

 DeepPotPT::~DeepPotPT() {
-  if (profiler_enabled) {
-    auto result = torch::autograd::profiler::disableProfiler();
-    if (result) {
-      result->save(profiler_file);
-    }
-    std::cout << "PyTorch profiler result saved to " << profiler_file
-              << std::endl;
-  }
+  try {
+    if (profiler_enabled) {
+      auto result = torch::autograd::profiler::disableProfiler();
+      if (result) result->save(profiler_file);
+      std::cout << "PyTorch profiler result saved to " << profiler_file << std::endl;
+    }
+  } catch (const std::exception& e) {
+    std::cerr << "Profiler teardown error: " << e.what() << std::endl;
+  } catch (...) {
+    std::cerr << "Profiler teardown: unknown error" << std::endl;
+  }
 }

source/api_cc/src/DeepPotPT.cc (1)

73-76: Construct the device after capability check to avoid kCUDA construction when CUDA is unavailable

Minor cleanup; avoids creating a CUDA device object when not needed.

-  torch::Device device(torch::kCUDA, gpu_id);
-  if (!gpu_enabled) {
-    device = torch::Device(torch::kCPU);
+  torch::Device device = gpu_enabled ? torch::Device(torch::kCUDA, gpu_id)
+                                     : torch::Device(torch::kCPU);
+  if (!gpu_enabled) {
     std::cout << "load model from: " << model << " to cpu " << std::endl;

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

doc/env.md (5)

93-100: Clarify scope and path semantics of DP_PROFILER.

Add two short notes so users know where the file lands and which backend this affects.

Apply this diff:

 ::: {envvar} DP_PROFILER
@@
-{{ pytorch_icon }} Enable the built-in PyTorch Kineto profiler for the PyTorch C++ (inference) backend.
+{{ pytorch_icon }} Enable the built-in PyTorch Kineto profiler for the PyTorch C++ (inference) backend (DeepPotPT only).
@@
-**Type**: string (output file stem)
+**Type**: string (output file stem or path prefix)
@@
-**Default**: unset (disabled)
+**Default**: unset (disabled)
+
+If a relative path is given, the trace is written relative to the current working directory of the process; absolute paths are supported.

---

101-105: Prevent filename collisions; define GPU ID semantics.

On CPU multi-process runs all ranks would target the same <ENV_VALUE>.json. Also clarify that the GPU ID is the local CUDA ordinal after CUDA_VISIBLE_DEVICES.

Apply this diff:

 The final file name is constructed as:
 
-- `<ENV_VALUE>_gpu<ID>.json` if running on GPU
-`- `<ENV_VALUE>.json` if running on CPU
+- `<ENV_VALUE>_gpu<ID>.json` if running on GPU (where `<ID>` is the local CUDA ordinal after `CUDA_VISIBLE_DEVICES`)
+- `<ENV_VALUE>.json` if running on CPU
+
+Note: On CPU multi-process runs, all ranks would otherwise write the same filename. Ensure each process uses a unique stem (or path) to avoid clobbering.

---

106-110: Add a link to PyTorch profiler docs.

Small cross-reference helps readers interpret traces.

Apply this diff:

-The trace can be examined with [Chrome trace viewer](https://ui.perfetto.dev/) (alternatively chrome://tracing). It includes:
+The trace can be examined with [Chrome trace viewer](https://ui.perfetto.dev/) (alternatively chrome://tracing) and tooling in [PyTorch profiler docs](https://pytorch.org/docs/stable/profiler.html). It includes:

---

113-117: Show per-rank filename example (CPU and GPU).

Provide a concrete pattern to avoid clobbering in multi-process runs.

Apply this diff:

 ```bash
 export DP_PROFILER=result
 mpirun -np 4 lmp -in in.lammps
 # Produces result_gpuX.json, where X is the GPU id used by each MPI rank.
+
+# Example (OpenMPI) to ensure unique files per rank:
+# GPU:
+#   export DP_PROFILER="result_rank${OMPI_COMM_WORLD_RANK}"
+#   mpirun -np 4 lmp -in in.lammps
+#   # Produces result_rank0_gpuX.json, result_rank1_gpuY.json, ...
+# CPU:
+#   export DP_PROFILER="cpu_result_rank${OMPI_COMM_WORLD_RANK}"
+#   mpirun -np 4 lmp -in in.lammps
+#   # Produces cpu_result_rank0.json, cpu_result_rank1.json, ...

---

`119-123`: **Tighten wording; add explicit CPU multi-process limitation.**

Minor grammar fix and explicit limitation avoids surprises.

Apply this diff:

```diff
-Tips:
+Tips:
@@
-- Large runs can generate sizable JSON files; consider limiting numbers of MD steps, like 20.
-- Currently this feature only supports single process, or multi-process runs where each process uses a distinct GPU on the same node.
+- Large runs can generate sizable JSON files; consider limiting the number of MD steps (e.g., ~20).
+- This feature currently supports: (a) single-process runs; or (b) multi-process runs where each process uses a distinct GPU on the same node. CPU multi-process runs require unique `DP_PROFILER` stems per process to avoid file overwrites.

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro