Пример #1
0
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
Пример #2
0
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
        ]
Пример #3
0
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")
Пример #4
0
    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,
        )
Пример #5
0
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))
Пример #6
0
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
Пример #7
0
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)
Пример #8
0
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")
Пример #9
0
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")
Пример #10
0
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()
Пример #11
0
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)
Пример #12
0
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)
Пример #13
0
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
Пример #15
0
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
        ]