1"""
  2# Description
  3
  4Functions to handle Slurm calls, to run calculations in clusters.  
  5
  6
  7# Index
  8
  9| | |
 10| --- | --- |
 11| `sbatch()`           | Sbatch all calculations |
 12| `squeue()`           | Get a Pandas DataFrame with info about the submitted calculations |
 13| `scancel()`          | Scancel all calculations, or applying some filters |
 14| `scancel_errors()`   | Scancel calculations with specific errors |
 15| `scancel_here()`     | Scancel all calculations running from a specific folder |
 16| `get_running_here()` | Get a list with all jobs running in a given folder |
 17| `check_template()`   | Checks that the slurm template is OK, and provides an example if not |
 18
 19---
 20"""
 21
 22
 23import os
 24import pandas as pd
 25import aton.call as call
 26import aton.file as file
 27import aton.txt.find as find
 28import aton.txt.edit as edit
 29from aton._version import __version__
 30
 31
 32def sbatch(
 33        prefix:str='',
 34        template:str='template.slurm',
 35        in_ext:str='.in',
 36        out_ext:str='.out',
 37        folder=None,
 38        files:list=[],
 39        testing:bool=False,
 40    ) -> None:
 41    """Sbatch all the calculations at once.
 42
 43    Calculation names should follow `prefix_ID.ext`,
 44    with `prefix` as the common name across calculations,
 45    followed by the calculation ID, used as JOB_NAME.
 46    The extensions from `in_ext` and `out_ext` ('.in' and '.out' by default)
 47    will be used for the INPUT and OUTPUT filenames of the slurm template.
 48
 49    The slurm template, `template.slurm` by default,
 50    must contain the keywords `JOBNAME`, `INPUT` and `OUTPUT`:
 51    ```
 52    #SBATCH --job-name=JOBNAME
 53    srun --cpu_bind=cores pw.x -inp INPUT > OUTPUT
 54    ```
 55
 56    Runs from the specified `folder`, current working directory if empty.
 57
 58    If more control is required, a custom list of `files` can be specified for sbatching.
 59
 60    If `testing = True` it skips the final sbatching,
 61    just printing the commands on the screen.
 62    """
 63    print('Sbatching all calculations...\n')
 64    key_input = 'INPUT'
 65    key_output = 'OUTPUT'
 66    key_jobname = 'JOBNAME'
 67    slurm_folder = 'slurms'
 68    folder = call.here(folder)
 69    # Get input files and abort if not found
 70    if not files:
 71        inputs_raw = file.get_list(folder=folder, include=prefix, abspath=False)
 72    else:
 73        inputs_raw = files
 74    inputs = []
 75    for filename in inputs_raw:
 76        if filename.endswith(in_ext):
 77            inputs.append(filename)
 78    if len(inputs) == 0:
 79        raise FileNotFoundError(f"Input files were not found! Expected {prefix}ID.{in_ext}")
 80    # Make the folder for the sbatch'ed slurm files
 81    call.bash(f"mkdir {slurm_folder}", folder, True, True)
 82    # Get the template
 83    slurm_file = check_template(template, folder)
 84    if not slurm_file:
 85        print(f'Aborting... Please correct {template}\n')
 86        return None
 87    for filename in inputs:
 88        # Get the file ID
 89        basename: str = os.path.basename(filename)
 90        basename_out: str = basename.replace(in_ext, out_ext)
 91        calc_id = basename.replace(prefix, '')
 92        calc_id = calc_id.replace(in_ext, '')
 93        calc_id = calc_id.replace('_', '')
 94        calc_id = calc_id.replace('-', '')
 95        calc_id = calc_id.replace('.', '')
 96        # Create slurm file for this supercell
 97        slurm_id = prefix + calc_id + '.slurm'
 98        # fixing dictionary with the words to replace in the template
 99        fixing_dict = {
100            key_jobname: calc_id,
101            key_input: basename,
102            key_output: basename_out
103        }
104        edit.from_template(slurm_file, slurm_id, fixing_dict)
105        if testing:
106            call.bash(f"echo {slurm_id}", folder)
107        else:
108            call.bash(f"sbatch {slurm_id}", folder, True, False)
109        call.bash(f"mv {slurm_id} {slurm_folder}", folder, False, True)  # Do not raise error if we can't move the file
110    print(f'Done! Temporary slurm files were moved to ./{slurm_folder}/\n')
111
112
113def squeue(user) -> pd.DataFrame:
114    """Returns a Pandas DataFrame with the jobs from a specific `user`"""
115    result = call.bash(command=f'squeue -u {user}', verbose=False)
116    data = result.stdout
117    lines = data.strip().split('\n')
118    data_rows = [line.split() for line in lines[1:]]
119    df = pd.DataFrame(data_rows, columns=lines[0].split())
120    return df
121
122
123def scancel(
124        user:str,
125        text:str='',
126        status:str='',
127        jobs:list=[],
128        key_jobid:str='JOBID',
129        key_name:str='NAME',
130        key_status:str='ST',
131        testing:bool=False,
132        ) -> None:
133    """Cancel all `user` jobs.
134
135    If a particular `text` string is provided,
136    only the calculations containing said string in the name will be deleted.
137
138    If a particular `status` string is provided,
139    only the calculations with said status will be cancelled.
140
141    If a list of `jobs` is provided, those JOBIDs will be cancelled.
142
143    These filters can all be combined to provide strict control.
144
145    If `testing = True`, it shows the calculations that would be deleted.
146
147    if the slurm squeue titles are different in your cluster,
148    you can specify them with `key_jobid`, `key_status` and `key_name`.
149    """
150    df = squeue(user)
151    if testing:
152        print('aton.api.slurm.scancel(testing=True):')
153        print(f'The following calculations would be killed for the user {user}')
154        print(f'{key_jobid}   {key_status}   {key_name}')
155    jobid_list = df[key_jobid].tolist()
156    name_list = df[key_name].tolist()
157    status_list = df[key_status].tolist()
158    for i, jobid in enumerate(jobid_list):
159        name = name_list[i]
160        st = status_list[i]
161        job = jobid_list[i]
162        # Should we delete this process?
163        bool_1: bool = status == '' and text == ''
164        bool_2: bool = status == st and text == ''
165        bool_3: bool = status == '' and text in name
166        bool_4: bool = status == st and text in name
167        will_delete: bool = bool_1 or bool_2 or bool_3 or bool_4
168        if jobs:
169            jobs = [str(i).strip() for i in jobs]
170            will_delete = will_delete and job in jobs
171        if will_delete:
172            if testing:
173                print(f'{jobid}   {st}   {name}')
174            else:
175                call.bash(f'scancel {jobid}')
176
177
178def scancel_errors(
179        user:str='',
180        errors=['oom_killed', 'OOM Killed', 'Out Of Memory'],
181        folder=None,
182        prefix:str='slurm-',
183        sufix:str='.out',
184        key_jobid:str='JOBID',
185        testing=False,
186        ) -> None:
187    """Cancel all running jobs with any matching `errors` in a given `folder`.
188
189    By default it matches Out Of Memory errors (OOM).
190    The jobs will be detected from the `<prefix>JOBID<sufix>` files, `slurm-JOBID.out` by default.
191
192    If an `user` name is provided, it additionally checks if the jobs are actually scheduled
193    to avoid reporting or killing jobs that have already finished.
194    If the slurm squeue 'JOBID' title is different in your cluster,
195    change it with `key_jobid`.
196
197    If `testing=True`, it will only print the jobs that would be cancelled.
198    """
199    filenames = get_running_here(user=user, folder=folder, prefix=prefix, sufix=sufix, key_jobid=key_jobid, get_files=True)
200    filenames_to_cancel = []
201    jobs_to_cancel = []
202    # Ensure errors is a list
203    if not isinstance(errors, list):
204        errors = [str(errors)]
205    # Find errors in the slurm files
206    for f in filenames:
207        matches = []
208        for e in errors:
209            m = find.lines(filepath=f, key=e)
210            matches.extend(m)
211        if matches:
212            filenames_to_cancel.append(f)
213    # Convert filenames to JOBIDs
214    for f in filenames_to_cancel:
215        f = f.replace(prefix, '')
216        f = f.replace(sufix, '')
217        jobs_to_cancel.append(f)
218    # Report or scancel
219    if testing:
220        print('The following jobs would be cancelled:')
221        for job in jobs_to_cancel:
222            print(f'{job}')
223        return None
224    for job in jobs_to_cancel:
225        call.bash(f'scancel {job}', folder)
226    return None
227
228
229def scancel_here(
230        user:str='',
231        folder=None,
232        prefix:str='slurm-',
233        sufix:str='.out',
234        key_jobid:str='JOBID',
235        testing:bool=False,
236        ) -> None:
237    """Cancel all running jobs in a given `folder`.
238
239    The jobs will be detected from the `<prefix>JOBID<sufix>` files, `slurm-JOBID.out` by default.
240
241    If an `user` name is provided, it additionally checks if the jobs are actually scheduled
242    to avoid reporting or killing jobs that have already finished.
243    If the slurm squeue 'JOBID' title is different in your cluster,
244    change it with `key_jobid`.
245
246    If `testing=True`, it will only print the jobs that would be cancelled.
247    """
248    jobs_to_cancel = get_running_here(user=user, folder=folder, prefix=prefix, sufix=sufix, key_jobid=key_jobid, get_files=False)
249    if testing:
250        print('The following jobs would be cancelled:')
251        for job in jobs_to_cancel:
252            print(f'{job}')
253        return None
254    for job in jobs_to_cancel:
255        call.bash(f'scancel {job}', folder)
256    return None
257
258
259def get_running_here(
260        user:str='',
261        folder=None,
262        prefix:str='slurm-',
263        sufix:str='.out',
264        key_jobid:str='JOBID',
265        get_files=False
266        ) -> list:
267    """Return a list with all running `jobs` in a given `folder`, or current working directory if not specified.
268
269    The jobs will be detected from the `<prefix>JOBID<sufix>` files, `slurm-JOBID.out` by default.
270
271    If an `user` name is provided, it additionally checks if the jobs are actually scheduled
272    to avoid reporting or killing jobs that have already finished.
273    If the slurm squeue 'JOBID' title is different in your cluster,
274    change it with `key_jobid`.
275
276    If `get_files=True`, a list with the filenames will be returned instead of the JOBIDs.
277    """
278    filenames = file.get_list(folder=folder, include=prefix, abspath=False)
279    if not filenames:
280        raise FileNotFoundError(f'To detect the calculations, {prefix}JOBID{sufix} files are needed!\nConfigure the folder, as well as the prefix and sufix if necessary.')
281    jobs = []
282    for filename in filenames:
283        filename = filename.replace(prefix, '')
284        filename = filename.replace(sufix, '')
285        jobs.append(filename)
286    # Check if the jobs are actually scheduled
287    if user:
288        df = squeue(user)
289        jobid_list = df[key_jobid].tolist()
290        jobs = [job for job in jobs if job in jobid_list]
291    # Return the filenames if specified
292    if get_files:
293        if user:
294            filenames = [f for f in filenames if any(job in f for job in jobs)]
295        return filenames
296    return jobs
297
298
299def check_template(
300        template:str='template.slurm',
301        folder=None,
302    ) -> str:
303    """Check the slurm `template` inside `folder`, to be used by `sbatch()`.
304
305    The current working directory is used if `folder` is not provided.
306    If the file does not exist or is invalid, creates a `template_EXAMPLE.slurm` file for reference.
307    """
308    folder = call.here(folder)
309    slurm_example = 'template_EXAMPLE.slurm'
310    new_slurm_file = os.path.join(folder, slurm_example)
311    # Default slurm template
312    content =f"""# Automatic slurm template created with ATON {__version__}
313# https://pablogila.github.io/aton
314
315#!/bin/bash
316#SBATCH --partition=general
317#SBATCH --qos=regular
318#SBATCH --job-name=JOBNAME
319#SBATCH --ntasks=16
320#SBATCH --time=12:00:00
321#SBATCH --mem=64G
322
323module purge
324module load QuantumESPRESSO/7.3-foss-2023a
325
326srun --cpu_bind=cores pw.x -inp INPUT > OUTPUT
327"""
328    # If the slurm template does not exist, create one
329    slurm_file = file.get(folder, template, return_anyway=True)
330    if not slurm_file:
331        with open(new_slurm_file, 'w') as f:
332            f.write(content)
333        print(f'!!! WARNING:  Slurm template missing, an example was generated automatically:\n'
334              f'{slurm_example}\n'
335              f'PLEASE CHECK it, UPDATE it and RENAME it to {template}\n'
336              'before using aton.api.slurm.sbatch()\n')
337        return None
338    # Check that the slurm file contains the INPUT_FILE, OUTPUT_FILE and JOB_NAME keywords
339    key_input = find.lines(slurm_file, 'INPUT')
340    key_output = find.lines(slurm_file, 'OUTPUT')
341    key_jobname = find.lines(slurm_file, 'JOBNAME')
342    missing = []
343    if not key_input:
344        missing.append('INPUT')
345    if not key_output:
346        missing.append('OUTPUT')
347    if not key_jobname:
348        missing.append('JOBNAME')
349    if len(missing) > 0:
350        with open(new_slurm_file, 'w') as f:
351            f.write(content)
352        print('!!! WARNING:  Some keywords were missing from your slurm template,\n'
353              f'PLEASE CHECK the example at {slurm_example}\n'
354              'before using aton.api.slurm.sbatch()\n'
355              f'The following keywords were missing from your {template}:')
356        for key in missing:
357            print(key)
358        print('')
359        return None
360    return slurm_file  # Ready to use!