Feature/one step save baseline

.dockerignore

@@ -33,4 +33,5 @@ tox.ini
 vcm-ml-data
 dataflow/*/env
 external/vcm/venv
-Dockerfile
+Dockerfile
+outdir/

Makefile

@@ -3,8 +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://')

docker/prognostic_run/requirements.txt

@@ -4,4 +4,5 @@ joblib
 zarr
 scikit-image
 google-cloud-logging
+gitpython
 backoff

environment.yml

@@ -49,5 +49,6 @@ dependencies:
   - pip:
     - gsutil
     - nc-time-axis>=1.2.0
+    - gitpython
     - bump2version>=0.5.11
     - yq

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

@@ -1,3 +1,4 @@
+from typing import Dict
 import yaml
 import f90nml
 
@@ -10,11 +11,11 @@ class dotdict(dict):
     __delattr__ = dict.__delitem__
 
 
-def get_runfile_config():
+def get_config() -> Dict:
     with open("fv3config.yml") as f:
         config = yaml.safe_load(f)
-    return dotdict(config["scikit_learn"])
+    return config
 
 
-def get_namelist():
+def get_namelist() -> f90nml.Namelist:
     return f90nml.read("input.nml")

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 $yaml   $image $output -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)