Adding word_rank and word_size env options

Author: bgopesh

projects/rocprofiler-sdk/source/bin/rocprofv3.py

@@ -123,29 +123,59 @@ def strtobool(val):
         raise ValueError(f"invalid truth value {val} (type={val_type})")
 
 
-def get_mpi_rank():
-    """Detect MPI rank from various MPI implementation environment variables.
-    Returns the rank as an integer, or None if not running under MPI.
+def get_mpi_rank_and_size(custom_rank_env=None, custom_size_env=None):
+    """Detect MPI rank and size from the same MPI implementation's environment variables.
+
+    This ensures that rank and size come from the same source (e.g., both from OpenMPI,
+    not mixing OpenMPI rank with SLURM size).
+
+    Args:
+        custom_rank_env: Optional custom environment variable name for rank detection.
+        custom_size_env: Optional custom environment variable name for size detection.
+
+    Returns:
+        Tuple of (rank, size) as integers, or (None, None) if not running under MPI.
     """
+    # If custom environment variables are specified, use them exclusively
+    if custom_rank_env is not None and custom_size_env is not None:
+        rank = int(os.environ[custom_rank_env]) if custom_rank_env in os.environ else None
+        size = int(os.environ[custom_size_env]) if custom_size_env in os.environ else None
+        return (rank, size)
+
+    # Try each MPI implementation in order, ensuring both rank and size come from the same source
     # OpenMPI
-    if "OMPI_COMM_WORLD_RANK" in os.environ:
-        return int(os.environ["OMPI_COMM_WORLD_RANK"])
+    if "OMPI_COMM_WORLD_RANK" in os.environ and "OMPI_COMM_WORLD_SIZE" in os.environ:
+        return (
+            int(os.environ["OMPI_COMM_WORLD_RANK"]),
+            int(os.environ["OMPI_COMM_WORLD_SIZE"]),
+        )
+
     # MVAPICH2
-    elif "MV2_COMM_WORLD_RANK" in os.environ:
-        return int(os.environ["MV2_COMM_WORLD_RANK"])
-    # SLURM
-    elif "SLURM_PROCID" in os.environ:
-        return int(os.environ["SLURM_PROCID"])
+    if "MV2_COMM_WORLD_RANK" in os.environ and "MV2_COMM_WORLD_SIZE" in os.environ:
+        return (
+            int(os.environ["MV2_COMM_WORLD_RANK"]),
+            int(os.environ["MV2_COMM_WORLD_SIZE"]),
+        )
+
+    # SLURM (try SLURM_PROCID with both SLURM_NPROCS and SLURM_NTASKS)
+    if "SLURM_PROCID" in os.environ:
+        if "SLURM_NPROCS" in os.environ:
+            return (int(os.environ["SLURM_PROCID"]), int(os.environ["SLURM_NPROCS"]))
+        elif "SLURM_NTASKS" in os.environ:
+            return (int(os.environ["SLURM_PROCID"]), int(os.environ["SLURM_NTASKS"]))
+
     # PMI (used by some MPI implementations)
-    elif "PMI_RANK" in os.environ:
-        return int(os.environ["PMI_RANK"])
+    if "PMI_RANK" in os.environ and "PMI_SIZE" in os.environ:
+        return (int(os.environ["PMI_RANK"]), int(os.environ["PMI_SIZE"]))
+
     # Flux
-    elif "FLUX_TASK_RANK" in os.environ:
-        return int(os.environ["FLUX_TASK_RANK"])
-    # MPICH
-    elif "PMI_ID" in os.environ:
-        return int(os.environ["PMI_ID"])
-    return None
+    if "FLUX_TASK_RANK" in os.environ and "FLUX_JOB_SIZE" in os.environ:
+        return (int(os.environ["FLUX_TASK_RANK"]), int(os.environ["FLUX_JOB_SIZE"]))
+
+    # MPICH (PMI_ID is rank-like, but no corresponding size variable in this check)
+    # Skip this to avoid returning incomplete information
+
+    return (None, None)
 
 
 def parse_rank_specification(rank_spec):
@@ -181,19 +211,24 @@ def parse_rank_specification(rank_spec):
     return ranks
 
 
-def should_rank_provide_output(mpi_ranks_spec=None):
+def should_rank_provide_output(
+    mpi_ranks_spec=None, custom_rank_env=None, custom_size_env=None
+):
     """Check if the current MPI rank should provide profile/trace output.
     Args:
         mpi_ranks_spec: String specification of ranks (e.g., "0-3,8,10-15")
                        If None, all ranks provide output (default behavior).
+        custom_rank_env: Optional custom environment variable name for rank detection.
+        custom_size_env: Optional custom environment variable name for world size detection.
     Returns:
         True if this rank should provide output, False otherwise.
     """
     # If no specification provided, all ranks provide output (default)
     if not mpi_ranks_spec:
         return True
 
-    current_rank = get_mpi_rank()
+    # Get both rank and size from the same source
+    current_rank, world_size = get_mpi_rank_and_size(custom_rank_env, custom_size_env)
 
     # If we can't detect the rank, assume we should provide output
     if current_rank is None:
@@ -202,6 +237,15 @@ def should_rank_provide_output(mpi_ranks_spec=None):
     # Parse the rank specification
     selected_ranks = parse_rank_specification(mpi_ranks_spec)
 
+    # Validate that selected ranks are within the valid range
+    if world_size is not None and selected_ranks:
+        max_selected_rank = max(selected_ranks)
+        if max_selected_rank >= world_size:
+            fatal_error(
+                f"Invalid rank specification: rank {max_selected_rank} is out of range. "
+                f"MPI world size is {world_size} (valid ranks: 0-{world_size-1})"
+            )
+
     # Check if current rank is in the selected set
     return current_rank in selected_ranks
 
@@ -315,8 +359,8 @@ def parse_arguments(args=None):
 
 For MPI applications, select specific ranks to provide profile/trace output:
 
-    $ mpirun -n 16 rocprofv3 --hip-trace --mpi-ranks 0-3,8 -- ./mympiapp
-    $ srun -n 32 rocprofv3 --hip-trace --mpi-ranks 0 -- ./myapp
+    $ mpirun -n 16 rocprofv3 --hip-trace --profile-mpi-ranks 0-3,8 -- ./mympiapp
+    $ srun -n 32 rocprofv3 --hip-trace --profile-mpi-ranks 0 -- ./myapp
 
 For attachment profiling of running processes:
 
@@ -624,12 +668,26 @@ def add_parser_bool_argument(gparser, *args, **kwargs):
     filter_options = parser.add_argument_group("Filtering options")
 
     filter_options.add_argument(
-        "--mpi-ranks",
+        "--profile-mpi-ranks",
         help="Specify which MPI ranks should provide profile/trace output using comma-separated ranges and individual ranks (e.g., '0-3,8,10-15'). If not specified, all ranks provide output. The tool runs on all ranks but only selected ranks generate output files.",
         default=os.environ.get("ROCPROF_MPI_RANKS", None),
         type=str,
         metavar="RANK_SPECIFICATION",
     )
+    filter_options.add_argument(
+        "--mpi-world-rank-var",
+        help="Specify the environment variable to use for determining the MPI rank (e.g., 'MY_CUSTOM_RANK_VAR'). If not specified, the tool will automatically detect the rank from common MPI environment variables.",
+        default=None,
+        type=str,
+        metavar="ENVIRONMENT_VARIABLE",
+    )
+    filter_options.add_argument(
+        "--mpi-world-size-var",
+        help="Specify the environment variable to use for determining the MPI world size (e.g., 'MY_CUSTOM_SIZE_VAR'). If not specified, the tool will automatically detect the world size from common MPI environment variables.",
+        default=None,
+        type=str,
+        metavar="ENVIRONMENT_VARIABLE",
+    )
     filter_options.add_argument(
         "--kernel-include-regex",
         help="Include the kernels matching this filter from counter-collection and thread-trace data (non-matching kernels will be excluded)",
@@ -1167,15 +1225,43 @@ def run(app_args, args, **kwargs):
     use_execv = kwargs.get("use_execv", True)
     app_pass = kwargs.get("pass_id", None)
 
+    # Validate custom MPI environment variables
+    # If one custom variable is specified, both must be provided
+    custom_rank_env = (
+        args.mpi_world_rank_var if has_set_attr(args, "mpi_world_rank_var") else None
+    )
+    custom_size_env = (
+        args.mpi_world_size_var if has_set_attr(args, "mpi_world_size_var") else None
+    )
+
+    if (custom_rank_env is not None and custom_size_env is None) or (
+        custom_rank_env is None and custom_size_env is not None
+    ):
+        fatal_error(
+            "When using custom MPI environment variables, "
+            "both --mpi-world-rank-var and --mpi-world-size-var must be specified"
+        )
+
+    # Set MPI_RANK and MPI_SIZE environment variables for use by the C++ code
+    # These variables are used for %rank% and %size% expansion in output paths
+    # Use get_mpi_rank_and_size to ensure rank and size come from the same source
+    mpi_rank, mpi_size = get_mpi_rank_and_size(custom_rank_env, custom_size_env)
+    if mpi_rank is not None:
+        app_env["MPI_RANK"] = str(mpi_rank)
+    if mpi_size is not None:
+        app_env["MPI_SIZE"] = str(mpi_size)
+
     # Check if this MPI rank should provide profile/trace output
     # If not, run the application without profiling instrumentation
-    if has_set_attr(args, "mpi_ranks"):
-        if not should_rank_provide_output(args.mpi_ranks):
-            current_rank = get_mpi_rank()
-            if current_rank is not None and args.log_level in ("info", "trace"):
+    if has_set_attr(args, "profile_mpi_ranks"):
+        if not should_rank_provide_output(
+            args.profile_mpi_ranks, custom_rank_env, custom_size_env
+        ):
+            # We already have mpi_rank from above, just use it for logging
+            if mpi_rank is not None and args.log_level in ("info", "trace"):
                 sys.stderr.write(
-                    f"[rocprofv3] MPI rank {current_rank} not in selected ranks "
-                    f"({args.mpi_ranks}), running application without profiling\n"
+                    f"[rocprofv3] MPI rank {mpi_rank} not in selected ranks "
+                    f"({args.profile_mpi_ranks}), running application without profiling\n"
                 )
                 sys.stderr.flush()
             # Execute application without profiling

projects/rocprofiler-sdk/source/docs/how-to/using-rocprofv3-with-mpi.rst

@@ -172,26 +172,26 @@ Assuming the hostname as `ubuntu-latest`, the process IDs as 3000020 and 3000019
 Selective rank profiling
 ==========================
 
-When running large-scale MPI jobs, collecting profiling data from all ranks can generate excessive output and may not be necessary. The ``--mpi-ranks`` option allows you to specify which MPI ranks should provide profile and trace output, while the tool still runs on all ranks to maintain program correctness.
+When running large-scale MPI jobs, collecting profiling data from all ranks can generate excessive output and may not be necessary. The ``--profile-mpi-ranks`` option allows you to specify which MPI ranks should provide profile and trace output, while the tool still runs on all ranks to maintain program correctness.
 
 Specifying ranks to profile
 -----------------------------
 
-Use the ``--mpi-ranks`` option with a comma-separated list of ranks and ranges:
+Use the ``--profile-mpi-ranks`` option with a comma-separated list of ranks and ranges:
 
 .. code-block:: bash
 
     # Profile only rank 0
-    mpirun -n 16 rocprofv3 --hip-trace --mpi-ranks 0 -- <application_path>
+    mpirun -n 16 rocprofv3 --hip-trace --profile-mpi-ranks 0 -- <application_path>
 
     # Profile ranks 0-3 and rank 8
-    mpirun -n 16 rocprofv3 --hip-trace --mpi-ranks 0-3,8 -- <application_path>
+    mpirun -n 16 rocprofv3 --hip-trace --profile-mpi-ranks 0-3,8 -- <application_path>
 
     # Profile ranks 0, 4, 8, and 12
-    mpirun -n 16 rocprofv3 --hip-trace --mpi-ranks 0,4,8,12 -- <application_path>
+    mpirun -n 16 rocprofv3 --hip-trace --profile-mpi-ranks 0,4,8,12 -- <application_path>
 
     # Profile a range of ranks (10 through 15)
-    srun -n 32 rocprofv3 --kernel-trace --mpi-ranks 10-15 -- <application_path>
+    srun -n 32 rocprofv3 --kernel-trace --profile-mpi-ranks 10-15 -- <application_path>
 
 The rank specification syntax supports:
 
@@ -202,7 +202,7 @@ The rank specification syntax supports:
 Behavior
 ---------
 
-When using ``--mpi-ranks``:
+When using ``--profile-mpi-ranks``:
 
 - The ``rocprofv3`` tool runs on **all** MPI ranks to avoid disrupting the application's execution
 - Only the specified ranks collect and output profiling/trace data
@@ -212,15 +212,35 @@ When using ``--mpi-ranks``:
 Default behavior
 -----------------
 
-If ``--mpi-ranks`` is not specified, all ranks provide output (default behavior), which is equivalent to running without rank filtering.
+If ``--profile-mpi-ranks`` is not specified, all ranks provide output (default behavior), which is equivalent to running without rank filtering.
 
 Supported MPI implementations
 -------------------------------
 
-The rank detection supports the following MPI implementations and job launchers:
+The rank and world size detection automatically supports the following MPI implementations and job launchers:
 
-- OpenMPI (``OMPI_COMM_WORLD_RANK``)
-- MVAPICH2 (``MV2_COMM_WORLD_RANK``)
-- MPICH (``PMI_ID``, ``PMI_RANK``)
-- SLURM (``SLURM_PROCID``)
-- Flux (``FLUX_TASK_RANK``)
+- OpenMPI (``OMPI_COMM_WORLD_RANK``, ``OMPI_COMM_WORLD_SIZE``)
+- MVAPICH2 (``MV2_COMM_WORLD_RANK``, ``MV2_COMM_WORLD_SIZE``)
+- MPICH (``PMI_ID``/``PMI_RANK``, ``PMI_SIZE``)
+- SLURM (``SLURM_PROCID``, ``SLURM_NPROCS``/``SLURM_NTASKS``)
+- Flux (``FLUX_TASK_RANK``, ``FLUX_JOB_SIZE``)
+
+Custom MPI environment variables
+----------------------------------
+
+For mixed environments or non-standard MPI configurations (such as interactive SLURM jobs using mpirun), you can specify custom environment variables for rank and world size detection:
+
+.. code-block:: bash
+
+    # Use custom environment variables for rank and world size detection
+    mpirun -n 16 rocprofv3 --hip-trace --profile-mpi-ranks 0-3 \
+        --mpi-world-rank-var MY_CUSTOM_RANK \
+        --mpi-world-size-var MY_CUSTOM_SIZE -- <application_path>
+
+When ``--mpi-world-rank-var`` and/or ``--mpi-world-size-var`` are specified, they take precedence over automatic detection. These options are useful when:
+
+- The MPI implementation uses non-standard environment variable names
+- You want to ensure specific environment variables are used instead of relying on auto-detection
+- Working in mixed MPI environments where multiple MPI-related variables might be present
+
+The ``--mpi-world-size-var`` option is particularly useful when using ``--profile-mpi-ranks`` to ensure that the specified rank ranges are validated against the actual MPI world size, preventing out-of-range errors.

projects/rocprofiler-sdk/tests/rocprofv3/mpi-ranks/CMakeLists.txt

@@ -34,7 +34,8 @@ if(MPIRUN_EXECUTABLE)
             $<TARGET_FILE:rocprofiler-sdk::rocprofv3> --hip-trace --kernel-trace
             --output-format csv json -d
             ${CMAKE_CURRENT_BINARY_DIR}/mpi-ranks-trace/rank.%env{OMPI_COMM_WORLD_RANK}%
-            -o out --mpi-ranks 0 ${PRELOAD_ARGS} -- $<TARGET_FILE:simple-transpose>
+            -o out --profile-mpi-ranks 0 ${PRELOAD_ARGS} --
+            $<TARGET_FILE:simple-transpose>
         DEPENDS simple-transpose
         TIMEOUT 120
         LABELS "integration-tests" "mpi-tests"
@@ -50,7 +51,8 @@ if(MPIRUN_EXECUTABLE)
             $<TARGET_FILE:rocprofiler-sdk::rocprofv3> --hip-trace --kernel-trace
             --output-format csv json -d
             ${CMAKE_CURRENT_BINARY_DIR}/mpi-ranks-multiple-trace/rank.%env{OMPI_COMM_WORLD_RANK}%
-            -o out --mpi-ranks 0-1,3 ${PRELOAD_ARGS} -- $<TARGET_FILE:simple-transpose>
+            -o out --profile-mpi-ranks 0-1,3 ${PRELOAD_ARGS} --
+            $<TARGET_FILE:simple-transpose>
         DEPENDS simple-transpose
         TIMEOUT 120
         LABELS "integration-tests" "mpi-tests"
@@ -86,7 +88,7 @@ rocprofiler_add_integration_execute_test(
     COMMAND
         $<TARGET_FILE:rocprofiler-sdk::rocprofv3> --hip-trace --kernel-trace
         --output-format csv json -d ${CMAKE_CURRENT_BINARY_DIR}/no-mpi-trace -o out
-        --mpi-ranks 0 ${PRELOAD_ARGS} -- $<TARGET_FILE:simple-transpose>
+        --profile-mpi-ranks 0 ${PRELOAD_ARGS} -- $<TARGET_FILE:simple-transpose>
     DEPENDS simple-transpose
     TIMEOUT 120
     LABELS "integration-tests"