def test_user_defined_k8s_config_in_run_tags(kubeconfig_file):
# Construct a K8s run launcher in a fake k8s environment.
mock_k8s_client_batch_api = mock.MagicMock()
celery_k8s_run_launcher = CeleryK8sRunLauncher(
instance_config_map="dagster-instance",
postgres_password_secret="dagster-postgresql-secret",
dagster_home="/opt/dagster/dagster_home",
load_incluster_config=False,
kubeconfig_file=kubeconfig_file,
k8s_client_batch_api=mock_k8s_client_batch_api,
)
# Construct Dagster run tags with user defined k8s config.
expected_resources = {
"requests": {"cpu": "250m", "memory": "64Mi"},
"limits": {"cpu": "500m", "memory": "2560Mi"},
}
user_defined_k8s_config = UserDefinedDagsterK8sConfig(
container_config={"resources": expected_resources},
)
user_defined_k8s_config_json = json.dumps(user_defined_k8s_config.to_dict())
tags = {"dagster-k8s/config": user_defined_k8s_config_json}
# Create fake external pipeline.
recon_pipeline = reconstructable(fake_pipeline)
recon_repo = recon_pipeline.repository
location_origin = InProcessRepositoryLocationOrigin(recon_repo)
location_handle = location_origin.create_handle()
repo_def = recon_repo.get_definition()
repo_handle = RepositoryHandle(
repository_name=repo_def.name,
repository_location_handle=location_handle,
)
fake_external_pipeline = external_pipeline_from_recon_pipeline(
recon_pipeline,
solid_selection=None,
repository_handle=repo_handle,
)
# Launch the run in a fake Dagster instance.
with instance_for_test() as instance:
celery_k8s_run_launcher.initialize(instance)
pipeline_name = "demo_pipeline"
run_config = {"execution": {"celery-k8s": {"config": {"job_image": "fake-image-name"}}}}
run = create_run_for_test(
instance,
pipeline_name=pipeline_name,
run_config=run_config,
tags=tags,
)
celery_k8s_run_launcher.launch_run(instance, run, fake_external_pipeline)
# Check that user defined k8s config was passed down to the k8s job.
mock_method_calls = mock_k8s_client_batch_api.method_calls
assert len(mock_method_calls) > 0
method_name, _args, kwargs = mock_method_calls[0]
assert method_name == "create_namespaced_job"
job_resources = kwargs["body"].spec.template.spec.containers[0].resources
assert job_resources == expected_resources
def test_no_postgres(kubeconfig_file):
# Construct a K8s run launcher in a fake k8s environment.
mock_k8s_client_batch_api = mock.MagicMock()
k8s_run_launcher = K8sRunLauncher(
service_account_name="dagit-admin",
instance_config_map="dagster-instance",
dagster_home="/opt/dagster/dagster_home",
job_image="fake_job_image",
load_incluster_config=False,
kubeconfig_file=kubeconfig_file,
k8s_client_batch_api=mock_k8s_client_batch_api,
)
# Create fake external pipeline.
recon_pipeline = reconstructable(fake_pipeline)
recon_repo = recon_pipeline.repository
repo_def = recon_repo.get_definition()
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)
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"
assert DAGSTER_PG_PASSWORD_ENV_VAR not in [
env.name
for env in kwargs["body"].spec.template.spec.containers[0].env
]
def _get_unloadable_sensor_origin():
working_directory = os.path.dirname(__file__)
recon_repo = ReconstructableRepository.for_file(__file__, "doesnt_exist",
working_directory)
return ExternalRepositoryOrigin(
InProcessRepositoryLocationOrigin(recon_repo),
"fake_repository").get_job_origin("doesnt_exist")
def get_external_origin(self):
"""
Hack! Inject origin that the k8s images will use. The BK image uses a different directory
structure (/workdir/python_modules/dagster-test/dagster_test/test_project) than the images
inside the kind cluster (/dagster_test/test_project). As a result the normal origin won't
work, we need to inject this one.
"""
return ExternalJobOrigin(
external_repository_origin=ExternalRepositoryOrigin(
repository_location_origin=InProcessRepositoryLocationOrigin(
recon_repo=ReconstructableRepository(
pointer=FileCodePointer(
python_file=
"/dagster_test/test_project/test_pipelines/repo.py",
fn_name="define_demo_execution_repo",
),
container_image=self._container_image,
executable_path="python",
entry_point=DEFAULT_DAGSTER_ENTRY_POINT,
)),
repository_name="demo_execution_repo",
),
job_name=self.name,
)
def get_test_project_external_pipeline(pipeline_name):
return (InProcessRepositoryLocationOrigin(
ReconstructableRepository.for_file(
file_relative_path(__file__, "test_pipelines/repo.py"),
"define_demo_execution_repo",
)).create_handle().create_location().get_repository(
"demo_execution_repo").get_full_external_pipeline(pipeline_name))
def _repo_location_origins_from_config(workspace_config, yaml_path):
ensure_workspace_config(workspace_config, yaml_path)
if "repository" in workspace_config:
warnings.warn(
# link to docs once they exist
"You are using the legacy repository yaml format. Please update your file "
"to abide by the new workspace file format."
)
origin = InProcessRepositoryLocationOrigin(
ReconstructableRepository.from_legacy_repository_yaml(yaml_path)
)
return {origin.location_name: origin}
location_origins = OrderedDict()
for location_config in workspace_config["load_from"]:
origin = _location_origin_from_location_config(location_config, yaml_path)
check.invariant(
location_origins.get(origin.location_name) is None,
'Cannot have multiple locations with the same name, got multiple "{name}"'.format(
name=origin.location_name,
),
)
location_origins[origin.location_name] = origin
return location_origins
def legacy_get_bar_repo_handle():
recon_repo = ReconstructableRepository.from_legacy_repository_yaml(
file_relative_path(__file__, "legacy_repository_file.yaml"))
return (RepositoryLocation.from_handle(
RepositoryLocationHandle.create_from_repository_location_origin(
InProcessRepositoryLocationOrigin(recon_repo))).get_repository(
"bar_repo").handle)
def get_test_project_external_repo(container_image=None):
with InProcessRepositoryLocationOrigin(
ReconstructableRepository.for_file(
file_relative_path(__file__, "test_pipelines/repo.py"),
"define_demo_execution_repo",
container_image=container_image,
)).create_location() as location:
yield location.get_repository("demo_execution_repo")
def get_test_project_external_repo(container_image=None):
return RepositoryLocation.from_handle(
RepositoryLocationHandle.create_from_repository_location_origin(
InProcessRepositoryLocationOrigin(
ReconstructableRepository.for_file(
file_relative_path(__file__, "test_pipelines/repo.py"),
"define_demo_execution_repo",
container_image=container_image,
)))).get_repository("demo_execution_repo")
def define_in_process_context(python_file, fn_name, instance):
check.inst_param(instance, "instance", DagsterInstance)
return WorkspaceProcessContext(
workspace=Workspace([
InProcessRepositoryLocationOrigin(
ReconstructableRepository.for_file(python_file, fn_name))
]),
instance=instance,
).create_request_context()
def get_test_project_external_pipeline(pipeline_name):
with InProcessRepositoryLocationOrigin(
LoadableTargetOrigin(
executable_path=sys.executable,
python_file=file_relative_path(__file__,
"test_pipelines/repo.py"),
attribute="define_demo_execution_repo",
)).create_location() as location:
yield location.get_repository(
"demo_execution_repo").get_full_external_pipeline(pipeline_name)
def _get_unloadable_sensor_origin(name):
working_directory = os.path.dirname(__file__)
loadable_target_origin = LoadableTargetOrigin(
executable_path=sys.executable,
python_file=__file__,
working_directory=working_directory,
)
return ExternalRepositoryOrigin(
InProcessRepositoryLocationOrigin(loadable_target_origin),
"fake_repository").get_instigator_origin(name)
def _unloadable_partition_set_origin():
working_directory = os.path.dirname(__file__)
return ExternalRepositoryOrigin(
InProcessRepositoryLocationOrigin(
LoadableTargetOrigin(
executable_path=sys.executable,
python_file=__file__,
working_directory=working_directory,
)),
"fake_repository",
).get_partition_set_origin("doesnt_exist")
def _mgr_fn(recon_repo):
check.inst_param(recon_repo, "recon_repo", ReconstructableRepository)
with Workspace([InProcessRepositoryLocationOrigin(recon_repo)]) as workspace:
yield workspace
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
]