Running PySR on Clusters with SLURM

[zouzaxd]

I submitted a job to SLURM with the following output:

Info: Starting SLURM job julia-4064213: srun -J julia-4064213 -n 64 -D /home/agripinoneto/PYSR -o /home/agripinoneto/PYSR/./julia-4064213-17367861256-%4t.out /apps/local/opt/julia/1.11.2/bin/julia --heap-size=5939M --worker=6vgeZiAhFp11MoXB
[ Info: Worker 0 (after 0 s): No output file "/home/agripinoneto/PYSR/./julia-4064213-17367861256-0000.out" yet
[ Info: Worker 0 (after 1 s): No output file "/home/agripinoneto/PYSR/./julia-4064213-17367861256-0000.out" yet

The job is running, but I'm not getting any output files, and it keeps showing "No output file yet" for a long time.

I am using PySR with the following configuration:

model = PySRRegressor(
    procs=128,
    parallelism='multiprocessing',
    cluster_manager="slurm",
)

And the job was submitted with the following SLURM command:

sbatch --qos=short --cpus-per-task=128 --nodes=2 ./script.sh

Questions:
Is there any issue with the PySR configuration that could be causing this problem?
Could the issue be with SLURM's resource allocation or the distribution of nodes?
What should I check to ensure that the output files are generated correctly and that the job is using the requested resources properly?

[MilesCranmer]

So the tricky thing about the slurm manager is you are only supposed to run it a single task inside that slurm allocation. It's easiest if you set up an interactive allocation with 2 nodes, and then ssh into the first node and run the script a single time.

Btw, with #794 I'm hoping that this will make things a bit easier for "standard" slurm scripts. If you have some time to try it, I'd love to hear if that fixes it for you (run with slurm_native instead of slurm)

[zouzaxd]

I installed the SlurmClusterManager, but I encountered the following message:

WARNING: could not import ClusterManagers.addprocs_slurm_native into Main
Compiling Julia backend...
Traceback (most recent call last):
  File "/home/agripinoneto/PYSR/PySR_scrit.py", line 141, in <module>
    models, errors = train_and_store_models(X.reshape(-1,1),Y)
                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/agripinoneto/PYSR/PySR_scrit.py", line 134, in train_and_store_models
    model.fit(treino_x, treino_y)  # Ajuste com pesos
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/agripinoneto/.local/lib/python3.12/site-packages/pysr/sr.py", line 2279, in fit
    self._run(X, y, runtime_params, weights=weights, seed=seed, category=category)
  File "/home/agripinoneto/.local/lib/python3.12/site-packages/pysr/sr.py", line 1845, in _run
    cluster_manager = _load_cluster_manager(cluster_manager)
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/agripinoneto/.local/lib/python3.12/site-packages/pysr/julia_helpers.py", line 34, in _load_cluster_manager
    return jl.seval(f"addprocs_{cluster_manager}")
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/agripinoneto/.julia/packages/PythonCall/Nr75f/src/JlWrap/module.jl", line 27, in seval
    return self._jl_callmethod($(pyjl_methodnum(pyjlmodule_seval)), expr)

It seems there might be an issue with importing ClusterManagers.addprocs_slurm_native. I am using PySR for Python, and the error is arising in this context.

[MilesCranmer]

You have to install pysr from that branch, since it’s not yet part of the released version yet

[MilesCranmer]

Sorry it turned out it was my fault. I fixed the PR: #794.

However, this won't actually work yet, until these upstream PRs are merged:

• Add support for params[:env] (that is, the env kwarg to the addprocs() function) kleinhenz/SlurmClusterManager.jl#20
• By default, propagate Base.active_project(), Base.LOAD_PATH, and Base.DEPOT_PATH to the workers (but allow the user to override that behavior by specifying JULIA_{PROJECT,LOAD_PATH,DEPOT_PATH} in params[:env]) kleinhenz/SlurmClusterManager.jl#30

In the meantime I would recommend trying to get the existing "slurm" cluster manager to work. I recommend using an interactive slurm allocation for this.