def test_construct_dagster_k8s_job_with_invalid_key_raises_error():
with pytest.raises(
Exception,
match="Unexpected keys in model class V1JobSpec: {'nonExistantKey'}"
):
UserDefinedDagsterK8sConfig(
job_spec_config={"nonExistantKey": "nonExistantValue"}, )
def test_user_defined_k8s_config_tags():
@solid(
tags={
USER_DEFINED_K8S_CONFIG_KEY: {
"container_config": {
"resources": {
"requests": {"cpu": "250m", "memory": "64Mi"},
"limits": {"cpu": "500m", "memory": "2560Mi"},
}
}
}
}
)
def my_solid(_):
pass
user_defined_k8s_config = get_user_defined_k8s_config(my_solid.tags)
assert user_defined_k8s_config.container_config
assert user_defined_k8s_config.container_config["resources"]
resources = user_defined_k8s_config.container_config["resources"]
assert resources["requests"]["cpu"] == "250m"
assert resources["requests"]["memory"] == "64Mi"
assert resources["limits"]["cpu"] == "500m"
assert resources["limits"]["memory"] == "2560Mi"
@solid
def no_resource_tags_solid(_):
pass
user_defined_k8s_config = get_user_defined_k8s_config(no_resource_tags_solid.tags)
assert user_defined_k8s_config == UserDefinedDagsterK8sConfig()
def test_user_defined_k8s_config_serialization():
cfg = UserDefinedDagsterK8sConfig(
container_config={
"resouces": {
"requests": {"cpu": "250m", "memory": "64Mi"},
"limits": {"cpu": "500m", "memory": "2560Mi"},
}
},
pod_template_spec_metadata={"key": "value"},
pod_spec_config={"key": "value"},
job_config={"key": "value"},
job_metadata={"key": "value"},
job_spec_config={"key": "value"},
)
assert UserDefinedDagsterK8sConfig.from_dict(cfg.to_dict()) == cfg
def test_construct_dagster_k8s_job_with_sidecar_container():
cfg = DagsterK8sJobConfig(
job_image="test/foo:latest",
dagster_home="/opt/dagster/dagster_home",
instance_config_map="test",
)
job = construct_dagster_k8s_job(cfg, [], "job123").to_dict()
assert job["spec"][
"ttl_seconds_after_finished"] == DEFAULT_K8S_JOB_TTL_SECONDS_AFTER_FINISHED
user_defined_cfg = UserDefinedDagsterK8sConfig(pod_spec_config={
"containers": [{
"command": ["echo", "HI"],
"image": "sidecar:bar",
"name": "sidecar"
}]
}, )
job = construct_dagster_k8s_job(
cfg, [], "job123", user_defined_k8s_config=user_defined_cfg).to_dict()
containers = job["spec"]["template"]["spec"]["containers"]
assert len(containers) == 2
assert containers[0]["image"] == "test/foo:latest"
assert containers[1]["image"] == "sidecar:bar"
assert containers[1]["command"] == ["echo", "HI"]
assert containers[1]["name"] == "sidecar"
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_step_handler_user_defined_config(kubeconfig_file):
mock_k8s_client_batch_api = mock.MagicMock()
handler = K8sStepHandler(
job_config=DagsterK8sJobConfig(instance_config_map="foobar",
job_image="bizbuz"),
job_namespace="foo",
load_incluster_config=False,
kubeconfig_file=kubeconfig_file,
k8s_client_batch_api=mock_k8s_client_batch_api,
)
# Construct Dagster solid 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}
with instance_for_test() as instance:
run = create_run_for_test(
instance,
pipeline_name="bar",
)
handler.launch_step(
StepHandlerContext(
instance,
ExecuteStepArgs(
reconstructable(bar).get_python_origin(), run.run_id,
["foo_solid"]),
{"foo_solid": tags},
))
# 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 kwargs["body"].spec.template.spec.containers[
0].image == "bizbuz"
job_resources = kwargs["body"].spec.template.spec.containers[
0].resources
assert job_resources == expected_resources
def test_construct_dagster_k8s_job_with_ttl():
cfg = DagsterK8sJobConfig(
job_image="test/foo:latest",
dagster_home="/opt/dagster/dagster_home",
instance_config_map="test",
)
job = construct_dagster_k8s_job(cfg, [], "job123").to_dict()
assert job["spec"][
"ttl_seconds_after_finished"] == DEFAULT_K8S_JOB_TTL_SECONDS_AFTER_FINISHED
user_defined_cfg = UserDefinedDagsterK8sConfig(
job_spec_config={"ttl_seconds_after_finished": 0}, )
job = construct_dagster_k8s_job(
cfg, [], "job123", user_defined_k8s_config=user_defined_cfg).to_dict()
assert job["spec"]["ttl_seconds_after_finished"] == 0
def test_user_defined_k8s_config_serialization():
cfg = UserDefinedDagsterK8sConfig(
container_config={
"resources": {
"requests": {
"cpu": "250m",
"memory": "64Mi"
},
"limits": {
"cpu": "500m",
"memory": "2560Mi"
},
}
},
pod_template_spec_metadata={"namespace": "value"},
pod_spec_config={"dns_policy": "value"},
job_config={"status": {
"completed_indexes": "value"
}},
job_metadata={"namespace": "value"},
job_spec_config={"backoff_limit": 120},
)
assert UserDefinedDagsterK8sConfig.from_dict(cfg.to_dict()) == cfg
def test_user_defined_config_from_tags():
config_args = {
"container_config": {
"resources": {
"requests": {
"cpu": "500m",
"memory": "128Mi"
},
"limits": {
"cpu": "1000m",
"memory": "1Gi"
},
}
},
"pod_template_spec_metadata": {
"pod_template_spec_key": "pod_template_spec_value"
},
"pod_spec_config": {
"pod_spec_config_key": "pod_spec_config_value"
},
"job_config": {
"job_config_key": "job_config_value"
},
"job_metadata": {
"job_metadata_key": "job_metadata_value"
},
"job_spec_config": {
"job_spec_config_key": "job_spec_config_value"
},
}
@pipeline(tags={USER_DEFINED_K8S_CONFIG_KEY: config_args})
def my_pipeline():
pass
assert get_user_defined_k8s_config(
my_pipeline.tags) == UserDefinedDagsterK8sConfig(**config_args)
def test_user_defined_config_from_tags():
config_args = {
"container_config": {
"resources": {
"requests": {"cpu": "500m", "memory": "128Mi"},
"limits": {"cpu": "1000m", "memory": "1Gi"},
}
},
"pod_template_spec_metadata": {"namespace": "pod_template_spec_value"},
"pod_spec_config": {"dns_policy": "pod_spec_config_value"},
"job_config": {"status": {"completed_indexes": "job_config_value"}},
"job_metadata": {"namespace": "job_metadata_value"},
"job_spec_config": {"backoff_limit": 120},
}
@pipeline(tags={USER_DEFINED_K8S_CONFIG_KEY: config_args})
def my_pipeline():
pass
assert get_user_defined_k8s_config(my_pipeline.tags) == UserDefinedDagsterK8sConfig(
**config_args
)
def _execute_step_k8s_job(
self,
execute_step_args_packed,
job_config_dict,
job_namespace,
load_incluster_config,
user_defined_k8s_config_dict=None,
kubeconfig_file=None,
):
"""Run step execution in a K8s job pod."""
execute_step_args = unpack_value(
check.dict_param(
execute_step_args_packed,
"execute_step_args_packed",
))
check.inst_param(execute_step_args, "execute_step_args",
ExecuteStepArgs)
check.invariant(
len(execute_step_args.step_keys_to_execute) == 1,
"Celery K8s task executor can only execute 1 step at a time",
)
# Celery will serialize this as a list
job_config = DagsterK8sJobConfig.from_dict(job_config_dict)
check.inst_param(job_config, "job_config", DagsterK8sJobConfig)
check.str_param(job_namespace, "job_namespace")
check.bool_param(load_incluster_config, "load_incluster_config")
user_defined_k8s_config = UserDefinedDagsterK8sConfig.from_dict(
user_defined_k8s_config_dict)
check.opt_inst_param(
user_defined_k8s_config,
"user_defined_k8s_config",
UserDefinedDagsterK8sConfig,
)
check.opt_str_param(kubeconfig_file, "kubeconfig_file")
# For when launched via DinD or running the cluster
if load_incluster_config:
kubernetes.config.load_incluster_config()
else:
kubernetes.config.load_kube_config(kubeconfig_file)
instance = DagsterInstance.from_ref(execute_step_args.instance_ref)
pipeline_run = instance.get_run_by_id(
execute_step_args.pipeline_run_id)
check.inst(
pipeline_run,
PipelineRun,
"Could not load run {}".format(execute_step_args.pipeline_run_id),
)
step_key = execute_step_args.step_keys_to_execute[0]
celery_worker_name = self.request.hostname
celery_pod_name = os.environ.get("HOSTNAME")
instance.report_engine_event(
"Task for step {step_key} picked up by Celery".format(
step_key=step_key),
pipeline_run,
EngineEventData([
EventMetadataEntry.text(celery_worker_name,
"Celery worker name"),
EventMetadataEntry.text(celery_pod_name,
"Celery worker Kubernetes Pod name"),
]),
CeleryK8sJobExecutor,
step_key=step_key,
)
if pipeline_run.status != PipelineRunStatus.STARTED:
instance.report_engine_event(
"Not scheduling step because pipeline run status is not STARTED",
pipeline_run,
EngineEventData([
EventMetadataEntry.text(step_key, "Step key"),
]),
CeleryK8sJobExecutor,
step_key=step_key,
)
return []
# Ensure we stay below k8s name length limits
k8s_name_key = get_k8s_job_name(execute_step_args.pipeline_run_id,
step_key)
retry_state = execute_step_args.known_state.get_retry_state()
if retry_state.get_attempt_count(step_key):
attempt_number = retry_state.get_attempt_count(step_key)
job_name = "dagster-job-%s-%d" % (k8s_name_key, attempt_number)
pod_name = "dagster-job-%s-%d" % (k8s_name_key, attempt_number)
else:
job_name = "dagster-job-%s" % (k8s_name_key)
pod_name = "dagster-job-%s" % (k8s_name_key)
input_json = serialize_dagster_namedtuple(execute_step_args)
args = ["dagster", "api", "execute_step", input_json]
job = construct_dagster_k8s_job(job_config, args, job_name,
user_defined_k8s_config, pod_name)
# Running list of events generated from this task execution
events = []
# Post event for starting execution
job_name = job.metadata.name
engine_event = instance.report_engine_event(
"Executing step {} in Kubernetes job {}".format(
step_key, job_name),
pipeline_run,
EngineEventData(
[
EventMetadataEntry.text(step_key, "Step key"),
EventMetadataEntry.text(job_name, "Kubernetes Job name"),
EventMetadataEntry.text(job_config.job_image, "Job image"),
EventMetadataEntry.text(job_config.image_pull_policy,
"Image pull policy"),
EventMetadataEntry.text(str(job_config.image_pull_secrets),
"Image pull secrets"),
EventMetadataEntry.text(
str(job_config.service_account_name),
"Service account name"),
],
marker_end=DELEGATE_MARKER,
),
CeleryK8sJobExecutor,
# validated above that step_keys is length 1, and it is not possible to use ETH or
# execution plan in this function (Celery K8s workers should not access to user code)
step_key=step_key,
)
events.append(engine_event)
try:
kubernetes.client.BatchV1Api().create_namespaced_job(
body=job, namespace=job_namespace)
except kubernetes.client.rest.ApiException as e:
if e.reason == "Conflict":
# There is an existing job with the same name so proceed and see if the existing job succeeded
instance.report_engine_event(
"Did not create Kubernetes job {} for step {} since job name already "
"exists, proceeding with existing job.".format(
job_name, step_key),
pipeline_run,
EngineEventData(
[
EventMetadataEntry.text(step_key, "Step key"),
EventMetadataEntry.text(job_name,
"Kubernetes Job name"),
],
marker_end=DELEGATE_MARKER,
),
CeleryK8sJobExecutor,
step_key=step_key,
)
else:
instance.report_engine_event(
"Encountered unexpected error while creating Kubernetes job {} for step {}, "
"exiting.".format(job_name, step_key),
pipeline_run,
EngineEventData(
[
EventMetadataEntry.text(step_key, "Step key"),
],
error=serializable_error_info_from_exc_info(
sys.exc_info()),
),
CeleryK8sJobExecutor,
step_key=step_key,
)
return []
try:
wait_for_job_success(
job_name=job_name,
namespace=job_namespace,
instance=instance,
run_id=execute_step_args.pipeline_run_id,
)
except (DagsterK8sError, DagsterK8sTimeoutError) as err:
step_failure_event = construct_step_failure_event_and_handle(
pipeline_run, step_key, err, instance=instance)
events.append(step_failure_event)
except DagsterK8sPipelineStatusException:
instance.report_engine_event(
"Terminating Kubernetes Job because pipeline run status is not STARTED",
pipeline_run,
EngineEventData([
EventMetadataEntry.text(step_key, "Step key"),
EventMetadataEntry.text(job_name, "Kubernetes Job name"),
EventMetadataEntry.text(job_namespace,
"Kubernetes Job namespace"),
]),
CeleryK8sJobExecutor,
step_key=step_key,
)
delete_job(job_name=job_name, namespace=job_namespace)
return []
except (
DagsterK8sUnrecoverableAPIError,
DagsterK8sAPIRetryLimitExceeded,
# We shouldn't see unwrapped APIExceptions anymore, as they should all be wrapped in
# a retry boundary. We still catch it here just in case we missed one so that we can
# report it to the event log
kubernetes.client.rest.ApiException,
) as err:
instance.report_engine_event(
"Encountered unexpected error while waiting on Kubernetes job {} for step {}, "
"exiting.".format(job_name, step_key),
pipeline_run,
EngineEventData(
[
EventMetadataEntry.text(step_key, "Step key"),
],
error=serializable_error_info_from_exc_info(
sys.exc_info()),
),
CeleryK8sJobExecutor,
step_key=step_key,
)
return []
try:
pod_names = get_pod_names_in_job(job_name, namespace=job_namespace)
except kubernetes.client.rest.ApiException as e:
instance.report_engine_event(
"Encountered unexpected error retreiving Pods for Kubernetes job {} for step {}, "
"exiting.".format(job_name, step_key),
pipeline_run,
EngineEventData(
[
EventMetadataEntry.text(step_key, "Step key"),
],
error=serializable_error_info_from_exc_info(
sys.exc_info()),
),
CeleryK8sJobExecutor,
step_key=step_key,
)
return []
# Post engine event for log retrieval
engine_event = instance.report_engine_event(
"Retrieving logs from Kubernetes Job pods",
pipeline_run,
EngineEventData(
[EventMetadataEntry.text("\n".join(pod_names), "Pod names")]),
CeleryK8sJobExecutor,
step_key=step_key,
)
events.append(engine_event)
logs = []
for pod_name in pod_names:
try:
raw_logs = retrieve_pod_logs(pod_name, namespace=job_namespace)
logs += raw_logs.split("\n")
except kubernetes.client.rest.ApiException as e:
instance.report_engine_event(
"Encountered unexpected error while fetching pod logs for Kubernetes job {}, "
"Pod name {} for step {}. Will attempt to continue with other pods."
.format(job_name, pod_name, step_key),
pipeline_run,
EngineEventData(
[
EventMetadataEntry.text(step_key, "Step key"),
],
error=serializable_error_info_from_exc_info(
sys.exc_info()),
),
CeleryK8sJobExecutor,
step_key=step_key,
)
events += filter_dagster_events_from_pod_logs(logs)
serialized_events = [
serialize_dagster_namedtuple(event) for event in events
]
return serialized_events
def _execute_step_k8s_job(
self,
instance_ref_dict,
step_keys,
run_config,
mode,
repo_name,
repo_location_name,
run_id,
job_config_dict,
job_namespace,
load_incluster_config,
retries_dict,
pipeline_origin_packed,
user_defined_k8s_config_dict=None,
kubeconfig_file=None,
):
"""Run step execution in a K8s job pod.
"""
check.dict_param(instance_ref_dict, "instance_ref_dict")
check.list_param(step_keys, "step_keys", of_type=str)
check.invariant(
len(step_keys) == 1, "Celery K8s task executor can only execute 1 step at a time"
)
check.dict_param(run_config, "run_config")
check.str_param(mode, "mode")
check.str_param(repo_name, "repo_name")
check.str_param(repo_location_name, "repo_location_name")
check.str_param(run_id, "run_id")
# Celery will serialize this as a list
job_config = DagsterK8sJobConfig.from_dict(job_config_dict)
check.inst_param(job_config, "job_config", DagsterK8sJobConfig)
check.str_param(job_namespace, "job_namespace")
check.bool_param(load_incluster_config, "load_incluster_config")
check.dict_param(retries_dict, "retries_dict")
pipeline_origin = unpack_value(
check.dict_param(
pipeline_origin_packed, "pipeline_origin_packed"
) # TODO: make part of args
)
check.inst(pipeline_origin, PipelineOrigin)
user_defined_k8s_config = UserDefinedDagsterK8sConfig.from_dict(
user_defined_k8s_config_dict
)
check.opt_inst_param(
user_defined_k8s_config, "user_defined_k8s_config", UserDefinedDagsterK8sConfig,
)
check.opt_str_param(kubeconfig_file, "kubeconfig_file")
# For when launched via DinD or running the cluster
if load_incluster_config:
kubernetes.config.load_incluster_config()
else:
kubernetes.config.load_kube_config(kubeconfig_file)
instance_ref = InstanceRef.from_dict(instance_ref_dict)
instance = DagsterInstance.from_ref(instance_ref)
pipeline_run = instance.get_run_by_id(run_id)
check.invariant(pipeline_run, "Could not load run {}".format(run_id))
step_key = step_keys[0]
celery_worker_name = self.request.hostname
celery_pod_name = os.environ.get("HOSTNAME")
instance.report_engine_event(
"Task for step {step_key} picked up by Celery".format(step_key=step_key),
pipeline_run,
EngineEventData(
[
EventMetadataEntry.text(celery_worker_name, "Celery worker name"),
EventMetadataEntry.text(celery_pod_name, "Celery worker Kubernetes Pod name"),
]
),
CeleryK8sJobExecutor,
step_key=step_key,
)
if pipeline_run.status != PipelineRunStatus.STARTED:
instance.report_engine_event(
"Not scheduling step because pipeline run status is not STARTED",
pipeline_run,
EngineEventData([EventMetadataEntry.text(step_key, "Step keys"),]),
CeleryK8sJobExecutor,
step_key=step_key,
)
return
# Ensure we stay below k8s name length limits
k8s_name_key = get_k8s_job_name(run_id, step_key)
retries = Retries.from_config(retries_dict)
if retries.get_attempt_count(step_key):
attempt_number = retries.get_attempt_count(step_key)
job_name = "dagster-job-%s-%d" % (k8s_name_key, attempt_number)
pod_name = "dagster-job-%s-%d" % (k8s_name_key, attempt_number)
else:
job_name = "dagster-job-%s" % (k8s_name_key)
pod_name = "dagster-job-%s" % (k8s_name_key)
input_json = serialize_dagster_namedtuple(
ExecuteStepArgs(
pipeline_origin=pipeline_origin,
pipeline_run_id=run_id,
instance_ref=None,
mode=mode,
step_keys_to_execute=step_keys,
run_config=run_config,
retries_dict=retries_dict,
)
)
command = ["dagster"]
args = ["api", "execute_step_with_structured_logs", input_json]
job = construct_dagster_k8s_job(
job_config, command, args, job_name, user_defined_k8s_config, pod_name
)
# Running list of events generated from this task execution
events = []
# Post event for starting execution
job_name = job.metadata.name
engine_event = instance.report_engine_event(
"Executing step {} in Kubernetes job {}".format(step_key, job_name),
pipeline_run,
EngineEventData(
[
EventMetadataEntry.text(step_key, "Step keys"),
EventMetadataEntry.text(job_name, "Kubernetes Job name"),
EventMetadataEntry.text(pod_name, "Kubernetes Pod name"),
EventMetadataEntry.text(job_config.job_image, "Job image"),
EventMetadataEntry.text(job_config.image_pull_policy, "Image pull policy"),
EventMetadataEntry.text(
str(job_config.image_pull_secrets), "Image pull secrets"
),
EventMetadataEntry.text(
str(job_config.service_account_name), "Service account name"
),
],
marker_end=DELEGATE_MARKER,
),
CeleryK8sJobExecutor,
# validated above that step_keys is length 1, and it is not possible to use ETH or
# execution plan in this function (Celery K8s workers should not access to user code)
step_key=step_key,
)
events.append(engine_event)
try:
kubernetes.client.BatchV1Api().create_namespaced_job(body=job, namespace=job_namespace)
except kubernetes.client.rest.ApiException as e:
if e.reason == "Conflict":
# There is an existing job with the same name so do not procede.
instance.report_engine_event(
"Did not create Kubernetes job {} for step {} since job name already "
"exists, exiting.".format(job_name, step_key),
pipeline_run,
EngineEventData(
[
EventMetadataEntry.text(step_key, "Step keys"),
EventMetadataEntry.text(job_name, "Kubernetes Job name"),
EventMetadataEntry.text(pod_name, "Kubernetes Pod name"),
],
marker_end=DELEGATE_MARKER,
),
CeleryK8sJobExecutor,
step_key=step_key,
)
else:
instance.report_engine_event(
"Encountered unexpected error while creating Kubernetes job {} for step {}, "
"exiting.".format(job_name, step_key),
pipeline_run,
EngineEventData(
[
EventMetadataEntry.text(step_key, "Step keys"),
EventMetadataEntry.text(e, "Error"),
]
),
CeleryK8sJobExecutor,
step_key=step_key,
)
return
try:
wait_for_job_success(
job_name=job_name, namespace=job_namespace, instance=instance, run_id=run_id,
)
except DagsterK8sPipelineStatusException:
instance.report_engine_event(
"Terminating Kubernetes Job because pipeline run status is not STARTED",
pipeline_run,
EngineEventData(
[
EventMetadataEntry.text(step_key, "Step keys"),
EventMetadataEntry.text(job_name, "Kubernetes Job name"),
EventMetadataEntry.text(job_namespace, "Kubernetes Job namespace"),
]
),
CeleryK8sJobExecutor,
step_key=step_key,
)
delete_job(job_name=job_name, namespace=job_namespace)
return
pod_names = get_pod_names_in_job(job_name, namespace=job_namespace)
# Post engine event for log retrieval
engine_event = instance.report_engine_event(
"Retrieving logs from Kubernetes Job pods",
pipeline_run,
EngineEventData([EventMetadataEntry.text("\n".join(pod_names), "Pod names")]),
CeleryK8sJobExecutor,
step_key=step_key,
)
events.append(engine_event)
logs = []
for pod_name in pod_names:
raw_logs = retrieve_pod_logs(pod_name, namespace=job_namespace)
logs += raw_logs.split("\n")
events += filter_dagster_events_from_pod_logs(logs)
serialized_events = [serialize_dagster_namedtuple(event) for event in events]
return serialized_events
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
]
