Docker image for runner

.github/workflows/deploy.yaml

@@ -0,0 +1,46 @@
+# This is a basic workflow to help you get started with Actions
+
+name: deploy
+
+# Controls when the action will run.
+on:
+  # Triggers the workflow on push or pull request events but only for the master branch
+  push:
+    branches:
+      - main
+
+  # Allows you to run this workflow manually from the Actions tab
+  workflow_dispatch:
+
+# A workflow run is made up of one or more jobs that can run sequentially or in parallel
+jobs:
+  deploy-tsadar:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v2
+
+      - name: Configure AWS credentials
+        uses: aws-actions/configure-aws-credentials@v1
+        with:
+          aws-access-key-id: ${{ secrets.AWS_ACCESS_KEY_ID }}
+          aws-secret-access-key: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
+          aws-region: us-east-1
+
+      - name: Login to Amazon ECR
+        id: login-ecr
+        uses: aws-actions/amazon-ecr-login@v1
+
+      - name: Build and tag api request image to Amazon ECR
+        env:
+          ECR_REGISTRY: ${{ steps.login-ecr.outputs.registry }}
+          ECR_REPOSITORY: continuum
+          IMAGE_TAG: latest-tsadar-runner
+        run: |
+          docker build -f docker/Dockerfile -t $ECR_REGISTRY/$ECR_REPOSITORY:$IMAGE_TAG .
+
+      - name: Push all images to Amazon ECR
+        env:
+          ECR_REGISTRY: ${{ steps.login-ecr.outputs.registry }}
+          ECR_REPOSITORY: continuum
+        run: |
+          docker push --all-tags $ECR_REGISTRY/$ECR_REPOSITORY

docker/Dockerfile

@@ -0,0 +1,11 @@
+# base image
+FROM mambaorg/micromamba:latest
+
+WORKDIR /app
+
+COPY --chown=$MAMBA_USER:$MAMBA_USER tsadar /app/tsadar
+COPY --chown=$MAMBA_USER:$MAMBA_USER env_gpu.yml docker/fit_for_app.py docker/run.sh /app/
+RUN chmod +x run.sh
+
+RUN micromamba install -y -n base -f /app/env_gpu.yml && \
+    micromamba clean --all --yes

docker/fit_for_app.py

@@ -0,0 +1,13 @@
+import sys
+from jax import config
+import jax
+
+config.update("jax_enable_x64", True)
+
+from tsadar.runner import run_for_app
+
+if __name__ == "__main__":
+    # print jax devices
+    print("jax devices: ", jax.devices())
+
+    run_for_app(sys.argv[1])

docker/run.sh

@@ -0,0 +1,2 @@
+#!/bin/bash
+micromamba run -n base python3 fit_for_app.py $1

env_gpu.yml

@@ -1,30 +1,27 @@
-name: tsadar-gpu
 channels:
   - conda-forge
 dependencies:
-  - python=3.10
-  - ipython
-  - ipykernel
+  - python=3.12
+  - pyhdf
   - pip
+  - jax[cuda12]
+  - jaxopt
+  - numpy<2
+  - scipy
+  - matplotlib
+  - pyyaml
+  - mlflow
+  - boto3
+  - flatten-dict
+  - typing-extensions
+  - optax
+  - tqdm
+  - xarray
+  - pandas
+  - xlrd
   - pip:
-      # works for regular pip packages
-      - --find-links https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
-      - jaxlib==0.4.26+cuda12.cudnn89
-      - jax==0.4.26
-      - jaxopt
-      - numpy
-      - scipy
-      - matplotlib
-      - pyhdf
-      - xlrd
-      - pyyaml
-      - mlflow
-      - boto3
-      - flatten-dict
-      - typing-extensions
-      - optax
-      - tqdm
-      - xarray
-      - mlflow_export_import
-      - pandas
-      - interpax
+    - mlflow_export_import
+    - interpax
+    - psutil
+    - pynvml
+

tsadar/data_handleing/load_ts_data.py

@@ -9,7 +9,7 @@
 BASE_FILES_PATH = os.path.join(os.path.dirname(__file__), "..", "aux")
 
 
-def loadData(sNum, sDay, loadspecs):
+def loadData(sNum, sDay, loadspecs, custom_path=False):
     """
     This function loads the appropriate data based off the provided shot number (sNum) automatically determining the
     type of data in the file. The flag sDay changes the default path to the temporary archive on the redwood server and
@@ -35,6 +35,8 @@ def loadData(sNum, sDay, loadspecs):
     """
     if sDay:
         folder = r"\\redwood\archive\tmp\thomson"
+    elif custom_path:
+        folder = custom_path
     else:
         folder = join(BASE_FILES_PATH, "data")
 

tsadar/process/prepare.py

@@ -1,6 +1,8 @@
 from typing import Dict
 
 import numpy as np
+import os
+
 from tsadar.process.evaluate_background import get_shot_bg
 from tsadar.data_handleing.load_ts_data import loadData
 from tsadar.process.correct_throughput import correctThroughput
@@ -20,8 +22,16 @@ def prepare_data(config: Dict) -> Dict:
 
     """
     # load data
+    custom_path = None
+    if "filenames" in config["data"].keys():
+        if config["data"]["filenames"]["epw"] is not None:
+            custom_path = os.path.dirname(config["data"]["filenames"]["epw-local"])
+
+        if config["data"]["filenames"]["iaw"] is not None:
+            custom_path = os.path.dirname(config["data"]["filenames"]["iaw-local"])
+
     [elecData, ionData, xlab, config["other"]["extraoptions"]["spectype"]] = loadData(
-        config["data"]["shotnum"], config["data"]["shotDay"], config["other"]["extraoptions"]
+        config["data"]["shotnum"], config["data"]["shotDay"], config["other"]["extraoptions"], custom_path=custom_path
     )
 
     # get scattering angles and weights

tsadar/runner.py

@@ -81,9 +81,27 @@ def run(cfg_path: str, mode: str) -> str:
     return run_id
 
 
-def run_for_app(cfg: Dict, mode: str) -> str:
-    with mlflow.start_run(run_name=cfg["mlflow"]["run"], log_system_metrics=True) as mlflow_run:
-        _run_(cfg, mode=mode)
+def run_for_app(run_id: str) -> str:
+    with mlflow.start_run(run_id=run_id, log_system_metrics=True) as mlflow_run:
+        # download config
+        with tempfile.TemporaryDirectory(dir=BASE_TEMPDIR) as temp_path:
+
+            dest_file_path = utils.download_file(f"config.yaml", mlflow_run.info.artifact_uri, temp_path)
+            with open(dest_file_path, "r") as fi:
+                config = yaml.safe_load(fi)
+
+            if config["data"]["filenames"]["epw"] is not None:
+                config["data"]["filenames"]["epw-local"] = utils.download_file(
+                    config["data"]["filenames"]["epw"], mlflow_run.info.artifact_uri, temp_path
+                )
+
+            if config["data"]["filenames"]["iaw"] is not None:
+                config["data"]["filenames"]["iaw-local"] = utils.download_file(
+                    config["data"]["filenames"]["iaw"], mlflow_run.info.artifact_uri, temp_path
+                )
+
+            _run_(config, mode="fit")
+
     return mlflow_run.info.run_id