def test_raise_on_error(kubeconfig_file):
# Construct a K8s run launcher in a fake k8s environment.
mock_k8s_client_batch_api = mock.MagicMock()
k8s_run_launcher = K8sRunLauncher(
service_account_name="dagit-admin",
instance_config_map="dagster-instance",
postgres_password_secret="dagster-postgresql-secret",
dagster_home="/opt/dagster/dagster_home",
job_image="fake_job_image",
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
repo_def = recon_repo.get_definition()
loadable_target_origin = LoadableTargetOrigin(python_file=__file__)
with instance_for_test() as instance:
with in_process_test_workspace(instance, loadable_target_origin) as workspace:
location = workspace.get_repository_location(workspace.repository_location_names[0])
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,
)
# Launch the run in a fake Dagster instance.
pipeline_name = "demo_pipeline"
run = create_run_for_test(
instance,
pipeline_name=pipeline_name,
external_pipeline_origin=fake_external_pipeline.get_external_origin(),
pipeline_code_origin=fake_external_pipeline.get_python_origin(),
)
k8s_run_launcher.register_instance(instance)
k8s_run_launcher.launch_run(LaunchRunContext(run, workspace))
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):
# 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 test_no_postgres(kubeconfig_file):
# Construct a K8s run launcher in a fake k8s environment.
mock_k8s_client_batch_api = mock.MagicMock()
k8s_run_launcher = K8sRunLauncher(
service_account_name="dagit-admin",
instance_config_map="dagster-instance",
dagster_home="/opt/dagster/dagster_home",
job_image="fake_job_image",
load_incluster_config=False,
kubeconfig_file=kubeconfig_file,
k8s_client_batch_api=mock_k8s_client_batch_api,
)
# Create fake external pipeline.
recon_pipeline = reconstructable(fake_pipeline)
recon_repo = recon_pipeline.repository
repo_def = recon_repo.get_definition()
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_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,
)
# Launch the run in a fake Dagster instance.
pipeline_name = "demo_pipeline"
run = create_run_for_test(
instance,
pipeline_name=pipeline_name,
external_pipeline_origin=fake_external_pipeline.
get_external_origin(),
pipeline_code_origin=fake_external_pipeline.get_python_origin(
),
)
k8s_run_launcher.register_instance(instance)
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_job_image"
# 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"
assert DAGSTER_PG_PASSWORD_ENV_VAR not in [
env.name
for env in kwargs["body"].spec.template.spec.containers[0].env
]
def test_check_run_health(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.Mock(spec_set=["read_namespaced_job"])
mock_k8s_client_batch_api.read_namespaced_job.side_effect = [
V1Job(status=V1JobStatus(failed=0, succeeded=0)),
V1Job(status=V1JobStatus(failed=0, succeeded=1)),
V1Job(status=V1JobStatus(failed=1, succeeded=0)),
]
k8s_run_launcher = K8sRunLauncher(
service_account_name="dagit-admin",
instance_config_map="dagster-instance",
postgres_password_secret="dagster-postgresql-secret",
dagster_home="/opt/dagster/dagster_home",
job_image="fake_job_image",
load_incluster_config=False,
kubeconfig_file=kubeconfig_file,
k8s_client_batch_api=mock_k8s_client_batch_api,
labels=labels,
)
# Create fake external pipeline.
recon_pipeline = reconstructable(fake_pipeline)
recon_repo = recon_pipeline.repository
repo_def = recon_repo.get_definition()
loadable_target_origin = LoadableTargetOrigin(python_file=__file__)
with instance_for_test() as instance:
with in_process_test_workspace(instance, loadable_target_origin) as workspace:
location = workspace.get_repository_location(workspace.repository_location_names[0])
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,
)
# Launch the run in a fake Dagster instance.
pipeline_name = "demo_pipeline"
run = create_run_for_test(
instance,
pipeline_name=pipeline_name,
external_pipeline_origin=fake_external_pipeline.get_external_origin(),
pipeline_code_origin=fake_external_pipeline.get_python_origin(),
)
k8s_run_launcher.register_instance(instance)
# same order as side effects
assert k8s_run_launcher.check_run_worker_health(run).status == WorkerStatus.RUNNING
assert k8s_run_launcher.check_run_worker_health(run).status == WorkerStatus.SUCCESS
assert k8s_run_launcher.check_run_worker_health(run).status == WorkerStatus.FAILED
def _external_pipeline_from_def(pipeline_def, solid_subset=None):
recon_pipeline = reconstructable(pipeline_def)
recon_repo = recon_pipeline.repository
repo_def = recon_repo.get_definition()
location_handle = LocationHandle('test', recon_repo.pointer)
repository_handle = RepositoryHandle(repo_def.name, location_handle)
return external_pipeline_from_recon_pipeline(
reconstructable(pipeline_def),
solid_subset=solid_subset,
repository_handle=repository_handle,
)
def _external_pipeline_from_def(pipeline_def, solid_selection=None):
recon_pipeline = reconstructable(pipeline_def)
recon_repo = recon_pipeline.repository
repo_def = recon_repo.get_definition()
location_handle = RepositoryLocationHandle.create_in_process_location(recon_repo.pointer)
repository_handle = RepositoryHandle(
repository_name=repo_def.name, repository_location_handle=location_handle,
)
return external_pipeline_from_recon_pipeline(
reconstructable(pipeline_def),
solid_selection=solid_selection,
repository_handle=repository_handle,
)
def get_external_pipeline(self, selector):
check.inst_param(selector, 'selector', PipelineSelector)
check.invariant(
selector.location_name == self.name,
'PipelineSelector location_name mismatch, got {selector.location_name} expected {self.name}'
.format(self=self, selector=selector),
)
return external_pipeline_from_recon_pipeline(
recon_pipeline=self.get_reconstructable_pipeline(
selector.pipeline_name),
solid_subset=selector.solid_subset,
repository_handle=self._external_repo.handle,
)
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())
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()
k8s_run_launcher = K8sRunLauncher(
service_account_name="dagit-admin",
instance_config_map="dagster-instance",
postgres_password_secret="dagster-postgresql-secret",
dagster_home="/opt/dagster/dagster_home",
job_image="fake_job_image",
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
repo_def = recon_repo.get_definition()
location_origin = InProcessRepositoryLocationOrigin(recon_repo)
with location_origin.create_location() as location:
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,
)
# Launch the run in a fake Dagster instance.
with instance_for_test() as instance:
pipeline_name = "demo_pipeline"
run = create_run_for_test(instance,
pipeline_name=pipeline_name,
tags=tags)
k8s_run_launcher.register_instance(instance)
run = k8s_run_launcher.launch_run(run, fake_external_pipeline)
updated_run = instance.get_run_by_id(run.run_id)
assert updated_run.tags[DOCKER_IMAGE_TAG] == "fake_job_image"
# 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
assert DAGSTER_PG_PASSWORD_ENV_VAR in [
env.name
for env in kwargs["body"].spec.template.spec.containers[0].env
]
def test_launcher_with_container_context(kubeconfig_file):
# Construct a K8s run launcher in a fake k8s environment.
mock_k8s_client_batch_api = mock.MagicMock()
k8s_run_launcher = K8sRunLauncher(
service_account_name="dagit-admin",
instance_config_map="dagster-instance",
postgres_password_secret="dagster-postgresql-secret",
dagster_home="/opt/dagster/dagster_home",
job_image="fake_job_image",
load_incluster_config=False,
kubeconfig_file=kubeconfig_file,
k8s_client_batch_api=mock_k8s_client_batch_api,
env_vars=["FOO_TEST"],
)
container_context_config = {"k8s": {"env_vars": ["BAR_TEST"]}}
# Create fake external pipeline.
recon_pipeline = reconstructable(fake_pipeline)
recon_repo = recon_pipeline.repository
repo_def = recon_repo.get_definition()
python_origin = recon_pipeline.get_python_origin()
python_origin = python_origin._replace(
repository_origin=python_origin.repository_origin._replace(
container_context=container_context_config,
)
)
loadable_target_origin = LoadableTargetOrigin(python_file=__file__)
with instance_for_test() as instance:
with in_process_test_workspace(instance, loadable_target_origin) as workspace:
location = workspace.get_repository_location(workspace.repository_location_names[0])
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,
)
# Launch the run in a fake Dagster instance.
pipeline_name = "demo_pipeline"
run = create_run_for_test(
instance,
pipeline_name=pipeline_name,
external_pipeline_origin=fake_external_pipeline.get_external_origin(),
pipeline_code_origin=python_origin,
)
k8s_run_launcher.register_instance(instance)
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_job_image"
# 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]
env_names = [env.name for env in container.env]
assert "BAR_TEST" in env_names
assert "FOO_TEST" in env_names
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()
)
