Modify resources of not-yet-running jobs (i.e. in READY or slurm-SUBMITTED state)

[janosh]

posting a script in case it's useful as a starting point for others.
this script modifies resource allocations on all jobs that aren't already running or finished, including jobs already submitted to slurm. needs to be run on the cluster in order to interact with slurm.

this minimal example changes the number of allocated processes to n_processes = 8

# %%
import subprocess

import pandas as pd
from jobflow_remote import JobController
from jobflow_remote.jobs.state import JobState


__author__ = "Janosh Riebesell"
__date__ = "2024-06-16"


start_time_key, end_time_key = "start_time", "end_time"
state_key, slurm_id_key = "state", "process_id"


job_controller = JobController.from_project_name("proj_name")
job_coll = job_controller.jobs


# %% get value counts of jobs by state
state_counts = pd.DataFrame(
    job_coll.aggregate([{"$group": {"_id": "$state", "count": {"$sum": 1}}}])
).set_index("_id")["count"]


# %%
not_yet_running_states = [
    JobState.WAITING,
    JobState.READY,
    JobState.CHECKED_OUT,
    JobState.UPLOADED,
    JobState.SUBMITTED,
]

not_yet_running_states_vals = [x.value for x in not_yet_running_states]

state_counts.loc[[*set(not_yet_running_states_vals) & set(state_counts.index)]]

df_changeable = pd.DataFrame(
    job_coll.find({state_key: {"$in": not_yet_running_states_vals}})
)
df_changeable[slurm_id_key] = [dct[slurm_id_key] for dct in df_changeable.remote]


# %% modify resources of jobs that aren't queued yet
n_processes = 8

for state in not_yet_running_states:
    # skip SUBMITTED/CHECKED_OUT, raises: ValueError action cannot be performed
    if state in (JobState.SUBMITTED, JobState.CHECKED_OUT):
        continue
    print(f"{state.value=}")
    job_controller.set_job_run_properties(
        resources={"processes": n_processes}, state=state
    )


# %% modify resources of jobs already in a slurm queue
changeable_ids = df_changeable[slurm_id_key].dropna().tolist()
slurm_ids = ",".join(changeable_ids)

scontrol_cmd = (
    f"scontrol update jobid={slurm_ids} NumCPUs={n_processes} NumTasks={n_processes}"
)
subprocess.run(scontrol_cmd.split(), check=True)  # noqa: S603


# %% since jobflow_remote.cli.job.resources() is slow, we could modify job docs in DB
# directly. seems less safe though as the resources() implementation seems to have
# builtin checks and propagation to child jobs
update_in_place: bool = False

if update_in_place:
    job_coll.update_many(
        {f"remote.{slurm_id_key}": {"$in": changeable_ids}},
        {"$set": {"resources": {"processes": n_processes}}},
    )

[gpetretto]

Hi @janosh,

thanks for providing the script. As a general comment, if you don't need/want to modify the jobs already submitted (so you don't need to interact with slurm and you can run the script on the Runner machine) or if you are running in the all-in-one configuration, a convenient way to have access to the DB without manually including the connection details is to use the JobController object:

jc = JobController.from_project_name()
job_coll = jc.jobs

In this way it is also possible to directly use the JobController to apply the changes instead of the CLI functions (i.e. JobController.set_job_run_properties, which is the backend for the CLI resources function). In this PR I also just made the set_job_doc_properties public, that can be used to set any property of the Job document in the DB.

I have seen that you mention that these methods are slow. I suppose the reason is because it does not apply a bulk update, but locks each document and updates them individually to prevent race conditions with the runner. In principle I could also add a method for bulk updates that will not be safe if the runner is active, but could be way faster than the one currently available.

[janosh]

thanks @gpetretto, that's very helpful! i updated the script accordingly