cleanup, get timeouts working, fix error propagation

hail/python/hailtop/aiocloud/aiogoogle/client/storage_client.py

@@ -71,10 +71,8 @@ def __init__(self,
         self._value = None
 
     async def cleanup_task(self, task: asyncio.Task):
-        print(f'cleaning up {task}')
         if task.done() and not task.cancelled():
             if exc := task.exception():
-                print(f'raising {exc}')
                 raise exc
         else:
             task.cancel()
@@ -85,6 +83,7 @@ async def cleanup_task(self, task: asyncio.Task):
 
     async def write(self, b):
         assert not self.closed
+        assert not self._request_task.done()
 
         fut = asyncio.ensure_future(self._it.feed(b))
         try:
@@ -104,9 +103,11 @@ async def write(self, b):
     async def _wait_closed(self):
         fut = asyncio.ensure_future(self._it.stop())
         self._exit_stack.push_async_callback(self.cleanup_task, fut)
-        async with await self._request_task as resp:
-            self._value = await resp.json()
-        await self._exit_stack.aclose()
+        try:
+            async with await self._request_task as resp:
+                self._value = await resp.json()
+        finally:
+            await self._exit_stack.aclose()
 
 
 class _TaskManager:
@@ -371,7 +372,6 @@ async def insert_object(self, bucket: str, name: str, **kwargs) -> WritableStrea
                 f'https://storage.googleapis.com/upload/storage/v1/b/{bucket}/o',
                 retry=False,
                 **kwargs))
-            print(f'InsertObjectStream {bucket}/{name}')
             return InsertObjectStream(it, request_task)
 
         # Write using resumable uploads.  See:
@@ -384,7 +384,6 @@ async def insert_object(self, bucket: str, name: str, **kwargs) -> WritableStrea
             **kwargs
         ) as resp:
             session_url = resp.headers['Location']
-        print(f'ResumableInsertObjectStream {bucket}/{name}')
         return ResumableInsertObjectStream(self._session, session_url, chunk_size)
 
     async def get_object(self, bucket: str, name: str, **kwargs) -> GetObjectStream:

hail/python/hailtop/aiocloud/common/session.py

@@ -1,6 +1,6 @@
 from contextlib import AsyncExitStack
 from types import TracebackType
-from typing import Optional, Type, TypeVar, Mapping, Union
+from typing import Optional, Type, TypeVar, Mapping, Union, Dict, Any
 import time
 import aiohttp
 import abc
@@ -67,6 +67,12 @@ async def close(self) -> None:
             del self._session
 
 
+def coerce_timeout(kwargs: Dict[str, Any]):
+    timeout = kwargs.get('timeout')
+    if timeout and isinstance(timeout, float) or isinstance(timeout, int):
+        kwargs['timeout'] = aiohttp.ClientTimeout(timeout)
+
+
 class Session(BaseSession):
     def __init__(self,
                  *,
@@ -82,6 +88,7 @@ def __init__(self,
             self._http_session = http_session
         else:
             self._owns_http_session = True
+            coerce_timeout(kwargs)
             self._http_session = httpx.ClientSession(**kwargs)
         self._credentials = credentials
 
@@ -96,6 +103,8 @@ async def request(self, method: str, url: str, **kwargs) -> aiohttp.ClientRespon
                 if k not in request_params:
                     request_params[k] = v
 
+        coerce_timeout(kwargs)
+
         # retry by default
         retry = kwargs.pop('retry', True)
         if retry:

hail/python/hailtop/aiotools/copy.py

@@ -81,7 +81,7 @@ async def copy(*,
                s3_kwargs: Optional[dict] = None,
                transfers: List[Transfer],
                verbose: bool = False,
-               totals: Optional[Tuple[int, int]] = None
+               totals: Optional[Tuple[int, int]] = None,
                ) -> None:
     with ThreadPoolExecutor() as thread_pool:
         if max_simultaneous_transfers is None:
@@ -173,6 +173,8 @@ async def main() -> None:
     parser.add_argument('-v', '--verbose', action='store_const',
                         const=True, default=False,
                         help='show logging information')
+    parser.add_argument('--timeout', type=str, default=None,
+                        help='show logging information')
     args = parser.parse_args()
 
     if args.verbose:
@@ -183,11 +185,27 @@ async def main() -> None:
     if args.files is None or args.files == '-':
         args.files = sys.stdin.read()
     files = json.loads(args.files)
-    gcs_kwargs = {'gcs_requester_pays_configuration': requester_pays_project}
+
+    timeout = args.timeout
+    if timeout:
+        timeout = float(timeout)
+    print(timeout)
+    gcs_kwargs = {
+        'gcs_requester_pays_configuration': requester_pays_project,
+        'timeout': timeout,
+    }
+    azure_kwargs = {
+        'timeout': timeout,
+    }
+    s3_kwargs = {
+        'timeout': timeout,
+    }
 
     await copy_from_dict(
         max_simultaneous_transfers=args.max_simultaneous_transfers,
         gcs_kwargs=gcs_kwargs,
+        azure_kwargs=azure_kwargs,
+        s3_kwargs=s3_kwargs,
         files=files,
         verbose=args.verbose
     )

hail/python/hailtop/aiotools/fs/copier.py

@@ -227,6 +227,7 @@ async def _copy_part(self,
                          this_part_size: int,
                          part_creator: MultiPartCreate,
                          return_exceptions: bool) -> None:
+        total_written = 0
         try:
             async with self.xfer_sema.acquire_manager(min(Copier.BUFFER_SIZE, this_part_size)):
                 async with await self.router_fs.open_from(srcfile, part_number * part_size, length=this_part_size) as srcf:
@@ -238,10 +239,10 @@ async def _copy_part(self,
                                 raise UnexpectedEOFError()
                             written = await destf.write(b)
                             assert written == len(b)
-                            source_report.finish_bytes(written)
+                            total_written += written
                             n -= len(b)
+            source_report.finish_bytes(total_written)
         except Exception as e:
-            print(f'exception {return_exceptions} {e}')
             if return_exceptions:
                 source_report.set_exception(e)
             else:

hail/python/hailtop/aiotools/fs/stream.py

@@ -110,16 +110,12 @@ def closed(self) -> bool:
         return self._closed
 
     async def __aenter__(self) -> 'WritableStream':
-        print(f'aenter {self}')
         return self
 
     async def __aexit__(
             self, exc_type: Optional[Type[BaseException]] = None,
             exc_value: Optional[BaseException] = None,
             exc_traceback: Optional[TracebackType] = None) -> None:
-        print(f'aexit {self}')
-        import traceback
-        traceback.print_stack()
         await self.wait_closed()
 
 

hail/python/hailtop/utils/utils.py

@@ -476,7 +476,6 @@ async def run_with_sema(pf: Callable[[], Awaitable[T]]):
             try:
                 return await pf()
             except Exception as exc:
-                print(f'bounded gather {exc}')
                 raise exc
 
     tasks = [asyncio.create_task(run_with_sema(pf)) for pf in pfs]