Feature/pickle rasterio #2131

jhamman · 2018-05-14T23:38:59Z

Closes rasterio backend should use DataStorePickleMixin (or something similar) #2121
Tests added
Tests passed
Fully documented, including whats-new.rst for all changes and api.rst for new API

…_rasterio

shoyer · 2018-05-15T02:03:13Z

xarray/backends/common.py

+
+        self.opener = opener
+        self.open_args = args
+        self.open_kwargs = kwargs


For simplicity, let's pool these into a single functools.partial object. Note that you can get out keyword arguments if desired from partial.keywords.

should we just have the backends pass in a partial function then?

If we don't anticipate adding any argument, then I think it's just as easy to make the partial function here. I suppose there's some potential overlap in needs with auto-closing a file, so I'm OK either way.

shoyer · 2018-05-15T02:08:05Z

xarray/backends/common.py

+        self._isopen = False
+
+    @property
+    def value(self):


I'm a little concerned about the complexity of making file opening lazy. State makes things harder to reason about. For example, what happens with the mode argument if never actually accessed the file?

For now can we stick with constructing value eagerly?

I'm actually not a big fan of dropping the lazy opening. For one, there will be a non-negligible performance penalty. Perhaps a better argument is that we've been using this logic successfully in xarray for a while now. That said, its not clear to me what the potential problems with this approach so I'm open to convincing otherwise.

Don't we always open files immediately to pull out metadata? (e.g., variables and dimensions)

I think you are right. I've removed most of this in 6669035

shoyer · 2018-05-15T02:09:24Z

xarray/backends/common.py

+    def __getitem__(self, key):
+        return self.value[key]
+
+    def __getattr__(self, name):


This magic is convenient but I think it could possibly be a source of bugs. I would stick to requiring pulling out .value.

Yeah, I was thinking someone would call me on this :). It was the path of least resistance for getting this working but I think a lot of care would be required not to trample on attributes. Do you think I should drop both getitem/getattr?

Removed in 097e264

shoyer · 2018-05-15T02:16:57Z

xarray/backends/common.py

@@ -503,3 +503,43 @@ def assert_open(self):
        if not self._isopen:
            raise AssertionError('internal failure: file must be open '
                                 'if `autoclose=True` is used.')
+
+
+class PickleByReconstructionWrapper(object):


It would be nice to add a unit test verifies that this works properly independently of any concrete datastore.

Maybe something simple with open()?

Done in 8b286c0

shoyer · 2018-05-15T02:18:09Z

xarray/tests/test_backends.py

+        with create_tmp_geotiff() as (tmp_file, expected):
+            with xr.open_rasterio(tmp_file) as rioda:
+                temp = pickle.dumps(rioda)
+                actual = pickle.loads(temp)


Something to think about here -- I'm pretty sure this leaks a file descriptor. I'm not quite sure what the best way to deal with this is (maybe it's not worth worrying about yet), but in the past this has caused issues with tests on various platforms (especially Windows).

And indeed, it's causing a test failure on Windows :).

I think it should work to explicitly open actual in a context manager, e.g.,

with pickle.loads(temp) as actual: ...

fmaussion

Thanks for digging into this! I don't fully understand what's going on here but I wondered if we shouldn't close the file instead of just dereferencing it.

fmaussion · 2018-05-15T10:35:15Z

xarray/backends/common.py

+
+    def __getstate__(self):
+        state = self.__dict__.copy()
+        del state['_ds']


shouldn't we close the file here? state['_ds'].close()

It might be used after it is pickled.

Ah OK, maybe I misunderstood. I thought the purpose of this wrapper was to dereference (and close) the file before 'dump()' and re-open it after pickle.load().

The basic idea is that most open file handles can't be pickled so we need to provide a mechanism to remove the existing handle and generate a new one in the dump/load steps in the pickling. In most cases, we do want to keep the original file open.

fmaussion · 2018-05-15T10:35:24Z

xarray/backends/common.py

+
+    def __setstate__(self, state):
+        self.__dict__.update(state)
+        self._ds = None


mrocklin

Glad to see work in this direction.

mrocklin · 2018-05-15T11:21:16Z

xarray/backends/common.py

+
+    def __getstate__(self):
+        state = self.__dict__.copy()
+        del state['_ds']


It might be used after it is pickled.

mrocklin · 2018-05-15T11:22:18Z

xarray/tests/test_backends.py

+            with xr.open_rasterio(tmp_file) as rioda:
+                temp = pickle.dumps(rioda)
+                with pickle.loads(temp) as actual:
+                    assert_equal(actual, rioda)


Does it make sense to also add an integration test reading rasterio data with dask.distributed?

Done in 55a3abc.

rsignell-usgs · 2018-05-15T19:02:06Z

@jhamman should I test this out on my original workflow or wait a bit?

jhamman · 2018-05-15T19:04:18Z

@rsignell-usgs - I think this should work in its current form. I'm writing some additional tests now and have some cleanup tasks as indicated in the reviews but I do not expect the behavior to change.

shoyer · 2018-05-15T20:14:06Z

Thinking about this a little more, maybe we should consolidate pickling with autoclosing into a single, composable wrapper class. Usage would look something like:

wrapper = FileWrapper(netCDF4.Dataset, path, mode='r', autoclose=False, kwargs=kwargs)
with wrapper:
    # will close/open wrapped file
    with wrapper.acquire() as ds:
        # context manager is dummy, simply returning the open file object
        ... # do stuff with ds, a netCDF4.Dataset instance
# or could write wrapper.close() instead of using the context manager

wrapper = FileWrapper(netCDF4.Dataset, path, mode='r', autoclose=True, kwargs=kwargs)
with wrapper:
    # explicit opening/closing the wrapper is a no-op if autoclose=True
    with wrapper.acquire() as ds:
        # context manager opens/closes file
        # if file is already open, it's a no-op
        ... # do stuff with ds, a netCDF4.Dataset instance

We could then expose xarray.backends.FileWrapper as part of our external API for backend authors (without coupling it into the Datastore class directly).

Eventually, we might even change the default for autoclose to automatically use LRU caching (but that obviously would come later).

rsignell-usgs · 2018-05-15T22:15:22Z

It's working for me!
https://gist.github.com/rsignell-usgs/ef81fb4306dac3a2406d0adb575b340f

jhamman · 2018-05-15T22:20:21Z

@shoyer - I like the idea of a wrapper object that achieves both of these tasks. I'm incrementally working that way and have just updated the Pickle Wrapper for now.

rsignell-usgs · 2018-06-05T23:00:51Z

@jhamman , still very much interested in this -- could the existing functionality be merged and enhanced later?

shoyer · 2018-06-06T21:34:03Z

xarray/tests/test_backends.py

+
+    lines = ['foo bar spam eggs']
+
+    with create_tmp_file() as tmp:


This needs to add allow_cleanup_failure=ON_WINDOWS for the test to pass on windows

Great, that did fix the windows failure.

…_rasterio

jhamman · 2018-06-07T13:08:10Z

All the tests are passing here. Its been a while so perhaps it would be good to get another review from @pydata/xarray.

@rsignell-usgs - can you give this another whirl on your use case?

fmaussion

LGTM, thanks!

rsignell-usgs · 2018-06-07T14:46:21Z

@jhamman , although I'm getting distributed workers to compute the mean from a bunch of images, I'm getting a "Failed to Serialize" error in cell [23] of this notebook:
https://gist.github.com/rsignell-usgs/90f15e2da918e3c6ba6ee5bb6095d594
If this is a bug, I think it was there before the recent updates.

You should be able to run this notebook without modification.

shoyer

There's one bug to fix here, but otherwise I agree that we can merge this soon and clean it up later.

shoyer · 2018-06-07T15:24:02Z

xarray/backends/common.py

+        return self._ds
+
+    def __getstate__(self):
+        del self._ds


You definitely need to delete this from state rather than from the class, like what we do in DataStorePickleMixin:

xarray/xarray/backends/common.py

Lines 461 to 468 in bc52f8a

def __getstate__(self):

state = self.__dict__.copy()

del state['_ds']

del state['_isopen']

if self._mode == 'w':

# file has already been created, don't override when restoring

state['_mode'] = 'a'

return state

Otherwise, this will break when you try to pickle this more than once, which is the error that shows up in @rsignell-usgs's notebook.

Thanks, I had just noticed this problem and it should be fixed/tested it in 35520c0.

rsignell-usgs · 2018-06-07T16:07:14Z

@jhamman woohoo! Cell [20] completes nicely now:
https://gist.github.com/rsignell-usgs/90f15e2da918e3c6ba6ee5bb6095d594
I'm getting some errors in Cell [20], but I think those are unrelated and didn't affect the successful completion of the tasks, right? (this is on an HPC system)

shoyer

LGTM

rsignell-usgs · 2018-06-07T18:45:42Z

Might this PR warrant a new minor release?

shoyer · 2018-06-07T19:18:55Z

@rsignell-usgs Sure, though interp() (#2104) is almost ready, too, so we might wait for that.

rsignell-usgs · 2018-06-07T19:20:24Z

Sounds good. Thanks @shoyer!

shoyer · 2018-06-08T05:00:59Z

v0.10.7 is out

Joseph Hamman added 5 commits May 12, 2018 12:56

add regression test

4319aac

Merge branch 'master' of github.com:pydata/xarray into feature/pickle…

fb8b3d7

…_rasterio

add PickleByReconstructionWrapper

2b7176b

docs

c8de3a5

Merge branch 'master' of github.com:pydata/xarray into feature/pickle…

45f0aeb

…_rasterio

jhamman mentioned this pull request May 15, 2018

add rasterio to docker image? pangeo-data/pangeo#249

Closed

shoyer reviewed May 15, 2018

View reviewed changes

load in context manager

99c9661

fmaussion reviewed May 15, 2018

View reviewed changes

mrocklin reviewed May 15, 2018

View reviewed changes

Joseph Hamman added 2 commits May 15, 2018 12:17

add distributed integration test

55a3abc

add test_pickle_reconstructor

8b286c0

drop lazy opening/caching and use partial function for open

6669035

stop using clever getattr hack

097e264

jhamman mentioned this pull request May 17, 2018

updated dockerfiles pangeo-data/pangeo#261

Merged

shoyer reviewed Jun 6, 2018

View reviewed changes

Joseph Hamman added 3 commits June 6, 2018 22:32

Merge branch 'master' of github.com:pydata/xarray into feature/pickle…

822a080

…_rasterio

allow_cleanup_failure=ON_WINDOWS in tests for windows

40c15a6

whats new fix

78ffcb4

fmaussion approved these changes Jun 7, 2018

View reviewed changes

shoyer reviewed Jun 7, 2018

View reviewed changes

fix bug in multiple pickles

35520c0

shoyer approved these changes Jun 7, 2018

View reviewed changes

fix for windows

8d7c768

jhamman merged commit 21a9f3d into pydata:master Jun 7, 2018

jhamman deleted the feature/pickle_rasterio branch June 7, 2018 18:03

rsignell-usgs mentioned this pull request Jun 14, 2018

Update notebook and worker containers with jhamman's xarray fix for rasterio ESIPFed/esiphub-dev#12

Closed

	def __getstate__(self):
	state = self.__dict__.copy()
	del state['_ds']
	del state['_isopen']
	if self._mode == 'w':
	# file has already been created, don't override when restoring
	state['_mode'] = 'a'
	return state

Feature/pickle rasterio #2131

Feature/pickle rasterio #2131

Conversation

jhamman commented May 14, 2018 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

fmaussion left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

mrocklin left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

rsignell-usgs commented May 15, 2018

jhamman commented May 15, 2018

shoyer commented May 15, 2018

rsignell-usgs commented May 15, 2018

jhamman commented May 15, 2018

rsignell-usgs commented Jun 5, 2018 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jhamman commented Jun 7, 2018

fmaussion left a comment

Choose a reason for hiding this comment

rsignell-usgs commented Jun 7, 2018 • edited Loading

shoyer left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

rsignell-usgs commented Jun 7, 2018 • edited Loading

shoyer left a comment

Choose a reason for hiding this comment

rsignell-usgs commented Jun 7, 2018

shoyer commented Jun 7, 2018

rsignell-usgs commented Jun 7, 2018

shoyer commented Jun 8, 2018

jhamman commented May 14, 2018 •

edited

Loading

rsignell-usgs commented Jun 5, 2018 •

edited

Loading

rsignell-usgs commented Jun 7, 2018 •

edited

Loading

rsignell-usgs commented Jun 7, 2018 •

edited

Loading