Cleanup jlaunch with multi-backend support

kylevedder · Apr 11, 2024 · 8ba3695 · 8ba3695
1 parent 2f82587
commit 8ba3695
Show file tree

Hide file tree

Showing 13 changed files with 401 additions and 617 deletions.
diff --git a/.gitignore b/.gitignore
@@ -23,4 +23,5 @@ cluster_tb_logs/
 dataset_count_compute_save_dir/
 checkpoint_eval_launch_dir/
 screenshots/
-eval_results/
+eval_results/
+launch_files/
diff --git a/data_prep_scripts/check_nsfp_status.py b/data_prep_scripts/check_nsfp_status.py
diff --git a/data_prep_scripts/compute_total_runtime.py b/data_prep_scripts/compute_total_runtime.py
diff --git a/data_prep_scripts/investigate_bucket_sizes.py b/data_prep_scripts/investigate_bucket_sizes.py
diff --git a/data_prep_scripts/investigate_pointclouds.py b/data_prep_scripts/investigate_pointclouds.py
diff --git a/data_prep_scripts/jlaunch.py b/data_prep_scripts/jlaunch.py
@@ -0,0 +1,20 @@
+#!/usr/bin/env python3
+import sys
+
+sys.path.insert(0, ".")
+
+import click
+
+from data_prep_scripts.jlaunch_backends import backends_list
+
+
+@click.group()
+def cli():
+    pass
+
+
+for backend in backends_list:
+    cli.add_command(backend)
+
+if __name__ == "__main__":
+    cli()
diff --git a/data_prep_scripts/jlaunch_backends/__init__.py b/data_prep_scripts/jlaunch_backends/__init__.py
@@ -0,0 +1,5 @@
+from .slurm import backend_slurm
+from .ngc import backend_ngc
+from .shared_utils import run_cmd
+
+backends_list = [backend_slurm, backend_ngc]