def test_step_handler(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,
)
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": {}},
)
)
# 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"
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 _get_step_handler_context(self, pipeline_context, steps,
active_execution) -> StepHandlerContext:
return StepHandlerContext(
instance=pipeline_context.plan_data.instance,
execute_step_args=ExecuteStepArgs(
pipeline_origin=pipeline_context.reconstructable_pipeline.
get_python_origin(),
pipeline_run_id=pipeline_context.pipeline_run.run_id,
step_keys_to_execute=[step.key for step in steps],
instance_ref=pipeline_context.plan_data.instance.get_ref(),
retry_mode=self.retries,
known_state=active_execution.get_known_state(),
),
step_tags={step.key: step.tags
for step in steps},
pipeline_run=pipeline_context.pipeline_run,
)