Merge pull request #1589 from AllenInstitute/GH-1569/bugfix/ecephys-image-download

Bugfix: Asynchronous Http Engine initialization in EcephysProjectCache

Author: kschelonka

allensdk/brain_observatory/ecephys/ecephys_project_api/ecephys_project_lims_api.py

@@ -568,7 +568,7 @@ def get_natural_scene_template(self, number: int) -> Iterable[bytes]:
 
     @classmethod
     def default(cls, lims_credentials: Optional[DbCredentials] = None, 
-                app_kwargs=None, asynchronous=True):
+                app_kwargs=None, asynchronous=False):
         """ Construct a "straightforward" lims api that can fetch data from 
         lims2.
 

allensdk/brain_observatory/ecephys/ecephys_project_api/ecephys_project_warehouse_api.py

@@ -307,7 +307,7 @@ def get_unit_analysis_metrics(self, unit_ids=None, ecephys_session_ids=None, ses
 
 
     @classmethod
-    def default(cls, asynchronous=True, **rma_kwargs):
+    def default(cls, asynchronous=False, **rma_kwargs):
         _rma_kwargs = {"scheme": "http", "host": "api.brain-map.org"}
         _rma_kwargs.update(rma_kwargs)
 

allensdk/brain_observatory/ecephys/ecephys_project_api/http_engine.py

@@ -79,6 +79,10 @@ def stream(self, route):
             if elapsed > self.timeout:
                 raise requests.Timeout(f"Download took {elapsed} seconds, but timeout was set to {self.timeout}")
 
+    @staticmethod
+    def write_bytes(path: str, stream: Iterable[bytes]):
+        write_from_stream(path, stream)
+
 
 AsyncStreamCallbackType = Callable[[AsyncIterator[bytes]], Awaitable[None]]
 
@@ -166,7 +170,13 @@ def __del__(self):
             nest_asyncio.apply()
             loop = asyncio.get_event_loop()
             loop.run_until_complete(self.session.close())
-        
+
+    @staticmethod
+    def write_bytes(
+            path: str,
+            coroutine: Callable[[AsyncStreamCallbackType], Awaitable[None]]):
+        write_bytes_from_coroutine(path, coroutine)
+
 
 def write_bytes_from_coroutine(
     path: str, 
@@ -217,7 +227,6 @@ def write_from_stream(path: str, stream: Iterable[bytes]):
         iterable yielding bytes to be written
 
     """
-
     with open(path, "wb") as fil:
         for chunk in stream:
             fil.write(chunk)

allensdk/brain_observatory/ecephys/ecephys_project_cache.py

@@ -74,7 +74,7 @@ def __init__(
             self,
             fetch_api: EcephysProjectApi = EcephysProjectWarehouseApi.default(), 
             fetch_tries: int = 2,
-            stream_writer: Callable = write_from_stream,
+            stream_writer: Optional[Callable] = None,
             manifest: Optional[Union[str, Path]] = None,
             version: Optional[str] = None,
             cache: bool = True):
@@ -83,6 +83,9 @@ def __init__(
         summaries of sessionwise data and provides tools for downloading detailed 
         sessionwise data (such as spike times).
 
+        To ensure correct configuration, it is recommended to use one of the
+        class constructors rather than to initialize this class directly.
+
         Parameters
         ==========
         fetch_api :
@@ -100,6 +103,14 @@ def __init__(
         fetch_tries : int
             Maximum number of times to attempt a download before giving up and
             raising an exception. Note that this is total tries, not retries
+        stream_writer: Callable
+            The method used to write from stream. Depends on whether the
+            engine is synchronous or asynchronous. If not set, will use the
+            `write_bytes` method native to the `fetch_api`'s `rma_engine`.
+            If the method is incompatible with the `fetch_api`'s `rma_engine`,
+            will likely encounter errors. For this reason it is recommended
+            to leave this field unspecified, or to use one of the class
+            constructors.
         manifest : str or Path
             full path at which manifest json will be stored (default =
             "ecephys_project_manifest.json" in the local directory.)
@@ -108,6 +119,43 @@ def __init__(
             recorded in the file at manifest, an error will be raised.
         cache: bool
             Whether to write to the cache (default=True)
+
+        Notes
+        =====
+        It is highly recommended to construct an instance of this class
+        using one of the following constructor methods:
+
+        from_warehouse(scheme: Optional[str] = None,
+                       host: Optional[str] = None,
+                       asynchronous: bool = True,
+                       manifest: Optional[Union[str, Path]] = None,
+                       version: Optional[str] = None,
+                       cache: bool = True,
+                       fetch_tries: int = 2)
+            Create an instance of EcephysProjectCache with an
+            EcephysProjectWarehouseApi. Retrieves released data stored
+            in the warehouse. Suitable for all users downloading
+            published Allen Institute data.
+        from_lims(lims_credentials: Optional[DbCredentials] = None,
+                  scheme: Optional[str] = None,
+                  host: Optional[str] = None,
+                  asynchronous: bool = True,
+                  manifest: Optional[str] = None,
+                  version: Optional[str] = None,
+                  cache: bool = True,
+                  fetch_tries: int = 2)
+            Create an instance of EcephysProjectCache with an
+            EcephysProjectLimsApi. Retrieves bleeding-edge data stored
+            locally on Allen Institute servers. Suitable for internal
+            users on-site at the Allen Institute or using the corporate
+            vpn. Requires Allen Institute database credentials.
+        fixed(manifest: Optional[Union[str, Path]] = None,
+              version: Optional[str] = None)
+            Create an instance of EcephysProjectCache that will only
+            use locally stored data, downloaded previously from LIMS
+            or warehouse using the EcephysProjectCache.
+            Suitable for users who want to analyze a fixed dataset they
+            have previously downloaded using EcephysProjectCache.
         """
         manifest_ = manifest or "ecephys_project_manifest.json"
         version_ = version or self.MANIFEST_VERSION
@@ -117,7 +165,24 @@ def __init__(
                                                   cache=cache)
         self.fetch_api = fetch_api
         self.fetch_tries = fetch_tries
-        self.stream_writer = stream_writer
+        self.stream_writer = (stream_writer
+                              or self.fetch_api.rma_engine.write_bytes)
+        if stream_writer is not None:
+            self.stream_writer = stream_writer
+        else:
+            if hasattr(self.fetch_api, "rma_engine"):    # EcephysProjectWarehouseApi    # noqa
+                self.stream_writer = self.fetch_api.rma_engine.write_bytes
+            # TODO: Make these names consistent in the different fetch apis
+            elif hasattr(self.fetch_api, "app_engine"):    # EcephysProjectLimsApi    # noqa
+                self.stream_writer = self.fetch_api.app_engine.write_bytes
+            else:
+                raise ValueError(
+                    "Must either set value for `stream_writer`, or use a "
+                    "`fetch_api` with an rma_engine or app_engine attribute "
+                    "that implements `write_bytes`. See `HttpEngine` and "
+                    "`AsyncHttpEngine` from "
+                    "allensdk.brain_observatory.ecephys.ecephys_project_api."
+                    "http_engine for examples.")
 
     def _get_sessions(self):
         path = self.get_cache_path(None, self.SESSIONS_KEY)
@@ -531,8 +596,8 @@ def _from_http_source_default(cls, fetch_api_cls, fetch_api_kwargs, **kwargs):
     def from_lims(cls, lims_credentials: Optional[DbCredentials] = None,
                   scheme: Optional[str] = None,
                   host: Optional[str] = None, 
-                  asynchronous: bool = True,
-                  manifest: Optional[str] = None, 
+                  asynchronous: bool = False,
+                  manifest: Optional[Union[str, Path]] = None, 
                   version: Optional[str] = None, 
                   cache: bool = True,
                   fetch_tries: int = 2):
@@ -557,7 +622,7 @@ def from_lims(cls, lims_credentials: Optional[DbCredentials] = None,
             value if unspecified. Will not be used unless `scheme` is
             also specified.
         asynchronous : bool
-            Whether to fetch file asynchronously. Defaults to True.
+            Whether to fetch file asynchronously. Defaults to False.
         manifest : str or Path
             full path at which manifest json will be stored
         version : str
@@ -586,7 +651,7 @@ def from_lims(cls, lims_credentials: Optional[DbCredentials] = None,
     def from_warehouse(cls, 
                        scheme: Optional[str] = None, 
                        host: Optional[str] = None,
-                       asynchronous: bool = True,
+                       asynchronous: bool = False,
                        manifest: Optional[Union[str, Path]] = None, 
                        version: Optional[str] = None,
                        cache: bool = True, 
@@ -607,7 +672,7 @@ def from_warehouse(cls,
             value if unspecified. Will not be used unless `scheme` is also
             specified.
         asynchronous : bool
-            Whether to fetch file asynchronously. Defaults to True.
+            Whether to fetch file asynchronously. Defaults to False.
         manifest : str or Path
             full path at which manifest json will be stored
         version : str
@@ -623,14 +688,15 @@ def from_warehouse(cls,
             app_kwargs = {"scheme": scheme, "host": host, 
                           "asynchronous": asynchronous}
         else:
-            app_kwargs = None
+            app_kwargs = {"asynchronous": asynchronous}
         return cls._from_http_source_default(
             EcephysProjectWarehouseApi, app_kwargs, manifest=manifest,
              version=version, cache=cache, fetch_tries=fetch_tries
         )
 
     @classmethod
-    def fixed(cls, manifest=None, version=None):
+    def fixed(cls, manifest: Optional[Union[str, Path]] = None, 
+              version: Optional[str] = None):
         """
         Creates a EcephysProjectCache that refuses to fetch any data 
         - only the existing local cache is accessible. Useful if you
@@ -701,5 +767,4 @@ def read_nwb(path):
     reader = pynwb.NWBHDF5IO(str(path), 'r')
     nwbfile = reader.read()
     nwbfile.identifier  # if the file is corrupt, make sure an exception gets raised during read
-    return nwbfile
-
+    return nwbfile

allensdk/test/brain_observatory/ecephys/test_ecephys_project_cache.py

@@ -11,7 +11,7 @@
 import allensdk.brain_observatory.ecephys.ecephys_project_cache as epc
 import allensdk.brain_observatory.ecephys.write_nwb.__main__ as write_nwb
 from allensdk.brain_observatory.ecephys.ecephys_project_api.http_engine import (
-    write_bytes_from_coroutine, AsyncHttpEngine
+    write_from_stream, write_bytes_from_coroutine, AsyncHttpEngine, HttpEngine
 )
 
 
@@ -109,12 +109,18 @@ def shared_tmpdir(tmpdir_factory):
     return str(tmpdir_factory.mktemp('test_ecephys_project_cache'))
 
 
+class MockEngine:
+    def __init__(self):
+        self.write_bytes = write_from_stream
+
+
 @pytest.fixture
 def mock_api(shared_tmpdir, raw_sessions, units, channels, raw_probes, analysis_metrics):
     class MockApi:
 
         def __init__(self, **kwargs):
             self.accesses = collections.defaultdict(lambda: 1)
+            self.rma_engine = MockEngine()
 
         def __getattr__(self, name):
             self.accesses[name] += 1
@@ -375,5 +381,83 @@ def test_from_lims_default(tmpdir_factory):
     cache = epc.EcephysProjectCache.from_lims(
         manifest=os.path.join(tmpdir, "manifest.json")
     )
-    assert isinstance(cache.fetch_api.app_engine, AsyncHttpEngine)
-    assert cache.stream_writer is epc.write_bytes_from_coroutine
+    assert isinstance(cache.fetch_api.app_engine, HttpEngine)
+    assert cache.stream_writer is write_from_stream
+    assert cache.fetch_api.app_engine.scheme == "http"
+    assert cache.fetch_api.app_engine.host == "lims2"
+
+
+def test_from_warehouse_default(tmpdir_factory):
+    tmpdir = str(tmpdir_factory.mktemp("test_from_warehouse_default"))
+
+    cache = epc.EcephysProjectCache.from_warehouse(
+        manifest=os.path.join(tmpdir, "manifest.json")
+    )
+    assert isinstance(cache.fetch_api.rma_engine, HttpEngine)
+    assert cache.stream_writer is write_from_stream
+    assert cache.fetch_api.rma_engine.scheme == "http"
+    assert cache.fetch_api.rma_engine.host == "api.brain-map.org"
+
+
+def test_init_default(tmpdir_factory):
+    tmpdir = str(tmpdir_factory.mktemp("test_init_default"))
+    cache = epc.EcephysProjectCache(
+        manifest=os.path.join(tmpdir, "manifest.json")
+    )
+    assert isinstance(cache.fetch_api.rma_engine, HttpEngine)
+    assert cache.stream_writer is cache.fetch_api.rma_engine.write_bytes
+    assert cache.fetch_api.rma_engine.scheme == "http"
+    assert cache.fetch_api.rma_engine.host == "api.brain-map.org"
+
+
+@pytest.mark.parametrize(
+    ("cache_constructor, asynchronous, engine_attr, expected_engine,"
+     "expected_scheme, expected_host, expected_stream_writer"), [
+        (
+            epc.EcephysProjectCache.from_lims, True,
+            "app_engine", AsyncHttpEngine, "http", "lims2",
+            write_bytes_from_coroutine
+        ),
+        (
+            epc.EcephysProjectCache.from_warehouse, True,
+            "rma_engine", AsyncHttpEngine, "http", "api.brain-map.org",
+            write_bytes_from_coroutine
+        ),
+        (
+            epc.EcephysProjectCache.from_warehouse, False,
+            "rma_engine", HttpEngine, "http", "api.brain-map.org",
+            write_from_stream
+        ),
+        (
+            epc.EcephysProjectCache.from_lims, False,
+            "app_engine", HttpEngine, "http", "lims2",
+            write_from_stream
+        ),
+    ])
+def test_stream_asynchronous_arg(
+        cache_constructor, asynchronous, engine_attr, expected_engine,
+        expected_scheme, expected_host, expected_stream_writer,
+        tmpdir_factory):
+    """ Ensure the proper stream engine is chosen from the `asynchronous`
+    argument in the EcephysProjectCache constructors (using other default
+    values)."""
+    tmpdir = str(tmpdir_factory.mktemp("test_stream_async_args"))
+    cache = cache_constructor(
+        asynchronous=asynchronous,
+        manifest=os.path.join(tmpdir, "manifest.json")
+    )
+    engine = getattr(cache.fetch_api, engine_attr)
+    assert isinstance(engine, expected_engine)
+    assert cache.stream_writer is expected_stream_writer
+    assert engine.scheme == expected_scheme
+    assert engine.host == expected_host
+
+
+def test_stream_writer_method_default_correct(tmpdir_factory):
+    """Checks that the stream_writer contained in the rma engine is used
+    when one is not supplied to the __init__ method.
+    """
+    tmpdir = str(tmpdir_factory.mktemp("test_from_warehouse_default"))
+    manifest = os.path.join(tmpdir, "manifest.json")
+    cache = epc.EcephysProjectCache(stream_writer=None, manifest=manifest)
+    assert cache.stream_writer == cache.fetch_api.rma_engine.write_bytes