Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add Optional "Name" to Agents #1612

Merged
merged 8 commits into from
Oct 15, 2019
Merged
Show file tree
Hide file tree
Changes from 7 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@ These changes are available in the [master branch](https://github.com/PrefectHQ/
- Additional debug logs to `Docker Container` and `Docker Image` tasks - [#920](https://github.com/PrefectHQ/prefect/issues/920)
- Changes to Fargate agent to support temporary credentials and IAM role based credentials within AWS compute such as a container or ec2 instance. [#1607](https://github.com/PrefectHQ/prefect/pull/1607)
- Local Secrets set through environment variable now retain their casing - [#1601](https://github.com/PrefectHQ/prefect/issues/1601)
- Agents can accept an optional `name` for logging and debugging - [#1612](https://github.com/PrefectHQ/prefect/pull/1612)
- Added AWS configuration options for Fargate Agent (task_role_arn, execution_role_arn) - [#1614](https://github.com/PrefectHQ/prefect/pull/1614)

### Task Library
Expand All @@ -37,7 +38,7 @@ These changes are available in the [master branch](https://github.com/PrefectHQ/

### Contributors

- None
- [Mark McDonald](https://github.com/mhmcdonal)

## 0.6.6 <Badge text="beta" type="success"/>

Expand Down
9 changes: 7 additions & 2 deletions src/prefect/agent/agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -34,16 +34,21 @@ class Agent:

In order for this to operate `PREFECT__CLOUD__AGENT__AUTH_TOKEN` must be set as an
environment variable or in your user configuration file.

Args:
- name (str, optional): An optional name to give this agent. Can also be set through
the environment variable `PREFECT__CLOUD__AGENT__NAME`. Defaults to "agent"
"""

def __init__(self) -> None:
def __init__(self, name: str = None) -> None:
self.name = name or config.cloud.agent.get("name", "agent")

token = config.cloud.agent.get("auth_token")

self.client = Client(api_token=token)
self._verify_token(token)

logger = logging.getLogger("agent")
logger = logging.getLogger(self.name)
logger.setLevel(config.cloud.agent.get("level"))
ch = logging.StreamHandler()
formatter = logging.Formatter(
Expand Down
5 changes: 4 additions & 1 deletion src/prefect/agent/fargate/agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,8 @@ class FargateAgent(Agent):
Fargate Agent can be found at https://docs.prefect.io/cloud/agent/fargate.html

Args:
- name (str, optional): An optional name to give this agent. Can also be set through
the environment variable `PREFECT__CLOUD__AGENT__NAME`. Defaults to "agent"
- aws_access_key_id (str, optional): AWS access key id for connecting the boto3
client. Defaults to the value set in the environment variable
`AWS_ACCESS_KEY_ID`.
Expand Down Expand Up @@ -56,6 +58,7 @@ class FargateAgent(Agent):

def __init__(
self,
name: str = None,
aws_access_key_id: str = None,
aws_secret_access_key: str = None,
aws_session_token: str = None,
Expand All @@ -70,7 +73,7 @@ def __init__(
task_cpu: str = None,
task_memory: str = None,
) -> None:
super().__init__()
super().__init__(name=name)

from boto3 import client as boto3_client

Expand Down
8 changes: 6 additions & 2 deletions src/prefect/agent/kubernetes/agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -34,10 +34,14 @@ class KubernetesAgent(Agent):
run on a k8s cluster or on a local machine where the kube_config is pointing at the
desired cluster. Information on using the Kubernetes Agent can be found at
https://docs.prefect.io/cloud/agent/kubernetes.html

Args:
- name (str, optional): An optional name to give this agent. Can also be set through
the environment variable `PREFECT__CLOUD__AGENT__NAME`. Defaults to "agent"
"""

def __init__(self) -> None:
super().__init__()
def __init__(self, name: str = None) -> None:
super().__init__(name=name)

from kubernetes import client, config

Expand Down
8 changes: 6 additions & 2 deletions src/prefect/agent/local/agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,15 +15,19 @@ class LocalAgent(Agent):
Local Agent can be found at https://docs.prefect.io/cloud/agent/local.html

Args:
- name (str, optional): An optional name to give this agent. Can also be set through
the environment variable `PREFECT__CLOUD__AGENT__NAME`. Defaults to "agent"
- base_url (str, optional): URL for a Docker daemon server. Defaults to
`unix:///var/run/docker.sock` however other hosts such as
`tcp://0.0.0.0:2375` can be provided
- no_pull (bool, optional): Flag on whether or not to pull flow images.
Defaults to `False` if not provided here or in context.
"""

def __init__(self, base_url: str = None, no_pull: bool = None) -> None:
super().__init__()
def __init__(
self, name: str = None, base_url: str = None, no_pull: bool = None
) -> None:
super().__init__(name=name)

if platform == "win32":
default_url = "npipe:////./pipe/docker_engine"
Expand Down
8 changes: 6 additions & 2 deletions src/prefect/agent/nomad/agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,10 +16,14 @@ class NomadAgent(Agent):
"""
Agent which deploys flow runs as Nomad jobs to a Nomad cluster based on the
`NOMAD_HOST` environment variable.

Args:
- name (str, optional): An optional name to give this agent. Can also be set through
the environment variable `PREFECT__CLOUD__AGENT__NAME`. Defaults to "agent"
"""

def __init__(self) -> None:
super().__init__()
def __init__(self, name: str = None) -> None:
super().__init__(name=name)

def deploy_flows(self, flow_runs: list) -> None:
"""
Expand Down
21 changes: 15 additions & 6 deletions src/prefect/cli/agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -43,27 +43,36 @@ def agent():


@agent.command(hidden=True)
@click.argument("name", default="local")
@click.argument("agent-option", default="local")
@click.option(
"--token", "-t", required=False, help="A Prefect Cloud API token.", hidden=True
)
@click.option(
"--verbose", "-v", is_flag=True, help="Enable verbose agent logs.", hidden=True
)
@click.option(
"--name",
"-n",
required=False,
help="A name to use for the agent",
hidden=True,
default=None,
)
@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(name, token, verbose, no_pull, base_url):
def start(agent_option, token, verbose, name, no_pull, base_url):
"""
Start an agent.

\b
Arguments:
name TEXT The name of an agent to start (e.g. `local`, `kubernetes`, `fargate`, `nomad`)
agent-option TEXT The name of an agent to start (e.g. `local`, `kubernetes`, `fargate`, `nomad`)
Defaults to `local`

\b
Options:
--token, -t TEXT A Prefect Cloud API token with RUNNER scope
--name, -n TEXT A name to use for the agent
--verbose, -v Enable verbose agent DEBUG logs
Defaults to INFO level logging

Expand All @@ -79,14 +88,14 @@ def start(name, token, verbose, no_pull, base_url):
"cloud.agent.level": "INFO" if not verbose else "DEBUG",
}
):
retrieved_agent = _agents.get(name, None)
retrieved_agent = _agents.get(agent_option, None)

if not retrieved_agent:
click.secho("{} is not a valid agent".format(name), fg="red")
click.secho("{} is not a valid agent".format(agent_option), fg="red")
return

with context(no_pull=no_pull, base_url=base_url):
from_qualified_name(retrieved_agent)().start()
from_qualified_name(retrieved_agent)(name=name).start()


@agent.command(hidden=True)
Expand Down
1 change: 1 addition & 0 deletions src/prefect/config.toml
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@ queue_interval = 30.0
# Agents require different API tokens
auth_token = ""
level = "INFO"
name = "agent"

[cloud.agent.resource_manager]
# Separate loop interval for resource managers
Expand Down
20 changes: 20 additions & 0 deletions tests/agent/test_agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,27 @@ def test_agent_config_options(runner_token):
with set_temporary_config({"cloud.agent.auth_token": "TEST_TOKEN"}):
agent = Agent()
assert agent.client.get_auth_token() == "TEST_TOKEN"
assert agent.name == "agent"
assert agent.logger
assert agent.logger.name == "agent"


def test_agent_name_set_options(runner_token, monkeypatch):
# Default
agent = Agent()
assert agent.name == "agent"
assert agent.logger.name == "agent"

# Init arg
agent = Agent(name="test1")
assert agent.name == "test1"
assert agent.logger.name == "test1"

# Config
with set_temporary_config({"cloud.agent.name": "test2"}):
agent = Agent()
assert agent.name == "test2"
assert agent.logger.name == "test2"


def test_agent_log_level(runner_token):
Expand Down
3 changes: 3 additions & 0 deletions tests/agent/test_fargate_agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@ def test_ecs_agent_config_options_default(monkeypatch, runner_token):

agent = FargateAgent()
assert agent
assert agent.name == "agent"
assert agent.cluster == "default"
assert not agent.subnets
assert not agent.security_groups
Expand All @@ -44,6 +45,7 @@ def test_ecs_agent_config_options_init(monkeypatch, runner_token):
monkeypatch.setattr("boto3.client", boto3_client)

agent = FargateAgent(
name="test",
aws_access_key_id="id",
aws_secret_access_key="secret",
aws_session_token="token",
Expand All @@ -59,6 +61,7 @@ def test_ecs_agent_config_options_init(monkeypatch, runner_token):
task_memory="2",
)
assert agent
assert agent.name == "test"
assert agent.task_role_arn == "task_role_arn"
assert agent.execution_role_arn == "execution_role_arn"
assert agent.cluster == "cluster"
Expand Down
4 changes: 3 additions & 1 deletion tests/agent/test_k8s_agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@ def test_k8s_agent_init(monkeypatch, runner_token):

agent = KubernetesAgent()
assert agent
assert agent.name == "agent"
assert agent.batch_client


Expand All @@ -30,8 +31,9 @@ def test_k8s_agent_config_options(monkeypatch, runner_token):
monkeypatch.setattr("kubernetes.config", k8s_config)

with set_temporary_config({"cloud.agent.auth_token": "TEST_TOKEN"}):
agent = KubernetesAgent()
agent = KubernetesAgent(name="test")
assert agent
assert agent.name == "test"
assert agent.client.get_auth_token() == "TEST_TOKEN"
assert agent.logger
assert agent.batch_client
Expand Down
4 changes: 3 additions & 1 deletion tests/agent/test_local_agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@ def test_local_agent_init(monkeypatch, runner_token):

agent = LocalAgent()
assert agent
assert agent.name == "agent"


def test_local_agent_config_options(monkeypatch, runner_token):
Expand All @@ -23,7 +24,8 @@ def test_local_agent_config_options(monkeypatch, runner_token):
monkeypatch.setattr("prefect.agent.local.agent.platform", "osx")

with set_temporary_config({"cloud.agent.auth_token": "TEST_TOKEN"}):
agent = LocalAgent()
agent = LocalAgent(name="test")
assert agent.name == "test"
assert agent.client.get_auth_token() == "TEST_TOKEN"
assert agent.logger
assert not agent.no_pull
Expand Down
4 changes: 3 additions & 1 deletion tests/agent/test_nomad_agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,12 +11,14 @@
def test_nomad_agent_init(runner_token):
agent = NomadAgent()
assert agent
assert agent.name == "agent"


def test_nomad_agent_config_options(runner_token):
with set_temporary_config({"cloud.agent.auth_token": "TEST_TOKEN"}):
agent = NomadAgent()
agent = NomadAgent(name="test")
assert agent
assert agent.name == "test"
assert agent.client.get_auth_token() == "TEST_TOKEN"
assert agent.logger

Expand Down
12 changes: 12 additions & 0 deletions tests/cli/test_agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,18 @@ def test_agent_start_verbose(monkeypatch, runner_token):
assert result.exit_code == 0


def test_agent_start_name(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", "--name", "test_agent"])
assert result.exit_code == 0


def test_agent_start_local_context_vars(monkeypatch, runner_token):
start = MagicMock()
monkeypatch.setattr("prefect.agent.local.LocalAgent.start", start)
Expand Down