def test_get_flow_image_raises_on_missing_info():
flow = Flow(
"test",
environment=LocalEnvironment(),
storage=Local(),
)
with pytest.raises(ValueError):
get_flow_image(flow=flow)
def test_get_flow_image_raises_on_missing_info():
flow = Flow(
"test",
run_config=UniversalRun(),
storage=Local(),
)
with pytest.raises(ValueError):
get_flow_image(flow=flow)
def execute(self, flow: "Flow", **kwargs: Any) -> None: # type: ignore
"""
Create a single Kubernetes job that runs the flow.
Args:
- flow (Flow): the Flow object
- **kwargs (Any): additional keyword arguments to pass to the runner
Raises:
- Exception: if the environment is unable to create the Kubernetes job
"""
docker_name = get_flow_image(flow)
from kubernetes import client, config
# Verify environment is running in cluster
try:
config.load_incluster_config()
except config.config_exception.ConfigException:
self.logger.error("Environment not currently running inside a cluster")
raise EnvironmentError("Environment not currently inside a cluster")
batch_client = client.BatchV1Api()
job = self._populate_run_time_job_spec_details(docker_name=docker_name)
# Create Job
try:
batch_client.create_namespaced_job(
namespace=prefect.context.get("namespace"), body=job
)
except Exception as exc:
self.logger.critical("Failed to create Kubernetes job: {}".format(exc))
raise exc
def test_get_flow_image_env_metadata():
flow = Flow(
"test",
environment=LocalEnvironment(metadata={"image": "repo/name:tag"}),
storage=Local(),
)
image = get_flow_image(flow=flow)
assert image == "repo/name:tag"
def test_get_flow_image_docker_storage():
flow = Flow(
"test",
environment=LocalEnvironment(),
storage=Docker(registry_url="test", image_name="name", image_tag="tag"),
)
image = get_flow_image(flow=flow)
assert image == "test/name:tag"
def test_get_flow_image_run_config():
flow = Flow(
"test",
run_config=DockerRun(image="repo/name:tag"),
storage=Local(),
)
image = get_flow_image(flow=flow)
assert image == "repo/name:tag"
def _render_task_definition_kwargs(self, flow: "Flow") -> dict:
task_definition_kwargs = self.task_definition_kwargs.copy()
env_values = [
{"name": "PREFECT__CLOUD__GRAPHQL", "value": config.cloud.graphql},
{"name": "PREFECT__CLOUD__USE_LOCAL_SECRETS", "value": "false"},
{
"name": "PREFECT__ENGINE__FLOW_RUNNER__DEFAULT_CLASS",
"value": "prefect.engine.cloud.CloudFlowRunner",
},
{
"name": "PREFECT__ENGINE__TASK_RUNNER__DEFAULT_CLASS",
"value": "prefect.engine.cloud.CloudTaskRunner",
},
{"name": "PREFECT__LOGGING__LOG_TO_CLOUD", "value": "true"},
{
"name": "PREFECT__LOGGING__EXTRA_LOGGERS",
"value": str(config.logging.extra_loggers),
},
]
# create containerDefinitions if they do not exist
if not task_definition_kwargs.get("containerDefinitions"):
task_definition_kwargs["containerDefinitions"] = []
task_definition_kwargs["containerDefinitions"].append({})
# set environment variables for all containers
for definition in task_definition_kwargs["containerDefinitions"]:
if not definition.get("environment"):
definition["environment"] = []
definition["environment"].extend(env_values)
# set name on first container
if not task_definition_kwargs["containerDefinitions"][0].get("name"):
task_definition_kwargs["containerDefinitions"][0]["name"] = ""
task_definition_kwargs.get("containerDefinitions")[0]["name"] = "flow-container"
# set image on first container
if not task_definition_kwargs["containerDefinitions"][0].get("image"):
task_definition_kwargs["containerDefinitions"][0]["image"] = ""
task_definition_kwargs.get("containerDefinitions")[0]["image"] = get_flow_image(
flow
)
# set command on first container
if not task_definition_kwargs["containerDefinitions"][0].get("command"):
task_definition_kwargs["containerDefinitions"][0]["command"] = []
task_definition_kwargs.get("containerDefinitions")[0]["command"] = [
"/bin/sh",
"-c",
"python -c 'import prefect; prefect.environments.execution.load_and_run_flow()'",
]
return task_definition_kwargs
def test_get_flow_image_docker_storage():
flow = Flow(
"test",
run_config=UniversalRun(),
storage=Docker(registry_url="test", image_name="name",
image_tag="tag"),
)
image = get_flow_image(flow=flow)
assert image == "test/name:tag"
def execute( # type: ignore
self, flow: "Flow", **kwargs: Any) -> None:
"""
Create a single Kubernetes job that spins up a dask scheduler, dynamically
creates worker pods, and runs the flow.
Args:
- flow (Flow): the Flow object
- **kwargs (Any): additional keyword arguments to pass to the runner
Raises:
- Exception: if the environment is unable to create the Kubernetes job
"""
docker_name = get_flow_image(flow)
from kubernetes import client, config
# Verify environment is running in cluster
try:
config.load_incluster_config()
except config.config_exception.ConfigException:
self.logger.error(
"Environment not currently running inside a cluster")
raise EnvironmentError(
"Environment not currently inside a cluster")
batch_client = client.BatchV1Api()
if self._scheduler_spec:
job = self._scheduler_spec
job = self._populate_scheduler_spec_yaml(yaml_obj=job,
docker_name=docker_name)
else:
with open(path.join(path.dirname(__file__),
"job.yaml")) as job_file:
job = yaml.safe_load(job_file)
job = self._populate_job_yaml(yaml_obj=job,
docker_name=docker_name)
# Create Job
try:
batch_client.create_namespaced_job(
namespace=prefect.context.get("namespace"), body=job)
except Exception as exc:
self.logger.critical(
"Failed to create Kubernetes job: {}".format(exc))
raise exc
def execute( # type: ignore
self, flow: "Flow", **kwargs: Any) -> None:
"""
Run the Fargate task that was defined for this flow.
Args:
- flow (Flow): the Flow object
- **kwargs (Any): additional keyword arguments to pass to the runner
"""
from boto3 import client as boto3_client
flow_run_id = prefect.context.get("flow_run_id", "unknown")
container_overrides = [{
"name":
"flow-container",
"environment": [
{
"name":
"PREFECT__CLOUD__AUTH_TOKEN",
"value":
config.cloud.agent.auth_token or config.cloud.auth_token,
},
{
"name": "PREFECT__CONTEXT__FLOW_RUN_ID",
"value": flow_run_id
},
{
"name": "PREFECT__CONTEXT__IMAGE",
"value": get_flow_image(flow)
},
],
}]
boto3_c = boto3_client(
"ecs",
aws_access_key_id=self.aws_access_key_id,
aws_secret_access_key=self.aws_secret_access_key,
aws_session_token=self.aws_session_token,
region_name=self.region_name,
)
boto3_c.run_task(
overrides={"containerOverrides": container_overrides},
launchType=self.launch_type,
**self.task_run_kwargs,
)
def execute(self, flow: "Flow") -> None: # type: ignore
"""
Create a single Kubernetes job that spins up a dask scheduler, dynamically
creates worker pods, and runs the flow.
Args:
- flow (Flow): the Flow object
Raises:
- Exception: if the environment is unable to create the Kubernetes job
"""
docker_name = get_flow_image(flow)
from kubernetes import client, config
# Verify environment is running in cluster
try:
config.load_incluster_config()
def setup(self, flow: "Flow") -> None: # type: ignore
"""
Register the task definition if it does not already exist.
Args:
- flow (Flow): the Flow object
"""
from boto3 import client as boto3_client
from botocore.exceptions import ClientError
boto3_c = boto3_client(
"ecs",
aws_access_key_id=self.aws_access_key_id,
aws_secret_access_key=self.aws_secret_access_key,
aws_session_token=self.aws_session_token,
region_name=self.region_name,
)
definition_exists = True
try:
boto3_c.describe_task_definition(
taskDefinition=self.task_definition_kwargs.get("family")
)
except ClientError:
definition_exists = False
if not definition_exists:
env_values = [
{"name": "PREFECT__CLOUD__GRAPHQL", "value": config.cloud.graphql},
{"name": "PREFECT__CLOUD__USE_LOCAL_SECRETS", "value": "false"},
{
"name": "PREFECT__ENGINE__FLOW_RUNNER__DEFAULT_CLASS",
"value": "prefect.engine.cloud.CloudFlowRunner",
},
{
"name": "PREFECT__ENGINE__TASK_RUNNER__DEFAULT_CLASS",
"value": "prefect.engine.cloud.CloudTaskRunner",
},
{"name": "PREFECT__LOGGING__LOG_TO_CLOUD", "value": "true"},
{
"name": "PREFECT__LOGGING__EXTRA_LOGGERS",
"value": str(config.logging.extra_loggers),
},
]
# create containerDefinitions if they do not exist
if not self.task_definition_kwargs.get("containerDefinitions"):
self.task_definition_kwargs["containerDefinitions"] = []
self.task_definition_kwargs["containerDefinitions"].append({})
# set environment variables for all containers
for definition in self.task_definition_kwargs["containerDefinitions"]:
if not definition.get("environment"):
definition["environment"] = []
definition["environment"].extend(env_values)
# set name on first container
if not self.task_definition_kwargs["containerDefinitions"][0].get("name"):
self.task_definition_kwargs["containerDefinitions"][0]["name"] = ""
self.task_definition_kwargs.get("containerDefinitions")[0][
"name"
] = "flow-container"
# set image on first container
if not self.task_definition_kwargs["containerDefinitions"][0].get("image"):
self.task_definition_kwargs["containerDefinitions"][0]["image"] = ""
self.task_definition_kwargs.get("containerDefinitions")[0][
"image"
] = get_flow_image(flow)
# set command on first container
if not self.task_definition_kwargs["containerDefinitions"][0].get(
"command"
):
self.task_definition_kwargs["containerDefinitions"][0]["command"] = []
self.task_definition_kwargs.get("containerDefinitions")[0]["command"] = [
"/bin/sh",
"-c",
"python -c 'import prefect; prefect.environments.execution.load_and_run_flow()'",
]
boto3_c.register_task_definition(**self.task_definition_kwargs)
