Skip to content

Commit

Permalink
Add logic to run_deepvariant to control postprocess cpus
Browse files Browse the repository at this point in the history
PiperOrigin-RevId: 575883359
  • Loading branch information
kishwarshafin authored and copybara-github committed Oct 23, 2023
1 parent 69e565d commit 6b16701
Showing 1 changed file with 12 additions and 0 deletions.
12 changes: 12 additions & 0 deletions scripts/run_deepvariant.py
Original file line number Diff line number Diff line change
Expand Up @@ -210,6 +210,14 @@
' has to be flag_name=true or flag_name=false.'
),
)
# Optional flag for postprocess variants
_POSTPROCESS_CPUS = flags.DEFINE_integer(
'postprocess_cpus',
None,
'Optional. Number of cpus to use during'
' postprocess_variants. Set to 0 to disable multiprocessing. Default is'
' None which sets to num_shards.',
)
_POSTPROCESS_VARIANTS_EXTRA_ARGS = flags.DEFINE_string(
'postprocess_variants_extra_args',
None,
Expand Down Expand Up @@ -436,10 +444,14 @@ def postprocess_variants_command(
ref: str, infile: str, outfile: str, extra_args: str, **kwargs
) -> Tuple[str, Optional[str]]:
"""Returns a postprocess_variants (command, logfile) for subprocess."""
cpus = _POSTPROCESS_CPUS.value
if not cpus:
cpus = _NUM_SHARDS.value
command = ['time', '/opt/deepvariant/bin/postprocess_variants']
command.extend(['--ref', '"{}"'.format(ref)])
command.extend(['--infile', '"{}"'.format(infile)])
command.extend(['--outfile', '"{}"'.format(outfile)])
command.extend(['--cpus', '"{}"'.format(cpus)])
# Extend the command with all items in kwargs and extra_args.
kwargs = _update_kwargs_with_warning(kwargs, _extra_args_to_dict(extra_args))
command = _extend_command_by_args_dict(command, kwargs)
Expand Down

0 comments on commit 6b16701

Please sign in to comment.