Part 2 of updates to get the workflow running on klone.

dirac-institute · Jul 23, 2024 · d494b97 · d494b97
1 parent 1264318
commit d494b97
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 9 deletions.
diff --git a/example_runtime_config.toml b/example_runtime_config.toml
@@ -19,6 +19,8 @@ staging_directory = "/home/drew/code/kbmod-wf/dev_staging"
 search_config_filepath = "/home/drew/code/kbmod-wf/dev_staging/search_config.yaml"
 
 [apps.reproject_wu]
+# Number of processors to use for parallelizing the reprojection
+n_workers = 8
 # The name of the observation site to use for reflex correction
 observation_site = "ctio"
 

diff --git a/src/kbmod_wf/resource_configs/klone_configuration.py b/src/kbmod_wf/resource_configs/klone_configuration.py
@@ -7,7 +7,7 @@
 
 walltimes = {
     "compute_bigmem": "01:00:00",
-    "large_mem": "12:00:00",
+    "large_mem": "04:00:00",
     "gpu_max": "12:00:00",
 }
 
@@ -28,11 +28,11 @@ def klone_resource_config():
                     account="astro",
                     min_blocks=0,
                     max_blocks=4,
-                    init_blocks=1,
+                    init_blocks=0,
                     parallelism=1,
                     nodes_per_block=1,
                     cores_per_node=1,  # perhaps should be 8???
-                    mem_per_node=64,  # In GB
+                    mem_per_node=256,  # In GB
                     exclusive=False,
                     walltime=walltimes["compute_bigmem"],
                     # Command to run before starting worker - i.e. conda activate <special_env>
@@ -42,11 +42,11 @@ def klone_resource_config():
             HighThroughputExecutor(
                 label="large_mem",
                 provider=SlurmProvider(
-                    partition="compute-bigmem",
+                    partition="ckpt-g2",
                     account="astro",
                     min_blocks=0,
                     max_blocks=2,
-                    init_blocks=1,
+                    init_blocks=0,
                     parallelism=1,
                     nodes_per_block=1,
                     cores_per_node=8,
@@ -59,13 +59,13 @@ def klone_resource_config():
             ),
             HighThroughputExecutor(
                 label="gpu",
-                available_accelerators=2,
+                available_accelerators=1,
                 provider=SlurmProvider(
                     partition="gpu-a40",
                     account="escience",
                     min_blocks=0,
                     max_blocks=2,
-                    init_blocks=1,
+                    init_blocks=0,
                     parallelism=1,
                     nodes_per_block=1,
                     cores_per_node=4,  # perhaps should be 8???

diff --git a/src/kbmod_wf/task_impls/reproject_wu.py b/src/kbmod_wf/task_impls/reproject_wu.py
@@ -85,9 +85,9 @@ def __init__(
         self.search_config = self.runtime_config.get("search_config", None)
 
         # Default to 8 workers if not in the config. Value must be 0<num workers<65.
-        self.n_workers = np.max(1, np.min(self.runtime_config.get("n_workers", 8), 64))
+        self.n_workers = max(1, min(self.runtime_config.get("n_workers", 8), 64))
 
-        self.uri_params = self._get_params_from_uri_file(uri_file=self.uri_file)
+        self.uri_params = self._get_params_from_uri_file(uri_file=self.uri_filepath)
         self.patch_size = self.uri_params["patch_size"]
         self.pixel_scale = self.uri_params["pixel_scale"]
         self.guess_dist = self.uri_params["dist_au"]  # ! Let's update the terminology here to be consistent.