def test_k8s_run_launcher_default( dagster_instance_for_k8s_run_launcher, helm_namespace_for_k8s_run_launcher, dagit_url_for_k8s_run_launcher, ): 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 = load_yaml_from_path( os.path.join(get_test_project_environments_path(), "env.yaml")) pipeline_name = "demo_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 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 _helm_namespace_helper(helm_chart_fn, request): """If an existing Helm chart namespace is specified via pytest CLI with the argument --existing-helm-namespace, we will use that chart. Otherwise, provision a test namespace and install Helm chart into that namespace. Yields the Helm chart namespace. """ existing_helm_namespace = request.config.getoption("--existing-helm-namespace") if existing_helm_namespace: yield existing_helm_namespace else: # Never bother cleaning up on Buildkite if IS_BUILDKITE: should_cleanup = False # Otherwise, always clean up unless --no-cleanup specified else: should_cleanup = not request.config.getoption("--no-cleanup") with test_namespace(should_cleanup) as namespace: with helm_test_resources(namespace, should_cleanup): docker_image = get_test_project_docker_image() with helm_chart_fn(namespace, docker_image, should_cleanup): print("Helm chart successfully installed in namespace %s" % namespace) yield namespace
def _launch_executor_run( dagit_url, run_config, dagster_instance_for_k8s_run_launcher, helm_namespace_for_k8s_run_launcher, pipeline_name="demo_k8s_executor_pipeline", num_steps=2, mode="default", ): run_id = launch_run_over_graphql(dagit_url, run_config=run_config, pipeline_name=pipeline_name, mode=mode) 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( ) events = dagster_instance_for_k8s_run_launcher.all_logs(run_id) assert len(_get_step_execution_events(events)) == num_steps return run_id
def test_valid_job_format(run_launcher): docker_image = get_test_project_docker_image() run_config = load_yaml_from_path( os.path.join(get_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.get_static_job_config(), args=["dagster", "api", "execute_run"], 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="", env_from=ENV_FROM, ).strip())
def test_terminate_launched_docker_run(): docker_image = get_test_project_docker_image() launcher_config = { "env_vars": ["AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY",], "network": "container:test-postgres-db-docker", } if IS_BUILDKITE: launcher_config["registry"] = get_buildkite_registry_config() else: find_local_test_image(docker_image) run_config = merge_yamls([os.path.join(get_test_project_environments_path(), "env_s3.yaml"),]) with docker_postgres_instance( overrides={ "run_launcher": { "class": "DockerRunLauncher", "module": "dagster_docker", "config": launcher_config, } } ) as instance: recon_pipeline = get_test_project_recon_pipeline("hanging_pipeline", docker_image) run = instance.create_run_for_pipeline( pipeline_def=recon_pipeline.get_definition(), run_config=run_config, ) run_id = run.run_id external_pipeline = ReOriginatedExternalPipelineForTest( get_test_project_external_pipeline("hanging_pipeline", container_image=docker_image), container_image=docker_image, ) instance.launch_run(run_id, external_pipeline) poll_for_step_start(instance, run_id) assert instance.run_launcher.can_terminate(run_id) assert instance.run_launcher.terminate(run_id) terminated_pipeline_run = poll_for_finished_run(instance, run_id, timeout=30) terminated_pipeline_run = instance.get_run_by_id(run_id) assert terminated_pipeline_run.status == PipelineRunStatus.CANCELED run_logs = instance.all_logs(run_id) _check_event_log_contains( run_logs, [ ("PIPELINE_CANCELING", "Sending pipeline termination request"), ("STEP_FAILURE", 'Execution of step "hanging_solid" failed.'), ("PIPELINE_CANCELED", 'Execution of pipeline "hanging_pipeline" canceled.'), ("ENGINE_EVENT", "Pipeline execution terminated by interrupt"), ("ENGINE_EVENT", "Process for pipeline exited"), ], )
def test_launch_docker_image_on_pipeline_config(): # Docker image name to use for launch specified as part of the pipeline origin # rather than in the run launcher instance config docker_image = get_test_project_docker_image() launcher_config = { "env_vars": [ "AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", ], "network": "container:test-postgres-db-docker", "container_kwargs": { "auto_remove": True, }, } if IS_BUILDKITE: launcher_config["registry"] = get_buildkite_registry_config() else: find_local_test_image(docker_image) run_config = merge_yamls([ os.path.join(get_test_project_environments_path(), "env.yaml"), os.path.join(get_test_project_environments_path(), "env_s3.yaml"), ]) with docker_postgres_instance( overrides={ "run_launcher": { "class": "DockerRunLauncher", "module": "dagster_docker", "config": launcher_config, } }) as instance: recon_pipeline = get_test_project_recon_pipeline( "demo_pipeline", docker_image) run = instance.create_run_for_pipeline( pipeline_def=recon_pipeline.get_definition(), run_config=run_config, ) with get_test_project_external_pipeline( "demo_pipeline", container_image=docker_image) as orig_pipeline: external_pipeline = ReOriginatedExternalPipelineForTest( orig_pipeline, container_image=docker_image, ) instance.launch_run(run.run_id, external_pipeline) poll_for_finished_run(instance, run.run_id, timeout=60) run = instance.get_run_by_id(run.run_id) assert run.status == PipelineRunStatus.SUCCESS assert run.tags[DOCKER_IMAGE_TAG] == docker_image
def _cluster_provider(request): from .kind import kind_cluster_exists, kind_cluster, kind_load_images if IS_BUILDKITE: print("Installing ECR credentials...") check_output( "aws ecr get-login --no-include-email --region us-west-2 | sh", shell=True) provider = request.config.getoption("--cluster-provider") # Use a kind cluster if provider == "kind": cluster_name = request.config.getoption("--kind-cluster") # Cluster will be deleted afterwards unless this is set. # This is to allow users to reuse an existing cluster in local test by running # `pytest --kind-cluster my-cluster --no-cleanup` -- this avoids the per-test run # overhead of cluster setup and teardown should_cleanup = True if IS_BUILDKITE else not request.config.getoption( "--no-cleanup") existing_cluster = kind_cluster_exists(cluster_name) with kind_cluster(cluster_name, should_cleanup=should_cleanup) as cluster_config: if not IS_BUILDKITE and not existing_cluster: docker_image = get_test_project_docker_image() try: client = docker.from_env() client.images.get(docker_image) print( # pylint: disable=print-call "Found existing image tagged {image}, skipping image build. To rebuild, first run: " "docker rmi {image}".format(image=docker_image)) except docker.errors.ImageNotFound: build_and_tag_test_image(docker_image) kind_load_images( cluster_name=cluster_config.name, local_dagster_test_image=docker_image, additional_images=additional_kind_images, ) yield cluster_config # Use cluster from kubeconfig elif provider == "kubeconfig": kubeconfig_file = os.getenv( "KUBECONFIG", os.path.expandvars("${HOME}/.kube/config")) kubernetes.config.load_kube_config(config_file=kubeconfig_file) yield ClusterConfig(name="from_system_kubeconfig", kubeconfig_file=kubeconfig_file) else: raise Exception("unknown cluster provider %s" % provider)
def test_terminate(): docker_image = get_test_project_docker_image() launcher_config = { "env_vars": [ "AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", ], "network": "container:test-postgres-db-docker", "image": docker_image, } _test_launch(docker_image, launcher_config, terminate=True)
def test_launch_docker_image_on_instance_config(): docker_image = get_test_project_docker_image() launcher_config = { "env_vars": [ "AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", ], "network": "container:test-postgres-db-docker", "image": docker_image, } if IS_BUILDKITE: launcher_config["registry"] = get_buildkite_registry_config() else: find_local_test_image(docker_image) run_config = merge_yamls([ os.path.join(get_test_project_environments_path(), "env.yaml"), os.path.join(get_test_project_environments_path(), "env_s3.yaml"), ]) with docker_postgres_instance( overrides={ "run_launcher": { "class": "DockerRunLauncher", "module": "dagster_docker", "config": launcher_config, } }) as instance: recon_pipeline = get_test_project_recon_pipeline("demo_pipeline") with get_test_project_workspace_and_external_pipeline( instance, "demo_pipeline") as ( workspace, orig_pipeline, ): external_pipeline = ReOriginatedExternalPipelineForTest( orig_pipeline) run = instance.create_run_for_pipeline( pipeline_def=recon_pipeline.get_definition(), run_config=run_config, external_pipeline_origin=external_pipeline.get_external_origin( ), pipeline_code_origin=external_pipeline.get_python_origin(), ) instance.launch_run(run.run_id, workspace) poll_for_finished_run(instance, run.run_id, timeout=60) assert instance.get_run_by_id( run.run_id).status == PipelineRunStatus.SUCCESS
def _launch_executor_run( run_config, dagster_instance_for_k8s_run_launcher, helm_namespace_for_k8s_run_launcher, ): pipeline_name = "demo_k8s_executor_pipeline" tags = {"key": "value"} 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="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, 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) 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() events = dagster_instance_for_k8s_run_launcher.all_logs(run.run_id) assert (len([ event for event in events if ("Executing step" in event.message and "in Kubernetes job" in event.message) ]) == 2)
def test_execute_celery_docker_image_on_pipeline_config(aws_creds): docker_image = get_test_project_docker_image() docker_config = { "network": "container:test-postgres-db-celery-docker", "container_kwargs": { "environment": [ "FIND_ME=here!", f"AWS_ACCESS_KEY_ID={aws_creds['aws_access_key_id']}", f"AWS_SECRET_ACCESS_KEY={aws_creds['aws_secret_access_key']}", ], # "auto_remove": False # uncomment when debugging to view container logs after execution }, } if IS_BUILDKITE: docker_config["registry"] = get_buildkite_registry_config() else: find_local_test_image(docker_image) run_config = merge_dicts( merge_yamls([ os.path.join(get_test_project_environments_path(), "env.yaml"), os.path.join(get_test_project_environments_path(), "env_s3.yaml"), os.path.join(get_test_project_environments_path(), "env_environment_vars.yaml"), ]), { "execution": { "celery-docker": { "config": { "docker": docker_config, "config_source": { "task_always_eager": True }, } } }, }, ) with celery_docker_postgres_instance() as instance: result = execute_pipeline( get_test_project_recon_pipeline("docker_celery_pipeline", docker_image), run_config=run_config, instance=instance, ) assert result.success assert result.result_for_solid("get_environment_solid").output_value( "result") == "here!"
def dagster_docker_image(): docker_image = get_test_project_docker_image() if not IS_BUILDKITE: try: client = docker.from_env() client.images.get(docker_image) print( # pylint: disable=print-call "Found existing image tagged {image}, skipping image build. To rebuild, first run: " "docker rmi {image}".format(image=docker_image)) except docker.errors.ImageNotFound: build_and_tag_test_image(docker_image) return docker_image
def test_launch_docker_image_multiple_networks(): docker_image = get_test_project_docker_image() launcher_config = { "env_vars": [ "AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", ], "networks": [ "container:test-postgres-db-docker", "postgres", ], "image": docker_image, } _test_launch(docker_image, launcher_config)
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 test_valid_job_format_with_backcompat_resources(run_launcher): docker_image = get_test_project_docker_image() run_config = load_yaml_from_path( os.path.join(get_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.get_static_job_config(), args=["dagster", "api", "execute_run"], 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, env_from=ENV_FROM, resources=""" resources: limits: cpu: 500m memory: 2560Mi requests: cpu: 250m memory: 64Mi""", ).strip())
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 test_k8s_run_launcher_default( dagster_instance_for_k8s_run_launcher, helm_namespace_for_k8s_run_launcher ): # 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 = load_yaml_from_path(os.path.join(get_test_project_environments_path(), "env.yaml")) pipeline_name = "demo_pipeline" tags = {"key": "value"} 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="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, 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) 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 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 test_docker_executor(): """ Note that this test relies on having AWS credentials in the environment. """ executor_config = { "execution": { "docker": { "config": { "networks": ["container:test-postgres-db-docker"], "env_vars": [ "AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", ], } } } } docker_image = get_test_project_docker_image() if IS_BUILDKITE: executor_config["execution"]["docker"]["config"][ "registry" ] = get_buildkite_registry_config() else: find_local_test_image(docker_image) run_config = merge_dicts( merge_yamls( [ os.path.join(get_test_project_environments_path(), "env.yaml"), os.path.join(get_test_project_environments_path(), "env_s3.yaml"), ] ), executor_config, ) with environ({"DOCKER_LAUNCHER_NETWORK": "container:test-postgres-db-docker"}): with docker_postgres_instance() as instance: recon_pipeline = get_test_project_recon_pipeline("demo_pipeline_docker", docker_image) assert execute_pipeline( recon_pipeline, run_config=run_config, instance=instance ).success
def test_execute_celery_docker_image_on_pipeline_config(): docker_image = get_test_project_docker_image() docker_config = { "env_vars": [ "AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", ], "network": "container:test-postgres-db-celery-docker", } if IS_BUILDKITE: docker_config["registry"] = get_buildkite_registry_config() else: find_local_test_image(docker_image) run_config = merge_dicts( merge_yamls([ os.path.join(get_test_project_environments_path(), "env.yaml"), os.path.join(get_test_project_environments_path(), "env_s3.yaml"), ]), { "execution": { "celery-docker": { "config": { "docker": docker_config, "config_source": { "task_always_eager": True }, } } }, }, ) with celery_docker_postgres_instance() as instance: result = execute_pipeline( get_test_project_recon_pipeline("docker_celery_pipeline", docker_image), run_config=run_config, instance=instance, ) assert result.success
def test_launch_docker_config_on_container_context(): docker_image = get_test_project_docker_image() launcher_config = {} _test_launch( docker_image, launcher_config, container_image=docker_image, container_context={ "docker": { "env_vars": [ "AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", ], "networks": [ "container:test-postgres-db-docker", "postgres", ], } }, )
def test_cant_combine_network_and_networks(): docker_image = get_test_project_docker_image() launcher_config = { "env_vars": [ "AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", ], "network": "container:test-postgres-db-docker", "networks": [ "postgres", ], "image": docker_image, } with pytest.raises(Exception, match="cannot set both `network` and `networks`"): with docker_postgres_instance( overrides={ "run_launcher": { "class": "DockerRunLauncher", "module": "dagster_docker", "config": launcher_config, } } ): pass
def test_valid_job_format_with_user_defined_k8s_config(run_launcher): docker_image = get_test_project_docker_image() run_config = load_yaml_from_path( os.path.join(get_test_project_environments_path(), "env.yaml")) pipeline_name = "demo_pipeline" run = PipelineRun(pipeline_name=pipeline_name, run_config=run_config) backoff_limit = 1234 ttl_seconds_after_finished = 5678 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" }, }, "job_spec_config": { "ttl_seconds_after_finished": ttl_seconds_after_finished, "backoff_limit": backoff_limit, }, "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.get_static_job_config(), args=["dagster", "api", "execute_run"], 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'", env_from=ENV_FROM, backoff_limit=backoff_limit, ttl_seconds_after_finished=ttl_seconds_after_finished, 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_launch_docker_no_network(): docker_image = get_test_project_docker_image() launcher_config = { "env_vars": [ "AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", ], } if IS_BUILDKITE: launcher_config["registry"] = get_buildkite_registry_config() else: find_local_test_image(docker_image) run_config = merge_yamls( [ os.path.join(get_test_project_environments_path(), "env.yaml"), os.path.join(get_test_project_environments_path(), "env_s3.yaml"), ] ) with docker_postgres_instance( overrides={ "run_launcher": { "class": "DockerRunLauncher", "module": "dagster_docker", "config": launcher_config, } }, # Ensure the container will time out and fail quickly conn_args={ "params": {"connect_timeout": 2}, }, ) as instance: recon_pipeline = get_test_project_recon_pipeline("demo_pipeline_s3", docker_image) with get_test_project_workspace_and_external_pipeline( instance, "demo_pipeline_s3", container_image=docker_image ) as (workspace, orig_pipeline): external_pipeline = ReOriginatedExternalPipelineForTest( orig_pipeline, container_image=docker_image, ) run = instance.create_run_for_pipeline( pipeline_def=recon_pipeline.get_definition(), run_config=run_config, external_pipeline_origin=external_pipeline.get_external_origin(), pipeline_code_origin=external_pipeline.get_python_origin(), ) instance.launch_run(run.run_id, workspace) # Container launches, but run is stuck in STARTING state # due to not being able to access the network run = instance.get_run_by_id(run.run_id) assert run.tags[DOCKER_IMAGE_TAG] == docker_image container_id = run.tags[DOCKER_CONTAINER_ID_TAG] run = instance.get_run_by_id(run.run_id) assert run.status == PipelineRunStatus.STARTING assert run.tags[DOCKER_IMAGE_TAG] == docker_image client = docker.client.from_env() container = None try: start_time = time.time() while True: container = client.containers.get(container_id) if time.time() - start_time > 60: raise Exception("Timed out waiting for container to exit") if container.status == "exited": break time.sleep(3) finally: if container: container.remove(force=True)
def test_docker_monitoring(): docker_image = get_test_project_docker_image() launcher_config = { "env_vars": [ "AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", ], "networks": ["container:test-postgres-db-docker"], "container_kwargs": { # "auto_remove": True, "volumes": ["/var/run/docker.sock:/var/run/docker.sock"], }, } if IS_BUILDKITE: launcher_config["registry"] = get_buildkite_registry_config() else: find_local_test_image(docker_image) run_config = merge_dicts( load_yaml_from_path( os.path.join(get_test_project_environments_path(), "env_s3.yaml")), { "solids": { "multiply_the_word_slow": { "inputs": { "word": "bar" }, "config": { "factor": 2, "sleep_time": 20 }, } }, "execution": { "docker": { "config": {} } }, }, ) with docker_postgres_instance({ "run_monitoring": { "enabled": True }, "run_launcher": { "class": "DockerRunLauncher", "module": "dagster_docker", "config": launcher_config, }, }) as instance: recon_pipeline = get_test_project_recon_pipeline( "demo_pipeline_docker_slow", docker_image) with get_test_project_workspace_and_external_pipeline( instance, "demo_pipeline_docker_slow", container_image=docker_image) as ( workspace, orig_pipeline, ): with start_daemon(): external_pipeline = ReOriginatedExternalPipelineForTest( orig_pipeline, container_image=docker_image) run = instance.create_run_for_pipeline( pipeline_def=recon_pipeline.get_definition(), run_config=run_config, external_pipeline_origin=external_pipeline. get_external_origin(), pipeline_code_origin=external_pipeline.get_python_origin(), ) with log_run_events(instance, run.run_id): instance.launch_run(run.run_id, workspace) start_time = time.time() while time.time() - start_time < 60: run = instance.get_run_by_id(run.run_id) if run.status == PipelineRunStatus.STARTED: break assert run.status == PipelineRunStatus.STARTING time.sleep(1) time.sleep(3) instance.run_launcher._get_container( # pylint:disable=protected-access instance.get_run_by_id(run.run_id)).stop() # daemon resumes the run poll_for_finished_run(instance, run.run_id, timeout=90) assert instance.get_run_by_id( run.run_id).status == PipelineRunStatus.SUCCESS
def test_execute_on_celery_k8s( # pylint: disable=redefined-outer-name,unused-argument dagster_instance_for_user_deployments, helm_namespace_for_user_deployments, ): namespace = helm_namespace_for_user_deployments 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, "--run-id", "foobar", ] 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) updated_run = dagster_instance_for_user_deployments.get_run_by_id("foobar") assert updated_run.tags[DOCKER_IMAGE_TAG] == get_test_project_docker_image()
def test_container_context_on_pipeline(): docker_image = get_test_project_docker_image() launcher_config = {} if IS_BUILDKITE: launcher_config["registry"] = get_buildkite_registry_config() else: find_local_test_image(docker_image) executor_config = { "execution": { "docker": { "config": {} } }, } run_config = merge_dicts( merge_yamls([ os.path.join(get_test_project_environments_path(), "env.yaml"), os.path.join(get_test_project_environments_path(), "env_s3.yaml"), ]), executor_config, ) with docker_postgres_instance( overrides={ "run_launcher": { "class": "DockerRunLauncher", "module": "dagster_docker", "config": launcher_config, } }) as instance: recon_pipeline = get_test_project_recon_pipeline( "demo_pipeline_docker", docker_image, container_context={ "docker": { "env_vars": [ "AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", ], "networks": ["container:test-postgres-db-docker"], "container_kwargs": { "auto_remove": True, "volumes": ["/var/run/docker.sock:/var/run/docker.sock"], }, } }, ) with get_test_project_workspace_and_external_pipeline( instance, "demo_pipeline_docker", container_image=docker_image) as ( workspace, orig_pipeline, ): external_pipeline = ReOriginatedExternalPipelineForTest( orig_pipeline, container_image=docker_image) run = instance.create_run_for_pipeline( pipeline_def=recon_pipeline.get_definition(), run_config=run_config, external_pipeline_origin=external_pipeline.get_external_origin( ), pipeline_code_origin=recon_pipeline.get_python_origin(), ) instance.launch_run(run.run_id, workspace) poll_for_finished_run(instance, run.run_id, timeout=60) for log in instance.all_logs(run.run_id): print(log) # pylint: disable=print-call assert instance.get_run_by_id( run.run_id).status == PipelineRunStatus.SUCCESS