From d63dc7128c310e9cb3d32521360caa3eac2380ef Mon Sep 17 00:00:00 2001 From: evilsocket Date: Mon, 10 Feb 2025 19:30:17 +0100 Subject: [PATCH 1/6] new: implemented run context --- dreadnode_cli/agent/cli.py | 48 +++++++++++++++++++++++++++++++++++++- dreadnode_cli/api.py | 8 +++++++ poetry.lock | 13 ++++++++++- pyproject.toml | 1 + 4 files changed, 68 insertions(+), 2 deletions(-) diff --git a/dreadnode_cli/agent/cli.py b/dreadnode_cli/agent/cli.py index b42c3a4..637a2b6 100644 --- a/dreadnode_cli/agent/cli.py +++ b/dreadnode_cli/agent/cli.py @@ -4,6 +4,7 @@ import time import typing as t +import toml import typer from rich import box, print from rich.live import Live @@ -318,6 +319,31 @@ def push( print(":tada: Agent pushed. use [bold]dreadnode agent deploy[/] to start a new run.") +def prepare_run_context( + env_vars: list[str] | None, parameters: list[str] | None, command: str | None +) -> Client.StrikeRunContext | None: + if not env_vars and not parameters and not command: + return None + + context = Client.StrikeRunContext() + + if env_vars: + context.environment = {env_var.split("=")[0]: env_var.split("=")[1] for env_var in env_vars} + + if parameters: + context.parameters = {} + for param in parameters: + if param.startswith("@"): + context.parameters.update(toml.load(open(param[1:]))) + else: + context.parameters.update(toml.loads(param)) + + if command: + context.command = command + + return context + + @cli.command(help="Start a new run using the latest active agent version") @pretty_cli def deploy( @@ -328,6 +354,22 @@ def deploy( pathlib.Path, typer.Option("--dir", "-d", help="The agent directory", file_okay=False, resolve_path=True), ] = pathlib.Path("."), + env_vars: t.Annotated[ + list[str] | None, + typer.Option("--env-var", "-e", help="Environment vars to override for this run (key=value)"), + ] = None, + parameters: t.Annotated[ + list[str] | None, + typer.Option( + "--param", + "-p", + help="Define custom parameters for this run (key = value in toml syntax or @filename.toml for multiple values)", + ), + ] = None, + command: t.Annotated[ + str | None, + typer.Option("--command", "-c", help="Override the container command for this run."), + ] = None, strike: t.Annotated[str | None, typer.Option("--strike", "-s", help="The strike to use for this run")] = None, watch: t.Annotated[bool, typer.Option("--watch", "-w", help="Watch the run status")] = True, ) -> None: @@ -346,6 +388,8 @@ def deploy( if strike is None: raise Exception("No strike specified, use -s/--strike or set the strike in the agent config") + context = prepare_run_context(env_vars, parameters, command) + user_models = UserModels.read() user_model: Client.UserModel | None = None @@ -376,7 +420,9 @@ def deploy( f"Model '{model}' is not user-defined nor is it available in strike '{strike_response.name}'" ) - run = client.start_strike_run(agent.latest_version.id, strike=strike, model=model, user_model=user_model) + run = client.start_strike_run( + agent.latest_version.id, strike=strike, model=model, user_model=user_model, context=context + ) agent_config.add_run(run.id).write(directory) formatted = format_run(run, server_url=server_config.url) diff --git a/dreadnode_cli/api.py b/dreadnode_cli/api.py index abf9335..1027872 100644 --- a/dreadnode_cli/api.py +++ b/dreadnode_cli/api.py @@ -351,6 +351,11 @@ class StrikeRunZone(_StrikeRunZone): outputs: list["Client.StrikeRunOutput"] inferences: list[dict[str, t.Any]] + class StrikeRunContext(BaseModel): + environment: dict[str, str] | None = None + parameters: dict[str, t.Any] | None = None + command: str | None = None + class _StrikeRun(BaseModel): id: UUID strike_id: UUID @@ -364,6 +369,7 @@ class _StrikeRun(BaseModel): agent_name: str | None = None agent_revision: int agent_version: "Client.StrikeAgentVersion" + context: "Client.StrikeRunContext" | None = None status: "Client.StrikeRunStatus" start: datetime | None end: datetime | None @@ -440,6 +446,7 @@ def start_strike_run( *, model: str | None = None, user_model: UserModel | None = None, + context: StrikeRunContext | None = None, strike: UUID | str | None = None, ) -> StrikeRunResponse: response = self.request( @@ -450,6 +457,7 @@ def start_strike_run( "model": model, "user_model": user_model.model_dump(mode="json") if user_model else None, "strike": str(strike) if strike else None, + "context": context.model_dump(mode="json") if context else None, }, ) return self.StrikeRunResponse(**response.json()) diff --git a/poetry.lock b/poetry.lock index f1c3292..70732ca 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1061,6 +1061,17 @@ files = [ {file = "sniffio-1.3.1.tar.gz", hash = "sha256:f4324edc670a0f49750a81b895f35c3adb843cca46f0530f79fc1babb23789dc"}, ] +[[package]] +name = "toml" +version = "0.10.2" +description = "Python Library for Tom's Obvious, Minimal Language" +optional = false +python-versions = ">=2.6, !=3.0.*, !=3.1.*, !=3.2.*" +files = [ + {file = "toml-0.10.2-py2.py3-none-any.whl", hash = "sha256:806143ae5bfb6a3c6e736a764057db0e6a0e05e338b5630894a5f779cabb4f9b"}, + {file = "toml-0.10.2.tar.gz", hash = "sha256:b3bda1d108d5dd99f4a20d24d9c348e91c4db7ab1b749200bded2f839ccbe68f"}, +] + [[package]] name = "tomli" version = "2.1.0" @@ -1154,4 +1165,4 @@ test = ["covdefaults (>=2.3)", "coverage (>=7.2.7)", "coverage-enable-subprocess [metadata] lock-version = "2.0" python-versions = "^3.10" -content-hash = "513aed5633093d6a74ed27b7117995eadee3e1c511e7835f0362f4de5246494e" +content-hash = "8f1e6d10643b00b626e29844fcea31b0487b4734f3e078d59f7ab55680df2550" diff --git a/pyproject.toml b/pyproject.toml index 97c4138..196556a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -20,6 +20,7 @@ httpx = "^0.27.2" ruamel-yaml = "^0.18.6" docker = "^7.1.0" pydantic-yaml = "^1.4.0" +toml = "^0.10.2" [tool.pytest.ini_options] asyncio_mode = "auto" From 6dd03ae0f0bc78045663602155319c629d92482c Mon Sep 17 00:00:00 2001 From: evilsocket Date: Mon, 10 Feb 2025 19:56:37 +0100 Subject: [PATCH 2/6] fix: fixed typing bug --- dreadnode_cli/api.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dreadnode_cli/api.py b/dreadnode_cli/api.py index 1027872..6be6d6b 100644 --- a/dreadnode_cli/api.py +++ b/dreadnode_cli/api.py @@ -369,7 +369,7 @@ class _StrikeRun(BaseModel): agent_name: str | None = None agent_revision: int agent_version: "Client.StrikeAgentVersion" - context: "Client.StrikeRunContext" | None = None + context: "Client.StrikeRunContext | None" = None status: "Client.StrikeRunStatus" start: datetime | None end: datetime | None From 684d93c25f5519012d1551b505692016f6bd9c1f Mon Sep 17 00:00:00 2001 From: evilsocket Date: Tue, 11 Feb 2025 11:00:40 +0100 Subject: [PATCH 3/6] new: showing run context if set (ENG-989) --- dreadnode_cli/agent/format.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/dreadnode_cli/agent/format.py b/dreadnode_cli/agent/format.py index 9ab9524..1f53b43 100644 --- a/dreadnode_cli/agent/format.py +++ b/dreadnode_cli/agent/format.py @@ -291,6 +291,15 @@ def format_run( table.add_row("start", format_time(run.start)) table.add_row("end", format_time(run.end)) + if run.context and (run.context.environment or run.context.parameters or run.context.command): + table.add_row("", "") + if run.context.environment: + table.add_row("environment", Pretty(run.context.environment)) + if run.context.parameters: + table.add_row("parameters", Pretty(run.context.parameters)) + if run.context.command: + table.add_row("command", run.context.command) + components: list[RenderableType] = [ table, format_zones_verbose(run.zones, include_logs=include_logs) if verbose else format_zones_summary(run.zones), From c79ee3ac60fa56efc563aecc53ff58fbd7d31678 Mon Sep 17 00:00:00 2001 From: evilsocket Date: Tue, 11 Feb 2025 11:08:59 +0100 Subject: [PATCH 4/6] improved run context format --- dreadnode_cli/agent/format.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/dreadnode_cli/agent/format.py b/dreadnode_cli/agent/format.py index 1f53b43..54ac073 100644 --- a/dreadnode_cli/agent/format.py +++ b/dreadnode_cli/agent/format.py @@ -294,11 +294,15 @@ def format_run( if run.context and (run.context.environment or run.context.parameters or run.context.command): table.add_row("", "") if run.context.environment: - table.add_row("environment", Pretty(run.context.environment)) + table.add_row( + "environment", " ".join(f"[magenta]{k}[/]=[yellow]{v}[/]" for k, v in run.context.environment.items()) + ) if run.context.parameters: - table.add_row("parameters", Pretty(run.context.parameters)) + table.add_row( + "parameters", " ".join(f"[magenta]{k}[/]=[yellow]{v}[/]" for k, v in run.context.parameters.items()) + ) if run.context.command: - table.add_row("command", run.context.command) + table.add_row("command", f"[bold][red]{run.context.command}[/red][/bold]") components: list[RenderableType] = [ table, From fa43c19bc7aae0c84d0271a57f97d28fbdc0b5ab Mon Sep 17 00:00:00 2001 From: evilsocket Date: Tue, 11 Feb 2025 11:11:55 +0100 Subject: [PATCH 5/6] docs: documented run context data --- CLI.md | 3 +++ README.md | 12 ++++++++++++ 2 files changed, 15 insertions(+) diff --git a/CLI.md b/CLI.md index 508037c..b9dfb60 100644 --- a/CLI.md +++ b/CLI.md @@ -85,6 +85,9 @@ $ dreadnode agent deploy [OPTIONS] * `-m, --model TEXT`: The inference model to use for this run * `-d, --dir DIRECTORY`: The agent directory [default: .] +* `-e, --env-var TEXT`: Environment vars to override for this run (key=value) +* `-p, --param TEXT`: Define custom parameters for this run (key = value in toml syntax or @filename.toml for multiple values) +* `-c, --command TEXT`: Override the container command for this run. * `-s, --strike TEXT`: The strike to use for this run * `-w, --watch`: Watch the run status [default: True] * `--help`: Show this message and exit. diff --git a/README.md b/README.md index e617bc1..b818aa0 100644 --- a/README.md +++ b/README.md @@ -165,6 +165,18 @@ dreadnode agent push # start a new run using the latest agent version. dreadnode agent deploy +# start a new run using the latest agent version with custom environment variables +dreadnode agent deploy --env-var TEST_ENV=test --env-var ANOTHER_ENV=another_value + +# start a new run using the latest agent version with custom parameters (using toml syntax) +dreadnode agent deploy --param "foo = 'bar'" --param "baz = 123.0" + +# start a new run using the latest agent version with custom parameters from a toml file +dreadnode agent deploy --param @parameters.toml + +# start a new run using the latest agent version and override the container command +dreadnode agent deploy --command "echo 'Hello, world!'" + # show the latest run of the currently active agent dreadnode agent latest From 905357cbce87b19702814b532791bc7d8e4e2084 Mon Sep 17 00:00:00 2001 From: evilsocket Date: Tue, 11 Feb 2025 11:17:01 +0100 Subject: [PATCH 6/6] fix: lint driven fix --- poetry.lock | 13 ++++++++++++- pyproject.toml | 1 + 2 files changed, 13 insertions(+), 1 deletion(-) diff --git a/poetry.lock b/poetry.lock index 70732ca..e7c53fc 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1114,6 +1114,17 @@ files = [ [package.dependencies] urllib3 = ">=2" +[[package]] +name = "types-toml" +version = "0.10.8.20240310" +description = "Typing stubs for toml" +optional = false +python-versions = ">=3.8" +files = [ + {file = "types-toml-0.10.8.20240310.tar.gz", hash = "sha256:3d41501302972436a6b8b239c850b26689657e25281b48ff0ec06345b8830331"}, + {file = "types_toml-0.10.8.20240310-py3-none-any.whl", hash = "sha256:627b47775d25fa29977d9c70dc0cbab3f314f32c8d8d0c012f2ef5de7aaec05d"}, +] + [[package]] name = "typing-extensions" version = "4.12.2" @@ -1165,4 +1176,4 @@ test = ["covdefaults (>=2.3)", "coverage (>=7.2.7)", "coverage-enable-subprocess [metadata] lock-version = "2.0" python-versions = "^3.10" -content-hash = "8f1e6d10643b00b626e29844fcea31b0487b4734f3e078d59f7ab55680df2550" +content-hash = "4dd62ed90a1469dd758b514a8e08d1397999ddb6bdd7eb72c15828a6ea0eca79" diff --git a/pyproject.toml b/pyproject.toml index 196556a..3a0511c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -21,6 +21,7 @@ ruamel-yaml = "^0.18.6" docker = "^7.1.0" pydantic-yaml = "^1.4.0" toml = "^0.10.2" +types-toml = "^0.10.8.20240310" [tool.pytest.ini_options] asyncio_mode = "auto"