diff --git a/jax/_src/profiler.py b/jax/_src/profiler.py index 912c90182977..f1016d262056 100644 --- a/jax/_src/profiler.py +++ b/jax/_src/profiler.py @@ -444,8 +444,7 @@ def trace(cls, runner: PGLEProfiler | None): if runner.fdo_profiles[-1] == b'': warnings.warn( "PGLE collected an empty trace, may be due to contention with " - "another tool that subscribes to CUPTI, such as Nsight Systems - check " - "for CUPTI_ERROR_MULTIPLE_SUBSCRIBERS_NOT_SUPPORTED from XLA. " + "another tool that subscribes to rocprofiler-sdk - check " "Consider populating a persistent compilation cache with PGLE enabled, " "and then profiling a second run that has the " "JAX_COMPILATION_CACHE_EXPECT_PGLE option enabled.", diff --git a/jaxlib/tools/gpu_version_script.lds b/jaxlib/tools/gpu_version_script.lds index 8e46b2c590b2..c40297736cda 100644 --- a/jaxlib/tools/gpu_version_script.lds +++ b/jaxlib/tools/gpu_version_script.lds @@ -4,6 +4,7 @@ VERS_1.0 { GetPjrtApi; MosaicGpuCompile; MosaicGpuUnload; + rocprofiler_configure; }; local: diff --git a/tests/profiler_test.py b/tests/profiler_test.py index 215e363e446d..65bab7acf77a 100644 --- a/tests/profiler_test.py +++ b/tests/profiler_test.py @@ -244,8 +244,7 @@ def _check_xspace_pb_exist(self, logdir): path = os.path.join(logdir, 'plugins', 'profile', '*', '*.xplane.pb') self.assertEqual(1, len(glob.glob(path)), 'Expected one path match: ' + path) - - @unittest.skip("Test causes OOMs") + @unittest.skipIf(not (portpicker and profiler_client and tf_profiler), "Test requires tensorflow.profiler and portpicker") def testSingleWorkerSamplingMode(self, delay_ms=None):