Slurm add --segment argument (#186)

* Support --segment; add cuda-event-trace=false

Signed-off-by: Guyue Huang <[email protected]>

* Remove irrelevant changes

Signed-off-by: Guyue Huang <[email protected]>

---------

Signed-off-by: Guyue Huang <[email protected]>
Signed-off-by: Guyue Huang <[email protected]>
Co-authored-by: Guyue Huang <[email protected]>

Author: guyueh1

nemo_run/core/execution/slurm.py

@@ -241,6 +241,7 @@ def your_slurm_executor() -> run.SlurmExecutor:
         "wait_all_nodes",
         "wckey",
         "wrap",
+        "segment",
     ]
 
     SRUN_ARGS = [
@@ -338,6 +339,7 @@ class ResourceRequest:
     monitor_group_job_wait_time: int = 60
     setup_lines: Optional[str] = None
     het_group_indices: Optional[list[int]] = None
+    segment: Optional[int] = None
 
     #: Set by the executor; cannot be initialized
     job_name: str = field(init=False, default="nemo-job")