Use shared async client to limit simultaneous connections

Author: BenGalewsky

examples/bigger_uproot.py

@@ -34,7 +34,7 @@
 
 spec = ServiceXSpec(
     General=General(
-        ServiceX="testing1",
+        ServiceX="servicex-uc-af",
         Codegen="uproot",
         OutputFormat="parquet",
         Delivery="LocalCache"

examples/single_file_uproot.py

@@ -43,17 +43,15 @@
 """
 query_ast = ast.parse(qstr)
 qastle_query = qastle.python_ast_to_text_ast(qastle.insert_linq_nodes(query_ast))
-print("From str", qastle_query)
 q2 = FuncADLQuery()
 q2.set_provided_qastle(qastle_query)
 print(q2.generate_selection_string())
-print("From python", query.generate_selection_string())
 spec = ServiceXSpec(
     General=General(
-        ServiceX="testing1",
+        ServiceX="servicex-uc-af",
         Codegen="uproot",
         OutputFormat="parquet",
-        Delivery="LocalCache"
+        Delivery="SignedURLs"
     ),
     Sample=[
         Sample(

servicex/query.py

@@ -407,9 +407,20 @@ async def get_signed_url(
             if progress:
                 progress.advance(task_id=download_progress, task_type="Download")
 
-        while True:
+        transform_is_complete = False
+        while not transform_is_complete:
             if not cached_record:
                 await asyncio.sleep(self.minio_polling_interval)
+
+            # Once the transform is complete we can stop polling since all the files
+            # are guaranteed to be in the bucket. Also, if we are just downloading or
+            # signing urls for a previous transform then we know it is complete as well
+            if cached_record or (
+                self.current_status and self.current_status.status == Status.complete
+            ):
+                transform_is_complete = True
+                await asyncio.sleep(10)
+
             if self.minio:
                 async with download_semaphore:
                     files = await self.minio.list_bucket()
@@ -443,13 +454,6 @@ async def get_signed_url(
                             )  # NOQA 501
                         files_seen.add(file.filename)
 
-            # Once the transform is complete we can stop polling since all the files
-            # are guaranteed to be in the bucket. Also, if we are just downloading or
-            # signing urls for a previous transform then we know it is complete as well
-            if cached_record or (
-                self.current_status and self.current_status.status == Status.complete
-            ):
-                break
 
         # Now just wait until all of our tasks complete
         await asyncio.gather(*download_tasks)

servicex/servicex_client.py

@@ -58,8 +58,8 @@ def get_codegen(_sample: Sample, _general: General):
             return _general.Codegen
 
 
-    async with httpx.AsyncClient(limits=httpx.Limits(max_connections=4)) as client:
-        download_semaphore = asyncio.Semaphore(10)
+    async with httpx.AsyncClient(limits=httpx.Limits(max_connections=4), timeout=None) as client:
+        download_semaphore = asyncio.Semaphore(20)
         servicex_semaphore = asyncio.Semaphore(4)
 
         sx = ServiceXClient(backend=config.General.ServiceX,
@@ -111,11 +111,11 @@ def get_codegen(_sample: Sample, _general: General):
         group = DatasetGroup(datasets, download_semaphore=download_semaphore)
 
         if config.General.Delivery == General.DeliveryEnum.LocalCache:
-            results = group.as_files()
-            return {obj.title: obj.signed_url_list for obj in results}
-        elif config.General.Delivery == General.DeliveryEnum.SignedURLs:
-            results = group.as_signed_urls()
+            results = await group.as_files_async()
             return {obj.title: obj.file_list for obj in results}
+        elif config.General.Delivery == General.DeliveryEnum.SignedURLs:
+            results = await group.as_signed_urls_async()
+            return {obj.title: obj.signed_url_list for obj in results}
 
 deliver = make_sync(deliver_async)
 

tests/test_servicex_adapter.py

@@ -57,7 +57,7 @@ def test_result_formats():
 @patch('servicex.servicex_adapter.httpx.AsyncClient.get')
 async def test_get_transforms(get, servicex, transform_status_response):
     get.return_value = httpx.Response(200, json=transform_status_response)
-    t = await servicex.get_transforms()
+    t = await servicex.get_transforms(self.servicex_semaphore)
     assert len(t) == 1
     assert t[0].request_id == "b8c508d0-ccf2-4deb-a1f7-65c839eebabf"
     get.assert_called_with(url='https://servicex.org/servicex/transformation', headers={})
@@ -68,7 +68,7 @@ async def test_get_transforms(get, servicex, transform_status_response):
 async def test_get_transforms_auth_error(get, servicex):
     with pytest.raises(AuthorizationError):
         get.return_value = httpx.Response(401)
-        await servicex.get_transforms()
+        await servicex.get_transforms(self.servicex_semaphore)
 
 
 @pytest.mark.asyncio
@@ -85,12 +85,12 @@ async def test_get_transforms_wlcg_bearer_token(decode, post, get, servicex,
 
     get.return_value = httpx.Response(200, json=transform_status_response)
     decode.return_value = {'exp': math.inf}
-    await servicex.get_transforms()
+    await servicex.get_transforms(self.servicex_semaphore)
 
     # Try with an expired token
     with pytest.raises(AuthorizationError):
         decode.return_value = {'exp': 0.0}
-        await servicex.get_transforms()
+        await servicex.get_transforms(self.servicex_semaphore)
 
     os.remove(token_file.name)
     del os.environ['BEARER_TOKEN_FILE']
@@ -103,7 +103,7 @@ async def test_get_transforms_with_refresh(get, post, transform_status_response)
     servicex = ServiceXAdapter(url="https://servicex.org", refresh_token="refrescas")
     post.return_value = httpx.Response(200, json={"access_token": "luckycharms"})
     get.return_value = httpx.Response(200, json=transform_status_response)
-    await servicex.get_transforms()
+    await servicex.get_transforms(self.servicex_semaphore)
 
     post.assert_called_with('https://servicex.org/token/refresh',
                             headers={'Authorization': 'Bearer refrescas'}, json=None)
@@ -118,7 +118,7 @@ def test_get_codegens(get, servicex):
         "uproot": "http://uproot-codegen",
         "xaod": "http://xaod-codegen"
     })
-    c = servicex.get_code_generators()
+    c = servicex.get_code_generators(self.servicex_semaphore)
     assert len(c) == 2
     assert c["uproot"] == "http://uproot-codegen"