Feature/one step save baseline

Makefile

@@ -3,7 +3,7 @@
 #################################################################################
 # GLOBALS                                                                       #
 #################################################################################
-VERSION = v0.1.0
+VERSION = v0.1.0-a1
 ENVIRONMENT_SCRIPTS = .environment-scripts
 PROJECT_DIR := $(shell dirname $(realpath $(lastword $(MAKEFILE_LIST))))
 BUCKET = [OPTIONAL] your-bucket-for-syncing-data (do not include 's3://')
@@ -24,21 +24,26 @@ endif
 # COMMANDS                                                                      #
 #################################################################################
 .PHONY: wheels build_images push_image
-wheels:
-	pip wheel --no-deps .
-	pip wheel --no-deps external/vcm
+# wheels:
+# 	pip wheel --no-deps .
+# 	pip wheel --no-deps external/vcm
+# 	pip wheel --no-deps external/fv3config
 
 # pattern rule for building docker images
 build_image_%:
 	docker build -f docker/$*/Dockerfile . -t us.gcr.io/vcm-ml/$*:$(VERSION)
 
-build_image_prognostic_run: wheels
+enter_%:
+	docker run -ti -w /fv3net -v $(shell pwd):/fv3net us.gcr.io/vcm-ml/$*:$(VERSION) bash
+
+build_image_prognostic_run:
 
 build_images: build_image_fv3net build_image_prognostic_run
 
-push_image:
-	docker push us.gcr.io/vcm-ml/fv3net:$(VERSION)
-	docker push us.gcr.io/vcm-ml/prognostic_run:$(VERSION)
+push_images: push_image_prognostic_run push_image_fv3net
+
+push_image_%:
+	docker push us.gcr.io/vcm-ml/$*:$(VERSION)
 
 enter: build_image
 	docker run -it -v $(shell pwd):/code \

docker/prognostic_run/Dockerfile

@@ -1,8 +1,14 @@
-FROM us.gcr.io/vcm-ml/fv3gfs-python:v0.2.1
+FROM us.gcr.io/vcm-ml/fv3gfs-python:v0.3.1
 
 
 COPY docker/prognostic_run/requirements.txt /tmp/requirements.txt
 RUN pip3 install -r /tmp/requirements.txt
-COPY fv3net-0.1.0-py3-none-any.whl /wheels/fv3net-0.1.0-py3-none-any.whl
-COPY vcm-0.1.0-py3-none-any.whl /wheels/vcm-0.1.0-py3-none-any.whl
-RUN pip3 install --no-deps /wheels/fv3net-0.1.0-py3-none-any.whl  && pip3 install /wheels/vcm-0.1.0-py3-none-any.whl
+RUN pip3 install wheel
+
+# cache external package installation
+COPY external/fv3config /fv3net/external/fv3config
+COPY external/vcm /fv3net/external/vcm
+RUN pip3 install -e /fv3net/external/vcm -e /fv3net/external/fv3config
+
+COPY . /fv3net
+RUN pip3 install --no-deps -e /fv3net

docker/prognostic_run/requirements.txt

@@ -2,4 +2,5 @@ scikit-learn==0.22.1
 dask
 joblib
 zarr
-scikit-image
+scikit-image
+google-cloud-logging

fv3net/pipelines/kube_jobs/one_step.py

@@ -1,12 +1,13 @@
 import logging
 import os
 import fsspec
+from toolz import assoc
 import uuid
 import yaml
 import re
 from copy import deepcopy
 from multiprocessing import Pool
-from typing import List, Tuple
+from typing import List, Dict
 
 import fv3config
 from . import utils
@@ -197,20 +198,16 @@ def _update_config(
     workflow_name: str,
     base_config_version: str,
     user_model_config: dict,
-    user_kubernetes_config: dict,
     input_url: str,
     config_url: str,
     timestep: str,
-) -> Tuple[dict]:
+) -> Dict:
     """
     Update kubernetes and fv3 configurations with user inputs
     to prepare for fv3gfs one-step runs.
     """
     base_model_config = utils.get_base_fv3config(base_config_version)
     model_config = utils.update_nested_dict(base_model_config, user_model_config)
-    kubernetes_config = utils.update_nested_dict(
-        deepcopy(KUBERNETES_CONFIG_DEFAULT), user_kubernetes_config
-    )
 
     model_config = fv3config.enable_restart(model_config)
     model_config["experiment_name"] = _get_experiment_name(workflow_name, timestep)
@@ -225,27 +222,19 @@ def _update_config(
         }
     )
 
-    return model_config, kubernetes_config
+    return model_config
 
 
 def _upload_config_files(
     model_config: dict,
-    kubernetes_config: dict,
     config_url: str,
-    local_vertical_grid_file=None,
+    local_vertical_grid_file,
     upload_config_filename="fv3config.yml",
-) -> Tuple[dict]:
+) -> str:
     """
     Upload any files to remote paths necessary for fv3config and the
     fv3gfs one-step runs.
     """
-
-    if "runfile" in kubernetes_config:
-        runfile_path = kubernetes_config["runfile"]
-        kubernetes_config["runfile"] = utils.transfer_local_to_remote(
-            runfile_path, config_url
-        )
-
     model_config["diag_table"] = utils.transfer_local_to_remote(
         model_config["diag_table"], config_url
     )
@@ -259,38 +248,22 @@ def _upload_config_files(
     with fsspec.open(config_path, "w") as config_file:
         config_file.write(yaml.dump(model_config))
 
-    return model_config, kubernetes_config
+    return config_path
 
 
-def prepare_and_upload_config(
-    workflow_name: str,
-    input_url: str,
-    config_url: str,
-    timestep: str,
-    one_step_config: dict,
-    base_config_version: str,
-    **kwargs,
-) -> Tuple[dict]:
-    """Update model and kubernetes configurations for this particular
-    timestep and upload necessary files to GCS"""
-
-    user_model_config = one_step_config["fv3config"]
-    user_kubernetes_config = one_step_config["kubernetes"]
-
-    model_config, kube_config = _update_config(
-        workflow_name,
-        base_config_version,
-        user_model_config,
-        user_kubernetes_config,
-        input_url,
-        config_url,
-        timestep,
-    )
-    model_config, kube_config = _upload_config_files(
-        model_config, kube_config, config_url, **kwargs
+def get_run_kubernetes_kwargs(user_kubernetes_config, config_url):
+
+    kubernetes_config = utils.update_nested_dict(
+        deepcopy(KUBERNETES_CONFIG_DEFAULT), user_kubernetes_config
     )
 
-    return model_config, kube_config
+    if "runfile" in kubernetes_config:
+        runfile_path = kubernetes_config["runfile"]
+        kubernetes_config["runfile"] = utils.transfer_local_to_remote(
+            runfile_path, config_url
+        )
+
+    return kubernetes_config
 
 
 def submit_jobs(
@@ -305,28 +278,51 @@ def submit_jobs(
     local_vertical_grid_file=None,
 ) -> None:
     """Submit one-step job for all timesteps in timestep_list"""
-    for timestep in timestep_list:
 
+    zarr_url = os.path.join(output_url, "big.zarr")
+    # kube kwargs are shared by all jobs
+    kube_kwargs = get_run_kubernetes_kwargs(one_step_config["kubernetes"], config_url)
+
+    def config_factory(**kwargs):
+        timestep = timestep_list[kwargs["index"]]
         curr_input_url = os.path.join(input_url, timestep)
-        curr_output_url = os.path.join(output_url, timestep)
         curr_config_url = os.path.join(config_url, timestep)
 
-        model_config, kube_config = prepare_and_upload_config(
+        config = deepcopy(one_step_config)
+        kwargs["url"] = zarr_url
+        config["fv3config"]["one_step"] = kwargs
+
+        model_config = _update_config(
             workflow_name,
+            base_config_version,
+            config["fv3config"],
             curr_input_url,
             curr_config_url,
             timestep,
-            one_step_config,
-            base_config_version,
-            local_vertical_grid_file=local_vertical_grid_file,
         )
+        return _upload_config_files(
+            model_config, curr_config_url, local_vertical_grid_file
+        )
+
+    def run_job(wait=False, **kwargs):
+        """Run a run_kubernetes job
 
-        jobname = model_config["experiment_name"]
-        kube_config["jobname"] = jobname
+        kwargs are passed workflows/one_step_jobs/runfile.py:post_process
+
+        """
+        uid = str(uuid.uuid4())
+        labels = assoc(job_labels, "jobid", uid)
+        model_config_url = config_factory(**kwargs)
         fv3config.run_kubernetes(
-            os.path.join(curr_config_url, "fv3config.yml"),
-            curr_output_url,
-            job_labels=job_labels,
-            **kube_config,
+            model_config_url, "/tmp/null", job_labels=labels, **kube_kwargs
         )
-        logger.info(f"Submitted job {jobname}")
+        if wait:
+            utils.wait_for_complete(job_labels, sleep_interval=10)
+
+    for k, timestep in enumerate(timestep_list):
+        if k == 0:
+            logger.info("Running the first time step to initialize the zarr store")
+            run_job(index=k, init=True, wait=True, timesteps=timestep_list)
+        else:
+            logger.info(f"Submitting job for timestep {timestep}")
+            run_job(index=k, init=False)

fv3net/runtime/__init__.py

@@ -1,3 +1,3 @@
 from . import sklearn_interface as sklearn
 from .state_io import init_writers, append_to_writers, CF_TO_RESTART_MAP
-from .config import get_runfile_config, get_namelist
+from .config import get_namelist, get_config

fv3net/runtime/config.py

@@ -10,10 +10,10 @@ class dotdict(dict):
     __delattr__ = dict.__delitem__
 
 
-def get_runfile_config():
+def get_config():
     with open("fv3config.yml") as f:
         config = yaml.safe_load(f)
-    return dotdict(config["scikit_learn"])
+    return config
 
 
 def get_namelist():

workflows/one_step_jobs/_run_steps.sh

@@ -0,0 +1,14 @@
+workdir=$(pwd)
+ src=gs://vcm-ml-data/orchestration-testing/test-andrep/coarsen_restarts_source-resolution_384_target-resolution_48/
+ output=gs://vcm-ml-data/testing-noah/one-step
+ image=us.gcr.io/vcm-ml/prognostic_run:v0.1.0-a1
+ yaml=$PWD/deep-conv-off.yml
+
+ gsutil -m rm -r $output > /dev/null
+ (
+    cd ../../
+    python $workdir/orchestrate_submit_jobs.py \
+        $src $output $yaml $image -o  \
+	--config-version v0.3
+ )
+

workflows/one_step_jobs/deep-conv-off.yml

@@ -1,5 +1,6 @@
 kubernetes:
   docker_image: us.gcr.io/vcm-ml/fv3gfs-python:v0.2.1
+  runfile: workflows/one_step_jobs/runfile.py
 fv3config:
   diag_table: workflows/one_step_jobs/diag_table_one_step
   namelist:

workflows/one_step_jobs/orchestrate_submit_jobs.py

@@ -83,7 +83,10 @@ def _create_arg_parser():
         one_step_config = yaml.load(file, Loader=yaml.FullLoader)
     workflow_name = Path(args.one_step_yaml).with_suffix("").name
     short_id = get_alphanumeric_unique_tag(8)
-    job_label = {"orchestrator-jobs": f"{workflow_name}-{short_id}"}
+    job_label = {
+        "orchestrator-jobs": f"{workflow_name}-{short_id}",
+        "workflow": "one_step_jobs",
+    }
 
     if not args.config_url:
         config_url = os.path.join(args.output_url, CONFIG_DIRECTORY_NAME)