Skip to content

Commit

Permalink
Add dedicated Arguments type
Browse files Browse the repository at this point in the history
  • Loading branch information
@jwodder @asmacdo
jwodder authored and asmacdo committed Jun 10, 2024
1 parent d52e7ee commit 48d606a
Show file tree
Hide file tree
Showing 2 changed files with 107 additions and 84 deletions.
172 changes: 98 additions & 74 deletions src/duct.py
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
#!/usr/bin/env python3
from __future__ import annotations
import argparse
from collections import defaultdict
from dataclasses import dataclass
from datetime import datetime
import json
import os
Expand Down Expand Up @@ -191,6 +193,101 @@ def __repr__(self):
)


@dataclass
class Arguments:
command: str
command_args: list[str]
output_prefix: str
sample_interval: float
report_interval: float
capture_outputs: str
outputs: str
record_types: str

@classmethod
def from_argv(cls) -> Arguments:
parser = argparse.ArgumentParser(
description="Gathers metrics on a command and all its child processes.",
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
)
parser.add_argument(
"command",
metavar="command [command_args ...]",
help="The command to execute, along with its arguments.",
)
parser.add_argument(
"--version", action="version", version=f"%(prog)s {__version__}"
)
parser.add_argument(
"command_args", nargs=argparse.REMAINDER, help="Arguments for the command."
)
parser.add_argument(
"-p",
"--output-prefix",
type=str,
default=os.getenv(
"DUCT_OUTPUT_PREFIX", ".duct/logs/{datetime_filesafe}-{pid}_"
),
help="File string format to be used as a prefix for the files -- the captured "
"stdout and stderr and the resource usage logs. The understood variables are "
"{datetime}, {datetime_filesafe}, and {pid}. "
"Leading directories will be created if they do not exist. "
"You can also provide value via DUCT_OUTPUT_PREFIX env variable. ",
)
parser.add_argument(
"--sample-interval",
"--s-i",
type=float,
default=float(os.getenv("DUCT_SAMPLE_INTERVAL", "1.0")),
help="Interval in seconds between status checks of the running process. "
"Sample interval should be larger than the runtime of the process or `duct` may "
"underreport the number of processes started.",
)
parser.add_argument(
"--report-interval",
"--r-i",
type=float,
default=float(os.getenv("DUCT_REPORT_INTERVAL", "60.0")),
help="Interval in seconds at which to report aggregated data.",
)
parser.add_argument(
"-c",
"--capture-outputs",
type=str,
default=os.getenv("DUCT_CAPTURE_OUTPUTS", "all"),
choices=["all", "none", "stdout", "stderr"],
help="Record stdout, stderr, all, or none to log files. "
"You can also provide value via DUCT_CAPTURE_OUTPUTS env variable.",
)
parser.add_argument(
"-o",
"--outputs",
type=str,
default="all",
choices=["all", "none", "stdout", "stderr"],
help="Print stdout, stderr, all, or none to stdout/stderr respectively.",
)
parser.add_argument(
"-t",
"--record-types",
type=str,
default="all",
choices=["all", "system-summary", "processes-samples"],
help="Record system-summary, processes-samples, or all",
)
args = parser.parse_args()
return cls(
command=args.command,
command_args=args.command_args,
output_prefix=args.output_prefix,
sample_interval=args.sample_interval,
report_interval=args.report_interval,
capture_outputs=args.capture_outputs,
outputs=args.outputs,
record_types=args.record_types,
)


def monitor_process(report, process, report_interval, sample_interval, stop_event):
while not stop_event.wait(timeout=sample_interval):
while True:
Expand All @@ -208,79 +305,6 @@ def monitor_process(report, process, report_interval, sample_interval, stop_even
report.number += 1


def create_and_parse_args():
parser = argparse.ArgumentParser(
description="Gathers metrics on a command and all its child processes.",
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
)
parser.add_argument(
"command",
metavar="command [command_args ...]",
help="The command to execute, along with its arguments.",
)
parser.add_argument(
"--version", action="version", version=f"%(prog)s {__version__}"
)
parser.add_argument(
"command_args", nargs=argparse.REMAINDER, help="Arguments for the command."
)
parser.add_argument(
"-p",
"--output-prefix",
type=str,
default=os.getenv(
"DUCT_OUTPUT_PREFIX", ".duct/logs/{datetime_filesafe}-{pid}_"
),
help="File string format to be used as a prefix for the files -- the captured "
"stdout and stderr and the resource usage logs. The understood variables are "
"{datetime}, {datetime_filesafe}, and {pid}. "
"Leading directories will be created if they do not exist. "
"You can also provide value via DUCT_OUTPUT_PREFIX env variable. ",
)
parser.add_argument(
"--sample-interval",
"--s-i",
type=float,
default=float(os.getenv("DUCT_SAMPLE_INTERVAL", "1.0")),
help="Interval in seconds between status checks of the running process. "
"Sample interval should be larger than the runtime of the process or `duct` may "
"underreport the number of processes started.",
)
parser.add_argument(
"--report-interval",
"--r-i",
type=float,
default=float(os.getenv("DUCT_REPORT_INTERVAL", "60.0")),
help="Interval in seconds at which to report aggregated data.",
)
parser.add_argument(
"-c",
"--capture-outputs",
type=str,
default=os.getenv("DUCT_CAPTURE_OUTPUTS", "all"),
choices=["all", "none", "stdout", "stderr"],
help="Record stdout, stderr, all, or none to log files. "
"You can also provide value via DUCT_CAPTURE_OUTPUTS env variable.",
)
parser.add_argument(
"-o",
"--outputs",
type=str,
default="all",
choices=["all", "none", "stdout", "stderr"],
help="Print stdout, stderr, all, or none to stdout/stderr respectively.",
)
parser.add_argument(
"-t",
"--record-types",
type=str,
default="all",
choices=["all", "system-summary", "processes-samples"],
help="Record system-summary, processes-samples, or all",
)
return parser.parse_args()


class TailPipe:
"""TailPipe simultaneously streams to an output stream (stdout or stderr) and a specified file."""

Expand Down Expand Up @@ -375,7 +399,7 @@ def ensure_directories(path: str) -> None:


def main():
args = create_and_parse_args()
args = Arguments.from_argv()
execute(args)


Expand Down
19 changes: 9 additions & 10 deletions test/test_execution.py
View file
Original file line number Diff line number Diff line change
@@ -1,11 +1,10 @@
import argparse
import os
from pathlib import Path
import shutil
from unittest import mock
import pytest
from utils import assert_files
from duct import execute
from duct import Arguments, execute

TEST_SCRIPT = str(Path(__file__).with_name("data") / "test_script.py")

Expand All @@ -17,7 +16,7 @@ def temp_output_dir(tmp_path):


def test_sanity_green(temp_output_dir):
args = argparse.Namespace(
args = Arguments(
command="echo",
command_args=["hello", "world"],
output_prefix=temp_output_dir,
Expand All @@ -34,7 +33,7 @@ def test_sanity_green(temp_output_dir):


def test_sanity_red(temp_output_dir):
args = argparse.Namespace(
args = Arguments(
command="false",
command_args=[],
output_prefix=temp_output_dir,
Expand All @@ -57,7 +56,7 @@ def test_sanity_red(temp_output_dir):


def test_outputs_full(temp_output_dir):
args = argparse.Namespace(
args = Arguments(
command=TEST_SCRIPT,
command_args=["--duration", "1"],
output_prefix=temp_output_dir,
Expand All @@ -73,7 +72,7 @@ def test_outputs_full(temp_output_dir):


def test_outputs_passthrough(temp_output_dir):
args = argparse.Namespace(
args = Arguments(
command=TEST_SCRIPT,
command_args=["--duration", "1"],
output_prefix=temp_output_dir,
Expand All @@ -91,7 +90,7 @@ def test_outputs_passthrough(temp_output_dir):


def test_outputs_capture(temp_output_dir):
args = argparse.Namespace(
args = Arguments(
command=TEST_SCRIPT,
command_args=["--duration", "1"],
output_prefix=temp_output_dir,
Expand All @@ -109,7 +108,7 @@ def test_outputs_capture(temp_output_dir):


def test_outputs_none(temp_output_dir):
args = argparse.Namespace(
args = Arguments(
command=TEST_SCRIPT,
command_args=["--duration", "1"],
output_prefix=temp_output_dir,
Expand All @@ -130,7 +129,7 @@ def test_outputs_none(temp_output_dir):


def test_exit_before_first_sample(temp_output_dir):
args = argparse.Namespace(
args = Arguments(
command="ls",
command_args=[],
output_prefix=temp_output_dir,
Expand All @@ -148,7 +147,7 @@ def test_exit_before_first_sample(temp_output_dir):


def test_run_less_than_report_interval(temp_output_dir):
args = argparse.Namespace(
args = Arguments(
command="sleep",
command_args=["0.01"],
output_prefix=temp_output_dir,
Expand Down

0 comments on commit 48d606a

Please sign in to comment.