def test_job_serialization():
cfg = DagsterK8sJobConfig(
job_image='test/foo:latest',
dagster_home='/opt/dagster/dagster_home',
image_pull_policy='Always',
image_pull_secrets=None,
service_account_name=None,
instance_config_map='some-instance-configmap',
postgres_password_secret='some-secret-name',
env_config_maps=None,
env_secrets=None,
)
assert DagsterK8sJobConfig.from_dict(cfg.to_dict()) == cfg
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 celery_k8s_config():
from dagster_k8s.job import DagsterK8sJobConfig
# DagsterK8sJobConfig provides config schema for specifying Dagster K8s Jobs
job_config = DagsterK8sJobConfig.config_type()
additional_config = {
'load_incluster_config': Field(
bool,
is_required=False,
default_value=True,
description='''Set this value if you are running the launcher within a k8s cluster. If
``True``, we assume the launcher is running within the target cluster and load config
using ``kubernetes.config.load_incluster_config``. Otherwise, we will use the k8s config
specified in ``kubeconfig_file`` (using ``kubernetes.config.load_kube_config``) or fall
back to the default kubeconfig. Default: ``True``.''',
),
'kubeconfig_file': Field(
Noneable(str),
is_required=False,
description='Path to a kubeconfig file to use, if not using default kubeconfig.',
),
'job_namespace': Field(
StringSource,
is_required=False,
default_value='default',
description='The namespace into which to launch new jobs. Note that any '
'other Kubernetes resources the Job requires (such as the service account) must be '
'present in this namespace. Default: ``"default"``',
),
}
cfg = merge_dicts(CELERY_CONFIG, job_config)
cfg = merge_dicts(cfg, additional_config)
return cfg
def config_type(cls):
job_cfg = DagsterK8sJobConfig.config_type()
scheduler_extra_cfg = {
'scheduler_namespace': Field(StringSource, is_required=True),
'load_incluster_config': Field(bool, is_required=False, default_value=True),
'kubeconfig_file': Field(Noneable(StringSource), is_required=False, default_value=None),
}
return merge_dicts(job_cfg, scheduler_extra_cfg)
def __init__(
self,
dagster_home,
service_account_name,
instance_config_map,
postgres_password_secret,
job_image,
load_incluster_config=True,
scheduler_namespace="default",
image_pull_policy="Always",
image_pull_secrets=None,
kubeconfig_file=None,
inst_data=None,
env_config_maps=None,
env_secrets=None,
):
self._inst_data = check.opt_inst_param(inst_data, "inst_data",
ConfigurableClassData)
if load_incluster_config:
check.invariant(
kubeconfig_file is None,
"`kubeconfig_file` is set but `load_incluster_config` is True.",
)
kubernetes.config.load_incluster_config()
else:
check.opt_str_param(kubeconfig_file, "kubeconfig_file")
kubernetes.config.load_kube_config(kubeconfig_file)
self._api = kubernetes.client.BatchV1beta1Api()
self._namespace = check.str_param(scheduler_namespace,
"scheduler_namespace")
self.grace_period_seconds = 5 # This should be passed in via config
self.job_config = DagsterK8sJobConfig(
job_image=check.str_param(job_image, "job_image"),
dagster_home=check.str_param(dagster_home, "dagster_home"),
image_pull_policy=check.str_param(image_pull_policy,
"image_pull_policy"),
image_pull_secrets=check.opt_list_param(image_pull_secrets,
"image_pull_secrets",
of_type=dict),
service_account_name=check.str_param(service_account_name,
"service_account_name"),
instance_config_map=check.str_param(instance_config_map,
"instance_config_map"),
postgres_password_secret=check.str_param(
postgres_password_secret, "postgres_password_secret"),
env_config_maps=check.opt_list_param(env_config_maps,
"env_config_maps",
of_type=str),
env_secrets=check.opt_list_param(env_secrets,
"env_secrets",
of_type=str),
)
def __init__(
self,
dagster_home,
service_account_name,
instance_config_map,
postgres_password_secret,
job_image,
load_incluster_config=True,
scheduler_namespace='default',
image_pull_policy='Always',
image_pull_secrets=None,
kubeconfig_file=None,
inst_data=None,
env_config_maps=None,
env_secrets=None,
):
self._inst_data = check.opt_inst_param(inst_data, 'inst_data',
ConfigurableClassData)
if load_incluster_config:
check.invariant(
kubeconfig_file is None,
'`kubeconfig_file` is set but `load_incluster_config` is True.',
)
kubernetes.config.load_incluster_config()
else:
check.opt_str_param(kubeconfig_file, 'kubeconfig_file')
kubernetes.config.load_kube_config(kubeconfig_file)
self._api = kubernetes.client.BatchV1beta1Api()
self._namespace = check.str_param(scheduler_namespace,
'scheduler_namespace')
self.grace_period_seconds = 5 # This should be passed in via config
self.job_config = DagsterK8sJobConfig(
job_image=check.str_param(job_image, 'job_image'),
dagster_home=check.str_param(dagster_home, 'dagster_home'),
image_pull_policy=check.str_param(image_pull_policy,
'image_pull_policy'),
image_pull_secrets=check.opt_list_param(image_pull_secrets,
'image_pull_secrets',
of_type=dict),
service_account_name=check.str_param(service_account_name,
'service_account_name'),
instance_config_map=check.str_param(instance_config_map,
'instance_config_map'),
postgres_password_secret=check.str_param(
postgres_password_secret, 'postgres_password_secret'),
env_config_maps=check.opt_list_param(env_config_maps,
'env_config_maps',
of_type=str),
env_secrets=check.opt_list_param(env_secrets,
'env_secrets',
of_type=str),
)
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 config_type(cls):
"""Include all arguments required for DagsterK8sJobConfig along with additional arguments
needed for the RunLauncher itself.
"""
from dagster_celery.executor import CELERY_CONFIG
job_cfg = DagsterK8sJobConfig.config_type_run_launcher()
run_launcher_extra_cfg = {
"load_incluster_config": Field(bool, is_required=False, default_value=True),
"kubeconfig_file": Field(Noneable(str), is_required=False, default_value=None),
}
res = merge_dicts(job_cfg, run_launcher_extra_cfg)
return merge_dicts(res, CELERY_CONFIG)
def get_k8s_job_config(self, job_image, exc_config):
return DagsterK8sJobConfig(
dagster_home=self.dagster_home,
instance_config_map=self.instance_config_map,
postgres_password_secret=self.postgres_password_secret,
job_image=check.opt_str_param(job_image, "job_image"),
image_pull_policy=exc_config.get("image_pull_policy", self._image_pull_policy),
image_pull_secrets=exc_config.get("image_pull_secrets", []) + self._image_pull_secrets,
service_account_name=exc_config.get("service_account_name", self._service_account_name),
env_config_maps=exc_config.get("env_config_maps", []) + self._env_config_maps,
env_secrets=exc_config.get("env_secrets", []) + self._env_secrets,
volume_mounts=exc_config.get("volume_mounts", []) + self._volume_mounts,
volumes=exc_config.get("volumes", []) + self._volumes,
labels=merge_dicts(self._labels, exc_config.get("labels", {})),
)
def dagster_k8s_executor(init_context: InitExecutorContext) -> Executor:
run_launcher = init_context.instance.run_launcher
exc_cfg = init_context.executor_config
job_config = DagsterK8sJobConfig(
dagster_home=run_launcher.dagster_home,
instance_config_map=run_launcher.instance_config_map,
postgres_password_secret=run_launcher.postgres_password_secret,
job_image=exc_cfg.get("job_image") or os.getenv("DAGSTER_CURRENT_IMAGE"),
image_pull_policy=exc_cfg.get("image_pull_policy"),
image_pull_secrets=exc_cfg.get("image_pull_secrets"),
service_account_name=exc_cfg.get("service_account_name"),
env_config_maps=exc_cfg.get("env_config_maps"),
env_secrets=exc_cfg.get("env_secrets"),
)
return StepDelegatingExecutor(
K8sStepHandler(
retries=RetryMode.DISABLED, # Not currently supported
job_config=job_config,
job_namespace=exc_cfg.get("job_namespace"),
)
)
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 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 _execute_step_k8s_job(
_self,
instance_ref_dict,
step_keys,
environment_dict,
mode,
pipeline_name,
run_id,
job_config_dict,
job_namespace,
load_incluster_config,
kubeconfig_file=None,
):
'''Run step execution in a K8s job pod.
'''
from dagster_k8s.job import DagsterK8sJobConfig, construct_dagster_graphql_k8s_job
from dagster_k8s.utils import get_pod_names_in_job, retrieve_pod_logs, wait_for_job_success
import kubernetes
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(environment_dict, 'environment_dict')
check.str_param(mode, 'mode')
check.str_param(pipeline_name, 'pipeline_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.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_keys_str = ", ".join(step_keys)
# Ensure we stay below k8s name length limits
k8s_name_key = _get_k8s_name_key(run_id, step_keys)
job_name = 'dagster-stepjob-%s' % k8s_name_key
pod_name = 'dagster-stepjob-%s' % k8s_name_key
variables = construct_variables(mode, environment_dict, pipeline_name, run_id, step_keys)
args = ['-p', 'executePlan', '-v', seven.json.dumps(variables)]
job = construct_dagster_graphql_k8s_job(job_config, args, job_name, pod_name)
# Running list of events generated from this task execution
events = []
# Post event for starting execution
engine_event = instance.report_engine_event(
'Executing steps {} in Kubernetes job {}'.format(step_keys_str, job.metadata.name),
pipeline_run,
EngineEventData(
[
EventMetadataEntry.text(step_keys_str, 'Step keys'),
EventMetadataEntry.text(job.metadata.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,
),
CeleryK8sJobEngine,
# 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_keys[0],
)
events.append(engine_event)
kubernetes.client.BatchV1Api().create_namespaced_job(body=job, namespace=job_namespace)
wait_for_job_success(job.metadata.name, namespace=job_namespace)
pod_names = get_pod_names_in_job(job.metadata.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')]),
CeleryK8sJobEngine,
step_key=step_keys[0],
)
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')
res = parse_raw_log_lines(logs)
handle_execution_errors(res, 'executePlan')
step_events = handle_execute_plan_result(res)
events += step_events
serialized_events = [serialize_dagster_namedtuple(event) for event in events]
return serialized_events
def __init__(
self,
dagster_home,
service_account_name,
instance_config_map,
postgres_password_secret,
job_image,
load_incluster_config=True,
scheduler_namespace="default",
image_pull_policy="Always",
image_pull_secrets=None,
kubeconfig_file=None,
inst_data=None,
env_config_maps=None,
env_secrets=None,
):
warnings.warn(
"`K8sScheduler` is deprecated and will be removed in the 0.12.0 dagster release."
" We recommend that you use the Dagster native scheduler instead, which runs automatically "
" as part of the dagster-daemon process. You can configure this scheduler by removing "
" the `scheduler` key from your `dagster.yaml` file. See"
" https://docs.dagster.io/deployment/dagster-daemon for more information on how to deploy."
)
self._inst_data = check.opt_inst_param(inst_data, "inst_data",
ConfigurableClassData)
if load_incluster_config:
check.invariant(
kubeconfig_file is None,
"`kubeconfig_file` is set but `load_incluster_config` is True.",
)
kubernetes.config.load_incluster_config()
else:
check.opt_str_param(kubeconfig_file, "kubeconfig_file")
kubernetes.config.load_kube_config(kubeconfig_file)
self._api = kubernetes.client.BatchV1beta1Api()
self._namespace = check.str_param(scheduler_namespace,
"scheduler_namespace")
self.grace_period_seconds = 5 # This should be passed in via config
self.job_config = DagsterK8sJobConfig(
job_image=check.str_param(job_image, "job_image"),
dagster_home=check.str_param(dagster_home, "dagster_home"),
image_pull_policy=check.str_param(image_pull_policy,
"image_pull_policy"),
image_pull_secrets=check.opt_list_param(image_pull_secrets,
"image_pull_secrets",
of_type=dict),
service_account_name=check.str_param(service_account_name,
"service_account_name"),
instance_config_map=check.str_param(instance_config_map,
"instance_config_map"),
postgres_password_secret=check.str_param(
postgres_password_secret, "postgres_password_secret"),
env_config_maps=check.opt_list_param(env_config_maps,
"env_config_maps",
of_type=str),
env_secrets=check.opt_list_param(env_secrets,
"env_secrets",
of_type=str),
)
def celery_k8s_job_executor(init_context):
'''Celery-based executor which launches tasks as Kubernetes Jobs.
The Celery executor exposes config settings for the underlying Celery app under
the ``config_source`` key. This config corresponds to the "new lowercase settings" introduced
in Celery version 4.0 and the object constructed from config will be passed to the
:py:class:`celery.Celery` constructor as its ``config_source`` argument.
(See https://docs.celeryproject.org/en/latest/userguide/configuration.html for details.)
The executor also exposes the ``broker``, `backend`, and ``include`` arguments to the
:py:class:`celery.Celery` constructor.
In the most common case, you may want to modify the ``broker`` and ``backend`` (e.g., to use
Redis instead of RabbitMQ). We expect that ``config_source`` will be less frequently
modified, but that when solid executions are especially fast or slow, or when there are
different requirements around idempotence or retry, it may make sense to execute pipelines
with variations on these settings.
If you'd like to configure a Celery Kubernetes Job executor in addition to the
:py:class:`~dagster.default_executors`, you should add it to the ``executor_defs`` defined on a
:py:class:`~dagster.ModeDefinition` as follows:
.. code-block:: python
from dagster import ModeDefinition, default_executors, pipeline
from dagster_celery.executor_k8s import celery_k8s_job_executor
@pipeline(mode_defs=[
ModeDefinition(executor_defs=default_executors + [celery_k8s_job_executor])
])
def celery_enabled_pipeline():
pass
Then you can configure the executor as follows:
.. code-block:: YAML
execution:
celery-k8s:
config:
job_image: 'my_repo.com/image_name:latest'
job_namespace: 'some-namespace'
broker: 'pyamqp://guest@localhost//' # Optional[str]: The URL of the Celery broker
backend: 'rpc://' # Optional[str]: The URL of the Celery results backend
include: ['my_module'] # Optional[List[str]]: Modules every worker should import
config_source: # Dict[str, Any]: Any additional parameters to pass to the
#... # Celery workers. This dict will be passed as the `config_source`
#... # argument of celery.Celery().
Note that the YAML you provide here must align with the configuration with which the Celery
workers on which you hope to run were started. If, for example, you point the executor at a
different broker than the one your workers are listening to, the workers will never be able to
pick up tasks for execution.
'''
from dagster_k8s.job import DagsterK8sJobConfig
check_cross_process_constraints(init_context)
job_config = DagsterK8sJobConfig(
job_image=init_context.executor_config.get('job_image'),
dagster_home=init_context.executor_config.get('dagster_home'),
image_pull_policy=init_context.executor_config.get(
'image_pull_policy'),
image_pull_secrets=init_context.executor_config.get(
'image_pull_secrets'),
service_account_name=init_context.executor_config.get(
'service_account_name'),
instance_config_map=init_context.executor_config.get(
'instance_config_map'),
postgres_password_secret=init_context.executor_config.get(
'postgres_password_secret'),
env_config_maps=init_context.executor_config.get('env_config_maps'),
env_secrets=init_context.executor_config.get('env_secrets'),
)
return CeleryK8sJobConfig(
broker=init_context.executor_config.get('broker'),
backend=init_context.executor_config.get('backend'),
config_source=init_context.executor_config.get('config_source'),
include=init_context.executor_config.get('include'),
retries=Retries.from_config(init_context.executor_config['retries']),
job_config=job_config,
job_namespace=init_context.executor_config.get('job_namespace'),
load_incluster_config=init_context.executor_config.get(
'load_incluster_config'),
kubeconfig_file=init_context.executor_config.get('kubeconfig_file'),
)
def launch_run(self, instance, run, external_pipeline):
check.inst_param(run, 'run', PipelineRun)
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_config = DagsterK8sJobConfig(
dagster_home=self.dagster_home,
instance_config_map=self.instance_config_map,
postgres_password_secret=self.postgres_password_secret,
job_image=exc_config.get('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'),
)
resources = get_k8s_resource_requirements(
frozentags(external_pipeline.tags))
job = construct_dagster_graphql_k8s_job(
job_config,
args=[
'-p',
'executeRunInProcess',
'-v',
seven.json.dumps({
'runId':
run.run_id,
'repositoryName':
external_pipeline.handle.repository_name,
'repositoryLocationName':
external_pipeline.handle.location_name,
}),
'--remap-sigterm',
],
job_name=job_name,
pod_name=pod_name,
component='runmaster',
resources=resources,
)
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 runmaster 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=CeleryK8sRunLauncher,
)
return run
def launch_run(self, instance, run, external_pipeline):
check.inst_param(run, 'run', PipelineRun)
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'),
)
resources = get_k8s_resource_requirements(
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',
resources=resources,
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=CeleryK8sRunLauncher,
)
return run
step_failure_event = DagsterEvent.step_failure_event(
step_context=step_context,
step_failure_data=StepFailureData(error=None, user_failure_data=None),
)
return [step_failure_event]
return []
def terminate_steps(self, step_keys: List[str]):
raise NotImplementedError()
@executor(
name="k8s",
config_schema=merge_dicts(
DagsterK8sJobConfig.config_type_pipeline_run(),
{
"job_namespace": Field(
StringSource,
is_required=False,
default_value="default",
)
},
{"retries": get_retries_config()},
),
requirements=multiple_process_executor_requirements(),
)
def dagster_k8s_executor(init_context: InitExecutorContext) -> Executor:
run_launcher = init_context.instance.run_launcher
exc_cfg = init_context.executor_config
job_config = DagsterK8sJobConfig(
def config_type(cls):
from dagster_celery.executor import CELERY_CONFIG
return merge_dicts(DagsterK8sJobConfig.config_type_run_launcher(), CELERY_CONFIG)
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 launch_run(self, instance, run, external_pipeline):
check.inst_param(run, 'run', PipelineRun)
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_config = DagsterK8sJobConfig(
dagster_home=self.dagster_home,
instance_config_map=self.instance_config_map,
postgres_password_secret=self.postgres_password_secret,
job_image=exc_config.get('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'),
)
resources = get_k8s_resource_requirements(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=external_pipeline.get_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='runmaster',
resources=resources,
)
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 runmaster 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=CeleryK8sRunLauncher,
)
return run
