Allow >1 CPU per task, and some misc clean up

invest_processes/pyproject.toml

@@ -10,7 +10,6 @@ maintainers = [
 dynamic = ["version"] # the version is provided dynamically by setuptools_scm
 dependencies = [
     "flask",
-    "natcap.invest",
     "pygeoapi>=0.23.4",
     "google-cloud-storage"
 ]

invest_processes/src/invest_processes/execute.py

@@ -1,8 +1,8 @@
 import logging
+import multiprocessing
 from pathlib import Path
 import textwrap
 
-from invest_processes.utils import download_and_extract_datastack
 from pygeoapi.process.base import BaseProcessor, ProcessorExecuteError
 
 LOGGER = logging.getLogger(__name__)
@@ -46,6 +46,8 @@
     }
 }
 
+CPUS_PER_TASK = 2  # number of cpus per task to submit to sbatch
+
 
 class ExecuteProcessor(BaseProcessor):
     """InVEST execute process"""
@@ -76,11 +78,15 @@ def create_slurm_script(self, datastack_url, workspace_dir):
         json_path = f'{workspace_dir}/datastack/parameters.invest.json'
         return textwrap.dedent(f"""\
             #!/bin/sh
+            #SBATCH --cpus-per-task={CPUS_PER_TASK}
 
             curl -o datastack.tgz "{datastack_url}"
             mkdir {workspace_dir}/datastack
             tar -xzvf datastack.tgz -C {workspace_dir}/datastack
             rm datastack.tgz
+
+            eval "$(~/bin/micromamba shell hook -s posix)"
+            micromamba activate invest_env
             MODEL_ID=$(python -c "from natcap.invest import datastack; print(datastack.extract_parameter_set('{json_path}').model_id)")
             invest --debug --taskgraph-log-level=DEBUG run \
                 --datastack {json_path} \

invest_processes/src/invest_processes/slurm_manager.py

@@ -711,10 +711,12 @@ def submit_slurm_job(self, processor, data_dict):
                 f'Submitting slurm job with the following command:\n{args}')
             result = subprocess.run(
                 args, capture_output=True, text=True, check=True)
-            LOGGER.info(f'stdout from sbatch: {result.stdout}')
 
         except subprocess.CalledProcessError as e:
-            raise RuntimeError('Error when submitting slurm job') from e
+            LOGGER.error('error:')
+            LOGGER.error(str(e))
+            LOGGER.error(f'stdout from sbatch: {e.stderr}')
+            raise RuntimeError(f'Error when submitting slurm job: {e.stderr}') from e
 
         job_id = result.stdout.strip()
         LOGGER.info(f"Job submitted successfully with ID: {job_id}")

invest_processes/src/invest_processes/validate.py

@@ -4,7 +4,6 @@
 import textwrap
 import time
 
-from invest_processes.utils import download_and_extract_datastack
 from pygeoapi.process.base import BaseProcessor, ProcessorExecuteError
 
 LOGGER = logging.getLogger(__name__)

invest_processes/startup_script.sh

@@ -0,0 +1,20 @@
+#!/bin/bash
+
+scl enable gcc-toolset-12 bash
+
+# install micromamba
+curl -Ls https://micro.mamba.pm/api/micromamba/linux-64/latest | tar -xvj ~/bin/micromamba
+eval "$(~/bin/micromamba shell hook -s posix)"
+micromamba --version
+
+# create and activate the environment
+micromamba create -y -n invest_env python=3.13 gdal
+micromamba activate invest_env
+pip install setuptools setuptools_scm build cython babel
+pip install --no-build-isolation git+https://github.com/natcap/pygeoprocessing.git
+pip install --no-build-isolation git+https://github.com/emlys/invest.git@bugfix/2572
+invest --version
+
+# activate the environment in future shells
+echo 'echo "activating micromamba env"' >> /etc/bashrc
+echo 'eval "$(~/bin/micromamba shell hook -s posix)" && micromamba activate invest_env' >> ~/.bashrc

slurm_cluster_config/hpc-slurm/main.tf

@@ -58,18 +58,20 @@ module "homefs" {
   region            = var.region
   reserved_ip_range = module.private_service_access.reserved_ip_range
   zone              = var.zone
+  size_gb           = 2560
 }
 
 module "debug_nodeset" {
   source                  = "github.com/GoogleCloudPlatform/cluster-toolkit//community/modules/compute/schedmd-slurm-gcp-v6-nodeset?ref=v1.90.0"
   allow_automatic_updates = false
   labels                  = var.labels
-  machine_type            = "c2-standard-8"
+  machine_type            = "c2-standard-4"
   disk_type               = "pd-ssd"
   name                    = "debug_nodeset"
   node_count_dynamic_max  = 4
   project_id              = var.project_id
   region                  = var.region
+  startup_script          = "/home/bin/startup_script.sh"
   subnetwork_self_link    = module.network.subnetwork_self_link
   zone                    = var.zone
 }
@@ -87,12 +89,12 @@ module "compute_nodeset" {
   allow_automatic_updates = false
   bandwidth_tier          = "gvnic_enabled"
   labels                  = var.labels
-  machine_type            = "c2-standard-8"
+  machine_type            = "c2-standard-4"
   name                    = "compute_nodeset"
   node_count_dynamic_max  = 20
   project_id              = var.project_id
   region                  = var.region
-  startup_script          = "startup_script.sh"
+  startup_script          = "/home/bin/startup_script.sh"
   subnetwork_self_link    = module.network.subnetwork_self_link
   zone                    = var.zone
 }
@@ -128,7 +130,7 @@ module "slurm_login" {
   source                  = "github.com/GoogleCloudPlatform/cluster-toolkit//community/modules/scheduler/schedmd-slurm-gcp-v6-login?ref=v1.90.0"
   enable_login_public_ips = true
   labels                  = var.labels
-  machine_type            = "n2-standard-4"
+  machine_type            = "n2-standard-2"
   name_prefix             = "slurm_login"
   project_id              = var.project_id
   region                  = var.region
@@ -143,6 +145,7 @@ module "slurm_controller" {
   enable_controller_public_ips = true
   labels                       = var.labels
   login_nodes                  = flatten([module.slurm_login.login_nodes])
+  machine_type                 = "n2-standard-2"
   network_storage              = flatten([module.homefs.network_storage])
   nodeset                      = flatten([module.h3_partition.nodeset, flatten([module.compute_partition.nodeset, flatten([module.debug_partition.nodeset])])])
   nodeset_dyn                  = flatten([module.h3_partition.nodeset_dyn, flatten([module.compute_partition.nodeset_dyn, flatten([module.debug_partition.nodeset_dyn])])])

slurm_cluster_config/login_node_playbook.yml

@@ -42,20 +42,27 @@
           - "--exclude=slurm_cluster_config"
           - "--exclude=tests"
 
+    - name: Move startup script to shared location
+      become: true
+      ansible.builtin.shell: |
+        mkdir /home/bin
+        mv /home/esoth_stanford_edu/invest-compute/invest_processes/startup_script.sh /home/bin
+        chmod +x /home/bin/startup_script.sh
+
     - name: Install and launch server
       args:
         chdir: invest-compute/invest_processes
         executable: /bin/bash # Ensure bash is used for shell hooks
       ansible.builtin.shell: |
         # install micromamba
         rm -f bin/micromamba
-        curl -Ls https://micro.mamba.pm/api/micromamba/linux-64/latest | tar -xvj bin/micromamba
-        eval "$(./bin/micromamba shell hook -s posix)"
+        curl -Ls https://micro.mamba.pm/api/micromamba/linux-64/latest | tar -xvj ~/bin/micromamba
+        eval "$(~/bin/micromamba shell hook -s posix)"
         micromamba --version
 
         # create and activate the environment
-        micromamba create -y -n env python=3.13 gdal natcap.invest==3.17.2
-        micromamba activate env
+        micromamba create -y -n server_env python=3.13
+        micromamba activate server_env
         pip install --no-cache-dir --only-binary=:all: pygeoapi
         pygeoapi --version