Data uploader and job queueing

.github/workflows/deploy.yaml

@@ -8,6 +8,7 @@ on:
   push:
     branches:
       - main
+      - feature/upload-data
 
   # Allows you to run this workflow manually from the Actions tab
   workflow_dispatch:
@@ -34,7 +35,7 @@ jobs:
         env:
           ECR_REGISTRY: ${{ steps.login-ecr.outputs.registry }}
           ECR_REPOSITORY: continuum
-          IMAGE_TAG: latest-tsadar
+          IMAGE_TAG: tsadar-app-v0.0.1
         run: |
           docker build -f Dockerfile -t $ECR_REGISTRY/$ECR_REPOSITORY:$IMAGE_TAG .
 

mambaenv.yaml

@@ -2,11 +2,13 @@ name: base
 channels:
   - conda-forge
 dependencies:
-  - python=3.12
-  - git
+  - python=3.10
   - pip
-  - pyhdf
-  - streamlit
 # install pip requirements
   - pip:
-    - git+https://github.com/ergodicio/tsadar.git
+    - mlflow
+    - numpy
+    - streamlit
+    - boto3
+    - protobuf==3.20.*
+    # - git+https://github.com/ergodicio/tsadar.git

tsadar_app.py

@@ -1,13 +1,7 @@
-from jax import config
-
-config.update("jax_enable_x64", True)
-
 import streamlit as st
-import yaml, os, mlflow
-from flatten_dict import flatten, unflatten
+import yaml, os, mlflow, tempfile, boto3
 
 from tsadar_gui import config
-from tsadar import run_for_app
 
 
 def process_file(file_path):
@@ -26,39 +20,8 @@ def process_file(file_path):
     config_dir = os.path.join(os.getcwd(), "temp")
 
     os.makedirs(config_dir, exist_ok=True)
-    if mode == "file":
-        defaults = st.file_uploader("Upload the defaults file", type=["yaml", "yml"])
-        if defaults:
-            with open(os.path.join(config_dir, defaults.name), "wb") as f:
-                f.write(defaults.getvalue())
-
-        inputs = st.file_uploader("Upload the inputs file", type=["yaml", "yml"])
-        if inputs:
-            with open(os.path.join(config_dir, inputs.name), "wb") as f:
-                f.write(inputs.getvalue())
-
-        with open(os.path.join(config_dir, "defaults.yaml"), "r") as f:
-            defs = yaml.safe_load(f)
-
-        with open(os.path.join(config_dir, "inputs.yaml"), "r") as f:
-            inps = yaml.safe_load(f)
-
-        defaults = flatten(defs)
-        defaults.update(flatten(inps))
-        cfg = unflatten(defaults)
-
-        # epw_file = st.file_uploader("Upload the EPW data file", type=["hdf", "txt"])
-        # if epw_file:
-        #     epw_file_path = epw_file.name
-        #     with open(epw_file.name, "wb") as f:
-        #         f.write(epw_file.getvalue())
-
-        # iaw_file = st.file_uploader("Upload the IAW data file", type=["hdf", "txt"])
-        # if iaw_file:
-        #     iaw_file = iaw_file.name
-        #     with open(iaw_file.name, "wb") as f:
-        #         f.write(iaw_file.getvalue())
-    elif mode == "Home":
+
+    if mode == "Home":
         st.header("Welcome to TSADARapp")
         st.write(
             "This web app is a Streamlit implementation of the [Thomson Scattering Analysis](https://www.github.com/ergodicio/inverse-thomson-scattering/) software [1]."
@@ -68,7 +31,7 @@ def process_file(file_path):
         )
 
     elif mode == "Fit":
-        cfg = config.get_config()
+        cfg, files = config.get_config()
 
         c1, c2 = st.columns(2)
 
@@ -82,9 +45,43 @@ def process_file(file_path):
             if st.button("Fit"):
                 # run and wait for the results
                 mlflow.set_experiment(cfg["mlflow"]["experiment"])
-                run_id = run_for_app(cfg, mode="fit")
 
-                st.write(f"The results can be found at the mlflow run id {run_id}")
+                with tempfile.TemporaryDirectory() as tempdir:
+                    with mlflow.start_run(run_name=cfg["mlflow"]["run"]) as mlflow_run:
+                        # write config yaml to disk
+                        with open(os.path.join(tempdir, "config.yaml"), "w") as f:
+                            yaml.dump(cfg, f)
+
+                        # write uploaded streamlit file to disk
+                        for k, fl in files.items():
+                            with open(os.path.join(tempdir, fl.name), "wb") as f:
+                                f.write(fl.read())
+
+                        mlflow.log_artifacts(tempdir)
+                        run_id = mlflow_run.info.run_id
+
+                # requests.post(
+                #     "http://sciapi.continuum/queue-run",
+                #     data={"jq_nm": "gpu", "jd_nm": "gpu", "job_name": "thomson", "run_type": "dm", "run_id": run_id},
+                # )
+                client = boto3.client("batch", region_name="us-east-1")
+
+                job_template = {
+                    "jobQueue": "gpu",
+                    "jobDefinition": "tsadar-gpu",
+                    "jobName": "tsadar",
+                    "parameters": {"run_id": run_id},
+                    "retryStrategy": {
+                        "attempts": 10,
+                        "evaluateOnExit": [{"action": "RETRY", "onStatusReason": "Host EC2*"}],
+                    },
+                }
+                submissionResult = client.submit_job(**job_template)
+                st.write(
+                    f"The job is queued. The status and results can be found at the mlflow experiment {cfg['mlflow']['experiment']} and run id {run_id}. Email [email protected] if you have any trouble."
+                )
+
+                # run_for_app(run_id)
 
     st.sidebar.title("About")
     ## Add attribution

tsadar_gui/config.py

@@ -7,9 +7,9 @@ def get_config():
     )
     st.write("More details can be found in the documentation at < >")
 
-    config = create_default_config()
+    config, files = create_default_config()
     # inputs = get_inputs()
-    return config
+    return config, files
 
 
 def optimizer_section():
@@ -403,6 +403,15 @@ def create_default_config():
         c1, c2 = st.columns(2)
         with c1:
             shotnum = st.number_input("shotnum", value=101675)
+            # file upload
+            files = {}
+            epw_file = st.file_uploader("Upload the EPW file", type=["hdf"])
+            if epw_file:
+                files["epw"] = epw_file
+            iaw_file = st.file_uploader("Upload the IAW file", type=["hdf"])
+            if iaw_file:
+                files["iaw"] = iaw_file
+
             shotDay = st.checkbox("shotDay", value=False)
             st.divider()
             st.write("Which lineouts?")
@@ -447,6 +456,7 @@ def create_default_config():
         data = {
             "shotnum": shotnum,
             "shotDay": shotDay,
+            "filenames": {},
             "lineouts": {
                 "type": lineout_type,
                 "start": lineout_start,
@@ -577,11 +587,26 @@ def create_default_config():
 
     with st.expander("Optimizer"):
         optimizer = optimizer_section()
-    return {
+
+    config = {
         "parameters": parameters,
         "data": data,
         "other": other,
         "mlflow": mlflow,
         "optimizer": optimizer,
         "plotting": plotting,
     }
+
+    if epw_file is not None:
+        if str(config["data"]["shotnum"]) in epw_file.name:
+            config["data"]["filenames"]["epw"] = epw_file.name
+        else:
+            st.error("The EPW file name does not match the shot number")
+
+    if iaw_file is not None:
+        if str(config["data"]["shotnum"]) in iaw_file.name:
+            config["data"]["filenames"]["iaw"] = iaw_file.name
+        else:
+            st.error("The IAW file name does not match the shot number")
+
+    return config, files