Several Biaflows containers are making use of AI models for segmentation. Instead of including the models into the Docker image, it is more convenient to store models permanently outside the container inside a folder on the HPC. This would also make it possible to store different (fine-tuned) models other than the default model loaded by the particular AI tool.

Proposal

• Create an additional parameter in slurm-config.ini named slurm_aimodels_path . In the example slurm-config.ini slurm_aimodels_path=my-scratch/models . There is already a parameter "slurm_model_paths" so I would suggest to use aimodels .

• At
  

  biomero/biomero/slurm_client.py

  Lines 2032 to 2038 in 15acc42

  	sbatch_env = {
  	"DATA_PATH": f"\"{self.slurm_data_path}/{input_data}\"",
  	"IMAGE_PATH": f"\"{self.slurm_images_path}/{model_path}\"",
  	"IMAGE_VERSION": f"{workflow_version}",
  	"SINGULARITY_IMAGE": f"\"{image}_{workflow_version}.sif\"",
  	"SCRIPT_PATH": f"\"{self.slurm_script_path}\"",
  	}

  Add "AIMODEL_PATH": f"{self.slurm_aimodels_path}",
  So we create a env variable AIMODEL_PATH and pass this to the slurm job template.

• At the job_template.sh add a --bind parameter to singularity run
  e.g. --bind "$AIMODEL_PATH:/tmp/models"

• For every container that wants/needs to make use of this we add a line to the DockerFile to tell the tool where to save the models. For example, microsam uses this env variable MICROSAM_CACHEDIR to indicate where to store the models. Cellpose does this similarly.

ENV MICROSAM_CACHEDIR=/tmp/models/microsam_cache

Example: https://github.com/maartenpaul/W_Segmentation-micro-sam/blob/7e9c5af7b1a6a7d3267977b2cbad8cdc724a4abc/Dockerfile#L16

• With this implementation if no AIMODEL_PATH is set when running the container with singularity it will be saved in the /tmp folder on the HPC and likely not cleaned up, maybe this is not ideal? Use home or scratch folder instead? When running with docker it requires downloading the model again every time, unless the /tmp folder is mounted.