Merge pull request #62 from demml/fix-streaming

Fix streaming

Former-commit-id: d978bc8cd1296ace11aa73012b9f8fbee74a1dbd [formerly 9163a46]
Former-commit-id: dc841b8

opsml/app/routes/auth.py

@@ -52,7 +52,7 @@ class UserDeleted(BaseModel):
 router = APIRouter()
 
 
-async def get_current_user(
+def get_current_user(
     request: Request,
     token: Annotated[str, Depends(oauth2_scheme)],
 ) -> User:
@@ -98,7 +98,7 @@ async def get_current_user(
     return user
 
 
-async def get_current_active_user(
+def get_current_active_user(
     current_user: Annotated[User, Depends(get_current_user)],
 ) -> User:
     if not current_user.is_active:
@@ -107,7 +107,7 @@ async def get_current_active_user(
 
 
 @router.post("/auth/token")
-async def login_for_access_token(
+def login_for_access_token(
     request: Request,
     response: Response,
     form_data: Annotated[OAuth2PasswordRequestForm, Depends()],
@@ -149,7 +149,7 @@ async def login_for_access_token(
 
 
 @router.get("/auth/token/rotate")
-async def create_refresh_token(
+def create_refresh_token(
     request: Request,
     response: Response,
     refresh_token: Annotated[Union[str, None], Cookie()] = None,
@@ -175,7 +175,7 @@ async def create_refresh_token(
     try:
         # check user
         auth_db: ServerAuthRegistry = request.app.state.auth_db
-        user = await get_current_user(request, refresh_token)
+        user = get_current_user(request, refresh_token)
 
         # create new access token
         refresh_token = auth_db.create_access_token(user, minutes=60)
@@ -193,7 +193,7 @@ async def create_refresh_token(
 
 
 @router.get("/auth/token/refresh")
-async def get_refresh_from_cookie(
+def get_refresh_from_cookie(
     request: Request,
     response: Response,
     refresh_token: Annotated[Union[str, None], Cookie()] = None,
@@ -207,7 +207,7 @@ async def get_refresh_from_cookie(
             headers={"WWW-Authenticate": "Bearer"},
         )
 
-    user = await get_current_user(request, refresh_token)
+    user = get_current_user(request, refresh_token)
     logger.info("Refreshing token for user: {}", user.username)
 
     # create new access token

opsml/app/routes/data.py

@@ -30,17 +30,17 @@
 
 
 @router.get("/data/download", name="download_data")
-async def download_data(request: Request, uid: str) -> StreamingResponse:
+def download_data(request: Request, uid: str) -> StreamingResponse:
     """Downloads data associated with a datacard"""
 
     registry: CardRegistry = request.app.state.registries.data
     datacard = cast(DataCard, registry.load_card(uid=uid))
     load_path = Path(datacard.uri / SaveName.DATA.value).with_suffix(datacard.interface.data_suffix)
-    return await download_artifacts_ui(request, str(load_path))
+    return download_artifacts_ui(request, str(load_path))
 
 
 @router.get("/data/download/profile", name="download_data_profile")
-async def download_data_profile(
+def download_data_profile(
     request: Request,
     uid: str,
 ) -> StreamingResponse:
@@ -49,7 +49,7 @@ async def download_data_profile(
     registry: CardRegistry = request.app.state.registries.data
     datacard = cast(DataCard, registry.load_card(uid=uid))
     load_path = Path(datacard.uri / SaveName.DATA_PROFILE.value).with_suffix(Suffix.HTML.value)
-    return await download_file(request, str(load_path))
+    return download_file(request, str(load_path))
 
 
 @router.post("/data/card", name="data_card", response_model=DataCardMetadata)

opsml/app/routes/files.py

@@ -13,7 +13,6 @@
 import streaming_form_data
 from fastapi import APIRouter, Depends, HTTPException, Request, status
 from fastapi.responses import StreamingResponse
-from starlette.concurrency import run_in_threadpool
 from starlette.requests import ClientDisconnect
 from streaming_form_data import StreamingFormDataParser
 from streaming_form_data.validators import MaxSizeValidator
@@ -114,7 +113,7 @@ async def upload_file(request: Request) -> Dict[str, str]:  # pragma: no cover
 
 
 @router.get("/files/download", name="download_file")
-async def download_file(request: Request, path: str) -> StreamingResponse:
+def download_file(request: Request, path: str) -> StreamingResponse:
     """Downloads a file
 
     Args:
@@ -130,8 +129,11 @@ async def download_file(request: Request, path: str) -> StreamingResponse:
     storage_client: StorageClientBase = request.app.state.storage_client
     try:
         file_path = Path(swap_opsml_root(request, Path(path)))
-        file_iterator = await run_in_threadpool(storage_client.iterfile, file_path, config.download_chunk_size)
-        return StreamingResponse(file_iterator, media_type="application/octet-stream")
+
+        return StreamingResponse(
+            storage_client.iterfile(file_path, config.download_chunk_size),
+            media_type="application/octet-stream",
+        )
 
     except Exception as error:
         logger.error("Server: Error downloading file {}", path)
@@ -141,7 +143,7 @@ async def download_file(request: Request, path: str) -> StreamingResponse:
         ) from error
 
 
-async def download_dir(request: Request, path: Path) -> StreamingResponse:
+def download_dir(request: Request, path: Path) -> StreamingResponse:
     """Downloads a file
 
     Args:
@@ -170,14 +172,16 @@ async def download_dir(request: Request, path: Path) -> StreamingResponse:
                     curr_rpath = Path(file_)
                     curr_lpath = lpath / curr_rpath.relative_to(rpath)
                     logger.info("Server: Downloading {} to {}", curr_rpath, curr_lpath)
-                    await run_in_threadpool(storage_client.get, curr_rpath, curr_lpath)
+                    storage_client.get(curr_rpath, curr_lpath)
                     zip_filepath = zipfile / curr_rpath.relative_to(rpath)
                     temp_zip.write(curr_lpath, zip_filepath)
 
             logger.info("Server: Sending zip file for {}", path)
-            iter_buffer = await run_in_threadpool(storage_client.iterbuffer, zip_io, config.download_chunk_size)
 
-            return StreamingResponse(iter_buffer, media_type="application/x-zip-compressed")
+            return StreamingResponse(
+                storage_client.iterbuffer(zip_io, config.download_chunk_size),
+                media_type="application/x-zip-compressed",
+            )
 
     except Exception as error:
         raise HTTPException(
@@ -187,7 +191,7 @@ async def download_dir(request: Request, path: Path) -> StreamingResponse:
 
 
 @router.get("/files/download/ui", name="download_artifacts")
-async def download_artifacts_ui(request: Request, path: str) -> StreamingResponse:
+def download_artifacts_ui(request: Request, path: str) -> StreamingResponse:
     """Downloads a file
 
     Args:
@@ -200,8 +204,8 @@ async def download_artifacts_ui(request: Request, path: str) -> StreamingRespons
         Streaming file response
     """
     if Path(path).suffix == "":
-        return await download_dir(request, Path(path))
-    return await download_file(request, path)
+        return download_dir(request, Path(path))
+    return download_file(request, path)
 
 
 @router.get("/files/list", name="list_files")
@@ -395,7 +399,7 @@ def get_file_to_view(request: Request, path: str) -> FileViewResponse:
 
 
 @router.post("/files/readme", name="create_readme")
-async def create_readme(
+def create_readme(
     request: Request,
     payload: ReadMeRequest,
 ) -> bool:

opsml/app/routes/healthcheck.py

@@ -11,7 +11,7 @@
 
 
 @router.get("/healthcheck", response_model=HealthCheckResult, name="healthcheck")
-def get_healthcheck() -> HealthCheckResult:
+async def get_healthcheck() -> HealthCheckResult:
     return HealthCheckResult(is_alive=True)
 
 
@@ -21,12 +21,13 @@ async def debug() -> DebugResponse:
         url=config.opsml_tracking_uri,
         storage=config.opsml_storage_uri,
         app_env=config.app_env,
+        app_version=config.app_version,
     )
 
 
 @router.get(
     "/error",
     description="An endpoint that will return a 500 error for debugging and alert testing",
 )
-def get_error() -> None:
+async def get_error() -> None:
     raise HTTPException(status_code=500)

opsml/app/routes/pydantic_models.py

@@ -43,6 +43,7 @@ class DebugResponse(BaseModel):
     url: str
     storage: str
     app_env: str
+    app_version: str
 
 
 class StorageSettingsResponse(BaseModel):

opsml/app/routes/runs.py

@@ -43,7 +43,7 @@ def get_runcard(request: Request, payload: CardRequest) -> RunCard:
 
 
 @router.get("/runs/graphs", name="graphs")
-async def get_graph_plots(request: Request, repository: str, name: str, version: str) -> Dict[str, Any]:
+def get_graph_plots(request: Request, repository: str, name: str, version: str) -> Dict[str, Any]:
     """Method for loading plots for a run
 
     Args:

opsml/settings/config.py

@@ -14,6 +14,7 @@
 from pydantic_settings import BaseSettings
 
 from opsml.types import StorageSystem
+from opsml.version import __version__
 
 
 class OpsmlAuthSettings(BaseModel):
@@ -23,6 +24,7 @@ class OpsmlAuthSettings(BaseModel):
 class OpsmlConfig(BaseSettings):
     app_name: str = "opsml"
     app_env: str = "development"
+    app_version: str = __version__
 
     opsml_storage_uri: str = "./opsml_registries"
     opsml_tracking_uri: str = "sqlite:///opsml.db"