def test_k8s_run_launcher_default(dagster_instance_for_k8s_run_launcher, helm_namespace_for_k8s_run_launcher, dagster_docker_image): run_config = merge_dicts( load_yaml_from_path( os.path.join(get_test_project_environments_path(), "env.yaml")), load_yaml_from_path( os.path.join(get_test_project_environments_path(), "env_s3.yaml")), { "execution": { "k8s": { "config": { "job_namespace": helm_namespace_for_k8s_run_launcher, "job_image": dagster_docker_image, "image_pull_policy": image_pull_policy(), "env_config_maps": ["dagster-pipeline-env"] + ([TEST_AWS_CONFIGMAP_NAME] if not IS_BUILDKITE else []), } } }, }, ) _launch_executor_run( run_config, dagster_instance_for_k8s_run_launcher, helm_namespace_for_k8s_run_launcher, )
def test_valid_job_format(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) job_name = 'dagster-run-%s' % run.run_id pod_name = 'dagster-run-%s' % run.run_id job = construct_dagster_graphql_k8s_job( run_launcher.job_config, args=[ '-p', 'executeRunInProcess', '-v', seven.json.dumps({'runId': run.run_id}), ], job_name=job_name, pod_name=pod_name, component='runmaster', ) 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='', ).strip())
def test_valid_job_format(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) 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, 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="", ).strip())
def test_k8s_run_monitoring( dagster_instance_for_k8s_run_launcher, helm_namespace_for_k8s_run_launcher, dagit_url_for_k8s_run_launcher, ): run_config = merge_dicts( load_yaml_from_path( os.path.join(get_test_project_environments_path(), "env_s3.yaml")), { "execution": { "k8s": { "config": { "job_namespace": helm_namespace_for_k8s_run_launcher, "image_pull_policy": image_pull_policy(), } } }, }, ) _launch_run_and_wait_for_resume( dagit_url_for_k8s_run_launcher, run_config, dagster_instance_for_k8s_run_launcher, helm_namespace_for_k8s_run_launcher, )
def test_k8s_run_launcher_volume_mounts( dagster_instance_for_k8s_run_launcher, helm_namespace_for_k8s_run_launcher, dagster_docker_image, dagit_url_for_k8s_run_launcher, ): run_config = merge_dicts( load_yaml_from_path( os.path.join(get_test_project_environments_path(), "env_s3.yaml")), { "execution": { "k8s": { "config": { "job_namespace": helm_namespace_for_k8s_run_launcher, "job_image": dagster_docker_image, "image_pull_policy": image_pull_policy(), } } }, }, ) _launch_executor_run( dagit_url_for_k8s_run_launcher, run_config, dagster_instance_for_k8s_run_launcher, helm_namespace_for_k8s_run_launcher, pipeline_name="volume_mount_pipeline", num_steps=1, mode="k8s", )
def test_k8s_run_launcher_default( dagster_instance_for_k8s_run_launcher, helm_namespace_for_k8s_run_launcher, dagster_docker_image ): # sanity check that we have a K8sRunLauncher check.inst(dagster_instance_for_k8s_run_launcher.run_launcher, K8sRunLauncher) pods = DagsterKubernetesClient.production_client().core_api.list_namespaced_pod( namespace=helm_namespace_for_k8s_run_launcher ) celery_pod_names = [p.metadata.name for p in pods.items if "celery-workers" in p.metadata.name] check.invariant(not celery_pod_names) run_config = merge_dicts( load_yaml_from_path(os.path.join(get_test_project_environments_path(), "env.yaml")), load_yaml_from_path(os.path.join(get_test_project_environments_path(), "env_s3.yaml")), { "execution": { "k8s": { "config": { "job_namespace": helm_namespace_for_k8s_run_launcher, "job_image": dagster_docker_image, "image_pull_policy": image_pull_policy(), "env_config_maps": ["dagster-pipeline-env"] + ([TEST_AWS_CONFIGMAP_NAME] if not IS_BUILDKITE else []), } } }, }, ) pipeline_name = "demo_k8s_executor_pipeline" tags = {"key": "value"} with get_test_project_location_and_external_pipeline(pipeline_name) as ( location, external_pipeline, ): run = create_run_for_test( dagster_instance_for_k8s_run_launcher, pipeline_name=pipeline_name, run_config=run_config, tags=tags, mode="default", pipeline_snapshot=external_pipeline.pipeline_snapshot, execution_plan_snapshot=location.get_external_execution_plan( external_pipeline, run_config, "default", None, None ).execution_plan_snapshot, ) dagster_instance_for_k8s_run_launcher.launch_run( run.run_id, ReOriginatedExternalPipelineForTest(external_pipeline), ) result = wait_for_job_and_get_raw_logs( job_name="dagster-run-%s" % run.run_id, namespace=helm_namespace_for_k8s_run_launcher ) assert "PIPELINE_SUCCESS" in result, "no match, result: {}".format(result) updated_run = dagster_instance_for_k8s_run_launcher.get_run_by_id(run.run_id) assert updated_run.tags[DOCKER_IMAGE_TAG] == get_test_project_docker_image()
def test_k8s_run_launcher_terminate( dagster_instance_for_k8s_run_launcher, helm_namespace_for_k8s_run_launcher, dagster_docker_image, dagit_url_for_k8s_run_launcher, ): pipeline_name = "slow_pipeline" run_config = merge_dicts( load_yaml_from_path(os.path.join(get_test_project_environments_path(), "env_s3.yaml")), { "execution": { "k8s": { "config": { "job_namespace": helm_namespace_for_k8s_run_launcher, "job_image": dagster_docker_image, "image_pull_policy": image_pull_policy(), } } }, }, ) run_id = launch_run_over_graphql( dagit_url_for_k8s_run_launcher, run_config=run_config, pipeline_name=pipeline_name, mode="k8s", ) wait_for_job(job_name="dagster-run-%s" % run_id, namespace=helm_namespace_for_k8s_run_launcher) timeout = datetime.timedelta(0, 30) start_time = datetime.datetime.now() while True: assert datetime.datetime.now() < start_time + timeout, "Timed out waiting for can_terminate" if can_terminate_run_over_graphql(dagit_url_for_k8s_run_launcher, run_id): break time.sleep(5) terminate_run_over_graphql(dagit_url_for_k8s_run_launcher, run_id=run_id) start_time = datetime.datetime.now() pipeline_run = None while True: assert datetime.datetime.now() < start_time + timeout, "Timed out waiting for termination" pipeline_run = dagster_instance_for_k8s_run_launcher.get_run_by_id(run_id) if pipeline_run.status == PipelineRunStatus.CANCELED: break time.sleep(5) # useful to have logs here, because the worker pods get deleted print(dagster_instance_for_k8s_run_launcher.all_logs(run_id)) # pylint: disable=print-call assert pipeline_run.status == PipelineRunStatus.CANCELED assert not can_terminate_run_over_graphql(dagit_url_for_k8s_run_launcher, run_id)
def test_execute_on_k8s_retry_pipeline( # pylint: disable=redefined-outer-name dagster_instance_for_k8s_run_launcher, helm_namespace_for_k8s_run_launcher, dagster_docker_image, dagit_url_for_k8s_run_launcher, ): run_config = merge_dicts( load_yaml_from_path(os.path.join(get_test_project_environments_path(), "env_s3.yaml")), { "execution": { "k8s": { "config": { "job_namespace": helm_namespace_for_k8s_run_launcher, "job_image": dagster_docker_image, "image_pull_policy": image_pull_policy(), } } }, }, ) pipeline_name = "retry_pipeline" run_id = launch_run_over_graphql( dagit_url_for_k8s_run_launcher, run_config=run_config, pipeline_name=pipeline_name, mode="k8s", ) result = wait_for_job_and_get_raw_logs( job_name="dagster-run-%s" % run_id, namespace=helm_namespace_for_k8s_run_launcher ) assert "PIPELINE_SUCCESS" in result, "no match, result: {}".format(result) stats = dagster_instance_for_k8s_run_launcher.get_run_stats(run_id) assert stats.steps_succeeded == 1 all_logs = dagster_instance_for_k8s_run_launcher.all_logs(run_id) assert DagsterEventType.STEP_START in [ event.dagster_event.event_type for event in all_logs if event.is_dagster_event ] assert DagsterEventType.STEP_UP_FOR_RETRY in [ event.dagster_event.event_type for event in all_logs if event.is_dagster_event ] assert DagsterEventType.STEP_RESTARTED in [ event.dagster_event.event_type for event in all_logs if event.is_dagster_event ] assert DagsterEventType.STEP_SUCCESS in [ event.dagster_event.event_type for event in all_logs if event.is_dagster_event ]
def test_valid_job_format_with_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' }, }) }) resources = get_k8s_resource_requirements(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, resources=resources, pod_name=pod_name, component='runmaster', ) 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 get_celery_engine_config(dagster_docker_image, job_namespace): return { "execution": { "celery-k8s": { "config": { "job_image": dagster_docker_image, "job_namespace": job_namespace, "image_pull_policy": image_pull_policy(), } } }, }
def test_k8s_executor_resource_requirements( dagster_instance_for_k8s_run_launcher, helm_namespace_for_k8s_run_launcher, dagster_docker_image, dagit_url_for_k8s_run_launcher, ): # sanity check that we have a K8sRunLauncher pods = DagsterKubernetesClient.production_client( ).core_api.list_namespaced_pod( namespace=helm_namespace_for_k8s_run_launcher) celery_pod_names = [ p.metadata.name for p in pods.items if "celery-workers" in p.metadata.name ] check.invariant(not celery_pod_names) run_config = merge_dicts( load_yaml_from_path( os.path.join(get_test_project_environments_path(), "env_s3.yaml")), { "execution": { "k8s": { "config": { "job_namespace": helm_namespace_for_k8s_run_launcher, "job_image": dagster_docker_image, "image_pull_policy": image_pull_policy(), } } }, }, ) pipeline_name = "resources_limit_pipeline" run_id = launch_run_over_graphql( dagit_url_for_k8s_run_launcher, run_config=run_config, pipeline_name=pipeline_name, mode="k8s", ) result = wait_for_job_and_get_raw_logs( job_name="dagster-run-%s" % run_id, namespace=helm_namespace_for_k8s_run_launcher) assert "PIPELINE_SUCCESS" in result, "no match, result: {}".format(result) updated_run = dagster_instance_for_k8s_run_launcher.get_run_by_id(run_id) assert updated_run.tags[DOCKER_IMAGE_TAG] == get_test_project_docker_image( )
def get_celery_engine_config(dagster_docker_image, job_namespace): return { "execution": { "celery-k8s": { "config": { "job_image": dagster_docker_image, "job_namespace": job_namespace, "image_pull_policy": image_pull_policy(), "env_config_maps": ["dagster-pipeline-env"] + ([TEST_AWS_CONFIGMAP_NAME] if not IS_BUILDKITE else []), } } }, }
def get_celery_job_engine_config(dagster_docker_image, job_namespace): return { "execution": { "config": merge_dicts( ({ "job_image": dagster_docker_image, } if dagster_docker_image else {}), { "job_namespace": job_namespace, "image_pull_policy": image_pull_policy(), }, ) }, }
def test_k8s_run_launcher_image_from_origin( dagster_instance_for_k8s_run_launcher, helm_namespace_for_k8s_run_launcher, dagit_url_for_k8s_run_launcher, ): # Like the previous test, but the executor doesn't supply an image - it's pulled # from the origin on the run instead pods = DagsterKubernetesClient.production_client( ).core_api.list_namespaced_pod( namespace=helm_namespace_for_k8s_run_launcher) celery_pod_names = [ p.metadata.name for p in pods.items if "celery-workers" in p.metadata.name ] check.invariant(not celery_pod_names) run_config = merge_dicts( load_yaml_from_path( os.path.join(get_test_project_environments_path(), "env.yaml")), load_yaml_from_path( os.path.join(get_test_project_environments_path(), "env_s3.yaml")), { "execution": { "k8s": { "config": { "job_namespace": helm_namespace_for_k8s_run_launcher, "image_pull_policy": image_pull_policy(), } } }, }, ) pipeline_name = "demo_k8s_executor_pipeline" run_id = launch_run_over_graphql(dagit_url_for_k8s_run_launcher, run_config=run_config, pipeline_name=pipeline_name) result = wait_for_job_and_get_raw_logs( job_name="dagster-run-%s" % run_id, namespace=helm_namespace_for_k8s_run_launcher) assert "PIPELINE_SUCCESS" in result, "no match, result: {}".format(result) updated_run = dagster_instance_for_k8s_run_launcher.get_run_by_id(run_id) assert updated_run.tags[DOCKER_IMAGE_TAG] == get_test_project_docker_image( )
def get_failing_celery_job_engine_config(dagster_docker_image, job_namespace): return { "execution": { "config": merge_dicts( ({ "job_image": dagster_docker_image, } if dagster_docker_image else {}), { "job_namespace": job_namespace, "image_pull_policy": image_pull_policy(), "env_config_maps": ["non-existent-config-map"], }, ) }, }
def k8s_scheduler(cluster_provider, helm_namespace_for_k8s_run_launcher): # pylint: disable=redefined-outer-name,unused-argument return K8sScheduler( scheduler_namespace=helm_namespace_for_k8s_run_launcher, image_pull_secrets=[{ "name": "element-dev-key" }], service_account_name="dagit-admin", instance_config_map="dagster-instance", postgres_password_secret="dagster-postgresql-secret", dagster_home="/opt/dagster/dagster_home", job_image=get_test_project_docker_image(), load_incluster_config=False, kubeconfig_file=cluster_provider.kubeconfig_file, image_pull_policy=image_pull_policy(), env_config_maps=["dagster-pipeline-env", "test-env-configmap"], env_secrets=["test-env-secret"], )
def k8s_scheduler( cluster_provider, helm_namespace ): # pylint: disable=redefined-outer-name,unused-argument return K8sScheduler( scheduler_namespace=helm_namespace, image_pull_secrets=[{'name': 'element-dev-key'}], service_account_name='dagit-admin', instance_config_map='dagster-instance', postgres_password_secret='dagster-postgresql-secret', dagster_home='/opt/dagster/dagster_home', job_image=test_project_docker_image(), load_incluster_config=False, kubeconfig_file=cluster_provider.kubeconfig_file, image_pull_policy=image_pull_policy(), env_config_maps=['dagster-pipeline-env', 'test-env-configmap'], env_secrets=['test-env-secret'], )
def get_celery_job_engine_config(dagster_docker_image, job_namespace, include_dagster_pipeline_env=False): return { "execution": { "config": merge_dicts( ({ "job_image": dagster_docker_image, } if dagster_docker_image else {}), { "job_namespace": job_namespace, "image_pull_policy": image_pull_policy(), }, ({ "env_config_maps": ["dagster-pipeline-env"] } if include_dagster_pipeline_env else {}), ) }, }
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 test_execute_on_celery_k8s_subchart_disabled( # pylint: disable=redefined-outer-name,unused-argument dagster_instance_for_user_deployments_subchart_disabled, helm_namespace_for_user_deployments_subchart_disabled, ): namespace = helm_namespace_for_user_deployments_subchart_disabled pipeline_name = "demo_pipeline_celery" core_api = kubernetes.client.CoreV1Api() batch_api = kubernetes.client.BatchV1Api() # Get name for dagit pod pods = core_api.list_namespaced_pod(namespace=namespace) dagit_pod_list = list( filter(lambda item: "dagit" in item.metadata.name, pods.items)) assert len(dagit_pod_list) == 1 dagit_pod = dagit_pod_list[0] dagit_pod_name = dagit_pod.metadata.name # Check that there are no run master jobs jobs = batch_api.list_namespaced_job(namespace=namespace) runmaster_job_list = list( filter(lambda item: "dagster-run-" in item.metadata.name, jobs.items)) assert len(runmaster_job_list) == 0 run_config_dict = { "storage": { "s3": { "config": { "s3_bucket": "dagster-scratch-80542c2" } } }, "execution": { "celery-k8s": { "config": { "image_pull_policy": image_pull_policy(), "env_config_maps": ["dagster-pipeline-env"] + ([TEST_AWS_CONFIGMAP_NAME] if not IS_BUILDKITE else []), "job_namespace": namespace, } } }, "loggers": { "console": { "config": { "log_level": "DEBUG" } } }, "solids": { "multiply_the_word": { "inputs": { "word": "bar" }, "config": { "factor": 2 } } }, } run_config_json = json.dumps(run_config_dict) exec_command = [ "dagster", "pipeline", "launch", "--repository", "demo_execution_repo", "--pipeline", pipeline_name, "--workspace", "/dagster-workspace/workspace.yaml", "--location", "user-code-deployment-1", "--config-json", run_config_json, ] stream( core_api.connect_get_namespaced_pod_exec, name=dagit_pod_name, namespace=namespace, command=exec_command, stderr=True, stdin=False, stdout=True, tty=False, _preload_content=False, ) runmaster_job_name = None timeout = datetime.timedelta(0, 90) start_time = datetime.datetime.now() while datetime.datetime.now( ) < start_time + timeout and not runmaster_job_name: jobs = batch_api.list_namespaced_job(namespace=namespace) runmaster_job_list = list( filter(lambda item: "dagster-run-" in item.metadata.name, jobs.items)) if len(runmaster_job_list) > 0: runmaster_job_name = runmaster_job_list[0].metadata.name assert runmaster_job_name result = wait_for_job_and_get_raw_logs(job_name=runmaster_job_name, namespace=namespace, wait_timeout=450) assert "PIPELINE_SUCCESS" in result, "no match, result: {}".format(result)
def test_k8s_run_launcher_terminate( dagster_instance_for_k8s_run_launcher, helm_namespace_for_k8s_run_launcher, dagster_docker_image ): pipeline_name = "slow_pipeline" tags = {"key": "value"} run_config = merge_dicts( load_yaml_from_path(os.path.join(get_test_project_environments_path(), "env_s3.yaml")), { "execution": { "k8s": { "config": { "job_namespace": helm_namespace_for_k8s_run_launcher, "job_image": dagster_docker_image, "image_pull_policy": image_pull_policy(), "env_config_maps": ["dagster-pipeline-env"] + ([TEST_AWS_CONFIGMAP_NAME] if not IS_BUILDKITE else []), } } }, }, ) with get_test_project_external_pipeline_hierarchy( dagster_instance_for_k8s_run_launcher, pipeline_name ) as ( workspace, location, _repo, external_pipeline, ): reoriginated_pipeline = ReOriginatedExternalPipelineForTest(external_pipeline) run = create_run_for_test( dagster_instance_for_k8s_run_launcher, pipeline_name=pipeline_name, run_config=run_config, tags=tags, mode="k8s", pipeline_snapshot=external_pipeline.pipeline_snapshot, execution_plan_snapshot=location.get_external_execution_plan( external_pipeline, run_config, "k8s", None, None ).execution_plan_snapshot, external_pipeline_origin=reoriginated_pipeline.get_external_origin(), pipeline_code_origin=reoriginated_pipeline.get_python_origin(), ) dagster_instance_for_k8s_run_launcher.launch_run(run.run_id, workspace) wait_for_job( job_name="dagster-run-%s" % run.run_id, namespace=helm_namespace_for_k8s_run_launcher ) timeout = datetime.timedelta(0, 30) start_time = datetime.datetime.now() while datetime.datetime.now() < start_time + timeout: if dagster_instance_for_k8s_run_launcher.run_launcher.can_terminate(run_id=run.run_id): break time.sleep(5) assert dagster_instance_for_k8s_run_launcher.run_launcher.can_terminate(run_id=run.run_id) assert dagster_instance_for_k8s_run_launcher.run_launcher.terminate(run_id=run.run_id) start_time = datetime.datetime.now() pipeline_run = None while datetime.datetime.now() < start_time + timeout: pipeline_run = dagster_instance_for_k8s_run_launcher.get_run_by_id(run.run_id) if pipeline_run.status == PipelineRunStatus.CANCELED: break time.sleep(5) # useful to have logs here, because the worker pods get deleted print( # pylint: disable=print-call dagster_instance_for_k8s_run_launcher.all_logs(run.run_id) ) assert pipeline_run.status == PipelineRunStatus.CANCELED assert not dagster_instance_for_k8s_run_launcher.run_launcher.terminate(run_id=run.run_id)
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 test_memoization_k8s_executor( dagster_instance_for_k8s_run_launcher, helm_namespace_for_k8s_run_launcher, dagster_docker_image, dagit_url_for_k8s_run_launcher, ): ephemeral_path = str(uuid.uuid4()) run_config = deep_merge_dicts( load_yaml_from_path( os.path.join(get_test_project_environments_path(), "env_s3.yaml")), { "execution": { "k8s": { "config": { "job_namespace": helm_namespace_for_k8s_run_launcher, "job_image": dagster_docker_image, "image_pull_policy": image_pull_policy(), } } }, }, ) run_config = deep_merge_dicts( run_config, { "resources": { "io_manager": { "config": { "s3_prefix": ephemeral_path } } } }, ) # wrap in try-catch to ensure that memoized results are always cleaned from s3 bucket try: pipeline_name = "memoization_pipeline" run_ids = [] for _ in range(2): run_id = launch_run_over_graphql( dagit_url_for_k8s_run_launcher, run_config=run_config, pipeline_name=pipeline_name, mode="k8s", ) result = wait_for_job_and_get_raw_logs( job_name="dagster-run-%s" % run_id, namespace=helm_namespace_for_k8s_run_launcher, ) assert "PIPELINE_SUCCESS" in result, "no match, result: {}".format( result) run_ids.append(run_id) # We expect that first run should have to run the step, since it has not yet been # memoized. unmemoized_run_id = run_ids[0] events = dagster_instance_for_k8s_run_launcher.all_logs( unmemoized_run_id) assert len(_get_step_execution_events(events)) == 1 # We expect that second run should not have to run the step, since it has been memoized. memoized_run_id = run_ids[1] events = dagster_instance_for_k8s_run_launcher.all_logs( memoized_run_id) assert len(_get_step_execution_events(events)) == 0 finally: cleanup_memoized_results(define_memoization_pipeline(), "k8s", dagster_instance_for_k8s_run_launcher, run_config)