def helm_namespace(
cluster_provider, request
): # pylint: disable=unused-argument, redefined-outer-name
'''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 = test_project_docker_image()
with helm_chart(namespace, docker_image, should_cleanup):
print('Helm chart successfully installed in namespace %s' % namespace)
yield namespace
def test_valid_job_format(run_launcher):
docker_image = test_project_docker_image()
run_config = load_yaml_from_path(
os.path.join(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.job_config,
command=['dagster-graphql'],
args=[
'-p',
'executeRunInProcess',
'-v',
seven.json.dumps({'runId': run.run_id}),
],
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='',
).strip())
def test_valid_job_format(run_launcher):
docker_image = test_project_docker_image()
run_config = load_yaml_from_path(
os.path.join(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.job_config,
command=["dagster"],
args=["api", "execute_run_with_structured_logs"],
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="",
).strip())
def test_valid_job_format(run_launcher):
docker_image = test_project_docker_image()
environment_dict = load_yaml_from_path(
os.path.join(test_project_environments_path(), 'env.yaml'))
pipeline_name = 'demo_pipeline'
run = PipelineRun(pipeline_name=pipeline_name,
environment_dict=environment_dict)
job_name = 'dagster-run-%s' % run.run_id
pod_name = 'dagster-run-%s' % run.run_id
job = construct_dagster_graphql_k8s_job(
run_launcher.job_config,
args=[
'-p',
'startPipelineExecutionForCreatedRun',
'-v',
seven.json.dumps({'runId': run.run_id}),
],
job_name=job_name,
pod_name=pod_name,
component='runmaster',
)
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,
).strip())
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 = 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 test_execute_celery_docker():
docker_image = test_project_docker_image()
docker_config = {
"image": docker_image,
"env_vars": ["AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY"],
}
if IS_BUILDKITE:
ecr_client = boto3.client("ecr", region_name="us-west-1")
token = ecr_client.get_authorization_token()
username, password = (base64.b64decode(
token["authorizationData"][0]
["authorizationToken"]).decode().split(":"))
registry = token["authorizationData"][0]["proxyEndpoint"]
docker_config["registry"] = {
"url": registry,
"username": username,
"password": password,
}
else:
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)
with seven.TemporaryDirectory() as temp_dir:
run_config = merge_dicts(
merge_yamls([
os.path.join(test_project_environments_path(), "env.yaml"),
os.path.join(test_project_environments_path(), "env_s3.yaml"),
]),
{
"execution": {
"celery-docker": {
"config": {
"docker": docker_config,
"config_source": {
"task_always_eager": True
},
}
}
},
},
)
result = execute_pipeline(
get_test_project_recon_pipeline("docker_celery_pipeline"),
run_config=run_config,
instance=DagsterInstance.local_temp(temp_dir),
)
assert result.success
def test_execute_celery_docker():
docker_image = test_project_docker_image()
docker_config = {
'image': docker_image,
'env_vars': ['AWS_ACCESS_KEY_ID', 'AWS_SECRET_ACCESS_KEY'],
}
if IS_BUILDKITE:
ecr_client = boto3.client('ecr', region_name='us-west-1')
token = ecr_client.get_authorization_token()
username, password = (base64.b64decode(
token['authorizationData'][0]
['authorizationToken']).decode().split(':'))
registry = token['authorizationData'][0]['proxyEndpoint']
docker_config['registry'] = {
'url': registry,
'username': username,
'password': password,
}
else:
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)
with seven.TemporaryDirectory() as temp_dir:
run_config = merge_dicts(
merge_yamls([
os.path.join(test_project_environments_path(), 'env.yaml'),
os.path.join(test_project_environments_path(), 'env_s3.yaml'),
]),
{
'execution': {
'celery-docker': {
'config': {
'docker': docker_config,
'config_source': {
'task_always_eager': True
},
}
}
},
},
)
result = execute_pipeline(
get_test_project_recon_pipeline('docker_celery_pipeline'),
run_config=run_config,
instance=DagsterInstance.local_temp(temp_dir),
)
assert result.success
def test_valid_job_format_with_backcompat_resources(run_launcher):
docker_image = test_project_docker_image()
run_config = load_yaml_from_path(
os.path.join(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.job_config,
command=["dagster-graphql"],
args=[
"-p",
"executeRunInProcess",
"-v",
seven.json.dumps({"runId": run.run_id}),
],
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,
resources="""
resources:
limits:
cpu: 500m
memory: 2560Mi
requests:
cpu: 250m
memory: 64Mi""",
).strip())
def test_valid_job_format_with_resources(run_launcher):
docker_image = test_project_docker_image()
run_config = load_yaml_from_path(
os.path.join(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'
},
})
})
resources = get_k8s_resource_requirements(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.job_config,
command=['dagster-graphql'],
args=[
'-p',
'executeRunInProcess',
'-v',
seven.json.dumps({'runId': run.run_id}),
],
job_name=job_name,
resources=resources,
pod_name=pod_name,
component='runmaster',
)
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='''
resources:
limits:
cpu: 500m
memory: 2560Mi
requests:
cpu: 250m
memory: 64Mi''',
).strip())
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-1 | 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 = 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 dagster_docker_image():
docker_image = 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 _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-1 | 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 = test_project_docker_image()
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 k8s_scheduler(cluster_provider, helm_namespace): # pylint: disable=redefined-outer-name,unused-argument
return K8sScheduler(
scheduler_namespace=helm_namespace,
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=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 run_launcher(cluster_provider, helm_namespace): # pylint: disable=redefined-outer-name,unused-argument
return K8sRunLauncher(
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=test_project_docker_image(),
load_incluster_config=False,
kubeconfig_file=cluster_provider.kubeconfig_file,
image_pull_policy=image_pull_policy(),
job_namespace=helm_namespace,
env_config_maps=['dagster-pipeline-env', 'test-env-configmap'],
env_secrets=['test-env-secret'],
)
def test_valid_job_format_with_user_defined_k8s_config(run_launcher):
docker_image = test_project_docker_image()
run_config = load_yaml_from_path(
os.path.join(test_project_environments_path(), "env.yaml"))
pipeline_name = "demo_pipeline"
run = PipelineRun(pipeline_name=pipeline_name, run_config=run_config)
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"
},
},
"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.job_config,
command=["dagster"],
args=["api", "execute_run_with_structured_logs"],
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'",
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_valid_job_format_with_user_defined_k8s_config(run_launcher):
docker_image = test_project_docker_image()
run_config = load_yaml_from_path(
os.path.join(test_project_environments_path(), 'env.yaml'))
pipeline_name = 'demo_pipeline'
run = PipelineRun(pipeline_name=pipeline_name, run_config=run_config)
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"
}
},
'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.job_config,
command=['dagster-graphql'],
args=[
'-p',
'executeRunInProcess',
'-v',
seven.json.dumps({'runId': run.run_id}),
],
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,
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())
