def test_bad_user_defined_k8s_config_tags():
@pipeline(tags={USER_DEFINED_K8S_CONFIG_KEY: {"other": {}}})
def my_solid():
pass
with pytest.raises(DagsterInvalidConfigError):
get_user_defined_k8s_config(my_solid.tags)
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_bad_user_defined_k8s_config_tags():
@pipeline(tags={USER_DEFINED_K8S_CONFIG_KEY: {"other": {}}})
def my_pipeline():
pass
with pytest.raises(
DagsterInvalidConfigError,
match='Received unexpected config entry "other" at the root'):
get_user_defined_k8s_config(my_pipeline.tags)
def test_bad_deprecated_resource_tags():
@pipeline(
tags={
K8S_RESOURCE_REQUIREMENTS_KEY: {
"other": {"cpu": "250m", "memory": "64Mi"},
}
}
)
def resource_tags_pipeline():
pass
with pytest.raises(DagsterInvalidConfigError):
get_user_defined_k8s_config(resource_tags_pipeline.tags)
def _submit_task_k8s_job(app, pipeline_context, step, queue, priority,
known_state):
user_defined_k8s_config = get_user_defined_k8s_config(step.tags)
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],
instance_ref=pipeline_context.instance.get_ref(),
retry_mode=pipeline_context.executor.retries.for_inner_plan(),
known_state=known_state,
should_verify_step=True,
)
task = create_k8s_job_task(app)
task_signature = task.si(
execute_step_args_packed=pack_value(execute_step_args),
job_config_dict=pipeline_context.executor.job_config.to_dict(),
job_namespace=pipeline_context.executor.job_namespace,
user_defined_k8s_config_dict=user_defined_k8s_config.to_dict(),
load_incluster_config=pipeline_context.executor.load_incluster_config,
kubeconfig_file=pipeline_context.executor.kubeconfig_file,
)
return task_signature.apply_async(
priority=priority,
queue=queue,
routing_key="{queue}.execute_step_k8s_job".format(queue=queue),
)
def test_tags_to_plan():
@solid
def blank(_):
pass
@pipeline
def k8s_ready():
blank.tag(
{
USER_DEFINED_K8S_CONFIG_KEY: {
"container_config": {
"resources": {
"requests": {"cpu": "250m", "memory": "64Mi"},
"limits": {"cpu": "500m", "memory": "2560Mi"},
}
}
}
}
)()
plan = create_execution_plan(k8s_ready)
step = list(plan.step_dict.values())[0]
user_defined_k8s_config = get_user_defined_k8s_config(step.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"
def test_construct_dagster_k8s_job_with_user_defined_service_account_name():
@graph
def user_defined_k8s_service_account_name_tags_graph():
pass
user_defined_k8s_config = get_user_defined_k8s_config(
user_defined_k8s_service_account_name_tags_graph.to_job(tags={
USER_DEFINED_K8S_CONFIG_KEY: {
"pod_spec_config": {
"service_account_name": "this-should-take-precedence",
},
},
}, ).tags)
cfg = DagsterK8sJobConfig(
job_image="test/foo:latest",
dagster_home="/opt/dagster/dagster_home",
instance_config_map="some-instance-configmap",
service_account_name="this-should-be-overriden",
)
job = construct_dagster_k8s_job(
cfg, ["foo", "bar"],
"job",
user_defined_k8s_config=user_defined_k8s_config).to_dict()
service_account_name = job["spec"]["template"]["spec"][
"service_account_name"]
assert service_account_name == "this-should-take-precedence"
def test_construct_dagster_k8s_job_with_user_defined_env_from():
@graph
def user_defined_k8s_env_from_tags_graph():
pass
# These fields still work even when using underscore keys
user_defined_k8s_config = get_user_defined_k8s_config(
user_defined_k8s_env_from_tags_graph.to_job(
tags={
USER_DEFINED_K8S_CONFIG_KEY: {
"container_config": {
"envFrom": [
{
"configMapRef": {
"name": "user_config_map_ref",
"optional": "True",
}
},
{
"secretRef": {
"name": "user_secret_ref_one",
"optional": "True"
}
},
{
"secretRef": {
"name": "user_secret_ref_two",
"optional": "False",
},
"prefix": "with_prefix",
},
]
}
}
}).tags)
cfg = DagsterK8sJobConfig(
job_image="test/foo:latest",
dagster_home="/opt/dagster/dagster_home",
instance_config_map="some-instance-configmap",
env_config_maps=["config_map"],
env_secrets=["secret"],
)
job = construct_dagster_k8s_job(
cfg, ["foo", "bar"],
"job",
user_defined_k8s_config=user_defined_k8s_config).to_dict()
env_from = job["spec"]["template"]["spec"]["containers"][0]["env_from"]
env_from_mapping = {(env_var.get("config_map_ref")
or env_var.get("secret_ref")).get("name"): env_var
for env_var in env_from}
assert len(env_from_mapping) == 5
assert env_from_mapping["config_map"]
assert env_from_mapping["user_config_map_ref"]
assert env_from_mapping["secret"]
assert env_from_mapping["user_secret_ref_one"]
assert env_from_mapping["user_secret_ref_two"]
def _submit_task_k8s_job(app, pipeline_context, step, queue, priority):
user_defined_k8s_config = get_user_defined_k8s_config(step.tags)
task = create_k8s_job_task(app)
recon_repo = pipeline_context.pipeline.get_reconstructable_repository()
task_signature = task.si(
instance_ref_dict=pipeline_context.instance.get_ref().to_dict(),
step_keys=[step.key],
run_config=pipeline_context.pipeline_run.run_config,
mode=pipeline_context.pipeline_run.mode,
repo_name=recon_repo.get_definition().name,
repo_location_name=pipeline_context.executor.repo_location_name,
run_id=pipeline_context.pipeline_run.run_id,
job_config_dict=pipeline_context.executor.job_config.to_dict(),
job_namespace=pipeline_context.executor.job_namespace,
user_defined_k8s_config_dict=user_defined_k8s_config.to_dict(),
retries_dict=pipeline_context.executor.retries.for_inner_plan().to_config(),
pipeline_origin_packed=pack_value(pipeline_context.pipeline.get_origin()),
load_incluster_config=pipeline_context.executor.load_incluster_config,
kubeconfig_file=pipeline_context.executor.kubeconfig_file,
)
return task_signature.apply_async(
priority=priority,
queue=queue,
routing_key="{queue}.execute_step_k8s_job".format(queue=queue),
)
def test_backcompat_resource_tags():
@solid(
tags={
K8S_RESOURCE_REQUIREMENTS_KEY: {
"requests": {
"cpu": "250m",
"memory": "64Mi"
},
"limits": {
"cpu": "500m",
"memory": "2560Mi"
},
}
})
def resource_tags_solid(_):
pass
user_defined_k8s_config = get_user_defined_k8s_config(
resource_tags_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"
def test_construct_dagster_k8s_job_with_user_defined_env():
@graph
def user_defined_k8s_env_tags_graph():
pass
user_defined_k8s_config = get_user_defined_k8s_config(
user_defined_k8s_env_tags_graph.to_job(
tags={
USER_DEFINED_K8S_CONFIG_KEY: {
"container_config": {
"env": [
{
"name": "ENV_VAR_1",
"value": "one"
},
{
"name": "ENV_VAR_2",
"value": "two"
},
{
"name": "DD_AGENT_HOST",
"valueFrom": {
"fieldRef": {
"fieldPath": "status.hostIP"
}
},
},
]
}
}
}).tags)
cfg = DagsterK8sJobConfig(
job_image="test/foo:latest",
dagster_home="/opt/dagster/dagster_home",
instance_config_map="some-instance-configmap",
)
job = construct_dagster_k8s_job(
cfg, ["foo", "bar"],
"job",
user_defined_k8s_config=user_defined_k8s_config).to_dict()
env = job["spec"]["template"]["spec"]["containers"][0]["env"]
env_mapping = remove_none_recursively(
{env_var["name"]: env_var
for env_var in env})
# Has DAGSTER_HOME and three additional env vars
assert len(env_mapping) == 4
assert env_mapping["ENV_VAR_1"]["value"] == "one"
assert env_mapping["ENV_VAR_2"]["value"] == "two"
assert env_mapping["DD_AGENT_HOST"]["value_from"] == {
"field_ref": {
"field_path": "status.hostIP"
}
}
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_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_construct_dagster_k8s_job_with_user_defined_volume_mounts_snake_case(
):
@graph
def user_defined_k8s_volume_mounts_tags_graph():
pass
# volume_mounts still work even when using underscore keys
user_defined_k8s_config = get_user_defined_k8s_config(
user_defined_k8s_volume_mounts_tags_graph.to_job(
tags={
USER_DEFINED_K8S_CONFIG_KEY: {
"container_config": {
"volume_mounts": [
{
"mountPath": "mount_path",
"mountPropagation": "mount_propagation",
"name": "a_volume_mount_one",
"readOnly": "False",
"subPath": "path/",
},
{
"mountPath": "mount_path",
"mountPropagation": "mount_propagation",
"name": "a_volume_mount_two",
"readOnly": "False",
"subPathExpr": "path/",
},
]
}
}
}).tags)
cfg = DagsterK8sJobConfig(
job_image="test/foo:latest",
dagster_home="/opt/dagster/dagster_home",
instance_config_map="some-instance-configmap",
)
job = construct_dagster_k8s_job(
cfg, ["foo", "bar"],
"job",
user_defined_k8s_config=user_defined_k8s_config).to_dict()
volume_mounts = job["spec"]["template"]["spec"]["containers"][0][
"volume_mounts"]
volume_mounts_mapping = {
volume_mount["name"]: volume_mount
for volume_mount in volume_mounts
}
assert len(volume_mounts_mapping) == 3
assert volume_mounts_mapping["dagster-instance"]
assert volume_mounts_mapping["a_volume_mount_one"]
assert volume_mounts_mapping["a_volume_mount_two"]
def launch_steps(
self,
step_contexts: List[IStepContext],
known_state: KnownExecutionState,
):
assert len(
step_contexts
) == 1, "Launching multiple steps is not currently supported"
step_context = step_contexts[0]
k8s_name_key = get_k8s_job_name(
self.pipeline_context.plan_data.pipeline_run.run_id,
step_context.step.key,
)
job_name = "dagster-job-%s" % (k8s_name_key)
pod_name = "dagster-job-%s" % (k8s_name_key)
pipeline_origin = self.pipeline_context.reconstructable_pipeline.get_python_origin(
)
execute_step_args = ExecuteStepArgs(
pipeline_origin=pipeline_origin,
pipeline_run_id=self.pipeline_context.pipeline_run.run_id,
step_keys_to_execute=[step_context.step.key],
instance_ref=self.pipeline_context.instance.get_ref(),
retry_mode=self.retries.for_inner_plan(),
known_state=known_state,
should_verify_step=True,
)
input_json = serialize_dagster_namedtuple(execute_step_args)
args = ["dagster", "api", "execute_step", input_json]
job_config = self._job_config
if not job_config.job_image:
job_config = job_config.with_image(
pipeline_origin.repository_origin.container_image)
if not job_config.job_image:
raise Exception(
"No image included in either executor config or the pipeline")
job = construct_dagster_k8s_job(
job_config,
args,
job_name,
get_user_defined_k8s_config(frozentags()),
pod_name,
)
kubernetes.config.load_incluster_config()
kubernetes.client.BatchV1Api().create_namespaced_job(
body=job, namespace=self._job_namespace)
def test_k8s_tag_op():
assert my_op
user_defined_cfg = get_user_defined_k8s_config(my_op.tags)
cfg = DagsterK8sJobConfig(
job_image="test/foo:latest",
dagster_home="/opt/dagster/dagster_home",
instance_config_map="test",
)
job = construct_dagster_k8s_job(cfg, [],
"job123",
user_defined_k8s_config=user_defined_cfg)
assert job.to_dict(
)["spec"]["template"]["spec"]["containers"][0]["resources"] == {
"requests": {
"cpu": "200m",
"memory": "32Mi"
},
"limits": None,
}
def _submit_task_k8s_job(app, plan_context, step, queue, priority, known_state):
user_defined_k8s_config = get_user_defined_k8s_config(step.tags)
pipeline_origin = plan_context.reconstructable_pipeline.get_python_origin()
execute_step_args = ExecuteStepArgs(
pipeline_origin=pipeline_origin,
pipeline_run_id=plan_context.pipeline_run.run_id,
step_keys_to_execute=[step.key],
instance_ref=plan_context.instance.get_ref(),
retry_mode=plan_context.executor.retries.for_inner_plan(),
known_state=known_state,
should_verify_step=True,
)
job_config = plan_context.executor.job_config
if not job_config.job_image:
job_config = job_config.with_image(pipeline_origin.repository_origin.container_image)
if not job_config.job_image:
raise Exception("No image included in either executor config or the dagster job")
task = create_k8s_job_task(app)
task_signature = task.si(
execute_step_args_packed=pack_value(execute_step_args),
job_config_dict=job_config.to_dict(),
job_namespace=plan_context.executor.job_namespace,
user_defined_k8s_config_dict=user_defined_k8s_config.to_dict(),
load_incluster_config=plan_context.executor.load_incluster_config,
job_wait_timeout=plan_context.executor.job_wait_timeout,
kubeconfig_file=plan_context.executor.kubeconfig_file,
)
return task_signature.apply_async(
priority=priority,
queue=queue,
routing_key="{queue}.execute_step_k8s_job".format(queue=queue),
)
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 launch_run(self, instance, run, external_pipeline):
check.inst_param(instance, "instance", DagsterInstance)
check.inst_param(run, "run", PipelineRun)
check.inst_param(external_pipeline, "external_pipeline", ExternalPipeline)
job_name = get_job_name_from_run_id(run.run_id)
pod_name = job_name
exc_config = _get_validated_celery_k8s_executor_config(run.run_config)
job_image = None
pipeline_origin = None
env_vars = None
if isinstance(external_pipeline.get_origin(), PipelineGrpcServerOrigin):
if exc_config.get("job_image"):
raise DagsterInvariantViolationError(
"Cannot specify job_image in executor config when loading pipeline "
"from GRPC server."
)
repository_location_handle = (
external_pipeline.repository_handle.repository_location_handle
)
if not isinstance(repository_location_handle, GrpcServerRepositoryLocationHandle):
raise DagsterInvariantViolationError(
"Expected RepositoryLocationHandle to be of type "
"GrpcServerRepositoryLocationHandle but found type {}".format(
type(repository_location_handle)
)
)
job_image = repository_location_handle.get_current_image()
env_vars = {"DAGSTER_CURRENT_IMAGE": job_image}
repository_name = external_pipeline.repository_handle.repository_name
pipeline_origin = PipelinePythonOrigin(
pipeline_name=external_pipeline.name,
repository_origin=repository_location_handle.get_repository_python_origin(
repository_name
),
)
else:
job_image = exc_config.get("job_image")
if not job_image:
raise DagsterInvariantViolationError(
"Cannot find job_image in celery-k8s executor config."
)
pipeline_origin = external_pipeline.get_origin()
job_config = DagsterK8sJobConfig(
dagster_home=self.dagster_home,
instance_config_map=self.instance_config_map,
postgres_password_secret=self.postgres_password_secret,
job_image=check.str_param(job_image, "job_image"),
image_pull_policy=exc_config.get("image_pull_policy"),
image_pull_secrets=exc_config.get("image_pull_secrets"),
service_account_name=exc_config.get("service_account_name"),
env_config_maps=exc_config.get("env_config_maps"),
env_secrets=exc_config.get("env_secrets"),
)
user_defined_k8s_config = get_user_defined_k8s_config(frozentags(external_pipeline.tags))
from dagster.cli.api import ExecuteRunArgs
input_json = serialize_dagster_namedtuple(
# depends on DagsterInstance.get() returning the same instance
# https://github.com/dagster-io/dagster/issues/2757
ExecuteRunArgs(
pipeline_origin=pipeline_origin, pipeline_run_id=run.run_id, instance_ref=None,
)
)
job = construct_dagster_k8s_job(
job_config,
command=["dagster"],
args=["api", "execute_run_with_structured_logs", input_json],
job_name=job_name,
pod_name=pod_name,
component="run_coordinator",
user_defined_k8s_config=user_defined_k8s_config,
env_vars=env_vars,
)
job_namespace = exc_config.get("job_namespace")
api = kubernetes.client.BatchV1Api()
api.create_namespaced_job(body=job, namespace=job_namespace)
self._instance.report_engine_event(
"Kubernetes run_coordinator job launched",
run,
EngineEventData(
[
EventMetadataEntry.text(job_name, "Kubernetes Job name"),
EventMetadataEntry.text(pod_name, "Kubernetes Pod name"),
EventMetadataEntry.text(job_namespace, "Kubernetes Namespace"),
EventMetadataEntry.text(run.run_id, "Run ID"),
]
),
cls=self.__class__,
)
return run
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 launch_run(self, instance, run, external_pipeline):
check.inst_param(instance, "instance", DagsterInstance)
check.inst_param(run, "run", PipelineRun)
check.inst_param(external_pipeline, "external_pipeline",
ExternalPipeline)
job_name = get_job_name_from_run_id(run.run_id)
pod_name = job_name
exc_config = _get_validated_celery_k8s_executor_config(run.run_config)
job_image = None
pipeline_origin = None
env_vars = None
job_image_from_executor_config = exc_config.get("job_image")
# If the user is using user-code deployments, we grab the image from the gRPC server.
if isinstance(
external_pipeline.get_external_origin().
external_repository_origin.repository_location_origin,
GrpcServerRepositoryLocationOrigin,
):
repository_location_handle = (
external_pipeline.repository_handle.repository_location_handle)
if not isinstance(repository_location_handle,
GrpcServerRepositoryLocationHandle):
raise DagsterInvariantViolationError(
"Expected RepositoryLocationHandle to be of type "
"GrpcServerRepositoryLocationHandle but found type {}".
format(type(repository_location_handle)))
repository_name = external_pipeline.repository_handle.repository_name
repository_origin = repository_location_handle.reload_repository_python_origin(
repository_name)
pipeline_origin = PipelinePythonOrigin(
pipeline_name=external_pipeline.name,
repository_origin=repository_origin)
job_image = repository_origin.container_image
env_vars = {"DAGSTER_CURRENT_IMAGE": job_image}
if job_image_from_executor_config:
raise DagsterInvariantViolationError(
"You have specified a job_image {job_image_from_executor_config} in your executor configuration, "
"but also {job_image} in your user-code deployment. You cannot specify a job_image "
"in your executor config when using user-code deployments because the job image is "
"pulled from the deployment. To resolve this error, remove the job_image "
"configuration from your executor configuration (which is a part of your run configuration)"
)
else:
if not job_image_from_executor_config:
raise DagsterInvariantViolationError(
"You have not specified a job_image in your executor configuration. "
"To resolve this error, specify the job_image configuration in the executor "
"config section in your run config. \n"
"Note: You may also be seeing this error because you are using the configured API. "
"Using configured with the celery-k8s executor is not supported at this time, "
"and the job_image must be configured at the top-level executor config without "
"using configured.")
job_image = job_image_from_executor_config
pipeline_origin = external_pipeline.get_python_origin()
job_config = DagsterK8sJobConfig(
dagster_home=self.dagster_home,
instance_config_map=self.instance_config_map,
postgres_password_secret=self.postgres_password_secret,
job_image=check.str_param(job_image, "job_image"),
image_pull_policy=exc_config.get("image_pull_policy"),
image_pull_secrets=exc_config.get("image_pull_secrets"),
service_account_name=exc_config.get("service_account_name"),
env_config_maps=exc_config.get("env_config_maps"),
env_secrets=exc_config.get("env_secrets"),
)
user_defined_k8s_config = get_user_defined_k8s_config(
frozentags(run.tags))
from dagster.cli.api import ExecuteRunArgs
input_json = serialize_dagster_namedtuple(
# depends on DagsterInstance.get() returning the same instance
# https://github.com/dagster-io/dagster/issues/2757
ExecuteRunArgs(
pipeline_origin=pipeline_origin,
pipeline_run_id=run.run_id,
instance_ref=None,
))
job = construct_dagster_k8s_job(
job_config,
args=["dagster", "api", "execute_run", input_json],
job_name=job_name,
pod_name=pod_name,
component="run_coordinator",
user_defined_k8s_config=user_defined_k8s_config,
env_vars=env_vars,
)
job_namespace = exc_config.get("job_namespace")
self._batch_api.create_namespaced_job(body=job,
namespace=job_namespace)
self._instance.report_engine_event(
"Kubernetes run_coordinator job launched",
run,
EngineEventData([
EventMetadataEntry.text(job_name, "Kubernetes Job name"),
EventMetadataEntry.text(job_namespace, "Kubernetes Namespace"),
EventMetadataEntry.text(run.run_id, "Run ID"),
]),
cls=self.__class__,
)
return run
def launch_run(self, run, external_pipeline):
check.inst_param(run, "run", PipelineRun)
check.inst_param(external_pipeline, "external_pipeline", ExternalPipeline)
job_name = get_job_name_from_run_id(run.run_id)
pod_name = job_name
exc_config = _get_validated_celery_k8s_executor_config(run.run_config)
env_vars = None
job_image_from_executor_config = exc_config.get("job_image")
pipeline_origin = external_pipeline.get_python_origin()
repository_origin = pipeline_origin.repository_origin
job_image = repository_origin.container_image
if job_image:
if job_image_from_executor_config:
job_image = job_image_from_executor_config
self._instance.report_engine_event(
f"You have specified a job_image {job_image_from_executor_config} in your executor configuration, "
f"but also {job_image} in your user-code deployment. Using the job image {job_image_from_executor_config} "
f"from executor configuration as it takes precedence.",
run,
cls=self.__class__,
)
else:
if not job_image_from_executor_config:
raise DagsterInvariantViolationError(
"You have not specified a job_image in your executor configuration. "
"To resolve this error, specify the job_image configuration in the executor "
"config section in your run config. \n"
"Note: You may also be seeing this error because you are using the configured API. "
"Using configured with the celery-k8s executor is not supported at this time, "
"and the job_image must be configured at the top-level executor config without "
"using configured."
)
job_image = job_image_from_executor_config
job_config = DagsterK8sJobConfig(
dagster_home=self.dagster_home,
instance_config_map=self.instance_config_map,
postgres_password_secret=self.postgres_password_secret,
job_image=check.str_param(job_image, "job_image"),
image_pull_policy=exc_config.get("image_pull_policy"),
image_pull_secrets=exc_config.get("image_pull_secrets"),
service_account_name=exc_config.get("service_account_name"),
env_config_maps=exc_config.get("env_config_maps"),
env_secrets=exc_config.get("env_secrets"),
)
self._instance.add_run_tags(
run.run_id,
{DOCKER_IMAGE_TAG: job_config.job_image},
)
user_defined_k8s_config = get_user_defined_k8s_config(frozentags(run.tags))
from dagster.cli.api import ExecuteRunArgs
input_json = serialize_dagster_namedtuple(
# depends on DagsterInstance.get() returning the same instance
# https://github.com/dagster-io/dagster/issues/2757
ExecuteRunArgs(
pipeline_origin=pipeline_origin,
pipeline_run_id=run.run_id,
instance_ref=None,
)
)
job = construct_dagster_k8s_job(
job_config,
args=["dagster", "api", "execute_run", input_json],
job_name=job_name,
pod_name=pod_name,
component="run_coordinator",
user_defined_k8s_config=user_defined_k8s_config,
env_vars=env_vars,
)
job_namespace = exc_config.get("job_namespace")
self._instance.report_engine_event(
"Creating Kubernetes run worker job",
run,
EngineEventData(
[
EventMetadataEntry.text(job_name, "Kubernetes Job name"),
EventMetadataEntry.text(job_namespace, "Kubernetes Namespace"),
EventMetadataEntry.text(run.run_id, "Run ID"),
]
),
cls=self.__class__,
)
self._batch_api.create_namespaced_job(body=job, namespace=job_namespace)
self._instance.report_engine_event(
"Kubernetes run worker job created",
run,
EngineEventData(
[
EventMetadataEntry.text(job_name, "Kubernetes Job name"),
EventMetadataEntry.text(job_namespace, "Kubernetes Namespace"),
EventMetadataEntry.text(run.run_id, "Run ID"),
]
),
cls=self.__class__,
)
return run
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())
def launch_run(self, context: LaunchRunContext) -> None:
run = context.pipeline_run
job_name = get_job_name_from_run_id(run.run_id)
pod_name = job_name
exc_config = _get_validated_celery_k8s_executor_config(run.run_config)
env_vars = None
job_image_from_executor_config = exc_config.get("job_image")
pipeline_origin = context.pipeline_code_origin
repository_origin = pipeline_origin.repository_origin
job_image = repository_origin.container_image
if job_image:
if job_image_from_executor_config:
job_image = job_image_from_executor_config
self._instance.report_engine_event(
f"You have specified a job_image {job_image_from_executor_config} in your executor configuration, "
f"but also {job_image} in your user-code deployment. Using the job image {job_image_from_executor_config} "
f"from executor configuration as it takes precedence.",
run,
cls=self.__class__,
)
else:
if not job_image_from_executor_config:
raise DagsterInvariantViolationError(
"You have not specified a job_image in your executor configuration. "
"To resolve this error, specify the job_image configuration in the executor "
"config section in your run config. \n"
"Note: You may also be seeing this error because you are using the configured API. "
"Using configured with the celery-k8s executor is not supported at this time, "
"and the job_image must be configured at the top-level executor config without "
"using configured.")
job_image = job_image_from_executor_config
job_config = self.get_k8s_job_config(job_image, exc_config)
self._instance.add_run_tags(
run.run_id,
{DOCKER_IMAGE_TAG: job_config.job_image},
)
user_defined_k8s_config = get_user_defined_k8s_config(
frozentags(run.tags))
from dagster.cli.api import ExecuteRunArgs
run_args = ExecuteRunArgs(
pipeline_origin=pipeline_origin,
pipeline_run_id=run.run_id,
instance_ref=self._instance.get_ref(),
)
job = construct_dagster_k8s_job(
job_config,
args=run_args.get_command_args(),
job_name=job_name,
pod_name=pod_name,
component="run_worker",
user_defined_k8s_config=user_defined_k8s_config,
env_vars=env_vars,
labels={
"dagster/job": pipeline_origin.pipeline_name,
},
)
job_namespace = exc_config.get("job_namespace")
self._instance.report_engine_event(
"Creating Kubernetes run worker job",
run,
EngineEventData([
EventMetadataEntry.text(job_name, "Kubernetes Job name"),
EventMetadataEntry.text(job_namespace, "Kubernetes Namespace"),
EventMetadataEntry.text(run.run_id, "Run ID"),
]),
cls=self.__class__,
)
self._batch_api.create_namespaced_job(body=job,
namespace=job_namespace)
self._instance.report_engine_event(
"Kubernetes run worker job created",
run,
EngineEventData([
EventMetadataEntry.text(job_name, "Kubernetes Job name"),
EventMetadataEntry.text(job_namespace, "Kubernetes Namespace"),
EventMetadataEntry.text(run.run_id, "Run ID"),
]),
cls=self.__class__,
)
def test_tags_to_dynamic_plan():
@solid(
tags={
USER_DEFINED_K8S_CONFIG_KEY: {
"container_config": {
"resources": {
"requests": {
"cpu": "500m",
"memory": "128Mi"
},
"limits": {
"cpu": "1000m",
"memory": "1Gi"
},
}
}
}
})
def multiply_inputs(_, x):
return 2 * x
@solid(
tags={
USER_DEFINED_K8S_CONFIG_KEY: {
"container_config": {
"resources": {
"requests": {
"cpu": "250m",
"memory": "64Mi"
},
"limits": {
"cpu": "500m",
"memory": "2560Mi"
},
}
}
}
},
output_defs=[DynamicOutputDefinition()],
)
def emit(_):
for i in range(3):
yield DynamicOutput(value=i, mapping_key=str(i))
@pipeline
def k8s_ready():
return emit().map(multiply_inputs)
known_state = KnownExecutionState(
{},
{
emit.name: {
"result": ["0", "1", "2"]
},
},
)
plan = create_execution_plan(k8s_ready, known_state=known_state)
emit_step = plan.get_step_by_key(emit.name)
user_defined_k8s_config = get_user_defined_k8s_config(emit_step.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"
for mapping_key in range(3):
multiply_inputs_step = plan.get_step_by_key(
f"{multiply_inputs.name}[{mapping_key}]")
dynamic_step_user_defined_k8s_config = get_user_defined_k8s_config(
multiply_inputs_step.tags)
assert dynamic_step_user_defined_k8s_config.container_config
assert dynamic_step_user_defined_k8s_config.container_config[
"resources"]
resources = dynamic_step_user_defined_k8s_config.container_config[
"resources"]
assert resources["requests"]["cpu"] == "500m"
assert resources["requests"]["memory"] == "128Mi"
assert resources["limits"]["cpu"] == "1000m"
assert resources["limits"]["memory"] == "1Gi"
