Update USER_GUIDE.md and Python imports to use the latest PARAM branch

.github/workflows/end_to_end_tests.yml

@@ -22,8 +22,9 @@ jobs:
     - name: Install PARAM
       run: |
         git clone https://github.com/facebookresearch/param.git
-        cd param/train/compute/python/
-        git checkout c83ce8429110a86549c40fec5a01acbd9fbd54a4
+        cd param/et_replay
+        git checkout 884a1f0154a16e2c170e456f8027f2646c9108ae
+        sed -i '/param_bench/d' pyproject.toml
         pip install .
 
     - name: Test chakra_trace_link Without Arguments

.github/workflows/python_tests.yml

@@ -22,8 +22,9 @@ jobs:
     - name: Install PARAM
       run: |
         git clone https://github.com/facebookresearch/param.git
-        cd param/train/compute/python/
-        git checkout c83ce8429110a86549c40fec5a01acbd9fbd54a4
+        cd param/et_replay
+        git checkout 884a1f0154a16e2c170e456f8027f2646c9108ae
+        sed -i '/param_bench/d' pyproject.toml
         pip install .
 
     - name: Install Dependencies

USER_GUIDE.md

@@ -30,8 +30,9 @@ Installing PARAM is necessary for Chakra to function properly as it imports esse
 
 ```bash
 $ git clone git@github.com:facebookresearch/param.git
-$ cd param/train/compute/python/
-$ git checkout c83ce8429110a86549c40fec5a01acbd9fbd54a4
+$ cd param/et_replay
+$ git checkout 884a1f0154a16e2c170e456f8027f2646c9108ae
+$ sed -i '' '13d' pyproject.toml
 $ pip install .
 ```
 

src/trace_link/kineto_operator.py

@@ -1,6 +1,6 @@
 from typing import Any, Dict, Optional
 
-from param_bench.train.compute.python.tools.execution_trace import Node as PyTorchOperator
+from et_replay.lib.execution_trace import Node as PyTorchOperator
 
 
 class KinetoOperator:

src/trace_link/trace_linker.py

@@ -6,14 +6,14 @@
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from typing import Dict, List, Optional, Tuple
 
-from param_bench.train.compute.python.tools.execution_trace import (
+from et_replay.lib.execution_trace import (
     EXECUTION_TRACE_PROCESS_ANNOTATION,
     EXECUTION_TRACE_THREAD_ANNOTATION,
 )
-from param_bench.train.compute.python.tools.execution_trace import (
+from et_replay.lib.execution_trace import (
     Node as PyTorchOperator,
 )
-from param_bench.train.compute.python.tools.utility import (
+from et_replay.lib.utils import (
     load_execution_trace_file,
     read_dictionary_from_json_file,
 )
@@ -416,7 +416,9 @@ def enforce_inter_thread_order(
 
         with ThreadPoolExecutor() as executor:
             futures = {
-                executor.submit(self.process_thread_inter_thread_order, tid, ops, kineto_tid_cpu_ops_map, threshold): tid
+                executor.submit(
+                    self.process_thread_inter_thread_order, tid, ops, kineto_tid_cpu_ops_map, threshold
+                ): tid
                 for tid, ops in kineto_tid_cpu_ops_map.items()
             }
 

tests/trace_link/test_trace_linker.py

@@ -4,11 +4,11 @@
 from chakra.src.trace_link.kineto_operator import KinetoOperator
 from chakra.src.trace_link.trace_linker import TraceLinker
 from chakra.src.trace_link.unique_id_assigner import UniqueIdAssigner
-from param_bench.train.compute.python.tools.execution_trace import (
+from et_replay.lib.execution_trace import (
     EXECUTION_TRACE_PROCESS_ANNOTATION,
     EXECUTION_TRACE_THREAD_ANNOTATION,
 )
-from param_bench.train.compute.python.tools.execution_trace import Node as PyTorchOperator
+from et_replay.lib.execution_trace import Node as PyTorchOperator
 
 
 @pytest.fixture