def test_user_defined_k8s_config_in_run_tags(kubeconfig_file):
# Construct a K8s run launcher in a fake k8s environment.
mock_k8s_client_batch_api = mock.MagicMock()
celery_k8s_run_launcher = CeleryK8sRunLauncher(
instance_config_map="dagster-instance",
postgres_password_secret="dagster-postgresql-secret",
dagster_home="/opt/dagster/dagster_home",
load_incluster_config=False,
kubeconfig_file=kubeconfig_file,
k8s_client_batch_api=mock_k8s_client_batch_api,
)
# Construct Dagster run tags with user defined k8s config.
expected_resources = {
"requests": {"cpu": "250m", "memory": "64Mi"},
"limits": {"cpu": "500m", "memory": "2560Mi"},
}
user_defined_k8s_config = UserDefinedDagsterK8sConfig(
container_config={"resources": expected_resources},
)
user_defined_k8s_config_json = json.dumps(user_defined_k8s_config.to_dict())
tags = {"dagster-k8s/config": user_defined_k8s_config_json}
# Create fake external pipeline.
recon_pipeline = reconstructable(fake_pipeline)
recon_repo = recon_pipeline.repository
location_origin = InProcessRepositoryLocationOrigin(recon_repo)
location_handle = location_origin.create_handle()
repo_def = recon_repo.get_definition()
repo_handle = RepositoryHandle(
repository_name=repo_def.name,
repository_location_handle=location_handle,
)
fake_external_pipeline = external_pipeline_from_recon_pipeline(
recon_pipeline,
solid_selection=None,
repository_handle=repo_handle,
)
# Launch the run in a fake Dagster instance.
with instance_for_test() as instance:
celery_k8s_run_launcher.initialize(instance)
pipeline_name = "demo_pipeline"
run_config = {"execution": {"celery-k8s": {"config": {"job_image": "fake-image-name"}}}}
run = create_run_for_test(
instance,
pipeline_name=pipeline_name,
run_config=run_config,
tags=tags,
)
celery_k8s_run_launcher.launch_run(instance, run, fake_external_pipeline)
# Check that user defined k8s config was passed down to the k8s job.
mock_method_calls = mock_k8s_client_batch_api.method_calls
assert len(mock_method_calls) > 0
method_name, _args, kwargs = mock_method_calls[0]
assert method_name == "create_namespaced_job"
job_resources = kwargs["body"].spec.template.spec.containers[0].resources
assert job_resources == expected_resources
def run_launcher(cluster_provider): # pylint: disable=redefined-outer-name,unused-argument
return CeleryK8sRunLauncher(
instance_config_map="dagster-instance",
postgres_password_secret="dagster-postgresql-secret",
dagster_home="/opt/dagster/dagster_home",
load_incluster_config=False,
kubeconfig_file=cluster_provider.kubeconfig_file,
)
def test_k8s_executor_config_override(kubeconfig_file):
# Construct a K8s run launcher in a fake k8s environment.
mock_k8s_client_batch_api = mock.MagicMock()
celery_k8s_run_launcher = CeleryK8sRunLauncher(
instance_config_map="dagster-instance",
postgres_password_secret="dagster-postgresql-secret",
dagster_home="/opt/dagster/dagster_home",
load_incluster_config=False,
kubeconfig_file=kubeconfig_file,
k8s_client_batch_api=mock_k8s_client_batch_api,
)
external_pipeline = get_test_project_external_pipeline(
"demo_pipeline", "my_image:tag")
# Launch the run in a fake Dagster instance.
with instance_for_test() as instance:
celery_k8s_run_launcher.register_instance(instance)
pipeline_name = "demo_pipeline"
# Launch without custom job_image
run = create_run_for_test(
instance,
pipeline_name=pipeline_name,
run_config={"execution": {
"celery-k8s": {}
}},
)
celery_k8s_run_launcher.launch_run(run, external_pipeline)
# Launch with custom job_image
run = create_run_for_test(
instance,
pipeline_name=pipeline_name,
run_config={
"execution": {
"celery-k8s": {
"config": {
"job_image": "fake-image-name"
}
}
}
},
)
celery_k8s_run_launcher.launch_run(run, external_pipeline)
# Check that user defined k8s config was passed down to the k8s job.
mock_method_calls = mock_k8s_client_batch_api.method_calls
assert len(mock_method_calls) > 0
_, _args, kwargs = mock_method_calls[0]
assert kwargs["body"].spec.template.spec.containers[
0].image == "my_image:tag"
_, _args, kwargs = mock_method_calls[1]
assert kwargs["body"].spec.template.spec.containers[
0].image == "fake-image-name"
def test_raise_on_error(kubeconfig_file):
mock_k8s_client_batch_api = mock.MagicMock()
celery_k8s_run_launcher = CeleryK8sRunLauncher(
instance_config_map="dagster-instance",
postgres_password_secret="dagster-postgresql-secret",
dagster_home="/opt/dagster/dagster_home",
load_incluster_config=False,
kubeconfig_file=kubeconfig_file,
k8s_client_batch_api=mock_k8s_client_batch_api,
fail_pod_on_run_failure=True,
)
# Create fake external pipeline.
recon_pipeline = reconstructable(fake_pipeline)
recon_repo = recon_pipeline.repository
with instance_for_test() as instance:
with in_process_test_workspace(instance, recon_repo) as workspace:
location = workspace.get_repository_location(
workspace.repository_location_names[0])
repo_def = recon_repo.get_definition()
repo_handle = RepositoryHandle(
repository_name=repo_def.name,
repository_location=location,
)
fake_external_pipeline = external_pipeline_from_recon_pipeline(
recon_pipeline,
solid_selection=None,
repository_handle=repo_handle,
)
celery_k8s_run_launcher.register_instance(instance)
pipeline_name = "demo_pipeline"
run_config = {
"execution": {
"celery-k8s": {
"config": {
"job_image": "fake-image-name"
}
}
}
}
run = create_run_for_test(
instance,
pipeline_name=pipeline_name,
run_config=run_config,
external_pipeline_origin=fake_external_pipeline.
get_external_origin(),
pipeline_code_origin=fake_external_pipeline.get_python_origin(
),
)
celery_k8s_run_launcher.launch_run(LaunchRunContext(
run, workspace))
# Check that user defined k8s config was passed down to the k8s job.
mock_method_calls = mock_k8s_client_batch_api.method_calls
assert len(mock_method_calls) > 0
method_name, _args, kwargs = mock_method_calls[0]
assert method_name == "create_namespaced_job"
container = kwargs["body"].spec.template.spec.containers[0]
args = container.args
assert (args == ExecuteRunArgs(
pipeline_origin=run.pipeline_code_origin,
pipeline_run_id=run.run_id,
instance_ref=instance.get_ref(),
set_exit_code_on_failure=True,
).get_command_args())
def test_user_defined_k8s_config_in_run_tags(kubeconfig_file):
labels = {"foo_label_key": "bar_label_value"}
# Construct a K8s run launcher in a fake k8s environment.
mock_k8s_client_batch_api = mock.MagicMock()
celery_k8s_run_launcher = CeleryK8sRunLauncher(
instance_config_map="dagster-instance",
postgres_password_secret="dagster-postgresql-secret",
dagster_home="/opt/dagster/dagster_home",
load_incluster_config=False,
kubeconfig_file=kubeconfig_file,
k8s_client_batch_api=mock_k8s_client_batch_api,
labels=labels,
)
# Construct Dagster run tags with user defined k8s config.
expected_resources = {
"requests": {
"cpu": "250m",
"memory": "64Mi"
},
"limits": {
"cpu": "500m",
"memory": "2560Mi"
},
}
user_defined_k8s_config = UserDefinedDagsterK8sConfig(
container_config={"resources": expected_resources}, )
user_defined_k8s_config_json = json.dumps(
user_defined_k8s_config.to_dict())
tags = {"dagster-k8s/config": user_defined_k8s_config_json}
# Create fake external pipeline.
recon_pipeline = reconstructable(fake_pipeline)
recon_repo = recon_pipeline.repository
with instance_for_test() as instance:
with in_process_test_workspace(instance, recon_repo) as workspace:
location = workspace.get_repository_location(
workspace.repository_location_names[0])
repo_def = recon_repo.get_definition()
repo_handle = RepositoryHandle(
repository_name=repo_def.name,
repository_location=location,
)
fake_external_pipeline = external_pipeline_from_recon_pipeline(
recon_pipeline,
solid_selection=None,
repository_handle=repo_handle,
)
celery_k8s_run_launcher.register_instance(instance)
pipeline_name = "demo_pipeline"
run_config = {
"execution": {
"celery-k8s": {
"config": {
"job_image": "fake-image-name"
}
}
}
}
run = create_run_for_test(
instance,
pipeline_name=pipeline_name,
run_config=run_config,
tags=tags,
external_pipeline_origin=fake_external_pipeline.
get_external_origin(),
pipeline_code_origin=fake_external_pipeline.get_python_origin(
),
)
celery_k8s_run_launcher.launch_run(LaunchRunContext(
run, workspace))
updated_run = instance.get_run_by_id(run.run_id)
assert updated_run.tags[DOCKER_IMAGE_TAG] == "fake-image-name"
# Check that user defined k8s config was passed down to the k8s job.
mock_method_calls = mock_k8s_client_batch_api.method_calls
assert len(mock_method_calls) > 0
method_name, _args, kwargs = mock_method_calls[0]
assert method_name == "create_namespaced_job"
container = kwargs["body"].spec.template.spec.containers[0]
job_resources = container.resources
assert job_resources == expected_resources
labels = kwargs["body"].spec.template.metadata.labels
assert labels["foo_label_key"] == "bar_label_value"
args = container.args
assert (args == ExecuteRunArgs(
pipeline_origin=run.pipeline_code_origin,
pipeline_run_id=run.run_id,
instance_ref=instance.get_ref(),
set_exit_code_on_failure=None,
).get_command_args())