|
| 1 | +import argparse |
| 2 | +import sys |
| 3 | +import os |
| 4 | +from contextlib import contextmanager |
| 5 | + |
| 6 | +import luigi |
| 7 | +import luigi.cmdline |
| 8 | + |
| 9 | +from rnaseq_pipeline.tasks import SubmitExperimentToGemma, SubmitExperimentsFromGoogleSpreadsheetToGemma, \ |
| 10 | + SubmitExperimentBatchInfoToGemma |
| 11 | + |
| 12 | +@contextmanager |
| 13 | +def umask(umask): |
| 14 | + print(f'Setting umask to 0x{umask:03o}') |
| 15 | + prev_umask = os.umask(umask) |
| 16 | + try: |
| 17 | + yield None |
| 18 | + finally: |
| 19 | + print(f'Restoring umask to 0x{prev_umask:03o}') |
| 20 | + os.umask(prev_umask) |
| 21 | + |
| 22 | +def parse_octal(s): |
| 23 | + return int(s, 8) |
| 24 | + |
| 25 | +def run_luigi_task(task, args): |
| 26 | + with umask(args.umask): |
| 27 | + luigi.build([task], workers=args.workers, detailed_summary=True, local_scheduler=args.local_scheduler) |
| 28 | + |
| 29 | +def run(args): |
| 30 | + with umask(0o002): |
| 31 | + luigi.run(args) |
| 32 | + |
| 33 | +def submit_experiment(argv): |
| 34 | + parser = argparse.ArgumentParser() |
| 35 | + parser.add_argument('--experiment-id', required=True, help='Experiment ID to submit to Gemma') |
| 36 | + parser.add_argument('--rerun', action='store_true', default=False, help='Rerun the experiment') |
| 37 | + parser.add_argument('--priority', type=int, default=100) |
| 38 | + parser.add_argument('--umask', type=parse_octal, default='002', |
| 39 | + help='Set a umask (defaults to 002 to make created files group-writable)') |
| 40 | + parser.add_argument('--workers', type=int, default=30, help='Number of workers to use (defaults to 30)') |
| 41 | + parser.add_argument('--local-scheduler', action='store_true', default=False) |
| 42 | + args = parser.parse_args(argv) |
| 43 | + run_luigi_task(SubmitExperimentToGemma(experiment_id=args.experiment_id, rerun=args.rerun, priority=args.priority), |
| 44 | + args) |
| 45 | + |
| 46 | +def submit_experiment_batch_info(argv): |
| 47 | + parser = argparse.ArgumentParser() |
| 48 | + parser.add_argument('--experiment-id', required=True, help='Experiment ID to submit to Gemma') |
| 49 | + parser.add_argument('--ignored-samples', nargs='+', default=[]) |
| 50 | + parser.add_argument('--rerun', action='store_true', default=False, help='Rerun the experiment') |
| 51 | + parser.add_argument('--umask', type=parse_octal, default='002', |
| 52 | + help='Set a umask (defaults to 002 to make created files group-writable)') |
| 53 | + parser.add_argument('--workers', type=int, default=30, help='Number of workers to use (defaults to 30)') |
| 54 | + parser.add_argument('--local-scheduler', action='store_true', default=False) |
| 55 | + args = parser.parse_args(argv) |
| 56 | + print(args.ignored_samples) |
| 57 | + run_luigi_task( |
| 58 | + SubmitExperimentBatchInfoToGemma(experiment_id=args.experiment_id, ignored_samples=args.ignored_samples, |
| 59 | + rerun=args.rerun), args) |
| 60 | + |
| 61 | +def submit_experiments_from_gsheet(argv): |
| 62 | + parser = argparse.ArgumentParser() |
| 63 | + parser.add_argument('--spreadsheet-id', required=True, help='Spreadsheet ID') |
| 64 | + parser.add_argument('--sheet-name', required=True, help='Sheet name') |
| 65 | + parser.add_argument('--umask', type=parse_octal, default='002', |
| 66 | + help='Set a umask (defaults to 002 to make created files group-writable)') |
| 67 | + parser.add_argument('--workers', type=int, default=200, help='Number of workers to use (defaults to 200)') |
| 68 | + parser.add_argument('--ignore-priority', action='store_true', help='Ignore the priority column in the spreadsheet') |
| 69 | + parser.add_argument('--local-scheduler', action='store_true', default=False) |
| 70 | + args = parser.parse_args(argv) |
| 71 | + run_luigi_task(SubmitExperimentsFromGoogleSpreadsheetToGemma(args.spreadsheet_id, args.sheet_name, |
| 72 | + ignore_priority=args.ignore_priority), args) |
| 73 | + |
| 74 | +def main(): |
| 75 | + if len(sys.argv) < 2: |
| 76 | + print('Usage: rnaseq-pipeline-cli <command>') |
| 77 | + return 1 |
| 78 | + command = sys.argv[1] |
| 79 | + if command == 'run': |
| 80 | + return run(sys.argv[2:]) |
| 81 | + elif command == 'submit-experiment': |
| 82 | + return submit_experiment(sys.argv[2:]) |
| 83 | + elif command == 'submit-experiment-batch-info': |
| 84 | + return submit_experiment_batch_info(sys.argv[2:]) |
| 85 | + elif command == 'submit-experiments-from-gsheet': |
| 86 | + return submit_experiments_from_gsheet(sys.argv[2:]) |
| 87 | + else: |
| 88 | + print( |
| 89 | + f'Unknown command {command}. Possible values are: submit-experiment, submit-experiment-batch-info, submit-experiments-from-gsheet.') |
| 90 | + return 1 |
0 commit comments