def execute( # type: ignore
self,
storage: "Storage",
flow_location: str,
**kwargs: Any # type: ignore
) -> None:
flow_run_info = None
flow_run_id = prefect.context.get("flow_run_id")
if self._on_execute:
# If an on_execute Callable has been provided, retrieve the flow run parameters
# and then allow the Callable a chance to update _provider_kwargs. This allows
# better sizing of the cluster resources based on parameters for this Flow run.
try:
client = Client()
flow_run_info = client.get_flow_run_info(flow_run_id)
parameters = flow_run_info.parameters or {} # type: ignore
self._on_execute(parameters, self._provider_kwargs)
except Exception as exc:
self.logger.info(
"Failed to retrieve flow run info with error: {}".format(
repr(exc)))
if "image" not in self._provider_kwargs or not self._provider_kwargs.get(
"image"):
# If image is not specified, use the Flow's image so that dependencies are
# identical on all containers: Flow runner, Dask scheduler, and Dask workers
flow_id = prefect.context.get("flow_id")
try:
client = Client()
if not flow_id: # We've observed cases where flow_id is None
if not flow_run_info:
flow_run_info = client.get_flow_run_info(flow_run_id)
flow_id = flow_run_info.flow_id
flow_info = client.graphql("""query {
flow(where: {id: {_eq: "%s"}}) {
storage
}
}""" % flow_id)
storage_info = flow_info["data"]["flow"][0]["storage"]
image = "{}/{}:{}".format(
storage_info["registry_url"],
storage_info["image_name"],
storage_info["image_tag"],
)
self.logger.info(
"Using Flow's Docker image for Dask scheduler & workers: {}"
.format(image))
self._provider_kwargs["image"] = image
except Exception as exc:
self.logger.info(
"Failed to retrieve flow info with error: {}".format(
repr(exc)))
self._create_dask_cluster()
self.logger.info(
"Executing on dynamically created Dask Cluster with scheduler address: {}"
.format(self.executor_kwargs["address"]))
super().execute(storage, flow_location, **kwargs)
def list_tenants():
"""
List available tenants
"""
check_override_auth_token()
client = Client()
tenants = client.get_available_tenants()
active_tenant_id = client._active_tenant_id
output = []
for item in tenants:
active = None
if item.id == active_tenant_id:
active = "*"
output.append([item.name, item.slug, item.id, active])
click.echo(
tabulate(
output,
headers=["NAME", "SLUG", "ID", ""],
tablefmt="plain",
numalign="left",
stralign="left",
)
)
def __init__(self) -> None:
self.loop_interval = prefect_config.cloud.agent.resource_manager.get(
"loop_interval")
self.client = Client(
api_token=prefect_config.cloud.agent.get("auth_token"))
self.namespace = os.getenv("NAMESPACE", "default")
logger = logging.getLogger("resource-manager")
logger.setLevel(logging.DEBUG)
ch = logging.StreamHandler(sys.stdout)
ch.setLevel(logging.DEBUG)
formatter = logging.Formatter(context.config.logging.format)
formatter.converter = time.gmtime # type: ignore
ch.setFormatter(formatter)
logger.addHandler(ch)
self.logger = logger
from kubernetes import client, config
try:
config.load_incluster_config()
except config.config_exception.ConfigException as exc:
self.logger.warning(
"{} Using out of cluster configuration option.".format(exc))
config.load_kube_config()
self.k8s_client = client
def add(token, config_path):
"""
Add a new Prefect Cloud auth token to use for Cloud communication.
\b
Options:
--token, -t TEXT A Prefect Cloud auth token [required]
--config-path, -c TEXT Path to a Prefect config.toml, defaults to `~/.prefect/config.toml`
"""
abs_directory = os.path.abspath(os.path.expanduser(config_path))
if not os.path.exists(abs_directory):
click.secho("{} does not exist".format(config_path), fg="red")
return
config = toml.load(abs_directory)
if not config.get("cloud"):
config["cloud"] = {}
config["cloud"]["auth_token"] = token
with open(abs_directory, "w") as file:
toml.dump(config, file)
click.echo("Auth token added to Prefect config")
client = Client()
client.token = token
result = client.graphql(query={"query": "hello"})
if not result.data.hello:
click.secho(
"Error attempting to use Prefect auth token {}".format(result))
def list_tokens():
"""
DEPRECATED. Please use API keys instead.
List your available Prefect Cloud API tokens.
"""
click.secho(
"WARNING: API tokens are deprecated. Please consider removing your remaining "
"tokens and using API keys instead.",
fg="yellow",
err=True, # Write to stderr in case the user is piping
)
client = Client()
output = client.graphql(query={"query": {"api_token": {"id", "name"}}})
if not output.get("data", None):
click.secho("Unable to list API tokens", fg="red")
return
tokens = []
for item in output.data.api_token:
tokens.append([item.name, item.id])
click.echo(
tabulate(
tokens,
headers=["NAME", "ID"],
tablefmt="plain",
numalign="left",
stralign="left",
))
def logout(token):
"""
Log out of Prefect Cloud
This will remove your cached authentication from disk.
"""
client = Client()
# Log out of API keys unless given the token flag
if client.api_key and not token:
# Check the source of the API key
abort_on_config_api_key(
"To log out, remove the config key `prefect.cloud.api_key`")
click.confirm(
"Are you sure you want to log out of Prefect Cloud? "
"This will remove your API key from this machine.",
default=False,
abort=True,
)
# Clear the key and tenant id then write to the cache
client.api_key = ""
client._tenant_id = ""
client.save_auth_to_disk()
click.secho("Logged out of Prefect Cloud", fg="green")
else:
raise TerminalError("You are not logged in to Prefect Cloud. "
"Use `prefect auth login` to log in first.")
def list_tenants():
"""
List available tenants
"""
client = Client()
try:
tenants = client.get_available_tenants()
except AuthorizationError:
raise TerminalError(
"You are not authenticated. Use `prefect auth login` first.")
output = []
for item in tenants:
active = None
if item.id == client.tenant_id:
active = "*"
output.append([item.name, item.slug, item.id, active])
click.echo(
tabulate(
output,
headers=["NAME", "SLUG", "ID", ""],
tablefmt="plain",
numalign="left",
stralign="left",
))
def login(token):
"""
Login to Prefect Cloud with an api token to use for Cloud communication.
\b
Options:
--token, -t TEXT A Prefect Cloud api token [required]
"""
client = Client()
client.login(api_token=token)
# Verify login obtained a valid api token
try:
client.graphql(query={"query": "hello"})
except AuthorizationError:
click.secho(
"Error attempting to use Prefect API token {}".format(token),
color="RED")
return
except ClientError:
click.secho("Error attempting to communicate with Prefect Cloud",
color="RED")
return
click.echo("Login successful")
def revoke_token(id):
"""
DEPRECATED. Please use API keys instead.
Revote a Prefect Cloud API token
\b
Options:
--id, -i TEXT The id of a token to revoke
"""
check_override_auth_token()
client = Client()
output = client.graphql(
query={
"mutation($input: delete_api_token_input!)": {
"delete_api_token(input: $input)": {"success"}
}
},
variables=dict(input=dict(token_id=id)),
)
if not output.get("data",
None) or not output.data.delete_api_token.success:
click.secho("Unable to revoke token with ID {}".format(id), fg="red")
return
click.secho("Token successfully revoked", fg="green")
def list_keys():
"""
List available Prefect Cloud API keys.
"""
client = Client()
response = client.graphql(
query={
"query": {
"auth_api_key": {
"id": True,
"name": True,
"expires_at": True,
}
}
})
keys = response.get("data", {}).get("auth_api_key")
if keys is None:
raise TerminalError(
f"Unexpected response from Prefect Cloud: {response}")
if not keys:
click.secho("You have not created any API keys", fg="yellow")
else:
click.echo(
tabulate(
[(key.name, key.id, key.expires_at or "NEVER")
for key in keys],
headers=["NAME", "ID", "EXPIRES AT"],
tablefmt="plain",
numalign="left",
stralign="left",
))
def status():
"""
Get the current Prefect authentication status
"""
client = Client()
click.echo(f"You are connecting to {client.api_server}")
if client.api_key:
click.echo("You are authenticating with an API key")
try:
click.echo(
f"You are logged in to tenant {client.get_default_tenant()}")
except Exception as exc:
click.echo(f"Your authentication is not working: {exc}")
if client._api_token:
click.secho(
"You are logged in with an API token. These have been deprecated in favor "
"of API keys." +
(" Since you have set an API key as well, this will be ignored."
if client.api_key else ""),
fg="yellow",
)
if not client._api_token and not client.api_key:
click.secho("You are not logged in!", fg="yellow")
def list_tokens():
"""
List your available Prefect Cloud API tokens.
"""
check_override_auth_token()
client = Client()
output = client.graphql(query={"query": {"api_token": {"id", "name"}}})
if not output.get("data", None):
click.secho("Unable to list API tokens", fg="red")
return
tokens = []
for item in output.data.api_token:
tokens.append([item.name, item.id])
click.echo(
tabulate(
tokens,
headers=["NAME", "ID"],
tablefmt="plain",
numalign="left",
stralign="left",
)
)
def create_token(name, scope):
"""
Create a Prefect Cloud API token.
For more info on API tokens visit https://docs.prefect.io/orchestration/concepts/api.html
\b
Options:
--name, -n TEXT A name to give the generated token
--scope, -s TEXT A scope for the token
"""
check_override_auth_token()
client = Client()
output = client.graphql(
query={
"mutation($input: create_api_token_input!)": {
"create_api_token(input: $input)": {"token"}
}
},
variables=dict(input=dict(name=name, scope=scope)),
)
if not output.get("data", None):
click.secho("Issue creating API token", fg="red")
return
click.echo(output.data.create_api_token.token)
def from_current_tenant(cls) -> "TenantView":
"""
Get an instance of this class filled with information by querying for the
tenant id set in the Prefect Client
Returns:
A populated `TenantView` instance
"""
client = Client()
return cls.from_tenant_id(client.tenant_id)
def _query_for_task_runs(
where: dict,
order_by: dict = None,
error_on_empty: bool = True,
) -> List[dict]:
"""
Query for task run data necessary to initialize `TaskRunView` instances
with `TaskRunView.from_task_run_data`.
Args:
- where (required): The Hasura `where` clause to filter by
- order_by (optional): An optional Hasura `order_by` clause to order results
by.
- error_on_empty (optional): If `True` and no tasks are found, a `ValueError`
will be raised.
Returns:
A list of dicts containing task run data
"""
client = Client()
query_args = {"where": where}
if order_by is not None:
query_args["order_by"] = order_by
query = {
"query": {
with_args("task_run", query_args): {
"id": True,
"name": True,
"task": {
"id": True,
"slug": True
},
"map_index": True,
"serialized_state": True,
"flow_run_id": True,
}
}
}
result = client.graphql(query)
task_runs = result.get("data", {}).get("task_run", None)
if task_runs is None:
raise ValueError(
f"Received bad result while querying for task runs where {where}: "
f"{result}")
if not task_runs and error_on_empty:
raise ValueError(
f"No task runs found while querying for task runs where {where}"
)
return task_runs
def delete_artifact(task_run_artifact_id: str) -> None:
"""
Delete an existing artifact
Args:
- task_run_artifact_id (str): the ID of an existing task run artifact
"""
if not _running_with_backend():
return
client = Client()
client.delete_task_run_artifact(task_run_artifact_id=task_run_artifact_id)
def _query_for_tenants(
where: dict,
order_by: dict = None,
error_on_empty: bool = True,
) -> List[dict]:
"""
Query for tenant data necessary to initialize `TenantView` instances with
`TenantView._from_tenant_data`.
Args:
- where (required): The Hasura `where` clause to filter by
- order_by (optional): An optional Hasura `order_by` clause to order
results by
- error_on_empty (optional): If `True` and no tenants are found, a
`ValueError` will be raised
Returns:
A list of dicts of tenant information
"""
client = Client()
query_args = {"where": where}
if order_by is not None:
query_args["order_by"] = order_by
tenant_query = {
"query": {
with_args("tenant", query_args): {
"id",
"slug",
"name",
}
}
}
result = client.graphql(tenant_query)
tenants = result.get("data", {}).get("tenant", None)
if tenants is None:
raise ValueError(
f"Received bad result while querying for tenants where {where}: "
f"{result}")
if not tenants: # Empty list
if error_on_empty:
raise ValueError(
f"No results found while querying for tenants where {where!r}"
)
return []
# Return a list
return tenants
def login(token):
"""
Log in to Prefect Cloud with an api token to use for Cloud communication.
\b
Options:
--token, -t TEXT A Prefect Cloud api token [required]
"""
check_override_auth_token()
client = Client(api_token=token)
# Verify login obtained a valid api token
try:
output = client.graphql(
query={"query": {
"user": {
"default_membership": "tenant_id"
}
}})
# Log into default membership
success_login = client.login_to_tenant(
tenant_id=output.data.user[0].default_membership.tenant_id)
if not success_login:
raise AuthorizationError
except AuthorizationError:
click.secho(
f"Error attempting to use Prefect API token {token}. "
"Please check that you are providing a USER scoped Personal Access Token.\n"
"For more information visit the documentation for USER tokens at "
"https://docs.prefect.io/orchestration/concepts/tokens.html#user",
fg="red",
)
return
except ClientError:
click.secho(
"Error attempting to communicate with Prefect Cloud. "
"Please check that you are providing a USER scoped Personal Access Token.\n"
"For more information visit the documentation for USER tokens at "
"https://docs.prefect.io/orchestration/concepts/tokens.html#user",
fg="red",
)
return
# save token
client.save_api_token()
click.secho("Login successful!", fg="green")
def switch_tenants(id, slug, default):
"""
Switch active tenant
\b
Options:
--id, -i TEXT A Prefect Cloud tenant id
--slug, -s TEXT A Prefect Cloud tenant slug
"""
# If the config specifies a tenant explicitly, it is used before this mechanism
if config.cloud.get("tenant_id"):
raise TerminalError(
"Your tenant id has been set in the Prefect config instead of with the "
"CLI. To switch tenants with the CLI, remove the config key "
" `prefect.cloud.tenant_id`")
client = Client()
# Deprecated API token check
if not client.api_key:
check_override_auth_token()
if default:
raise TerminalError(
"The default tenant flag can only be used with API keys.")
else: # Using an API key
if default:
# Clear the set tenant on disk
client.tenant_id = None
client.save_auth_to_disk()
click.secho(
"Tenant restored to the default tenant for your API key: "
f"{client.get_default_tenant()}",
fg="green",
)
return
login_success = client.login_to_tenant(tenant_slug=slug, tenant_id=id)
if not login_success:
raise TerminalError("Unable to switch tenant!")
# `login_to_tenant` will write to disk if using an API token, if using an API key
# we will write to disk manually here
if client.api_key:
client.save_auth_to_disk()
click.secho(f"Tenant switched to {client.tenant_id}", fg="green")
def teardown(self, worker=None):
from prefect import Client
msg = """
Lost communication with Dask worker:
{}
""".format(worker)
Client().write_run_logs([
dict(
flow_run_id=self.flow_run_id,
name="DaskWorkerPlugin",
message=msg,
level="ERROR",
)
])
def update_link(task_run_artifact_id: str, link: str) -> None:
"""
Update an existing link artifact. This function will replace the current link
artifact with the new link provided.
Args:
- task_run_artifact_id (str): the ID of an existing task run artifact
- link (str): the new link to update the artifact with
"""
if not _running_with_backend():
return
client = Client()
client.update_task_run_artifact(task_run_artifact_id=task_run_artifact_id,
data={"link": link})
def update_markdown(task_run_artifact_id: str, markdown: str) -> None:
"""
Update an existing markdown artifact. This function will replace the current markdown
artifact with the new markdown provided.
Args:
- task_run_artifact_id (str): the ID of an existing task run artifact
- markdown (str): the new markdown to update the artifact with
"""
if not _running_with_backend():
return
client = Client()
client.update_task_run_artifact(task_run_artifact_id=task_run_artifact_id,
data={"markdown": markdown})
def delete_artifact(task_run_artifact_id: str) -> None:
"""
Delete an existing artifact
Note: The functionality here is experimental, and may change between
versions without notice. Use at your own risk.
Args:
- task_run_artifact_id (str): the ID of an existing task run artifact
"""
if not _running_with_backend():
return
client = Client()
client.delete_task_run_artifact(task_run_artifact_id=task_run_artifact_id)
def register_workflow(
self,
environment: str,
prefect_execution_environment: str,
prefect_workflow_register_token: str,
workflow_cpu_configuration: int,
workflow_memory_configuration: int,
) -> None:
"""
Registers the workflow to Prefect Cloud
Parameters:
environment [string] -- environment the workflow should be pushed to
prefect_execution_environment [string] -- e.g: ecs_fargate, kubernetes
prefect_workflow_register_token [string]
-- Prefect API token to register workflows
workflow_cpu_configuration [int] -- e.g: 256,512,1024,2048,4096
workflow_memory_configuration [int] -- e.g: 512,30720
"""
# set flow properties
flow_module, flow_name = self.set_workflow_properties(
environment,
prefect_execution_environment,
workflow_cpu_configuration,
workflow_memory_configuration,
)
# Authenticate to ECR as the registration process pushes the image to AWS
self.aws_conn_helpers.ecr_authenticate()
# Create ECR repository
self.create_workflow_ecr_repository(flow_name=flow_name)
try:
# Instantiate the prefect client
prefect_client = Client(api_token=prefect_workflow_register_token)
# Register the Workflow
prefect_client.register(
flow=flow_module.flow,
project_name=f"{environment}_dataflow_automation")
except Exception as e:
raise e
else:
logging.info(f"Workflow '{flow_name}' registered successfully!")
def is_agent_up_prefect_cloud(pytestconfig):
"""
Check if the agent is up on Prefect Cloud
"""
# get prefect agent API token
prefect_agent_token = pytestconfig.getoption("prefect_agent_token")
# Instantiate the prefect client
prefect_client = Client(api_token=prefect_agent_token)
# query prefect cloud agents
query = """
query RunningFlows {
agent {
labels,
last_queried
}
}
"""
agents = prefect_client.graphql(query=query)["data"]["agent"]
environment_agent_info = [
item for item in agents
if f"{ENV}_dataflow_automation" in item["labels"]
]
if len(environment_agent_info) == 0:
# agent doesn't even appear on prefect cloud yet
return ResourceCheckStatus.RETRY
else:
last_queried_time = environment_agent_info[0]["last_queried"]
# date
last_queried_time = datetime.strptime(last_queried_time,
"%Y-%m-%dT%H:%M:%S.%f%z")
difference = datetime.utcnow().replace(
tzinfo=None) - last_queried_time.replace(tzinfo=None)
# If the agent has been queried recently (meaning it is up)
if int(difference.total_seconds()) < 15:
return ResourceCheckStatus.FINISHED
else:
# agent appears on prefect cloud but doesn't seem to be up
return ResourceCheckStatus.RETRY
def update_markdown(task_run_artifact_id: str, markdown: str) -> None:
"""
Update an existing markdown artifact. This function will replace the current markdown
artifact with the new markdown provided.
Note: The functionality here is experimental, and may change between
versions without notice. Use at your own risk.
Args:
- task_run_artifact_id (str): the ID of an existing task run artifact
- markdown (str): the new markdown to update the artifact with
"""
if not _running_with_backend():
return
client = Client()
client.update_task_run_artifact(task_run_artifact_id=task_run_artifact_id,
data={"markdown": markdown})
def create_link(link: str) -> Optional[str]:
"""
Create a link artifact
Args:
- link (str): the link to post
Returns:
- str: the task run artifact ID
"""
if not _running_with_backend():
return None
client = Client()
return client.create_task_run_artifact(
task_run_id=context.get("task_run_id"),
kind="link",
data={"link": link})
def login(key):
"""
Login to Prefect Cloud
Create an API key in the UI then login with it here:
$ prefect auth login -k YOUR-KEY
You will be switched to the default tenant associated with the key. After login,
your available tenants can be seen with `prefect auth list-tenants` and you can
change the default tenant on this machine using `prefect auth switch-tenants`.
The given key will be stored on disk for later access. Prefect will default to using
this key for all interaction with the API but frequently overrides can be passed to
individual commands or functions. To remove your key from disk, see
`prefect auth logout`.
"""
if not key:
raise TerminalError("You must supply an API key!")
abort_on_config_api_key(
"To log in with the CLI, remove the config key `prefect.cloud.api_key`"
)
# Ignore any tenant id that has been previously set via login
client = Client(api_key=key)
client._tenant_id = None
try:
tenant_id = client._get_auth_tenant()
except AuthorizationError:
raise TerminalError("Unauthorized. Invalid Prefect Cloud API key.")
except ClientError:
raise TerminalError(
"Error attempting to communicate with Prefect Cloud.")
else:
client.tenant_id = tenant_id
client.save_auth_to_disk()
tenant = TenantView.from_tenant_id(tenant_id)
click.secho(
f"Logged in to Prefect Cloud tenant {tenant.name!r} ({tenant.slug})",
fg="green",
)
return
def status():
"""
Get the current Prefect authentication status
"""
client = Client()
click.echo(f"You are connecting to {client.api_server}")
if client.api_key:
click.echo("You are authenticating with an API key")
try:
click.echo(f"You are logged in to tenant {client.tenant_id}")
except Exception as exc:
click.echo(f"Your authentication is not working: {exc}")
if AUTH_TOKEN_SETTINGS_PATH.exists():
click.secho(
"The authentication tokens settings path still exists. These have been "
"removed in favor of API keys. We recommend purging old tokens with "
"`prefect auth purge-tokens`",
fg="yellow",
)
if config.cloud.get("auth_token"):
if os.environ.get("PREFECT__CLOUD__AUTH_TOKEN"):
click.secho(
"An authentication token is set via environment variable. "
"These have been removed in favor of API keys and the variable will be "
"ignored. We recommend unsetting the 'PREFECT__CLOUD__AUTH_TOKEN' key",
fg="yellow",
)
else:
click.secho(
"An authentication token is set via the prefect config file. "
"These have been removed in favor of API keys and the setting will be "
"ignored. We recommend removing the 'prefect.cloud.auth_token' key",
fg="yellow",
)
if not client.api_key:
click.secho(
"You are not logged in! Use `prefect auth login` to login with an API key.",
fg="yellow",
)
def revoke_key(id):
"""
Revoke a Prefect Cloud API key.
"""
client = Client()
output = client.graphql(
query={
"mutation($input: delete_api_key_input!)": {
"delete_api_key(input: $input)": {"success"}
}
},
variables=dict(input=dict(key_id=id)),
)
if not output.get("data", None) or not output.data.delete_api_key.success:
raise TerminalError(f"Unable to revoke key {id!r}")
click.secho("Key successfully revoked!", fg="green")
