wip

Author: HansKallekleiv

backend/src/backend/primary/user_session_proxy.py

@@ -17,7 +17,9 @@
 class _RedisUserJobs:
     def __init__(self) -> None:
         # redis.Redis does not yet have namespace support - https://github.com/redis/redis-py/issues/12 - need to prefix manually.
-        self._redis_client = redis.Redis.from_url(config.REDIS_USER_SESSION_URL, decode_responses=True)
+        self._redis_client = redis.Redis.from_url(
+            config.REDIS_USER_SESSION_URL, decode_responses=True
+        )
 
     def get_job_name(self, user_id: str) -> Optional[str]:
         return self._redis_client.get("user-job-name:" + user_id)
@@ -51,7 +53,9 @@ async def _active_running_job(self, user_id: str) -> bool:
             return True
 
         async with httpx.AsyncClient() as client:
-            res = await client.get(f"http://{self._name}:{self._port}/api/v1/jobs/{existing_job_name}")
+            res = await client.get(
+                f"http://{self._name}:{self._port}/api/v1/jobs/{existing_job_name}"
+            )
 
         job = res.json()
 
@@ -88,8 +92,8 @@ async def _create_new_job(self, user_id: str) -> None:
                     # these could be dynamic based on e.g. the selected ensemble sizess by the user.
                     json={
                         "resources": {
-                            "limits": {"memory": "64GiB", "cpu": "4"},
-                            "requests": {"memory": "32GiB", "cpu": "2"},
+                            "limits": {"memory": "128GiB", "cpu": "16"},
+                            "requests": {"memory": "32GiB", "cpu": "4"},
                         }
                     },
                 )
@@ -115,7 +119,9 @@ async def get_base_url(self, user_id: str) -> str:
 RADIX_JOB_SCHEDULER_INSTANCE = RadixJobScheduler("backend-user-session", 8000)
 
 
-async def proxy_to_user_session(request: Request, authenticated_user: AuthenticatedUser) -> Any:
+async def proxy_to_user_session(
+    request: Request, authenticated_user: AuthenticatedUser
+) -> Any:
     # Ideally this function should probably be a starlette/FastAPI middleware, but it appears that
     # it is not yet possible to put middleware on single routes through decorator like in express.js.
 

backend/src/backend/user_session/routers/surface/router.py

@@ -14,18 +14,15 @@
 from fastapi import APIRouter, Depends, Request, Body
 
 from src.backend.auth.auth_helper import AuthenticatedUser, AuthHelper
-
-from src.services.sumo_access.surface_access import SurfaceAccess
 import asyncio
 from src.backend.primary.routers.surface import schemas
-from .test_async import async_get_cached_surf
-from .test_go import go_get_surface_blobs
+
 from sumo.wrapper import SumoClient
-from fmu.sumo.explorer.objects import CaseCollection, Case, SurfaceCollection
+from fmu.sumo.explorer.objects import CaseCollection
 from src.services.utils.perf_timer import PerfTimer
-import base64
+
 from io import BytesIO
-from azure.storage.blob.aio import BlobServiceClient, BlobClient, ContainerClient
+from azure.storage.blob.aio import BlobClient, ContainerClient
 import requests
 
 LOGGER = logging.getLogger(__name__)
@@ -42,169 +39,73 @@ async def well_intersection_reals_from_user_session(
 ) -> List[schemas.SurfaceIntersectionPoints]:
     body = await request.json()
     ensemble_ident = schemas.EnsembleIdent(**body.get("ensemble_ident"))
-    realization_surface_set_spec = schemas.RealizationsSurfaceSetSpec(**body.get("realizations_surface_set_spec"))
+    realization_surface_set_spec = schemas.RealizationsSurfaceSetSpec(
+        **body.get("realizations_surface_set_spec")
+    )
     surface_fence_spec = schemas.SurfaceFenceSpec(**body.get("surface_fence_spec"))
     timer = PerfTimer()
     # Config
     case_uuid = ensemble_ident.case_uuid
-    snames = realization_surface_set_spec.surface_names
-    sattr = realization_surface_set_spec.surface_attribute
     ensemble_name = ensemble_ident.ensemble_name
-    realization_nums = realization_surface_set_spec.realization_nums
     intersections = []
 
-    uuids = get_uuids(
-        case_uuid, ensemble_name, realization_nums, snames, sattr, authenticated_user.get_sumo_access_token()
+    uuids = get_surface_set_uuids(
+        case_uuid,
+        ensemble_name,
+        realization_surface_set_spec,
+        authenticated_user.get_sumo_access_token(),
     )
     base_uri, auth_token = get_base_uri_and_auth_token_for_case(
         case_uuid, "prod", authenticated_user.get_sumo_access_token()
     )
 
-    async with ContainerClient.from_container_url(container_url=base_uri, credential=auth_token) as container_client:
+    async with ContainerClient.from_container_url(
+        container_url=base_uri, credential=auth_token
+    ) as container_client:
         coro_array = []
         timer = PerfTimer()
 
         for uuid in uuids:
-            coro_array.append(my_download_to_file(container_client=container_client, sumo_surf_uuid=uuid))
+            coro_array.append(
+                download_blob(container_client=container_client, sumo_surf_uuid=uuid)
+            )
         res_array = await asyncio.gather(*coro_array)
-        dl_time_s = timer.lap_s()
-        print(f"download surfs: {dl_time_s:.2f}s", flush=True)
-        LOGGER.info(f"download surfs: {dl_time_s:.2f}s", extra={"download_surfs": dl_time_s})
+        elapsed_download = timer.lap_s()
 
         tot_mb = 0
         for res in res_array:
             tot_mb += len(res) / (1024 * 1024)
-        print(f"Total MB downloaded: {tot_mb:.2f}MB  =>  {tot_mb/dl_time_s:.2f}MB/s", flush=True)
-        LOGGER.info(
-            f"Total MB downloaded: {tot_mb:.2f}MB  =>  {tot_mb/dl_time_s:.2f}MB/s",
-            extra={"total_mb": tot_mb, "mb_per_s": tot_mb / dl_time_s},
-        )
 
-    # surfs = [xtgeo.surface_from_file(BytesIO(bytestr), fformat="irap_binary") for bytestr in res_array]
     surfaces = await load_xtgeo(res_array)
-    print(f"convert surfs: {timer.lap_s():.2f}s", flush=True)
-    fence_arr = np.array(
-        [
-            surface_fence_spec.x_points,
-            surface_fence_spec.y_points,
-            np.zeros(len(surface_fence_spec.y_points)),
-            surface_fence_spec.cum_length,
-        ]
-    ).T
-    intersections = await make_intersections(surfaces, fence_arr)
-
-    print(f"intersect surfs: {timer.lap_s():.2f}s", flush=True)
-    return ORJSONResponse([section.dict() for section in intersections])
-
+    elapsed_xtgeo = timer.lap_s()
 
-async def my_download_to_file(container_client: ContainerClient, sumo_surf_uuid):
-    blob_client: BlobClient = container_client.get_blob_client(blob=sumo_surf_uuid)
-    download_stream = await blob_client.download_blob(
-        max_concurrency=4,
-    )
-    return await download_stream.readall()
-
-
-@router.post("/well_intersection_reals_from_user_session")
-async def well_intersection_reals_from_user_session2(
-    request: Request,
-    authenticated_user: AuthenticatedUser = Depends(AuthHelper.get_authenticated_user),
-) -> List[schemas.SurfaceIntersectionPoints]:
-    body = await request.json()
-
-    ensemble_ident = schemas.EnsembleIdent(**body.get("ensemble_ident"))
-    realization_surface_set_spec = schemas.RealizationsSurfaceSetSpec(**body.get("realizations_surface_set_spec"))
-    surface_fence_spec = schemas.SurfaceFenceSpec(**body.get("surface_fence_spec"))
-    timer = PerfTimer()
-    # Config
-    case_uuid = ensemble_ident.case_uuid
-    snames = realization_surface_set_spec.surface_names
-    sattr = realization_surface_set_spec.surface_attribute
-    ensemble_name = ensemble_ident.ensemble_name
-    realization_nums = realization_surface_set_spec.realization_nums
-    # Get uuids
-    uuids = get_uuids(
-        case_uuid, ensemble_name, realization_nums, snames, sattr, authenticated_user.get_sumo_access_token()
-    )
-    elapsed_meta = timer.lap_ms()
-
-    # Check if cached
-    uuids_to_download = []
-    surfaces = []
-    for uuid in uuids:
-        possible_surface = await cache.get(f"{authenticated_user._user_id}-{uuid}")
-        if possible_surface is None:
-            uuids_to_download.append(uuid)
-        else:
-            surfaces.append(possible_surface)
-    elapsed_cache = timer.lap_ms()
-    if uuids_to_download:
-        # Download remaining
-        data_map_b64 = go_get_surface_blobs(authenticated_user.get_sumo_access_token(), case_uuid, uuids_to_download)
-        elapsed_download = timer.lap_ms()
-
-        # Convert to xtgeo
-        downloaded_surface_dict = await b64_to_xtgeo(data_map_b64)
-        elapsed_xtgeo = timer.lap_ms()
-
-        # Add to cache
-        for uuid, surface in downloaded_surface_dict.items():
-            await cache.set(f"{authenticated_user._user_id}-{uuid}", surface)
-            surfaces.append(surface)
-    else:
-        elapsed_download = 0
-        elapsed_xtgeo = 0
-    # Intersect
-    fence_arr = np.array(
-        [
-            surface_fence_spec.x_points,
-            surface_fence_spec.y_points,
-            np.zeros(len(surface_fence_spec.y_points)),
-            surface_fence_spec.cum_length,
-        ]
-    ).T
-    intersections = await make_intersections(surfaces, fence_arr)
-    elapsed_intersect = timer.lap_ms()
+    intersections = await make_intersections(surfaces, surface_fence_spec)
+    elapsed_intersect = timer.lap_s()
 
-    result = [intersection.dict() for intersection in intersections]
-    elapsed_response_format = timer.lap_ms()
     LOGGER.info(
         f"Got intersected surface set from Sumo: {timer.elapsed_ms()}ms ("
-        f"meta={elapsed_meta}ms, "
-        f"cache={elapsed_cache}ms, "
         f"download={elapsed_download}ms, "
         f"xtgeo={elapsed_xtgeo}ms, "
         f"intersect={elapsed_intersect}ms, "
-        f"response_format={elapsed_response_format}ms) ",
+        f"size={tot_mb:.2f}MB, "
+        f"speed={tot_mb/elapsed_download:.2f}MB/s)",
         extra={
-            "meta": elapsed_meta,
-            "cache": elapsed_cache,
             "download": elapsed_download,
             "xtgeo": elapsed_xtgeo,
             "intersect": elapsed_intersect,
-            "response_format": elapsed_response_format,
+            "size": tot_mb,
+            "speed": tot_mb / elapsed_download,
         },
     )
-    return ORJSONResponse(result)
-
+    return ORJSONResponse([section.dict() for section in intersections])
 
-async def b64_to_xtgeo(data_map_b64):
-    def to_xtgeo(object_id, b64_blob):
-        bytestr = base64.b64decode(b64_blob)
-        xtgeo_surface = xtgeo.surface_from_file(BytesIO(bytestr), fformat="irap_binary")
-        return {object_id: xtgeo_surface}
 
-    loop = asyncio.get_running_loop()
-    with ThreadPoolExecutor() as executor:
-        tasks = [
-            loop.run_in_executor(executor, to_xtgeo, object_id, b64_blob)
-            for object_id, b64_blob in data_map_b64.items()
-        ]
-        results = await asyncio.gather(*tasks)
-    downloaded_surface_dict = {}
-    for result in results:
-        downloaded_surface_dict.update(result)
-    return downloaded_surface_dict
+async def download_blob(container_client: ContainerClient, sumo_surf_uuid):
+    blob_client: BlobClient = container_client.get_blob_client(blob=sumo_surf_uuid)
+    download_stream = await blob_client.download_blob(
+        max_concurrency=4,
+    )
+    return await download_stream.readall()
 
 
 def make_intersection(surf, fence_arr):
@@ -217,11 +118,21 @@ def make_intersection(surf, fence_arr):
     return intersection
 
 
-async def make_intersections(surfaces, fence_arr):
+async def make_intersections(surfaces, surface_fence_spec):
     loop = asyncio.get_running_loop()
-
+    fence_arr = np.array(
+        [
+            surface_fence_spec.x_points,
+            surface_fence_spec.y_points,
+            np.zeros(len(surface_fence_spec.y_points)),
+            surface_fence_spec.cum_length,
+        ]
+    ).T
     with ProcessPoolExecutor() as executor:
-        tasks = [loop.run_in_executor(executor, make_intersection, surf, fence_arr) for surf in surfaces]
+        tasks = [
+            loop.run_in_executor(executor, make_intersection, surf, fence_arr)
+            for surf in surfaces
+        ]
         intersections = await asyncio.gather(*tasks)
     return intersections
 
@@ -233,20 +144,27 @@ def load_surf(bytestr) -> xtgeo.RegularSurface:
 async def load_xtgeo(res_array):
     loop = asyncio.get_running_loop()
     with ProcessPoolExecutor() as executor:
-        tasks = [loop.run_in_executor(executor, load_surf, bytestr) for bytestr in res_array]
+        tasks = [
+            loop.run_in_executor(executor, load_surf, bytestr) for bytestr in res_array
+        ]
         surfaces = await asyncio.gather(*tasks)
     return surfaces
 
 
-def get_uuids(case_uuid, ensemble_name, realization_nums, snames, sattr, bearer_token):
+def get_surface_set_uuids(
+    case_uuid,
+    ensemble_name,
+    realization_surface_set_spec: schemas.RealizationsSurfaceSetSpec,
+    bearer_token,
+):
     sumo_client = SumoClient(env="prod", token=bearer_token, interactive=False)
     case_collection = CaseCollection(sumo_client).filter(uuid=case_uuid)
     case = case_collection[0]
     surface_collection = case.surfaces.filter(
         iteration=ensemble_name,
-        name=snames,
-        tagname=sattr,
-        realization=realization_nums,
+        name=realization_surface_set_spec.surface_names,
+        tagname=realization_surface_set_spec.surface_attribute,
+        realization=realization_surface_set_spec.realization_nums,
     )
     return [surf.uuid for surf in surface_collection]