Test cupti fix (#54)

Summary:
## What does this PR do?
Fixes flaky test
https://app.circleci.com/pipelines/github/facebookresearch/HolisticTraceAnalysis/144/workflows/2bf5d461-8f9f-4b15-9f7e-c7df57f61d62/jobs/678

## Before submitting

- [ ] Was this discussed/approved via a Github issue? (no need for typos, doc improvements)
  - [x] N/A
- [ ] Did you write any new necessary tests?
  - [x] N/A
- [ ] Did you make sure to update the docs?
  - [x] N/A
- [ ] Did you update the [changelog](https://github.com/facebookresearch/HolisticTraceAnalysis/blob/main/CHANGELOG.md)?
  - [x] N/A

Pull Request resolved: #54

Reviewed By: briancoutinho

Differential Revision: D45995642

Pulled By: anupambhatnagar

fbshipit-source-id: afe9decabf0d7c3c73fe089344cb01411045b2f6

tests/test_trace_analysis.py

@@ -306,7 +306,17 @@ def test_get_cupti_counter_data_with_operators(self):
         cupti_profiler_t = TraceAnalysis(trace_dir=cupti_profiler_trace_dir)
         counters_df = cupti_profiler_t.get_cupti_counter_data_with_operators()[0]
 
-        test_row = counters_df.iloc[5].to_dict()
+        self.assertEqual(len(counters_df), 77)
+
+        # Example trace has CUPTI SASS FLOPS instruction counters
+        counter_names = set(CUDA_SASS_INSTRUCTION_COUNTER_FLOPS.keys())
+        self.assertEqual(
+            set(counters_df.columns.unique()) & counter_names, counter_names
+        )
+
+        # Pick 5th kernel that executed.
+        test_row = counters_df.sort_values(axis=0, by="ts").iloc[5].to_dict()
+
         self.assertEqual(test_row["cat"], "cuda_profiler_range")
         self.assertTrue("fft2d_r2c_32x32" in test_row["name"])
 
@@ -316,12 +326,6 @@ def test_get_cupti_counter_data_with_operators(self):
         self.assertEqual(test_row["top_level_op"], "aten::conv2d")
         self.assertEqual(test_row["bottom_level_op"], "aten::_convolution")
 
-        # Example trace has CUPTI SASS FLOPS instruction counters
-        counter_names = set(CUDA_SASS_INSTRUCTION_COUNTER_FLOPS.keys())
-        self.assertEqual(
-            set(counters_df.columns.unique()) & counter_names, counter_names
-        )
-
 
 if __name__ == "__main__":  # pragma: no cover
     unittest.main()