cmd-queue 0.1.20__py3-none-any.whl → 0.2.1__py3-none-any.whl

cmd_queue/__init__.py

@@ -306,7 +306,7 @@ Example:
 __mkinit__ = """
 mkinit -m cmd_queue
 """
-__version__ = '0.1.20'
+__version__ = '0.2.1'
 
 
 __submodules__ = {

cmd_queue/base_queue.py

@@ -135,6 +135,11 @@ class Queue(ub.NiceRepr):
             name = kwargs.get('name', None)
             if name is None:
                 name = kwargs['name'] = self.name + '-job-{}'.format(self.num_real_jobs)
+
+            # TODO: make sure name is path safe.
+            if ':' in name:
+                raise ValueError('Name must be path-safe')
+
             if self.all_depends:
                 depends = kwargs.get('depends', None)
                 if depends is None:

cmd_queue/main.py

@@ -94,12 +94,15 @@ class CommonShowRun(CommonConfig):
 
     backend = scfg.Value('tmux', help='the execution backend to use', choices=['tmux', 'slurm', 'serial', 'airflow'])
 
+    gpus = scfg.Value(None, help='a comma separated list of the gpu numbers to spread across. tmux backend only.')
+
     def _build_queue(config):
         import cmd_queue
         import json
         queue = cmd_queue.Queue.create(size=max(1, config['workers']),
                                        backend=config['backend'],
-                                       name=config['qname'])
+                                       name=config['qname'],
+                                       gpus=config['gpus'])
         # Run a new CLI queue
         data = json.loads(config.cli_queue_fpath.read_text())
         print('data = {}'.format(ub.urepr(data, nl=1)))

cmd_queue/serial_queue.py

@@ -9,30 +9,6 @@ from cmd_queue import base_queue
 from cmd_queue.util import util_tags
 
 
-def indent(text, prefix='    '):
-    r"""
-    Indents a block of text
-
-    Args:
-        text (str): text to indent
-        prefix (str, default = '    '): prefix to add to each line
-
-    Returns:
-        str: indented text
-
-        >>> from cmd_queue.serial_queue import *  # NOQA
-        >>> text = ['aaaa', 'bb', 'cc\n   dddd\n    ef\n']
-        >>> text = indent(text)
-        >>> print(text)
-        >>> text = indent(text)
-        >>> print(text)
-    """
-    if isinstance(text, (list, tuple)):
-        return indent('\n'.join(text), prefix)
-    else:
-        return prefix + text.replace('\n', '\n' + prefix)
-
-
 class BashJob(base_queue.Job):
     r"""
     A job meant to run inside of a larger bash file. Analog of SlurmJob
@@ -115,6 +91,21 @@ class BashJob(base_queue.Job):
         self.tags = util_tags.Tags.coerce(tags)
         self.allow_indent = allow_indent
 
+    def _test_bash_syntax_errors(self):
+        """
+        Check for bash syntax errors
+
+        Example:
+            >>> from cmd_queue.serial_queue import *  # NOQA
+            >>> # Demo full boilerplate for a job with dependencies
+            >>> self = BashJob('basd syhi(', name='job1')
+            >>> import pytest
+            >>> with pytest.raises(SyntaxError):
+            >>>     self._test_bash_syntax_errors()
+        """
+        bash_text = self.finalize_text()
+        _check_bash_text_for_syntax_errors(bash_text)
+
     def finalize_text(self, with_status=True, with_gaurds=True,
                       conditionals=None, **kwargs):
         script = []
@@ -575,6 +566,10 @@ class SerialQueue(base_queue.Queue):
         r"""
         Print info about the commands, optionally with rich
 
+        Args:
+            *args: see :func:`cmd_queue.base_queue.Queue.print_commands`.
+            **kwargs: see :func:`cmd_queue.base_queue.Queue.print_commands`.
+
         CommandLine:
             xdoctest -m cmd_queue.serial_queue SerialQueue.print_commands
 
@@ -713,3 +708,40 @@ def _bash_json_dump(json_fmt_parts, fpath):
     printf_part = 'printf ' +  printf_body + ' \\\n    ' + printf_args
     dump_code = printf_part + ' \\\n    ' + redirect_part
     return dump_code
+
+
+def indent(text, prefix='    '):
+    r"""
+    Indents a block of text
+
+    Args:
+        text (str): text to indent
+        prefix (str, default = '    '): prefix to add to each line
+
+    Returns:
+        str: indented text
+
+        >>> from cmd_queue.serial_queue import *  # NOQA
+        >>> text = ['aaaa', 'bb', 'cc\n   dddd\n    ef\n']
+        >>> text = indent(text)
+        >>> print(text)
+        >>> text = indent(text)
+        >>> print(text)
+    """
+    if isinstance(text, (list, tuple)):
+        return indent('\n'.join(text), prefix)
+    else:
+        return prefix + text.replace('\n', '\n' + prefix)
+
+
+def _check_bash_text_for_syntax_errors(bash_text):
+    import tempfile
+    tmpdir = tempfile.TemporaryDirectory()
+    with tmpdir:
+        dpath = ub.Path(tmpdir.name)
+        fpath = dpath / 'job_text.sh'
+        fpath.write_text(bash_text)
+        info = ub.cmd(['bash', '-nv', fpath])
+        if info.returncode != 0:
+            print(info.stderr)
+            raise SyntaxError('bash syntax error')

cmd_queue/slurm_queue.py

@@ -41,23 +41,51 @@ from cmd_queue import base_queue  # NOQA
 from cmd_queue.util import util_tags
 
 
-def _coerce_mem(mem):
+try:
+    from functools import cache  # Python 3.9+ only
+except ImportError:
+    from ubelt import memoize as cache
+
+
+@cache
+def _unit_registery():
+    import sys
+    if sys.version_info[0:2] == (3, 9):
+        # backwards compatability support for numpy 2.0 and pint on cp39
+        try:
+            import numpy as np
+        except ImportError:
+            ...
+        else:
+            if not np.__version__.startswith('1.'):
+                np.cumproduct = np.cumprod
+    import pint
+    reg = pint.UnitRegistry()
+    return reg
+
+
+def _coerce_mem_megabytes(mem):
     """
+    Transform input into an integer representing amount of megabytes.
+
     Args:
         mem (int | str): integer number of megabytes or a parseable string
 
+    Returns:
+        int: number of megabytes
+
     Example:
+        >>> # xdoctest: +REQUIRES(module:pint)
         >>> from cmd_queue.slurm_queue import *  # NOQA
-        >>> print(_coerce_mem(30602))
-        >>> print(_coerce_mem('4GB'))
-        >>> print(_coerce_mem('32GB'))
-        >>> print(_coerce_mem('300000000 bytes'))
+        >>> print(_coerce_mem_megabytes(30602))
+        >>> print(_coerce_mem_megabytes('4GB'))
+        >>> print(_coerce_mem_megabytes('32GB'))
+        >>> print(_coerce_mem_megabytes('300000000 bytes'))
     """
     if isinstance(mem, int):
         assert mem > 0
     elif isinstance(mem, str):
-        import pint
-        reg = pint.UnitRegistry()
+        reg = _unit_registery()
         mem = reg.parse_expression(mem)
         mem = int(mem.to('megabytes').m)
     else:
@@ -190,6 +218,7 @@ class SlurmJob(base_queue.Job):
     Represents a slurm job that hasn't been submitted yet
 
     Example:
+        >>> # xdoctest: +REQUIRES(module:pint)
         >>> from cmd_queue.slurm_queue import *  # NOQA
         >>> self = SlurmJob('python -c print("hello world")', 'hi', cpus=5, gpus=1, mem='10GB')
         >>> command = self._build_sbatch_args()
@@ -235,17 +264,13 @@ class SlurmJob(base_queue.Job):
         return ' \\\n    '.join(args)
 
     def _build_sbatch_args(self, jobname_to_varname=None):
-        # job_name = 'todo'
-        # output_fpath = '$HOME/.cache/slurm/logs/job-%j-%x.out'
-        # command = "python -c 'import sys; sys.exit(1)'"
-        # -c 2 -p priority --gres=gpu:1
         sbatch_args = ['sbatch']
         if self.name:
             sbatch_args.append(f'--job-name="{self.name}"')
         if self.cpus:
             sbatch_args.append(f'--cpus-per-task={self.cpus}')
         if self.mem:
-            mem = _coerce_mem(self.mem)
+            mem = _coerce_mem_megabytes(self.mem)
             sbatch_args.append(f'--mem={mem}')
         if self.gpus and 'gres' not in self._sbatch_kvargs:
             ub.schedule_deprecation(
@@ -277,7 +302,8 @@ class SlurmJob(base_queue.Job):
 
         for key, value in self._sbatch_kvargs.items():
             key = key.replace('_', '-')
-            sbatch_args.append(f'--{key}="{value}"')
+            if value is not None:
+                sbatch_args.append(f'--{key}="{value}"')
 
         for key, flag in self._sbatch_flags.items():
             if flag:
@@ -342,6 +368,19 @@ class SlurmQueue(base_queue.Queue):
     CommandLine:
        xdoctest -m cmd_queue.slurm_queue SlurmQueue
 
+    Example:
+        >>> from cmd_queue.slurm_queue import *  # NOQA
+        >>> self = SlurmQueue()
+        >>> job0 = self.submit('echo "hi from $SLURM_JOBID"')
+        >>> job1 = self.submit('echo "hi from $SLURM_JOBID"', depends=[job0])
+        >>> job2 = self.submit('echo "hi from $SLURM_JOBID"', depends=[job1])
+        >>> job3 = self.submit('echo "hi from $SLURM_JOBID"', depends=[job1, job2])
+        >>> self.write()
+        >>> self.print_commands()
+        >>> # xdoctest: +REQUIRES(--run)
+        >>> if not self.is_available():
+        >>>     self.run()
+
     Example:
         >>> from cmd_queue.slurm_queue import *  # NOQA
         >>> self = SlurmQueue()
@@ -384,6 +423,11 @@ class SlurmQueue(base_queue.Queue):
         self.unused_kwargs = kwargs
         self.queue_id = name + '-' + stamp + '-' + ub.hash_data(uuid.uuid4())[0:8]
         self.dpath = ub.Path.appdir('cmd_queue/slurm') / self.queue_id
+        if 0:
+            # hack for submission on different systems, probably dont want to
+            # do this.
+            self.dpath = self.dpath.shrinkuser(home='$HOME')
+
         self.log_dpath = self.dpath / 'logs'
         self.fpath = self.dpath / (self.queue_id + '.sh')
         self.shell = shell
@@ -395,6 +439,37 @@ class SlurmQueue(base_queue.Queue):
     def __nice__(self):
         return self.queue_id
 
+    @classmethod
+    def _slurm_checks(cls):
+        status = {}
+        info = {}
+        info['squeue_fpath'] = ub.find_exe('squeue')
+        status['has_squeue'] = bool(info['squeue_fpath'])
+        status['slurmd_running'] = False
+        import psutil
+        for p in psutil.process_iter():
+            if p.name() == 'slurmd':
+                status['slurmd_running'] = True
+                info['slurmd_info'] = {
+                    'pid': p.pid,
+                    'name': p.name(),
+                    'status': p.status(),
+                    'create_time': p.create_time(),
+                }
+                break
+        status['squeue_working'] = (ub.cmd('squeue')['ret'] == 0)
+
+        sinfo = ub.cmd('sinfo --json')
+        status['sinfo_working'] = False
+        if sinfo['ret'] == 0:
+            status['sinfo_working'] = True
+            import json
+            sinfo_out = json.loads(sinfo['out'])
+            has_working_nodes = not all(
+                node['state'] == 'down'
+                for node in sinfo_out['nodes'])
+            status['has_working_nodes'] = has_working_nodes
+
     @classmethod
     def is_available(cls):
         """
@@ -407,15 +482,23 @@ class SlurmQueue(base_queue.Queue):
                 squeue_working = (ub.cmd('squeue')['ret'] == 0)
                 if squeue_working:
                     # Check if nodes are available or down
-                    sinfo = ub.cmd('sinfo --json')
-                    if sinfo['ret'] == 0:
-                        import json
-                        sinfo_out = json.loads(sinfo['out'])
-                        has_working_nodes = not all(
-                            node['state'] == 'down'
-                            for node in sinfo_out['nodes'])
-                        if has_working_nodes:
-                            return True
+                    # note: the --json command is not available in
+                    # slurm-wlm 19.05.5, but it is in slurm-wlm 21.08.5
+                    sinfo_version_str = ub.cmd('sinfo --version').stdout.strip().split(' ')[1]
+                    sinfo_major_version = int(sinfo_version_str.split('.')[0])
+                    if sinfo_major_version < 21:
+                        # Dont check in this case
+                        return True
+                    else:
+                        sinfo = ub.cmd('sinfo --json')
+                        if sinfo['ret'] == 0:
+                            import json
+                            sinfo_out = json.loads(sinfo['out'])
+                            has_working_nodes = not all(
+                                node['state'] == 'down'
+                                for node in sinfo_out['nodes'])
+                            if has_working_nodes:
+                                return True
         return False
 
     def submit(self, command, **kwargs):
@@ -520,6 +603,10 @@ class SlurmQueue(base_queue.Queue):
             info = ub.cmd('squeue --format="%i %P %j %u %t %M %D %R"')
             stream = io.StringIO(info['out'])
             df = pd.read_csv(stream, sep=' ')
+            
+            # Only include job names that this queue created
+            job_names = [job.name for job in self.jobs]
+            df = df[df['NAME'].isin(job_names)]
             jobid_history.update(df['JOBID'])
 
             num_running = (df['ST'] == 'R').sum()
@@ -700,4 +787,9 @@ sbatch \
 squeue
 
 
+
+References:
+    https://stackoverflow.com/questions/74164136/slurm-accessing-stdout-stderr-location-of-a-completed-job
+
+
 """

cmd_queue/slurmify.py

@@ -0,0 +1,116 @@
+r"""
+Helper script to wrap a command with sbatch, but using a more srun like syntax.
+
+.. code:: bash
+
+    python -m cmd_queue.slurmify \
+        --jobname="my_job" \
+        --depends=None \
+        --gpus=1 \
+        --mem=16GB \
+        --cpus_per_task=5 \
+        --ntasks=1 \
+        --ntasks-per-node=1 \
+        --partition=community \
+        -- \
+            python -c 'import sys; print("hello world"); sys.exit(0)'
+"""
+#!/usr/bin/env python3
+import scriptconfig as scfg
+import ubelt as ub
+
+
+class SlurmifyCLI(scfg.DataConfig):
+    __command__ = 'slurmify'
+
+    jobname = scfg.Value(None, help='for submit, this is the name of the new job')
+    depends = scfg.Value(None, help='comma separated jobnames to depend on')
+
+    command = scfg.Value(None, type=str, position=1, nargs='*', help=ub.paragraph(
+        '''
+        Specifies the bash command to queue.
+        Care must be taken when specifying this argument.  If specifying as a
+        key/value pair argument, it is important to quote and escape the bash
+        command properly.  A more convinient way to specify this command is as
+        a positional argument. End all of the options to this CLI with `--` and
+        then specify your full command.
+        '''))
+
+    gpus = scfg.Value(None, help='a comma separated list of the gpu numbers to spread across. tmux backend only.')
+    workers = scfg.Value(1, help='number of concurrent queues for the tmux backend.')
+
+    mem = scfg.Value(None, help='')
+    partition = scfg.Value(1, help='slurm partition')
+
+    ntasks = scfg.Value(None, help='')
+    ntasks_per_node = scfg.Value(None, help='')
+    cpus_per_task = scfg.Value(None, help='')
+
+    @classmethod
+    def main(cls, cmdline=1, **kwargs):
+        """
+        Example:
+            >>> # xdoctest: +SKIP
+            >>> from cmd_queue.slurmify import *  # NOQA
+            >>> cmdline = 0
+            >>> kwargs = dict()
+            >>> cls = SlurmifyCLI
+            >>> cls.main(cmdline=cmdline, **kwargs)
+        """
+        import rich
+        from rich.markup import escape
+        config = cls.cli(cmdline=cmdline, data=kwargs, strict=True)
+        rich.print('config = ' + escape(ub.urepr(config, nl=1)))
+
+        # import json
+        # Run a new CLI queue
+        row = {'type': 'command', 'command': config['command']}
+        if config.jobname:
+            row['name'] = config.jobname
+        if config.depends:
+            row['depends'] = config.depends
+
+        import cmd_queue
+        queue = cmd_queue.Queue.create(
+            size=max(1, config['workers']),
+            backend='slurm',
+            name='slurmified',
+            gpus=config['gpus'],
+            mem=config['mem'],
+            partition=config['partition'],
+            ntasks=config['ntasks'],
+            ntasks_per_node=config['ntasks_per_node'],
+        )
+        try:
+            bash_command = row['command']
+            if isinstance(bash_command, list):
+                if len(bash_command) == 1:
+                    # hack
+                    import shlex
+                    if shlex.quote(bash_command[0]) == bash_command[0]:
+                        bash_command = bash_command[0]
+                    else:
+                        bash_command = shlex.quote(bash_command[0])
+                else:
+                    import shlex
+                    bash_command = ' '.join([shlex.quote(str(p)) for p in bash_command])
+            submitkw = ub.udict(row) & {'name', 'depends'}
+            queue.submit(bash_command, log=False, **submitkw)
+        except Exception:
+            print('row = {}'.format(ub.urepr(row, nl=1)))
+            raise
+        queue.print_commands()
+
+        # config.cli_queue_fpath.write_text(json.dumps(row))
+        # 'sbatch --job-name="test_job1" --output="$HOME/.cache/slurm/logs/job-%j-%x.out" --wrap=""
+
+__cli__ = SlurmifyCLI
+
+if __name__ == '__main__':
+    """
+
+    CommandLine:
+        python ~/code/cmd_queue/cmd_queue/slurmify.py
+        python -m cmd_queue.slurmify
+    """
+    __cli__.main()

cmd_queue/tmux_queue.py

@@ -1055,4 +1055,38 @@ if 0:
     tmux kill-session -t my_session_id
 
     tmux new-session -d -s my_session_id -e "MYVAR1" -- "bash"
+
+
+
+    #### to start a tmux session with 4 panes
+    tmux new-session -d -s my_session_id1 "bash"
+    tmux send -t my_session_id1 "tmux split-window -h -t 0" Enter
+    tmux send -t my_session_id1 "tmux split-window -v -t 0" Enter
+    tmux send -t my_session_id1 "tmux split-window -v -t 2" Enter
+
+    # Now send a command to each pane
+    tmux send -t my_session_id1 "tmux select-pane -t 0" Enter
+    tmux send -t my_session_id1 "echo pane0" Enter
+    tmux send -t my_session_id1 "tmux select-pane -t 1" Enter
+    tmux send -t my_session_id1 "echo pane1" Enter
+    tmux send -t my_session_id1 "tmux select-pane -t 2" Enter
+    tmux send -t my_session_id1 "echo pane2" Enter
+    tmux send -t my_session_id1 "tmux select-pane -t 3" Enter
+    tmux send -t my_session_id1 "echo pane3" Enter
+
+    # https://stackoverflow.com/questions/54954177/how-to-write-a-tmux-script-so-that-it-automatically-split-windows-and-opens-a-se
+    # https://tmuxcheatsheet.com/
+    # https://gist.github.com/Starefossen/5955406
+
+    # List the bindings
+    tmux list-keys
+
+    # Can arange the splits in a session via a preset layout
+    # Preset layouts are:
+    # even-horizontal, even-vertical, main-horizontal, main-vertical, or tiled.
+    tmux select-layout -t "${SESSION_NAME}" even-vertical
+
+    # switch to an existing session
+    tmux switch -t "${SESSION_NAME}"
+
     """