Function compile_command

pointersect/meta_script/meta_script_utils.py:44–102

Compose the command to execute on bolt. Args: script_filename: the .py filename. Note that it is assumed to be executed from the base folder. num_gpus: number of gpus params: the command-line arguments. They will be composed as -

(
        script_filename: str,
        num_gpus: int,
        params: T.Dict[str, T.Any],
        switch_names: T.List[str] = tuple(),
        prefix: str = '/venv',
        use_xvfb: bool = False,
)

def compile_command(
        script_filename: str,
        num_gpus: int,
        params: T.Dict[str, T.Any],
        switch_names: T.List[str] = tuple(),
        prefix: str = '/venv',
        use_xvfb: bool = False,
) -> str:
    """
    Compose the command to execute on bolt.

    Args:
        script_filename:
            the .py filename.  Note that it is assumed to be executed from the base folder.
        num_gpus:
            number of gpus
        params:
            the command-line arguments. They will be composed as --{key} {val}, except for
            those contained in `switch_names` (which will become --{key} if val = true).
        switch_names:
            a list of names of params to be treated as switches.
        use_xvfb:
            if true, use xvfb to simulate an x window

    Returns:

    """

    xvfb_prefix = 'xvfb-run -a -s "-screen 0 800x600x24" ' if use_xvfb else ''
    if num_gpus > 1:
        command = f'source /miniconda/etc/profile.d/conda.sh; ' \
                  f'conda activate {prefix}; ' \
                  f'CUDA_HOME="{prefix}/pkgs/cuda-toolkit" ' \
                  f'PYTHONPATH="." ' \
                  f'{xvfb_prefix}' \
                  f'torchrun --standalone --nnodes=1 --nproc_per_node={num_gpus} ' \
                  f'{script_filename} '
    else:
        command = f'source /miniconda/etc/profile.d/conda.sh; ' \
                  f'conda activate {prefix}; ' \
                  f'CUDA_HOME="{prefix}/pkgs/cuda-toolkit" ' \
                  f'PYTHONPATH="." ' \
                  f'{xvfb_prefix}' \
                  f'python ' \
                  f'{script_filename} '

    for key in params.keys():
        val = params[key]
        if key in switch_names:
            if val:
                command += '--{} '.format(key)
        else:
            if isinstance(val, (list, tuple)):
                command += '--{} '.format(key)
                for v in val:
                    command += '{} '.format(v)
            else:
                command += '--{} {} '.format(key, val)