def test_k8s_run_launcher_default(
dagster_instance_for_k8s_run_launcher,
helm_namespace_for_k8s_run_launcher,
dagit_url_for_k8s_run_launcher,
):
pods = DagsterKubernetesClient.production_client(
).core_api.list_namespaced_pod(
namespace=helm_namespace_for_k8s_run_launcher)
celery_pod_names = [
p.metadata.name for p in pods.items
if "celery-workers" in p.metadata.name
]
check.invariant(not celery_pod_names)
run_config = load_yaml_from_path(
os.path.join(get_test_project_environments_path(), "env.yaml"))
pipeline_name = "demo_pipeline"
run_id = launch_run_over_graphql(dagit_url_for_k8s_run_launcher,
run_config=run_config,
pipeline_name=pipeline_name)
result = wait_for_job_and_get_raw_logs(
job_name="dagster-run-%s" % run_id,
namespace=helm_namespace_for_k8s_run_launcher)
assert "PIPELINE_SUCCESS" in result, "no match, result: {}".format(result)
updated_run = dagster_instance_for_k8s_run_launcher.get_run_by_id(run_id)
assert updated_run.tags[DOCKER_IMAGE_TAG] == get_test_project_docker_image(
)
def test_k8s_run_launcher_default(
dagster_instance_for_k8s_run_launcher, helm_namespace_for_k8s_run_launcher, dagster_docker_image
):
# sanity check that we have a K8sRunLauncher
check.inst(dagster_instance_for_k8s_run_launcher.run_launcher, K8sRunLauncher)
pods = DagsterKubernetesClient.production_client().core_api.list_namespaced_pod(
namespace=helm_namespace_for_k8s_run_launcher
)
celery_pod_names = [p.metadata.name for p in pods.items if "celery-workers" in p.metadata.name]
check.invariant(not celery_pod_names)
run_config = merge_dicts(
load_yaml_from_path(os.path.join(get_test_project_environments_path(), "env.yaml")),
load_yaml_from_path(os.path.join(get_test_project_environments_path(), "env_s3.yaml")),
{
"execution": {
"k8s": {
"config": {
"job_namespace": helm_namespace_for_k8s_run_launcher,
"job_image": dagster_docker_image,
"image_pull_policy": image_pull_policy(),
"env_config_maps": ["dagster-pipeline-env"]
+ ([TEST_AWS_CONFIGMAP_NAME] if not IS_BUILDKITE else []),
}
}
},
},
)
pipeline_name = "demo_k8s_executor_pipeline"
tags = {"key": "value"}
with get_test_project_location_and_external_pipeline(pipeline_name) as (
location,
external_pipeline,
):
run = create_run_for_test(
dagster_instance_for_k8s_run_launcher,
pipeline_name=pipeline_name,
run_config=run_config,
tags=tags,
mode="default",
pipeline_snapshot=external_pipeline.pipeline_snapshot,
execution_plan_snapshot=location.get_external_execution_plan(
external_pipeline, run_config, "default", None, None
).execution_plan_snapshot,
)
dagster_instance_for_k8s_run_launcher.launch_run(
run.run_id,
ReOriginatedExternalPipelineForTest(external_pipeline),
)
result = wait_for_job_and_get_raw_logs(
job_name="dagster-run-%s" % run.run_id, namespace=helm_namespace_for_k8s_run_launcher
)
assert "PIPELINE_SUCCESS" in result, "no match, result: {}".format(result)
updated_run = dagster_instance_for_k8s_run_launcher.get_run_by_id(run.run_id)
assert updated_run.tags[DOCKER_IMAGE_TAG] == get_test_project_docker_image()
def _helm_namespace_helper(helm_chart_fn, request):
"""If an existing Helm chart namespace is specified via pytest CLI with the argument
--existing-helm-namespace, we will use that chart.
Otherwise, provision a test namespace and install Helm chart into that namespace.
Yields the Helm chart namespace.
"""
existing_helm_namespace = request.config.getoption("--existing-helm-namespace")
if existing_helm_namespace:
yield existing_helm_namespace
else:
# Never bother cleaning up on Buildkite
if IS_BUILDKITE:
should_cleanup = False
# Otherwise, always clean up unless --no-cleanup specified
else:
should_cleanup = not request.config.getoption("--no-cleanup")
with test_namespace(should_cleanup) as namespace:
with helm_test_resources(namespace, should_cleanup):
docker_image = get_test_project_docker_image()
with helm_chart_fn(namespace, docker_image, should_cleanup):
print("Helm chart successfully installed in namespace %s" % namespace)
yield namespace
def _launch_executor_run(
dagit_url,
run_config,
dagster_instance_for_k8s_run_launcher,
helm_namespace_for_k8s_run_launcher,
pipeline_name="demo_k8s_executor_pipeline",
num_steps=2,
mode="default",
):
run_id = launch_run_over_graphql(dagit_url,
run_config=run_config,
pipeline_name=pipeline_name,
mode=mode)
result = wait_for_job_and_get_raw_logs(
job_name="dagster-run-%s" % run_id,
namespace=helm_namespace_for_k8s_run_launcher)
assert "PIPELINE_SUCCESS" in result, "no match, result: {}".format(result)
updated_run = dagster_instance_for_k8s_run_launcher.get_run_by_id(run_id)
assert updated_run.tags[DOCKER_IMAGE_TAG] == get_test_project_docker_image(
)
events = dagster_instance_for_k8s_run_launcher.all_logs(run_id)
assert len(_get_step_execution_events(events)) == num_steps
return run_id
def test_valid_job_format(run_launcher):
docker_image = get_test_project_docker_image()
run_config = load_yaml_from_path(
os.path.join(get_test_project_environments_path(), "env.yaml"))
pipeline_name = "demo_pipeline"
run = PipelineRun(pipeline_name=pipeline_name, run_config=run_config)
job_name = "dagster-run-%s" % run.run_id
pod_name = "dagster-run-%s" % run.run_id
job = construct_dagster_k8s_job(
job_config=run_launcher.get_static_job_config(),
args=["dagster", "api", "execute_run"],
job_name=job_name,
pod_name=pod_name,
component="run_coordinator",
)
assert (yaml.dump(
remove_none_recursively(job.to_dict()),
default_flow_style=False).strip() == EXPECTED_JOB_SPEC.format(
run_id=run.run_id,
job_image=docker_image,
image_pull_policy=image_pull_policy(),
dagster_version=dagster_version,
resources="",
env_from=ENV_FROM,
).strip())
def test_terminate_launched_docker_run():
docker_image = get_test_project_docker_image()
launcher_config = {
"env_vars": ["AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY",],
"network": "container:test-postgres-db-docker",
}
if IS_BUILDKITE:
launcher_config["registry"] = get_buildkite_registry_config()
else:
find_local_test_image(docker_image)
run_config = merge_yamls([os.path.join(get_test_project_environments_path(), "env_s3.yaml"),])
with docker_postgres_instance(
overrides={
"run_launcher": {
"class": "DockerRunLauncher",
"module": "dagster_docker",
"config": launcher_config,
}
}
) as instance:
recon_pipeline = get_test_project_recon_pipeline("hanging_pipeline", docker_image)
run = instance.create_run_for_pipeline(
pipeline_def=recon_pipeline.get_definition(), run_config=run_config,
)
run_id = run.run_id
external_pipeline = ReOriginatedExternalPipelineForTest(
get_test_project_external_pipeline("hanging_pipeline", container_image=docker_image),
container_image=docker_image,
)
instance.launch_run(run_id, external_pipeline)
poll_for_step_start(instance, run_id)
assert instance.run_launcher.can_terminate(run_id)
assert instance.run_launcher.terminate(run_id)
terminated_pipeline_run = poll_for_finished_run(instance, run_id, timeout=30)
terminated_pipeline_run = instance.get_run_by_id(run_id)
assert terminated_pipeline_run.status == PipelineRunStatus.CANCELED
run_logs = instance.all_logs(run_id)
_check_event_log_contains(
run_logs,
[
("PIPELINE_CANCELING", "Sending pipeline termination request"),
("STEP_FAILURE", 'Execution of step "hanging_solid" failed.'),
("PIPELINE_CANCELED", 'Execution of pipeline "hanging_pipeline" canceled.'),
("ENGINE_EVENT", "Pipeline execution terminated by interrupt"),
("ENGINE_EVENT", "Process for pipeline exited"),
],
)
def test_launch_docker_image_on_pipeline_config():
# Docker image name to use for launch specified as part of the pipeline origin
# rather than in the run launcher instance config
docker_image = get_test_project_docker_image()
launcher_config = {
"env_vars": [
"AWS_ACCESS_KEY_ID",
"AWS_SECRET_ACCESS_KEY",
],
"network": "container:test-postgres-db-docker",
"container_kwargs": {
"auto_remove": True,
},
}
if IS_BUILDKITE:
launcher_config["registry"] = get_buildkite_registry_config()
else:
find_local_test_image(docker_image)
run_config = merge_yamls([
os.path.join(get_test_project_environments_path(), "env.yaml"),
os.path.join(get_test_project_environments_path(), "env_s3.yaml"),
])
with docker_postgres_instance(
overrides={
"run_launcher": {
"class": "DockerRunLauncher",
"module": "dagster_docker",
"config": launcher_config,
}
}) as instance:
recon_pipeline = get_test_project_recon_pipeline(
"demo_pipeline", docker_image)
run = instance.create_run_for_pipeline(
pipeline_def=recon_pipeline.get_definition(),
run_config=run_config,
)
with get_test_project_external_pipeline(
"demo_pipeline",
container_image=docker_image) as orig_pipeline:
external_pipeline = ReOriginatedExternalPipelineForTest(
orig_pipeline,
container_image=docker_image,
)
instance.launch_run(run.run_id, external_pipeline)
poll_for_finished_run(instance, run.run_id, timeout=60)
run = instance.get_run_by_id(run.run_id)
assert run.status == PipelineRunStatus.SUCCESS
assert run.tags[DOCKER_IMAGE_TAG] == docker_image
def _cluster_provider(request):
from .kind import kind_cluster_exists, kind_cluster, kind_load_images
if IS_BUILDKITE:
print("Installing ECR credentials...")
check_output(
"aws ecr get-login --no-include-email --region us-west-2 | sh",
shell=True)
provider = request.config.getoption("--cluster-provider")
# Use a kind cluster
if provider == "kind":
cluster_name = request.config.getoption("--kind-cluster")
# Cluster will be deleted afterwards unless this is set.
# This is to allow users to reuse an existing cluster in local test by running
# `pytest --kind-cluster my-cluster --no-cleanup` -- this avoids the per-test run
# overhead of cluster setup and teardown
should_cleanup = True if IS_BUILDKITE else not request.config.getoption(
"--no-cleanup")
existing_cluster = kind_cluster_exists(cluster_name)
with kind_cluster(cluster_name,
should_cleanup=should_cleanup) as cluster_config:
if not IS_BUILDKITE and not existing_cluster:
docker_image = get_test_project_docker_image()
try:
client = docker.from_env()
client.images.get(docker_image)
print( # pylint: disable=print-call
"Found existing image tagged {image}, skipping image build. To rebuild, first run: "
"docker rmi {image}".format(image=docker_image))
except docker.errors.ImageNotFound:
build_and_tag_test_image(docker_image)
kind_load_images(
cluster_name=cluster_config.name,
local_dagster_test_image=docker_image,
additional_images=additional_kind_images,
)
yield cluster_config
# Use cluster from kubeconfig
elif provider == "kubeconfig":
kubeconfig_file = os.getenv(
"KUBECONFIG", os.path.expandvars("${HOME}/.kube/config"))
kubernetes.config.load_kube_config(config_file=kubeconfig_file)
yield ClusterConfig(name="from_system_kubeconfig",
kubeconfig_file=kubeconfig_file)
else:
raise Exception("unknown cluster provider %s" % provider)
def test_terminate():
docker_image = get_test_project_docker_image()
launcher_config = {
"env_vars": [
"AWS_ACCESS_KEY_ID",
"AWS_SECRET_ACCESS_KEY",
],
"network": "container:test-postgres-db-docker",
"image": docker_image,
}
_test_launch(docker_image, launcher_config, terminate=True)
def test_launch_docker_image_on_instance_config():
docker_image = get_test_project_docker_image()
launcher_config = {
"env_vars": [
"AWS_ACCESS_KEY_ID",
"AWS_SECRET_ACCESS_KEY",
],
"network": "container:test-postgres-db-docker",
"image": docker_image,
}
if IS_BUILDKITE:
launcher_config["registry"] = get_buildkite_registry_config()
else:
find_local_test_image(docker_image)
run_config = merge_yamls([
os.path.join(get_test_project_environments_path(), "env.yaml"),
os.path.join(get_test_project_environments_path(), "env_s3.yaml"),
])
with docker_postgres_instance(
overrides={
"run_launcher": {
"class": "DockerRunLauncher",
"module": "dagster_docker",
"config": launcher_config,
}
}) as instance:
recon_pipeline = get_test_project_recon_pipeline("demo_pipeline")
with get_test_project_workspace_and_external_pipeline(
instance, "demo_pipeline") as (
workspace,
orig_pipeline,
):
external_pipeline = ReOriginatedExternalPipelineForTest(
orig_pipeline)
run = instance.create_run_for_pipeline(
pipeline_def=recon_pipeline.get_definition(),
run_config=run_config,
external_pipeline_origin=external_pipeline.get_external_origin(
),
pipeline_code_origin=external_pipeline.get_python_origin(),
)
instance.launch_run(run.run_id, workspace)
poll_for_finished_run(instance, run.run_id, timeout=60)
assert instance.get_run_by_id(
run.run_id).status == PipelineRunStatus.SUCCESS
def _launch_executor_run(
run_config,
dagster_instance_for_k8s_run_launcher,
helm_namespace_for_k8s_run_launcher,
):
pipeline_name = "demo_k8s_executor_pipeline"
tags = {"key": "value"}
with get_test_project_external_pipeline_hierarchy(
dagster_instance_for_k8s_run_launcher, pipeline_name) as (
workspace,
location,
_repo,
external_pipeline,
):
reoriginated_pipeline = ReOriginatedExternalPipelineForTest(
external_pipeline)
run = create_run_for_test(
dagster_instance_for_k8s_run_launcher,
pipeline_name=pipeline_name,
run_config=run_config,
tags=tags,
mode="default",
pipeline_snapshot=external_pipeline.pipeline_snapshot,
execution_plan_snapshot=location.get_external_execution_plan(
external_pipeline, run_config, "default", None,
None).execution_plan_snapshot,
external_pipeline_origin=reoriginated_pipeline.get_external_origin(
),
pipeline_code_origin=reoriginated_pipeline.get_python_origin(),
)
dagster_instance_for_k8s_run_launcher.launch_run(run.run_id, workspace)
result = wait_for_job_and_get_raw_logs(
job_name="dagster-run-%s" % run.run_id,
namespace=helm_namespace_for_k8s_run_launcher)
assert "PIPELINE_SUCCESS" in result, "no match, result: {}".format(
result)
updated_run = dagster_instance_for_k8s_run_launcher.get_run_by_id(
run.run_id)
assert updated_run.tags[
DOCKER_IMAGE_TAG] == get_test_project_docker_image()
events = dagster_instance_for_k8s_run_launcher.all_logs(run.run_id)
assert (len([
event for event in events
if ("Executing step" in event.message
and "in Kubernetes job" in event.message)
]) == 2)
def test_execute_celery_docker_image_on_pipeline_config(aws_creds):
docker_image = get_test_project_docker_image()
docker_config = {
"network": "container:test-postgres-db-celery-docker",
"container_kwargs": {
"environment": [
"FIND_ME=here!",
f"AWS_ACCESS_KEY_ID={aws_creds['aws_access_key_id']}",
f"AWS_SECRET_ACCESS_KEY={aws_creds['aws_secret_access_key']}",
],
# "auto_remove": False # uncomment when debugging to view container logs after execution
},
}
if IS_BUILDKITE:
docker_config["registry"] = get_buildkite_registry_config()
else:
find_local_test_image(docker_image)
run_config = merge_dicts(
merge_yamls([
os.path.join(get_test_project_environments_path(), "env.yaml"),
os.path.join(get_test_project_environments_path(), "env_s3.yaml"),
os.path.join(get_test_project_environments_path(),
"env_environment_vars.yaml"),
]),
{
"execution": {
"celery-docker": {
"config": {
"docker": docker_config,
"config_source": {
"task_always_eager": True
},
}
}
},
},
)
with celery_docker_postgres_instance() as instance:
result = execute_pipeline(
get_test_project_recon_pipeline("docker_celery_pipeline",
docker_image),
run_config=run_config,
instance=instance,
)
assert result.success
assert result.result_for_solid("get_environment_solid").output_value(
"result") == "here!"
def dagster_docker_image():
docker_image = get_test_project_docker_image()
if not IS_BUILDKITE:
try:
client = docker.from_env()
client.images.get(docker_image)
print( # pylint: disable=print-call
"Found existing image tagged {image}, skipping image build. To rebuild, first run: "
"docker rmi {image}".format(image=docker_image))
except docker.errors.ImageNotFound:
build_and_tag_test_image(docker_image)
return docker_image
def test_launch_docker_image_multiple_networks():
docker_image = get_test_project_docker_image()
launcher_config = {
"env_vars": [
"AWS_ACCESS_KEY_ID",
"AWS_SECRET_ACCESS_KEY",
],
"networks": [
"container:test-postgres-db-docker",
"postgres",
],
"image": docker_image,
}
_test_launch(docker_image, launcher_config)
def test_k8s_executor_resource_requirements(
dagster_instance_for_k8s_run_launcher,
helm_namespace_for_k8s_run_launcher,
dagster_docker_image,
dagit_url_for_k8s_run_launcher,
):
# sanity check that we have a K8sRunLauncher
pods = DagsterKubernetesClient.production_client(
).core_api.list_namespaced_pod(
namespace=helm_namespace_for_k8s_run_launcher)
celery_pod_names = [
p.metadata.name for p in pods.items
if "celery-workers" in p.metadata.name
]
check.invariant(not celery_pod_names)
run_config = merge_dicts(
load_yaml_from_path(
os.path.join(get_test_project_environments_path(), "env_s3.yaml")),
{
"execution": {
"k8s": {
"config": {
"job_namespace": helm_namespace_for_k8s_run_launcher,
"job_image": dagster_docker_image,
"image_pull_policy": image_pull_policy(),
}
}
},
},
)
pipeline_name = "resources_limit_pipeline"
run_id = launch_run_over_graphql(
dagit_url_for_k8s_run_launcher,
run_config=run_config,
pipeline_name=pipeline_name,
mode="k8s",
)
result = wait_for_job_and_get_raw_logs(
job_name="dagster-run-%s" % run_id,
namespace=helm_namespace_for_k8s_run_launcher)
assert "PIPELINE_SUCCESS" in result, "no match, result: {}".format(result)
updated_run = dagster_instance_for_k8s_run_launcher.get_run_by_id(run_id)
assert updated_run.tags[DOCKER_IMAGE_TAG] == get_test_project_docker_image(
)
def test_valid_job_format_with_backcompat_resources(run_launcher):
docker_image = get_test_project_docker_image()
run_config = load_yaml_from_path(
os.path.join(get_test_project_environments_path(), "env.yaml"))
pipeline_name = "demo_pipeline"
run = PipelineRun(pipeline_name=pipeline_name, run_config=run_config)
tags = validate_tags({
K8S_RESOURCE_REQUIREMENTS_KEY: ({
"requests": {
"cpu": "250m",
"memory": "64Mi"
},
"limits": {
"cpu": "500m",
"memory": "2560Mi"
},
})
})
user_defined_k8s_config = get_user_defined_k8s_config(tags)
job_name = "dagster-run-%s" % run.run_id
pod_name = "dagster-run-%s" % run.run_id
job = construct_dagster_k8s_job(
job_config=run_launcher.get_static_job_config(),
args=["dagster", "api", "execute_run"],
job_name=job_name,
user_defined_k8s_config=user_defined_k8s_config,
pod_name=pod_name,
component="run_coordinator",
)
assert (yaml.dump(
remove_none_recursively(job.to_dict()),
default_flow_style=False).strip() == EXPECTED_JOB_SPEC.format(
run_id=run.run_id,
job_image=docker_image,
image_pull_policy=image_pull_policy(),
dagster_version=dagster_version,
env_from=ENV_FROM,
resources="""
resources:
limits:
cpu: 500m
memory: 2560Mi
requests:
cpu: 250m
memory: 64Mi""",
).strip())
def test_k8s_run_launcher_image_from_origin(
dagster_instance_for_k8s_run_launcher,
helm_namespace_for_k8s_run_launcher,
dagit_url_for_k8s_run_launcher,
):
# Like the previous test, but the executor doesn't supply an image - it's pulled
# from the origin on the run instead
pods = DagsterKubernetesClient.production_client(
).core_api.list_namespaced_pod(
namespace=helm_namespace_for_k8s_run_launcher)
celery_pod_names = [
p.metadata.name for p in pods.items
if "celery-workers" in p.metadata.name
]
check.invariant(not celery_pod_names)
run_config = merge_dicts(
load_yaml_from_path(
os.path.join(get_test_project_environments_path(), "env.yaml")),
load_yaml_from_path(
os.path.join(get_test_project_environments_path(), "env_s3.yaml")),
{
"execution": {
"k8s": {
"config": {
"job_namespace": helm_namespace_for_k8s_run_launcher,
"image_pull_policy": image_pull_policy(),
}
}
},
},
)
pipeline_name = "demo_k8s_executor_pipeline"
run_id = launch_run_over_graphql(dagit_url_for_k8s_run_launcher,
run_config=run_config,
pipeline_name=pipeline_name)
result = wait_for_job_and_get_raw_logs(
job_name="dagster-run-%s" % run_id,
namespace=helm_namespace_for_k8s_run_launcher)
assert "PIPELINE_SUCCESS" in result, "no match, result: {}".format(result)
updated_run = dagster_instance_for_k8s_run_launcher.get_run_by_id(run_id)
assert updated_run.tags[DOCKER_IMAGE_TAG] == get_test_project_docker_image(
)
def test_k8s_run_launcher_default(
dagster_instance_for_k8s_run_launcher, helm_namespace_for_k8s_run_launcher
):
# sanity check that we have a K8sRunLauncher
check.inst(dagster_instance_for_k8s_run_launcher.run_launcher, K8sRunLauncher)
pods = DagsterKubernetesClient.production_client().core_api.list_namespaced_pod(
namespace=helm_namespace_for_k8s_run_launcher
)
celery_pod_names = [p.metadata.name for p in pods.items if "celery-workers" in p.metadata.name]
check.invariant(not celery_pod_names)
run_config = load_yaml_from_path(os.path.join(get_test_project_environments_path(), "env.yaml"))
pipeline_name = "demo_pipeline"
tags = {"key": "value"}
with get_test_project_external_pipeline_hierarchy(
dagster_instance_for_k8s_run_launcher, pipeline_name
) as (
workspace,
location,
_repo,
external_pipeline,
):
reoriginated_pipeline = ReOriginatedExternalPipelineForTest(external_pipeline)
run = create_run_for_test(
dagster_instance_for_k8s_run_launcher,
pipeline_name=pipeline_name,
run_config=run_config,
tags=tags,
mode="default",
pipeline_snapshot=external_pipeline.pipeline_snapshot,
execution_plan_snapshot=location.get_external_execution_plan(
external_pipeline, run_config, "default", None, None
).execution_plan_snapshot,
external_pipeline_origin=reoriginated_pipeline.get_external_origin(),
pipeline_code_origin=reoriginated_pipeline.get_python_origin(),
)
dagster_instance_for_k8s_run_launcher.launch_run(run.run_id, workspace)
result = wait_for_job_and_get_raw_logs(
job_name="dagster-run-%s" % run.run_id, namespace=helm_namespace_for_k8s_run_launcher
)
assert "PIPELINE_SUCCESS" in result, "no match, result: {}".format(result)
updated_run = dagster_instance_for_k8s_run_launcher.get_run_by_id(run.run_id)
assert updated_run.tags[DOCKER_IMAGE_TAG] == get_test_project_docker_image()
def k8s_scheduler(cluster_provider, helm_namespace_for_k8s_run_launcher): # pylint: disable=redefined-outer-name,unused-argument
return K8sScheduler(
scheduler_namespace=helm_namespace_for_k8s_run_launcher,
image_pull_secrets=[{
"name": "element-dev-key"
}],
service_account_name="dagit-admin",
instance_config_map="dagster-instance",
postgres_password_secret="dagster-postgresql-secret",
dagster_home="/opt/dagster/dagster_home",
job_image=get_test_project_docker_image(),
load_incluster_config=False,
kubeconfig_file=cluster_provider.kubeconfig_file,
image_pull_policy=image_pull_policy(),
env_config_maps=["dagster-pipeline-env", "test-env-configmap"],
env_secrets=["test-env-secret"],
)
def test_docker_executor():
"""
Note that this test relies on having AWS credentials in the environment.
"""
executor_config = {
"execution": {
"docker": {
"config": {
"networks": ["container:test-postgres-db-docker"],
"env_vars": [
"AWS_ACCESS_KEY_ID",
"AWS_SECRET_ACCESS_KEY",
],
}
}
}
}
docker_image = get_test_project_docker_image()
if IS_BUILDKITE:
executor_config["execution"]["docker"]["config"][
"registry"
] = get_buildkite_registry_config()
else:
find_local_test_image(docker_image)
run_config = merge_dicts(
merge_yamls(
[
os.path.join(get_test_project_environments_path(), "env.yaml"),
os.path.join(get_test_project_environments_path(), "env_s3.yaml"),
]
),
executor_config,
)
with environ({"DOCKER_LAUNCHER_NETWORK": "container:test-postgres-db-docker"}):
with docker_postgres_instance() as instance:
recon_pipeline = get_test_project_recon_pipeline("demo_pipeline_docker", docker_image)
assert execute_pipeline(
recon_pipeline, run_config=run_config, instance=instance
).success
def test_execute_celery_docker_image_on_pipeline_config():
docker_image = get_test_project_docker_image()
docker_config = {
"env_vars": [
"AWS_ACCESS_KEY_ID",
"AWS_SECRET_ACCESS_KEY",
],
"network": "container:test-postgres-db-celery-docker",
}
if IS_BUILDKITE:
docker_config["registry"] = get_buildkite_registry_config()
else:
find_local_test_image(docker_image)
run_config = merge_dicts(
merge_yamls([
os.path.join(get_test_project_environments_path(), "env.yaml"),
os.path.join(get_test_project_environments_path(), "env_s3.yaml"),
]),
{
"execution": {
"celery-docker": {
"config": {
"docker": docker_config,
"config_source": {
"task_always_eager": True
},
}
}
},
},
)
with celery_docker_postgres_instance() as instance:
result = execute_pipeline(
get_test_project_recon_pipeline("docker_celery_pipeline",
docker_image),
run_config=run_config,
instance=instance,
)
assert result.success
def test_launch_docker_config_on_container_context():
docker_image = get_test_project_docker_image()
launcher_config = {}
_test_launch(
docker_image,
launcher_config,
container_image=docker_image,
container_context={
"docker": {
"env_vars": [
"AWS_ACCESS_KEY_ID",
"AWS_SECRET_ACCESS_KEY",
],
"networks": [
"container:test-postgres-db-docker",
"postgres",
],
}
},
)
def test_cant_combine_network_and_networks():
docker_image = get_test_project_docker_image()
launcher_config = {
"env_vars": [
"AWS_ACCESS_KEY_ID",
"AWS_SECRET_ACCESS_KEY",
],
"network": "container:test-postgres-db-docker",
"networks": [
"postgres",
],
"image": docker_image,
}
with pytest.raises(Exception, match="cannot set both `network` and `networks`"):
with docker_postgres_instance(
overrides={
"run_launcher": {
"class": "DockerRunLauncher",
"module": "dagster_docker",
"config": launcher_config,
}
}
):
pass
def test_valid_job_format_with_user_defined_k8s_config(run_launcher):
docker_image = get_test_project_docker_image()
run_config = load_yaml_from_path(
os.path.join(get_test_project_environments_path(), "env.yaml"))
pipeline_name = "demo_pipeline"
run = PipelineRun(pipeline_name=pipeline_name, run_config=run_config)
backoff_limit = 1234
ttl_seconds_after_finished = 5678
tags = validate_tags({
USER_DEFINED_K8S_CONFIG_KEY: ({
"container_config": {
"resources": {
"requests": {
"cpu": "250m",
"memory": "64Mi"
},
"limits": {
"cpu": "500m",
"memory": "2560Mi"
},
}
},
"pod_template_spec_metadata": {
"annotations": {
"cluster-autoscaler.kubernetes.io/safe-to-evict": "true"
},
"labels": {
"spotinst.io/restrict-scale-down": "true"
},
},
"job_spec_config": {
"ttl_seconds_after_finished": ttl_seconds_after_finished,
"backoff_limit": backoff_limit,
},
"pod_spec_config": {
"affinity": {
"nodeAffinity": {
"requiredDuringSchedulingIgnoredDuringExecution": {
"nodeSelectorTerms": [{
"matchExpressions": [{
"key":
"kubernetes.io/e2e-az-name",
"operator":
"In",
"values": ["e2e-az1", "e2e-az2"],
}]
}]
}
}
}
},
})
})
user_defined_k8s_config = get_user_defined_k8s_config(tags)
job_name = "dagster-run-%s" % run.run_id
pod_name = "dagster-run-%s" % run.run_id
job = construct_dagster_k8s_job(
job_config=run_launcher.get_static_job_config(),
args=["dagster", "api", "execute_run"],
job_name=job_name,
user_defined_k8s_config=user_defined_k8s_config,
pod_name=pod_name,
component="run_coordinator",
)
assert (yaml.dump(remove_none_recursively(job.to_dict()),
default_flow_style=False).strip() ==
EXPECTED_CONFIGURED_JOB_SPEC.format(
run_id=run.run_id,
job_image=docker_image,
image_pull_policy=image_pull_policy(),
dagster_version=dagster_version,
labels="spotinst.io/restrict-scale-down: 'true'",
env_from=ENV_FROM,
backoff_limit=backoff_limit,
ttl_seconds_after_finished=ttl_seconds_after_finished,
resources="""
resources:
limits:
cpu: 500m
memory: 2560Mi
requests:
cpu: 250m
memory: 64Mi""",
annotations="""annotations:
cluster-autoscaler.kubernetes.io/safe-to-evict: \'true\'""",
affinity="""affinity:
nodeAffinity:
requiredDuringSchedulingIgnoredDuringExecution:
nodeSelectorTerms:
- matchExpressions:
- key: kubernetes.io/e2e-az-name
operator: In
values:
- e2e-az1
- e2e-az2""",
).strip())
def test_launch_docker_no_network():
docker_image = get_test_project_docker_image()
launcher_config = {
"env_vars": [
"AWS_ACCESS_KEY_ID",
"AWS_SECRET_ACCESS_KEY",
],
}
if IS_BUILDKITE:
launcher_config["registry"] = get_buildkite_registry_config()
else:
find_local_test_image(docker_image)
run_config = merge_yamls(
[
os.path.join(get_test_project_environments_path(), "env.yaml"),
os.path.join(get_test_project_environments_path(), "env_s3.yaml"),
]
)
with docker_postgres_instance(
overrides={
"run_launcher": {
"class": "DockerRunLauncher",
"module": "dagster_docker",
"config": launcher_config,
}
},
# Ensure the container will time out and fail quickly
conn_args={
"params": {"connect_timeout": 2},
},
) as instance:
recon_pipeline = get_test_project_recon_pipeline("demo_pipeline_s3", docker_image)
with get_test_project_workspace_and_external_pipeline(
instance, "demo_pipeline_s3", container_image=docker_image
) as (workspace, orig_pipeline):
external_pipeline = ReOriginatedExternalPipelineForTest(
orig_pipeline,
container_image=docker_image,
)
run = instance.create_run_for_pipeline(
pipeline_def=recon_pipeline.get_definition(),
run_config=run_config,
external_pipeline_origin=external_pipeline.get_external_origin(),
pipeline_code_origin=external_pipeline.get_python_origin(),
)
instance.launch_run(run.run_id, workspace)
# Container launches, but run is stuck in STARTING state
# due to not being able to access the network
run = instance.get_run_by_id(run.run_id)
assert run.tags[DOCKER_IMAGE_TAG] == docker_image
container_id = run.tags[DOCKER_CONTAINER_ID_TAG]
run = instance.get_run_by_id(run.run_id)
assert run.status == PipelineRunStatus.STARTING
assert run.tags[DOCKER_IMAGE_TAG] == docker_image
client = docker.client.from_env()
container = None
try:
start_time = time.time()
while True:
container = client.containers.get(container_id)
if time.time() - start_time > 60:
raise Exception("Timed out waiting for container to exit")
if container.status == "exited":
break
time.sleep(3)
finally:
if container:
container.remove(force=True)
def test_docker_monitoring():
docker_image = get_test_project_docker_image()
launcher_config = {
"env_vars": [
"AWS_ACCESS_KEY_ID",
"AWS_SECRET_ACCESS_KEY",
],
"networks": ["container:test-postgres-db-docker"],
"container_kwargs": {
# "auto_remove": True,
"volumes": ["/var/run/docker.sock:/var/run/docker.sock"],
},
}
if IS_BUILDKITE:
launcher_config["registry"] = get_buildkite_registry_config()
else:
find_local_test_image(docker_image)
run_config = merge_dicts(
load_yaml_from_path(
os.path.join(get_test_project_environments_path(), "env_s3.yaml")),
{
"solids": {
"multiply_the_word_slow": {
"inputs": {
"word": "bar"
},
"config": {
"factor": 2,
"sleep_time": 20
},
}
},
"execution": {
"docker": {
"config": {}
}
},
},
)
with docker_postgres_instance({
"run_monitoring": {
"enabled": True
},
"run_launcher": {
"class": "DockerRunLauncher",
"module": "dagster_docker",
"config": launcher_config,
},
}) as instance:
recon_pipeline = get_test_project_recon_pipeline(
"demo_pipeline_docker_slow", docker_image)
with get_test_project_workspace_and_external_pipeline(
instance,
"demo_pipeline_docker_slow",
container_image=docker_image) as (
workspace,
orig_pipeline,
):
with start_daemon():
external_pipeline = ReOriginatedExternalPipelineForTest(
orig_pipeline, container_image=docker_image)
run = instance.create_run_for_pipeline(
pipeline_def=recon_pipeline.get_definition(),
run_config=run_config,
external_pipeline_origin=external_pipeline.
get_external_origin(),
pipeline_code_origin=external_pipeline.get_python_origin(),
)
with log_run_events(instance, run.run_id):
instance.launch_run(run.run_id, workspace)
start_time = time.time()
while time.time() - start_time < 60:
run = instance.get_run_by_id(run.run_id)
if run.status == PipelineRunStatus.STARTED:
break
assert run.status == PipelineRunStatus.STARTING
time.sleep(1)
time.sleep(3)
instance.run_launcher._get_container( # pylint:disable=protected-access
instance.get_run_by_id(run.run_id)).stop()
# daemon resumes the run
poll_for_finished_run(instance, run.run_id, timeout=90)
assert instance.get_run_by_id(
run.run_id).status == PipelineRunStatus.SUCCESS
def test_execute_on_celery_k8s( # pylint: disable=redefined-outer-name,unused-argument
dagster_instance_for_user_deployments,
helm_namespace_for_user_deployments,
):
namespace = helm_namespace_for_user_deployments
pipeline_name = "demo_pipeline_celery"
core_api = kubernetes.client.CoreV1Api()
batch_api = kubernetes.client.BatchV1Api()
# Get name for dagit pod
pods = core_api.list_namespaced_pod(namespace=namespace)
dagit_pod_list = list(filter(lambda item: "dagit" in item.metadata.name, pods.items))
assert len(dagit_pod_list) == 1
dagit_pod = dagit_pod_list[0]
dagit_pod_name = dagit_pod.metadata.name
# Check that there are no run master jobs
jobs = batch_api.list_namespaced_job(namespace=namespace)
runmaster_job_list = list(filter(lambda item: "dagster-run-" in item.metadata.name, jobs.items))
assert len(runmaster_job_list) == 0
run_config_dict = {
"storage": {"s3": {"config": {"s3_bucket": "dagster-scratch-80542c2"}}},
"execution": {
"celery-k8s": {
"config": {
"image_pull_policy": image_pull_policy(),
"env_config_maps": ["dagster-pipeline-env"]
+ ([TEST_AWS_CONFIGMAP_NAME] if not IS_BUILDKITE else []),
"job_namespace": namespace,
}
}
},
"loggers": {"console": {"config": {"log_level": "DEBUG"}}},
"solids": {"multiply_the_word": {"inputs": {"word": "bar"}, "config": {"factor": 2}}},
}
run_config_json = json.dumps(run_config_dict)
exec_command = [
"dagster",
"pipeline",
"launch",
"--repository",
"demo_execution_repo",
"--pipeline",
pipeline_name,
"--workspace",
"/dagster-workspace/workspace.yaml",
"--location",
"user-code-deployment-1",
"--config-json",
run_config_json,
"--run-id",
"foobar",
]
stream(
core_api.connect_get_namespaced_pod_exec,
name=dagit_pod_name,
namespace=namespace,
command=exec_command,
stderr=True,
stdin=False,
stdout=True,
tty=False,
_preload_content=False,
)
runmaster_job_name = None
timeout = datetime.timedelta(0, 90)
start_time = datetime.datetime.now()
while datetime.datetime.now() < start_time + timeout and not runmaster_job_name:
jobs = batch_api.list_namespaced_job(namespace=namespace)
runmaster_job_list = list(
filter(lambda item: "dagster-run-" in item.metadata.name, jobs.items)
)
if len(runmaster_job_list) > 0:
runmaster_job_name = runmaster_job_list[0].metadata.name
assert runmaster_job_name
result = wait_for_job_and_get_raw_logs(
job_name=runmaster_job_name, namespace=namespace, wait_timeout=450
)
assert "PIPELINE_SUCCESS" in result, "no match, result: {}".format(result)
updated_run = dagster_instance_for_user_deployments.get_run_by_id("foobar")
assert updated_run.tags[DOCKER_IMAGE_TAG] == get_test_project_docker_image()
def test_container_context_on_pipeline():
docker_image = get_test_project_docker_image()
launcher_config = {}
if IS_BUILDKITE:
launcher_config["registry"] = get_buildkite_registry_config()
else:
find_local_test_image(docker_image)
executor_config = {
"execution": {
"docker": {
"config": {}
}
},
}
run_config = merge_dicts(
merge_yamls([
os.path.join(get_test_project_environments_path(), "env.yaml"),
os.path.join(get_test_project_environments_path(), "env_s3.yaml"),
]),
executor_config,
)
with docker_postgres_instance(
overrides={
"run_launcher": {
"class": "DockerRunLauncher",
"module": "dagster_docker",
"config": launcher_config,
}
}) as instance:
recon_pipeline = get_test_project_recon_pipeline(
"demo_pipeline_docker",
docker_image,
container_context={
"docker": {
"env_vars": [
"AWS_ACCESS_KEY_ID",
"AWS_SECRET_ACCESS_KEY",
],
"networks": ["container:test-postgres-db-docker"],
"container_kwargs": {
"auto_remove": True,
"volumes":
["/var/run/docker.sock:/var/run/docker.sock"],
},
}
},
)
with get_test_project_workspace_and_external_pipeline(
instance, "demo_pipeline_docker",
container_image=docker_image) as (
workspace,
orig_pipeline,
):
external_pipeline = ReOriginatedExternalPipelineForTest(
orig_pipeline, container_image=docker_image)
run = instance.create_run_for_pipeline(
pipeline_def=recon_pipeline.get_definition(),
run_config=run_config,
external_pipeline_origin=external_pipeline.get_external_origin(
),
pipeline_code_origin=recon_pipeline.get_python_origin(),
)
instance.launch_run(run.run_id, workspace)
poll_for_finished_run(instance, run.run_id, timeout=60)
for log in instance.all_logs(run.run_id):
print(log) # pylint: disable=print-call
assert instance.get_run_by_id(
run.run_id).status == PipelineRunStatus.SUCCESS
