xee dataflow error: [ee.ee_exception.EEException: Earth Engine client library not initialized. ]

[puzhao8]

Hi all, I am trying to use xee dataflow to download data from GEE by following this tutorial.

I built a container using the following Dockerfile, and submit the dataflow pipeline in the running container to ensure the match between local pipeline submission and google cloud runtime environment [ref]. Earth engine and service account key were set as ENV in the Dockerfile. DirectRunner works well, but DataflowRunner didn't work with the error: [ee.ee_exception.EEException: Earth Engine client library not initialized. ]

It is appreciated if anyone could help with this.

# xee-dataflow-run
# Use the Apache Beam Python SDK as the base image
FROM apache/beam_python3.10_sdk:2.62.0

# Set the working directory
WORKDIR /app

# Copy the requirements file
COPY requirements.txt ./

# Install system dependencies
RUN apt-get update && apt-get install -y \
    libffi-dev \
    libssl-dev \
    && rm -rf /var/lib/apt/lists/*

# Install Python dependencies
RUN pip install --no-cache-dir -r requirements.txt && \
    rm -rf /root/.cache/pip

# Verify Apache Beam installation
RUN pip list | grep apache-beam

# Copy application code
COPY ee_to_zarr_dataflow.py  ./
COPY .config /root/.config

RUN ls -l /root/.config/earthengine/credentials || echo "EE credentials not found!"

# Set environment variables
ENV EE_CONFIG="/root/.config/earthengine/credentials"
ENV GOOGLE_APPLICATION_CREDENTIALS="/root/.config/earthengine/service-account-key.json"
ENV GOOGLE_CLOUD_PROJECT='xx'
ENV CLOUDSDK_CORE_PROJECT='xx'

# Prevent Python from writing .pyc files & use unbuffered output
ENV PYTHONDONTWRITEBYTECODE=1
ENV PYTHONUNBUFFERED=1

# Set the entrypoint to Apache Beam SDK launcher
ENTRYPOINT ["/opt/apache/beam/boot"]

In the running container, I submitted the dataflow task using the following command line:

python ee_to_zarr_dataflow.py \
  --input NASA/GPM_L3/IMERG_V06 \
  --output gs://{BUCKET}/output/ \
  --target_chunks='time=1' \
  --runner DataflowRunner \
  --project {project} \
  --region us-central1 \
  --temp_location gs://{BUCKET}/tmp/ \
  --service_account_email xee-dataflow-controller@{project}.iam.gserviceaccount.com \
  --experiments=use_runner_v2 \
  --sdk_location=container \
  --sdk_container_image us-central1-docker.pkg.dev/{project}/xee-dataflow-run/beam-runner \
  --job_name xee-dataflow-run

But the following error happened.

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/usr/local/lib/python3.10/site-packages/apache_beam/runners/worker/sdk_worker.py", line 311, in _execute
    response = task()
  File "/usr/local/lib/python3.10/site-packages/apache_beam/runners/worker/sdk_worker.py", line 386, in <lambda>
    lambda: self.create_worker().do_instruction(request), request)
  File "/usr/local/lib/python3.10/site-packages/apache_beam/runners/worker/sdk_worker.py", line 656, in do_instruction
    return getattr(self, request_type)(
  File "/usr/local/lib/python3.10/site-packages/apache_beam/runners/worker/sdk_worker.py", line 694, in process_bundle
    bundle_processor.process_bundle(instruction_id))
  File "/usr/local/lib/python3.10/site-packages/apache_beam/runners/worker/bundle_processor.py", line 1274, in process_bundle
    input_op_by_transform_id[element.transform_id].process_encoded(
  File "/usr/local/lib/python3.10/site-packages/apache_beam/runners/worker/bundle_processor.py", line 237, in process_encoded
    self.output(decoded_value)
  File "apache_beam/runners/worker/operations.py", line 567, in apache_beam.runners.worker.operations.Operation.output
  File "apache_beam/runners/worker/operations.py", line 569, in apache_beam.runners.worker.operations.Operation.output
  File "apache_beam/runners/worker/operations.py", line 260, in apache_beam.runners.worker.operations.SingletonElementConsumerSet.receive
  File "apache_beam/runners/worker/operations.py", line 263, in apache_beam.runners.worker.operations.SingletonElementConsumerSet.receive
  File "apache_beam/runners/worker/operations.py", line 950, in apache_beam.runners.worker.operations.DoOperation.process
  File "apache_beam/runners/worker/operations.py", line 951, in apache_beam.runners.worker.operations.DoOperation.process
  File "apache_beam/runners/common.py", line 1503, in apache_beam.runners.common.DoFnRunner.process
  File "apache_beam/runners/common.py", line 1591, in apache_beam.runners.common.DoFnRunner._reraise_augmented
  File "apache_beam/runners/common.py", line 1501, in apache_beam.runners.common.DoFnRunner.process
  File "apache_beam/runners/common.py", line 689, in apache_beam.runners.common.SimpleInvoker.invoke_process
  File "apache_beam/runners/common.py", line 1686, in apache_beam.runners.common._OutputHandler.handle_process_outputs
  File "apache_beam/runners/common.py", line 1799, in apache_beam.runners.common._OutputHandler._write_value_to_tag
  File "apache_beam/runners/worker/operations.py", line 263, in apache_beam.runners.worker.operations.SingletonElementConsumerSet.receive
  File "apache_beam/runners/worker/operations.py", line 950, in apache_beam.runners.worker.operations.DoOperation.process
  File "apache_beam/runners/worker/operations.py", line 951, in apache_beam.runners.worker.operations.DoOperation.process
  File "apache_beam/runners/common.py", line 1503, in apache_beam.runners.common.DoFnRunner.process
  File "apache_beam/runners/common.py", line 1591, in apache_beam.runners.common.DoFnRunner._reraise_augmented
  File "apache_beam/runners/common.py", line 1501, in apache_beam.runners.common.DoFnRunner.process
  File "apache_beam/runners/common.py", line 689, in apache_beam.runners.common.SimpleInvoker.invoke_process
  File "apache_beam/runners/common.py", line 1686, in apache_beam.runners.common._OutputHandler.handle_process_outputs
  File "apache_beam/runners/common.py", line 1799, in apache_beam.runners.common._OutputHandler._write_value_to_tag
  File "apache_beam/runners/worker/operations.py", line 263, in apache_beam.runners.worker.operations.SingletonElementConsumerSet.receive
  File "apache_beam/runners/worker/operations.py", line 950, in apache_beam.runners.worker.operations.DoOperation.process
  File "apache_beam/runners/worker/operations.py", line 951, in apache_beam.runners.worker.operations.DoOperation.process
  File "apache_beam/runners/common.py", line 1503, in apache_beam.runners.common.DoFnRunner.process
  File "apache_beam/runners/common.py", line 1591, in apache_beam.runners.common.DoFnRunner._reraise_augmented
  File "apache_beam/runners/common.py", line 1501, in apache_beam.runners.common.DoFnRunner.process
  File "apache_beam/runners/common.py", line 689, in apache_beam.runners.common.SimpleInvoker.invoke_process
  File "apache_beam/runners/common.py", line 1686, in apache_beam.runners.common._OutputHandler.handle_process_outputs
  File "apache_beam/runners/common.py", line 1799, in apache_beam.runners.common._OutputHandler._write_value_to_tag
  File "apache_beam/runners/worker/operations.py", line 263, in apache_beam.runners.worker.operations.SingletonElementConsumerSet.receive
  File "apache_beam/runners/worker/operations.py", line 950, in apache_beam.runners.worker.operations.DoOperation.process
  File "apache_beam/runners/worker/operations.py", line 951, in apache_beam.runners.worker.operations.DoOperation.process
  File "apache_beam/runners/common.py", line 1503, in apache_beam.runners.common.DoFnRunner.process
  File "apache_beam/runners/common.py", line 1612, in apache_beam.runners.common.DoFnRunner._reraise_augmented
  File "apache_beam/runners/common.py", line 1501, in apache_beam.runners.common.DoFnRunner.process
  File "apache_beam/runners/common.py", line 689, in apache_beam.runners.common.SimpleInvoker.invoke_process
  File "apache_beam/runners/common.py", line 1676, in apache_beam.runners.common._OutputHandler.handle_process_outputs
  File "/usr/local/lib/python3.10/site-packages/xarray_beam/_src/core.py", line 448, in _key_to_chunks
    result = chunk.chunk().compute(num_workers=num_threads)
  File "/usr/local/lib/python3.10/site-packages/xarray/core/dataset.py", line 1073, in compute
    return new.load(**kwargs)
  File "/usr/local/lib/python3.10/site-packages/xarray/core/dataset.py", line 900, in load
    evaluated_data: tuple[np.ndarray[Any, Any], ...] = chunkmanager.compute(
  File "/usr/local/lib/python3.10/site-packages/xarray/namedarray/daskmanager.py", line 85, in compute
    return compute(*data, **kwargs)  # type: ignore[no-untyped-call, no-any-return]
  File "/usr/local/lib/python3.10/site-packages/dask/base.py", line 662, in compute
    results = schedule(dsk, keys, **kwargs)
  File "/usr/local/lib/python3.10/site-packages/xarray/core/indexing.py", line 575, in __array__
    return np.asarray(self.get_duck_array(), dtype=dtype)
  File "/usr/local/lib/python3.10/site-packages/xarray/core/indexing.py", line 578, in get_duck_array
    return self.array.get_duck_array()
  File "/usr/local/lib/python3.10/site-packages/xarray/core/indexing.py", line 835, in get_duck_array
    self._ensure_cached()
  File "/usr/local/lib/python3.10/site-packages/xarray/core/indexing.py", line 832, in _ensure_cached
    self.array = as_indexable(self.array.get_duck_array())
  File "/usr/local/lib/python3.10/site-packages/xarray/core/indexing.py", line 789, in get_duck_array
    return self.array.get_duck_array()
  File "/usr/local/lib/python3.10/site-packages/xarray/core/indexing.py", line 652, in get_duck_array
    array = self.array[self.key]
  File "/usr/local/lib/python3.10/site-packages/xee/ext.py", line 813, in __getitem__
    return indexing.explicit_indexing_adapter(
  File "/usr/local/lib/python3.10/site-packages/xarray/core/indexing.py", line 1013, in explicit_indexing_adapter
    result = raw_indexing_method(raw_key.tuple)
  File "/usr/local/lib/python3.10/site-packages/xee/ext.py", line 945, in _raw_indexing_method
    for (i, j, k), arr in pool.map(
  File "/usr/local/lib/python3.10/site-packages/xee/ext.py", line 962, in _make_tile
    target_image = self._slice_collection(slice(istart, iend))
  File "/usr/local/lib/python3.10/site-packages/xee/ext.py", line 872, in _slice_collection
    imgs = self.store.image_collection.toList(list_range, offset=start).slice(
  File "/usr/local/lib/python3.10/site-packages/ee/collection.py", line 890, in toList
    return apifunction.ApiFunction.call_(
  File "/usr/local/lib/python3.10/site-packages/ee/apifunction.py", line 84, in call_
    return cls.lookup(name).call(*args, **kwargs)
  File "/usr/local/lib/python3.10/site-packages/ee/apifunction.py", line 137, in lookup
    result = cls.lookupInternal(name)
  File "/usr/local/lib/python3.10/site-packages/ee/apifunction.py", line 154, in lookupInternal
    cls.initialize()
  File "/usr/local/lib/python3.10/site-packages/ee/apifunction.py", line 161, in initialize
    signatures = data.getAlgorithms()
  File "/usr/local/lib/python3.10/site-packages/ee/data.py", line 1481, in getAlgorithms
    _get_cloud_projects()
  File "/usr/local/lib/python3.10/site-packages/ee/data.py", line 352, in _get_cloud_projects
    raise ee_exception.EEException(_NOT_INITIALIZED_MESSAGE)
ee.ee_exception.EEException: Earth Engine client library not initialized. See http://goo.gle/ee-auth. [while running 'DatasetToChunks/KeyToChunks-ptransform-55']