add pyflexplot to automatic workflow

MeteoSwiss · Nov 22, 2024 · 417cf81 · 417cf81
1 parent 98e7dd5
commit 417cf81
Show file tree

Hide file tree

Showing 2 changed files with 51 additions and 9 deletions.
diff --git a/flex_container_orchestrator/config/.env b/flex_container_orchestrator/config/.env
@@ -36,7 +36,7 @@ S3_ACCESS_KEY=<secret data>
 S3_SECRET_KEY=<secret data>
 
 ## IMAGE
-FLEXPART_TAG=2411.b41e06a177edb101e546f6b180b342e7e32bff8f
+FLEXPART_TAG=2411.abb884c4add463ea64544488ed30acec79be41fe
 FLEXPART_ECR_REPO=493666016161.dkr.ecr.eu-central-2.amazonaws.com/numericalweatherpredictions/dispersionmodelling/flexpart-ifs/flexpart-containerize
 
 ## Time settings
@@ -46,3 +46,23 @@ TDELTA=6
 TFREQ_F=6
 # Frequency of IFS runs in hour(s)
 TFREQ=6
+
+#----------- Pyflexplot
+
+## Input/output S3 bucket settings
+MAIN__AWS__S3__INPUT__NAME=flexpart-output
+MAIN__AWS__S3__OUTPUT__NAME=pyflexplot-output
+MAIN__AWS__S3__INPUT__ENDPOINT_URL=https://object-store.os-api.cci1.ecmwf.int
+MAIN__AWS__S3__OUTPUT__ENDPOINT_URL=https://object-store.os-api.cci1.ecmwf.int
+## EWC buckets access and secret keys
+MAIN__AWS__S3__INPUT__S3_ACCESS_KEY=<secret data>
+MAIN__AWS__S3__INPUT__S3_SECRET_KEY=<secret data>
+MAIN__AWS__S3__OUTPUT__S3_ACCESS_KEY=<secret data>
+MAIN__AWS__S3__OUTPUT__S3_SECRET_KEY=<secret data>
+
+## Specific NWP model
+preset=opr/ifs-hres-eu/all_pdf
+
+## IMAGE
+PYFLEXPLOT_TAG=2411.35a30e3ba6ce01d8537b8497fff8ebc042bcc39c
+PYFLEXPLOT_ECR_REPO=493666016161.dkr.ecr.eu-central-2.amazonaws.com/numericalweatherpredictions/dispersionmodelling/flexpart-ifs/pyflexplot
diff --git a/flex_container_orchestrator/services/flexpart_service.py b/flex_container_orchestrator/services/flexpart_service.py
@@ -109,25 +109,47 @@ def launch_containers(date: str, location: str, time: str, step: str) -> None:
 
     logger.info("Aggregator launch script executed successfully.")
 
-    # ====== Third part: Run Flexpart ======
+    # ====== Third part: Run Flexpart and Pyflexplot ======
     try:
+        flexpart_image = f"{os.getenv('FLEXPART_ECR_REPO')}:{os.getenv('FLEXPART_TAG')}"
+        pyflexplot_image = f"{os.getenv('PYFLEXPLOT_ECR_REPO')}:{os.getenv('PYFLEXPLOT_TAG')}"
+
         # Loop through each configuration and execute Flexpart
-        docker_image = f"{os.getenv('FLEXPART_ECR_REPO')}:{os.getenv('FLEXPART_TAG')}"
         for config in configurations:
             env_vars = [["-e", f"{key.strip()}={value}"] for key, value in config.items()]
 
-            # Build the Docker command as a list
-            docker_command = [
+            # Docker command for Flexpart as a list
+            docker_command_flexpart = [
                 "docker", "run",
                 "--env-file", env_file_path,
                 *[item for sublist in env_vars for item in sublist],
                 "--rm",
-                docker_image,
+                flexpart_image,
+            ]
+
+            logger.info("Running: %s", " ".join(docker_command_flexpart))
+            run_command(docker_command_flexpart)
+
+            s3_dest_bucket = os.environ.get("MAIN__AWS__S3__OUTPUT__NAME", "pyflexplot-output")
+            s3_input_bucket = os.environ.get("MAIN__AWS__S3__INPUT__NAME", "flexpart-output")
+
+            # Docker command for Pyflexplot as a list
+            docker_command_pyflexplot = [
+                "docker", "run",
+                "--env-file", env_file_path,
+                "--rm",
+                pyflexplot_image,
+                "--preset", os.environ.get("preset", ""),
+                "--merge-pdfs",
+                f"--dest=s3://{s3_dest_bucket}",
+                "--setup", "infile",
+                f"s3://{s3_input_bucket}/{config['IBDATE']}{config['IBTIME']}/sandbox/grid_conc_{config['IBDATE']}{config['IBTIME']}0000.nc",
+                "--setup", "base_time", f"{config['IBDATE']}{config['IBTIME']}"
             ]
 
-            logger.info("Running: %s", " ".join(docker_command))
-            run_command(docker_command)
+            logger.info("Running: %s", " ".join(docker_command_pyflexplot))
+            run_command(docker_command_pyflexplot)
 
     except subprocess.CalledProcessError:
-        logger.error("Launch Flexpart script encountered an error.")
+        logger.error("Launch Flexpart and Pyflexplot encountered an error.")
         sys.exit(1)