From 38b43454a187bac5e88e4bbf54e98a2cccdd996f Mon Sep 17 00:00:00 2001 From: Christian Tabedzki <35670232+tabedzki@users.noreply.github.com> Date: Fri, 10 Jul 2026 12:46:36 -0400 Subject: [PATCH] feat(slurm): uv env prefetch + conda->uv for ephys, wire up spockmk2 ephys generator Compute nodes have no network, so packages must be present before sbatch. The submitter already SSHes to the head node, so sync the uv env there first: - prefetch_uv_env(): on the head node (which has network), run `uv python install 3.14 && uv sync --locked` before sbatch, reusing the existing ssh connection and mirroring the on-spock local-run branch. Raises so a dependency problem surfaces at submit time, not minutes later on a node. Gated to the BrainCogsEphysSorters repo; imaging is untouched. - Dispatch ephys (process_repository == BrainCogsEphysSorters) to the previously dead generate_slurm_spockmk2_ephys; imaging spock keeps generate_slurm_spock on conda. - generate_slurm_spockmk2_ephys: replace conda with `uv sync --frozen --offline` + `uv run --frozen --offline`; keep `module load matlab` for the KS2/KS3 shell-outs; add --cpus-per-task=8 (DREDge peak detection is the CPU-parallel bottleneck) and --time=16:00:00. Also fix a pre-existing bug where `source ~/.bashrc` was concatenated onto the first #SBATCH line and swallowed it. Companion to BrainCOGS/BrainCogsEphysSorters#5 (DREDge) and the uv migration. Closes #94 Closes #54 --- u19_pipeline/automatic_job/slurm_creator.py | 65 ++++++++++++++++++--- 1 file changed, 58 insertions(+), 7 deletions(-) diff --git a/u19_pipeline/automatic_job/slurm_creator.py b/u19_pipeline/automatic_job/slurm_creator.py index d1d91c1c..464e3638 100644 --- a/u19_pipeline/automatic_job/slurm_creator.py +++ b/u19_pipeline/automatic_job/slurm_creator.py @@ -40,7 +40,11 @@ def generate_slurm_file(job_id, program_selection_params): print('slurm_dict', slurm_dict) if program_selection_params['process_cluster'] == 'spock': - slurm_text = generate_slurm_spock(slurm_dict) + # Ephys (BrainCogsEphysSorters) runs on uv; imaging keeps the conda path. + if program_selection_params['process_repository'] == 'BrainCogsEphysSorters': + slurm_text = generate_slurm_spockmk2_ephys(slurm_dict) + else: + slurm_text = generate_slurm_spock(slurm_dict) else: slurm_text = generate_slurm_tiger(slurm_dict) @@ -67,6 +71,40 @@ def generate_slurm_file(job_id, program_selection_params): return status, slurm_destination +def prefetch_uv_env(program_selection_params, cluster_vars, repository_dir): + ''' + Sync the uv environment on the head node before submitting the job. + + Compute nodes have no network access, so all packages must be present in the uv cache + (on the mounted $HOME) before sbatch. The head node has network, so we run + `uv sync --locked` there. `uv sync --locked` is idempotent and near-instant on a warm + cache, so this is cheap on every submit and only does real work after a uv.lock change. + + Raises so a dependency problem surfaces at submit time with a real error message, + instead of as a job that dies minutes later on a compute node. + ''' + + # bash -lc so a login profile is sourced and `uv` is on PATH over non-interactive ssh. + remote = ("cd " + repository_dir + " && " + "uv python install 3.14 && " + "uv sync --locked") + + if program_selection_params['process_cluster'] == 'spock' and is_this_spock(): + command = ['bash', '-lc', remote] # already on spock: run locally + else: + command = ['ssh', cluster_vars['user']+"@"+cluster_vars['hostname'], 'bash', '-lc', remote] + + print('prefetch_uv_env', command) + p = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE) + p.wait() + stdout, stderr = p.communicate() + print(stdout) + print(stderr) + + if p.returncode != config.system_process['SUCCESS']: + raise Exception('uv env prefetch failed before sbatch: ' + stderr.decode('UTF-8')) + + def queue_slurm_file(job_id, program_selection_params, raw_directory, proc_directory, modality, slurm_location): id_slurm_job = -1 @@ -81,6 +119,11 @@ def queue_slurm_file(job_id, program_selection_params, raw_directory, proc_direc processing_repository = program_selection_params['process_repository'] repository_dir = pathlib.Path(cluster_vars[modality+'_process_dir'],processing_repository).as_posix() + # The BrainCogsEphysSorters repo runs on uv. Compute nodes have no network, so sync the + # environment on the head node (which does) before submitting; the job then runs offline. + if processing_repository == 'BrainCogsEphysSorters': + prefetch_uv_env(program_selection_params, cluster_vars, repository_dir) + command = ['ssh', cluster_vars['user']+"@"+cluster_vars['hostname'], 'sbatch', "--export=recording_process_id="+job_id+ ",raw_data_directory='"+raw_directory+ @@ -230,10 +273,20 @@ def module_defininition_text(): def generate_slurm_spockmk2_ephys(slurm_dict): + # Ephys-specific resources. DREDge peak detection is CPU-parallel and is the DREDge + # wall-time bottleneck, so give the (single) task multiple cores; keep a generous + # walltime since DREDge is additive on top of Kilosort. Copy so we don't mutate the + # shared default dict used by the imaging path. + slurm_dict = copy.deepcopy(slurm_dict) + slurm_dict['cpus-per-task'] = 8 + slurm_dict['time'] = '16:00:00' + + # #SBATCH directives must come first; put `source ~/.bashrc` in the body (after the + # directives) so it does not swallow the first #SBATCH line, and so `uv` is on PATH. slurm_text = '#!/bin/bash\n' - slurm_text += 'source ~/.bashrc' slurm_text += create_slurm_params_file(slurm_dict) slurm_text += ''' + source ~/.bashrc echo "SLURM_JOB_ID: ${SLURM_JOB_ID}" echo "SLURM_SUBMIT_DIR: ${SLURM_SUBMIT_DIR}" echo "RECORDING_PROCESS_ID: ${recording_process_id}" @@ -242,14 +295,12 @@ def generate_slurm_spockmk2_ephys(slurm_dict): echo "REPOSITORY_DIR: ${repository_dir}" echo "PROCESS_SCRIPT_PATH: ${process_script_path}" - module load anacondapy/2023.07-cuda -s + # matlab is still needed for the Kilosort2 / Kilosort3 shell-outs; uv replaces conda. module load matlab/R2024a -s - conda activate BraincogsEphysSorters_Env - cd ${repository_dir} - python -u ${process_script_path} - #python ${process_script_path} ${recording_process_id} + uv sync --frozen --offline + uv run --frozen --offline python -u ${process_script_path} ''' return slurm_text