format experiment server

src/ert/experiment_server/main.py

@@ -1,53 +1,81 @@
 import asyncio
 import multiprocessing as mp
 import os
-from concurrent.futures import ProcessPoolExecutor
+import queue
 import uuid
+from concurrent.futures import ProcessPoolExecutor
 from multiprocessing.queues import Queue
-import queue
+from typing import Dict, Tuple, Union
 
 from fastapi import BackgroundTasks, FastAPI, HTTPException, WebSocket
+from fastapi.encoders import jsonable_encoder
 from pydantic import BaseModel, Field
 
+from ert.config import ErtConfig, QueueSystem
 from ert.ensemble_evaluator import EvaluatorServerConfig
-from ert.run_models.model_factory import create_model
-from ert.run_models.base_run_model import BaseRunModel, StatusEvents
-from ert.gui.simulation.ensemble_experiment_panel import Arguments as EnsembleExperimentArguments
-from ert.gui.simulation.ensemble_smoother_panel import Arguments as EnsembleSmootherArguments
-from ert.gui.simulation.evaluate_ensemble_panel import Arguments as EvaluateEnsembleArguments
-from ert.gui.simulation.iterated_ensemble_smoother_panel import Arguments as IteratedEnsembleSmootherArguments
+from ert.ensemble_evaluator.event import EndEvent, _UpdateEvent
+from ert.gui.simulation.ensemble_experiment_panel import (
+    Arguments as EnsembleExperimentArguments,
+)
+from ert.gui.simulation.ensemble_smoother_panel import (
+    Arguments as EnsembleSmootherArguments,
+)
+from ert.gui.simulation.evaluate_ensemble_panel import (
+    Arguments as EvaluateEnsembleArguments,
+)
+from ert.gui.simulation.iterated_ensemble_smoother_panel import (
+    Arguments as IteratedEnsembleSmootherArguments,
+)
 from ert.gui.simulation.manual_update_panel import Arguments as ManualUpdateArguments
-from ert.gui.simulation.multiple_data_assimilation_panel import Arguments as MultipleDataAssimilationArguments
+from ert.gui.simulation.multiple_data_assimilation_panel import (
+    Arguments as MultipleDataAssimilationArguments,
+)
 from ert.gui.simulation.single_test_run_panel import Arguments as SingleTestRunArguments
+from ert.run_models.base_run_model import BaseRunModel, StatusEvents
+from ert.run_models.model_factory import create_model
 from ert.storage import open_storage
-from ert.ensemble_evaluator.event import _UpdateEvent, EndEvent
 
-from typing import Dict, Union, Tuple
-
-from ert.config import ErtConfig, QueueSystem
-from fastapi.encoders import jsonable_encoder
 
 class Experiment(BaseModel):
-    args: Union[EnsembleExperimentArguments, EnsembleSmootherArguments, EvaluateEnsembleArguments, IteratedEnsembleSmootherArguments, ManualUpdateArguments, MultipleDataAssimilationArguments, SingleTestRunArguments] = Field(..., discriminator='mode')
+    args: Union[
+        EnsembleExperimentArguments,
+        EnsembleSmootherArguments,
+        EvaluateEnsembleArguments,
+        IteratedEnsembleSmootherArguments,
+        ManualUpdateArguments,
+        MultipleDataAssimilationArguments,
+        SingleTestRunArguments,
+    ] = Field(..., discriminator="mode")
     ert_config: ErtConfig
 
-mp_ctx = mp.get_context('fork')
-process_pool = ProcessPoolExecutor(max_workers=max((os.cpu_count() or 1) - 2, 1), mp_context=mp_ctx)
+
+mp_ctx = mp.get_context("fork")
+process_pool = ProcessPoolExecutor(
+    max_workers=max((os.cpu_count() or 1) - 2, 1), mp_context=mp_ctx
+)
 app = FastAPI()
+experiments: Dict[str, Tuple[BaseRunModel, "Queue[StatusEvents]"]] = {}
 
 
 @app.get("/")
 async def root():
     return {"message": "ping"}
 
-experiments : Dict[str, Tuple[BaseRunModel, "Queue[StatusEvents]"]]= {}
 
-async def run_experiment(experiment_id:str, evaluator_server_config: EvaluatorServerConfig):
+async def run_experiment(
+    experiment_id: str, evaluator_server_config: EvaluatorServerConfig
+):
     loop = asyncio.get_running_loop()
     print(f"Starting experiment {experiment_id}")
-    await loop.run_in_executor(None, lambda: experiments[experiment_id][0].start_simulations_thread(evaluator_server_config))
+    await loop.run_in_executor(
+        None,
+        lambda: experiments[experiment_id][0].start_simulations_thread(
+            evaluator_server_config
+        ),
+    )
     print(f"Experiment {experiment_id} done")
 
+
 @app.post("/experiments/")
 async def submit_experiment(experiment: Experiment, background_tasks: BackgroundTasks):
     storage = open_storage(experiment.ert_config.ens_path, "w")
@@ -60,7 +88,10 @@ async def submit_experiment(experiment: Experiment, background_tasks: Background
             status_queue,
         )
     except ValueError as e:
-        return HTTPException(status_code=420, detail=f"{experiment.args.mode} was not valid, failed with: {e}")
+        return HTTPException(
+            status_code=420,
+            detail=f"{experiment.args.mode} was not valid, failed with: {e}",
+        )
 
     port_range = None
     if model.queue_system == QueueSystem.LOCAL:
@@ -70,9 +101,12 @@ async def submit_experiment(experiment: Experiment, background_tasks: Background
     experiment_id = str(uuid.uuid4())
     experiments[experiment_id] = (model, status_queue)
 
-    background_tasks.add_task(run_experiment, experiment_id, evaluator_server_config=evaluator_server_config)
+    background_tasks.add_task(
+        run_experiment, experiment_id, evaluator_server_config=evaluator_server_config
+    )
     return {"message": "Experiment Started", "experiment_id": experiment_id}
 
+
 @app.put("/experiments/{experiment_id}/cancel")
 async def cancel_experiment(experiment_id: str):
     if experiment_id in experiments:
@@ -92,7 +126,7 @@ async def websocket_endpoint(websocket: WebSocket, experiment_id: str):
         try:
             item: StatusEvents = q.get(block=False)
         except queue.Empty:
-            asyncio.sleep(0.01)
+            await asyncio.sleep(0.01)
             continue
 
         if isinstance(item, _UpdateEvent):
@@ -104,4 +138,3 @@ async def websocket_endpoint(websocket: WebSocket, experiment_id: str):
         await asyncio.sleep(0.1)
         if isinstance(item, EndEvent):
             break
-