From b5bfe29d9c2904e345d4753e41bb5266a5a6b55e Mon Sep 17 00:00:00 2001 From: Josh Meek Date: Wed, 13 Nov 2019 09:53:01 -0500 Subject: [PATCH 1/8] Add option for Fargate Agent kwargs to start CLI command --- src/prefect/cli/agent.py | 33 ++++++++++++++++++++++++++++++--- 1 file changed, 30 insertions(+), 3 deletions(-) diff --git a/src/prefect/cli/agent.py b/src/prefect/cli/agent.py index e0b9f793b1f1..5f09138cbdcb 100644 --- a/src/prefect/cli/agent.py +++ b/src/prefect/cli/agent.py @@ -1,5 +1,6 @@ import click +import prefect from prefect import config, context from prefect.utilities.configuration import set_temporary_config from prefect.utilities.serialization import from_qualified_name @@ -42,7 +43,10 @@ def agent(): pass -@agent.command(hidden=True) +@agent.command( + hidden=True, + context_settings=dict(ignore_unknown_options=True, allow_extra_args=True,), +) @click.argument("agent-option", default="local") @click.option( "--token", "-t", required=False, help="A Prefect Cloud API token.", hidden=True @@ -67,7 +71,8 @@ def agent(): ) @click.option("--no-pull", is_flag=True, help="Pull images flag.", hidden=True) @click.option("--base-url", "-b", help="Docker daemon base URL.", hidden=True) -def start(agent_option, token, name, verbose, label, no_pull, base_url): +@click.pass_context +def start(ctx, agent_option, token, name, verbose, label, no_pull, base_url): """ Start an agent. @@ -90,7 +95,19 @@ def start(agent_option, token, name, verbose, label, no_pull, base_url): --base-url, -b TEXT A Docker daemon host URL for a LocalAgent --no-pull Pull images for a LocalAgent Defaults to pulling if not provided + + \b + Fargate Agent Options: + Any of the configuration options outlined in the docs can be provided here + https://docs.prefect.io/cloud/agent/fargate.html#configuration """ + + # Split context + kwargs = dict() + for item in ctx.args: + item = item.replace("--", "") + kwargs.update([item.split("=")]) + tmp_config = {"cloud.agent.auth_token": token or config.cloud.agent.auth_token} if verbose: tmp_config["cloud.agent.level"] = "DEBUG" @@ -102,7 +119,17 @@ def start(agent_option, token, name, verbose, label, no_pull, base_url): click.secho("{} is not a valid agent".format(agent_option), fg="red") return - with context(no_pull=no_pull, base_url=base_url): + _agent = from_qualified_name(retrieved_agent) + + if agent_option == "local": + from_qualified_name(retrieved_agent)( + name=name, labels=list(label), base_url=base_url, no_pull=no_pull, + ).start() + elif agent_option == "fargate": + from_qualified_name(retrieved_agent)( + name=name, labels=list(label), **kwargs + ).start() + else: from_qualified_name(retrieved_agent)(name=name, labels=list(label)).start() From 13fc49d3728449f083a549144bc16f230f16f758 Mon Sep 17 00:00:00 2001 From: Josh Meek Date: Wed, 13 Nov 2019 10:01:19 -0500 Subject: [PATCH 2/8] Update Agent CLI tests --- tests/cli/test_agent.py | 84 ++++++++++++++++++++++++++++++++++++++++- 1 file changed, 83 insertions(+), 1 deletion(-) diff --git a/tests/cli/test_agent.py b/tests/cli/test_agent.py index ddeea3436f79..4ee1bd54e46c 100644 --- a/tests/cli/test_agent.py +++ b/tests/cli/test_agent.py @@ -55,6 +55,88 @@ def test_agent_start_verbose(monkeypatch, runner_token): assert result.exit_code == 0 +def test_agent_start_local(monkeypatch, runner_token): + start = MagicMock() + monkeypatch.setattr("prefect.agent.local.LocalAgent.start", start) + + docker_client = MagicMock() + monkeypatch.setattr("prefect.agent.local.agent.docker.APIClient", docker_client) + + runner = CliRunner() + result = runner.invoke(agent, ["start", "local"]) + assert result.exit_code == 0 + + +def test_agent_start_kubernetes(monkeypatch, runner_token): + start = MagicMock() + monkeypatch.setattr("prefect.agent.kubernetes.KubernetesAgent.start", start) + + k8s_config = MagicMock() + monkeypatch.setattr("kubernetes.config", k8s_config) + + runner = CliRunner() + result = runner.invoke(agent, ["start", "kubernetes"]) + assert result.exit_code == 0 + + +def test_agent_start_kubernetes_kwargs_ignored(monkeypatch, runner_token): + start = MagicMock() + monkeypatch.setattr("prefect.agent.kubernetes.KubernetesAgent.start", start) + + k8s_config = MagicMock() + monkeypatch.setattr("kubernetes.config", k8s_config) + + runner = CliRunner() + result = runner.invoke(agent, ["start", "kubernetes", "test_kwarg=ignored"]) + assert result.exit_code == 0 + + +def test_agent_start_fargate(monkeypatch, runner_token): + start = MagicMock() + monkeypatch.setattr("prefect.agent.fargate.FargateAgent.start", start) + + boto3_client = MagicMock() + monkeypatch.setattr("boto3.client", boto3_client) + + runner = CliRunner() + result = runner.invoke(agent, ["start", "fargate"]) + assert result.exit_code == 0 + + +def test_agent_start_fargate_kwargs(monkeypatch, runner_token): + start = MagicMock() + monkeypatch.setattr("prefect.agent.fargate.FargateAgent.start", start) + + boto3_client = MagicMock() + monkeypatch.setattr("boto3.client", boto3_client) + + runner = CliRunner() + result = runner.invoke(agent, ["start", "fargate", "taskRoleArn=test"]) + assert result.exit_code == 0 + + +def test_agent_start_fargate_kwargs_received(monkeypatch, runner_token): + start = MagicMock() + monkeypatch.setattr("prefect.agent.fargate.FargateAgent.start", start) + + fargate_agent = MagicMock() + monkeypatch.setattr("prefect.agent.fargate.FargateAgent", fargate_agent) + + boto3_client = MagicMock() + monkeypatch.setattr("boto3.client", boto3_client) + + runner = CliRunner() + result = runner.invoke( + agent, ["start", "fargate", "taskRoleArn=arn", "--volumes=vol"] + ) + assert result.exit_code == 0 + + assert fargate_agent.called + fargate_agent.assert_called_with( + labels=[], name=None, taskRoleArn="arn", volumes="vol" + ) + + def test_agent_start_name(monkeypatch, runner_token): start = MagicMock() monkeypatch.setattr("prefect.agent.local.LocalAgent.start", start) @@ -67,7 +149,7 @@ def test_agent_start_name(monkeypatch, runner_token): assert result.exit_code == 0 -def test_agent_start_local_context_vars(monkeypatch, runner_token): +def test_agent_start_local_vars(monkeypatch, runner_token): start = MagicMock() monkeypatch.setattr("prefect.agent.local.LocalAgent.start", start) From 6b986398b41090decee26a0c5067466c86c8daba Mon Sep 17 00:00:00 2001 From: Josh Meek Date: Wed, 13 Nov 2019 10:08:01 -0500 Subject: [PATCH 3/8] Update Fargate Agent documentation to note CLI kwarg options --- docs/cloud/agent/fargate.md | 25 +++++++++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/docs/cloud/agent/fargate.md b/docs/cloud/agent/fargate.md index c8d4d007860a..5363f7e8b366 100644 --- a/docs/cloud/agent/fargate.md +++ b/docs/cloud/agent/fargate.md @@ -201,3 +201,28 @@ $ prefect agent start fargate :::warning Outbound Traffic If you encounter issues with Fargate raising errors in cases of client timeouts or inability to pull containers then you may need to adjust your `networkConfiguration`. Visit [this discussion thread](https://github.com/aws/amazon-ecs-agent/issues/1128#issuecomment-351545461) for more information on configuring AWS security groups. ::: + +#### Prefect CLI Using Kwargs + +All configuration options for the Fargate Agent can also be provided to the `prefect agent start fargate` CLI command. They must match the camel casing used by boto3. + +```bash +$ export AWS_ACCESS_KEY_ID=... +$ export AWS_SECRET_ACCESS_KEY=... +$ export REGION_NAME=us-east-1 + +$ prefect agent start fargate cpu=256 memory=512 networkConfiguration="{'awsvpcConfiguration': {'assignPublicIp': 'ENABLED', 'subnets': ['my_subnet_id'], 'securityGroups': []}}" +``` + +Kwarg values can also be provided through environment variables. This is useful in situations where case sensitive environment variables are desired or when using templating tools like Terraform to deploy your Agent. + +```bash +$ export AWS_ACCESS_KEY_ID=... +$ export AWS_SECRET_ACCESS_KEY=... +$ export REGION_NAME=us-east-1 +$ export CPU=256 +$ export MEMORY=512 +$ export NETWORK_CONFIGURATION="{'awsvpcConfiguration': {'assignPublicIp': 'ENABLED', 'subnets': ['my_subnet_id'], 'securityGroups': []}}" + +$ prefect agent start fargate cpu=$CPU memory=$MEMORY networkConfiguration=$NETWORK_CONFIGURATION +``` From 25d10b333016bcf2cae78289d512b097036bd00c Mon Sep 17 00:00:00 2001 From: Josh Meek Date: Wed, 13 Nov 2019 10:14:55 -0500 Subject: [PATCH 4/8] Remove extra prefect import --- src/prefect/cli/agent.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/prefect/cli/agent.py b/src/prefect/cli/agent.py index 5f09138cbdcb..c7fe3d36e70d 100644 --- a/src/prefect/cli/agent.py +++ b/src/prefect/cli/agent.py @@ -1,7 +1,6 @@ import click -import prefect -from prefect import config, context +from prefect import config from prefect.utilities.configuration import set_temporary_config from prefect.utilities.serialization import from_qualified_name From 73f7a0904e16fb14fc988a02a890fbc39b8bdbc9 Mon Sep 17 00:00:00 2001 From: Josh Meek Date: Wed, 13 Nov 2019 10:15:32 -0500 Subject: [PATCH 5/8] Update CHANGELOG --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 48354d67b34b..be4ec16c48c9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -16,6 +16,7 @@ These changes are available in the [master branch](https://github.com/PrefectHQ/ - Loosen `containerDefinitions` requirements for `FargateTaskEnvironment` - [#1713](https://github.com/PrefectHQ/prefect/pull/1713) - Local Docker agent proactively fails flow runs if image cannot be pulled - [#1395](https://github.com/PrefectHQ/prefect/issues/1395) - Add graceful keyboard interrupt shutdown for all agents - [#1731](https://github.com/PrefectHQ/prefect/pull/1731) +- `agent start` CLI command now allows for Agent kwargs - [#1737](https://github.com/PrefectHQ/prefect/pull/1737) ### Task Library From 5c95785c618e423c226ef8e7a86ea39635ed65af Mon Sep 17 00:00:00 2001 From: Josh Meek Date: Wed, 13 Nov 2019 10:22:48 -0500 Subject: [PATCH 6/8] Skip agent cli tests on missing imports --- tests/cli/test_agent.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tests/cli/test_agent.py b/tests/cli/test_agent.py index 4ee1bd54e46c..cf41b029325d 100644 --- a/tests/cli/test_agent.py +++ b/tests/cli/test_agent.py @@ -1,9 +1,14 @@ from unittest.mock import MagicMock from click.testing import CliRunner +import pytest from prefect.cli.agent import agent +pytest.importorskip("boto3") +pytest.importorskip("botocore") +pytest.importorskip("kubernetes") + def test_agent_init(): runner = CliRunner() From 572435251261341061372d1d07709552cd706acd Mon Sep 17 00:00:00 2001 From: Josh Meek Date: Wed, 13 Nov 2019 12:44:09 -0500 Subject: [PATCH 7/8] Bound google cloud storage requirement to < 1.23.0 --- setup.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/setup.py b/setup.py index ee51accd8713..394520dd71ff 100644 --- a/setup.py +++ b/setup.py @@ -20,7 +20,7 @@ "dropbox": ["dropbox ~= 9.0"], "google": [ "google-cloud-bigquery >= 1.6.0, < 2.0", - "google-cloud-storage >= 1.13, < 2.0", + "google-cloud-storage >= 1.13, < 1.23.0", ], "kubernetes": ["kubernetes >= 9.0.0a1, < 10.0", "dask-kubernetes >= 0.8.0"], "rss": ["feedparser >= 5.0.1, < 6.0"], From 62ea50a6252a4ae227e643563c988902e6a487d3 Mon Sep 17 00:00:00 2001 From: Josh Meek Date: Wed, 13 Nov 2019 13:06:51 -0500 Subject: [PATCH 8/8] Update Fargate Agent CLI doc for note about double hyphen options --- docs/cloud/agent/fargate.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/cloud/agent/fargate.md b/docs/cloud/agent/fargate.md index 5363f7e8b366..7227b94db20d 100644 --- a/docs/cloud/agent/fargate.md +++ b/docs/cloud/agent/fargate.md @@ -204,7 +204,7 @@ If you encounter issues with Fargate raising errors in cases of client timeouts #### Prefect CLI Using Kwargs -All configuration options for the Fargate Agent can also be provided to the `prefect agent start fargate` CLI command. They must match the camel casing used by boto3. +All configuration options for the Fargate Agent can also be provided to the `prefect agent start fargate` CLI command. They must match the camel casing used by boto3 but both the single kwarg as well as with the standard prefix of `--` are accepted. This means that `taskRoleArn=""` is the same as `--taskRoleArn=""`. ```bash $ export AWS_ACCESS_KEY_ID=...