Fix in file position logic: file descriptor and Python-side handle (#20270)

philippslang · facebook-github-bot · commit f23fb66e6e67 · 2019-05-09T08:20:01.000-07:00
Summary: This addresses #18436 The logic replicates the essence of closing file descriptors in numpy: https://github.com/numpy/numpy/blob/bf20e3034085716c4559ec4bf31b23b6016f266c/numpy/core/include/numpy/npy_3kcompat.h#L278 This stores the position of the file descriptor before resetting it to the Python handle offset, then resets to the original position before exit. The Python-side handle is then updated to reflect the new position. Also added somewhat more demanding tests to cover this. Pull Request resolved: #20270 Differential Revision: D15275902 Pulled By: soumith fbshipit-source-id: 5ca8a52b61c7718d2e69571f72f80b1350b0acdb
diff --git a/test/test_torch.py b/test/test_torch.py
@@ -9287,33 +9287,45 @@ def test_serialization_gzip(self):
 
     def test_serialization_offset(self):
         a = torch.randn(5, 5)
-        i = 41
-        for use_name in (False, True):
-            # Passing filename to torch.save(...) will cause the file to be opened twice,
-            # which is not supported on Windows
-            if sys.platform == "win32" and use_name:
-                continue
-            with tempfile.NamedTemporaryFile() as f:
-                handle = f if not use_name else f.name
-                pickle.dump(i, f)
-                torch.save(a, f)
-                f.seek(0)
-                j = pickle.load(f)
-                b = torch.load(f)
-            self.assertTrue(torch.equal(a, b))
-            self.assertEqual(i, j)
+        b = torch.randn(2, 2)
+        m = torch.nn.Conv2d(1, 1, (1, 3))
+        i, j = 41, 43
+        with tempfile.NamedTemporaryFile() as f:
+            pickle.dump(i, f)
+            torch.save(a, f)
+            pickle.dump(j, f)
+            torch.save(b, f)
+            torch.save(m, f)
+            f.seek(0)
+            i_loaded = pickle.load(f)
+            a_loaded = torch.load(f)
+            j_loaded = pickle.load(f)
+            b_loaded = torch.load(f)
+            m_loaded = torch.load(f)
+        self.assertTrue(torch.equal(a, a_loaded))
+        self.assertTrue(torch.equal(b, b_loaded))
+        self.assertTrue(m.kernel_size == m_loaded.kernel_size)
+        self.assertEqual(i, i_loaded)
+        self.assertEqual(j, j_loaded)
 
     def test_serialization_offset_filelike(self):
         a = torch.randn(5, 5)
-        i = 41
+        b = torch.randn(2, 3)
+        i, j = 41, 43
         with BytesIOContext() as f:
             pickle.dump(i, f)
-            torch.save(a, f)
+            torch.save(a, f)            
+            pickle.dump(j, f)
+            torch.save(b, f)
             f.seek(0)
-            j = pickle.load(f)
-            b = torch.load(f)
-        self.assertTrue(torch.equal(a, b))
-        self.assertEqual(i, j)
+            i_loaded = pickle.load(f)
+            a_loaded = torch.load(f)
+            j_loaded = pickle.load(f)
+            b_loaded = torch.load(f)
+        self.assertTrue(torch.equal(a, a_loaded))
+        self.assertTrue(torch.equal(b, b_loaded))
+        self.assertEqual(i, i_loaded)
+        self.assertEqual(j, j_loaded)
 
     def test_serialization_offset_gzip(self):
         a = torch.randn(5, 5)
diff --git a/torch/csrc/generic/StorageMethods.cpp b/torch/csrc/generic/StorageMethods.cpp
@@ -261,7 +261,8 @@ static PyObject *THPStorage_(setFromFile)(THPStorage *self, PyObject *args)
   }
 
   // file is backed by a fd
-  int fd = PyObject_AsFileDescriptor(file);
+  const int fd = PyObject_AsFileDescriptor(file);
+  const auto fd_original_pos = lseek(fd, 0, SEEK_CUR);
   if (offset != Py_None) {
     lseek(fd, THPUtils_unpackLong(offset), SEEK_SET);
   }
@@ -272,6 +273,17 @@ static PyObject *THPStorage_(setFromFile)(THPStorage *self, PyObject *args)
     return nullptr;
   Py_INCREF(self);
 
+  // the file descriptor is returned to original position and
+  // the file handle at python call-site needs updating to the
+  // advanced postion
+  const auto fd_current_pos = lseek(fd, 0, SEEK_CUR);
+  lseek(fd, fd_original_pos, SEEK_SET);
+  const auto seek_return = PyObject_CallMethod(file, "seek", "li", fd_current_pos, 0);
+  if (seek_return == nullptr) {
+      return nullptr;
+  }
+  Py_DECREF(seek_return);
+
   return (PyObject *) self;
   END_HANDLE_TH_ERRORS
 }
diff --git a/torch/serialization.py b/torch/serialization.py
@@ -579,6 +579,7 @@ def persistent_load(saved_id):
     for key in deserialized_storage_keys:
         assert key in deserialized_objects
         deserialized_objects[key]._set_from_file(f, offset, f_should_read_directly)
-        offset = None
+        if offset is not None:
+            offset = f.tell()
 
     return result
