import subprocess
import os
from .. import filesystem
from ..config._serialization import load
from ._native import CONFIG_OUT_FILENAME, run_native
DOCKER_OUTDIR = "/outdir"
DOCKER_CONFIG_LOCATION = os.path.join("/", CONFIG_OUT_FILENAME)
DOCKER_RUNFILE = "/runfile.py"
DOCKER_COMMAND = ["docker", "run"]
DOCKER_KEYFILE = "/gcs_key.json"
FV3RUN_MODULE = "fv3config.fv3run"
[docs]def run_docker(
config_dict_or_location,
outdir,
docker_image,
runfile=None,
keyfile=None,
capture_output=True,
):
"""Run the FV3GFS model in a docker container with the given configuration.
Copies the resulting directory to a target location. Will use the Google cloud
storage key at ``$GOOGLE_APPLICATION_CREDENTIALS`` by default. Requires the
fv3gfs-python package and fv3config to be installed in the docker image.
Args:
config_dict_or_location (dict or str): a configuration dictionary, or a
location (local or on Google cloud storage) of a yaml file containing
a configuration dictionary
outdir (str): location to copy the resulting run directory
runfile (str, optional): Python model script to use in place of the default.
docker_image (str, optional): If given, run this command inside a container
using this docker image. Image must have this package and fv3gfs-python
installed.
keyfile (str, optional): location of a Google cloud storage key to use
inside the docker container
capture_output (bool, optional): If true, then the stderr and stdout
streams will be redirected to the files `outdir/stderr.log` and `outdir/stdout.log`
respectively.
"""
if keyfile is None:
keyfile = os.environ.get("GOOGLE_APPLICATION_CREDENTIALS", None)
if isinstance(config_dict_or_location, str):
config_dict = _load_yaml(config_dict_or_location)
else:
config_dict = config_dict_or_location
filesystem.get_fs(outdir).makedirs(outdir, exist_ok=True)
bind_mount_args = []
docker_args = []
_get_credentials_args(keyfile, docker_args, bind_mount_args)
_get_local_data_bind_mounts(config_dict, bind_mount_args)
_get_docker_args(docker_args, bind_mount_args, outdir)
runfile_in_docker = _get_runfile_args(runfile, bind_mount_args)
if filesystem.is_local_path(outdir):
docker_outdir = DOCKER_OUTDIR
else:
docker_outdir = outdir
python_command = run_native.command(
config_dict,
docker_outdir,
runfile=runfile_in_docker,
capture_output=capture_output,
)
subprocess.check_call(
DOCKER_COMMAND + bind_mount_args + docker_args + [docker_image] + python_command
)
def _load_yaml(url):
with filesystem.open(url) as f:
return load(f)
def _get_runfile_args(runfile, bind_mount_args) -> str:
if runfile is not None:
if filesystem.is_local_path(runfile):
bind_mount_args += ["-v", f"{os.path.abspath(runfile)}:{DOCKER_RUNFILE}"]
return DOCKER_RUNFILE
else:
return runfile
def _get_paths(config_dict):
"""Return a list of all paths referenced by the config dict."""
return_list = [
config_dict["diag_table"],
config_dict["data_table"],
config_dict["forcing"],
config_dict["initial_conditions"],
]
patch_files = config_dict.get("patch_files", [])
if isinstance(patch_files, list):
return_list.extend(patch_files)
else:
return_list.append(patch_files)
return return_list
def _is_local_path(maybe_path_or_object):
return (
isinstance(maybe_path_or_object, str)
and os.path.isabs(maybe_path_or_object)
and filesystem.is_local_path(maybe_path_or_object)
)
def _get_local_data_paths(config_dict):
"""Return a list of all local paths referenced by the config dict."""
local_paths = []
for potential_path in _get_paths(config_dict):
if _is_local_path(potential_path):
local_paths.append(potential_path)
elif isinstance(potential_path, list):
for asset in potential_path:
if filesystem.is_local_path(asset["source_location"]):
local_paths.append(
os.path.join(asset["source_location"], asset["source_name"])
)
elif isinstance(potential_path, dict):
asset = potential_path
if filesystem.is_local_path(asset["source_location"]):
local_paths.append(
os.path.join(asset["source_location"], asset["source_name"])
)
return local_paths
def _get_local_data_bind_mounts(config_dict, bind_mount_args):
for local_path in _get_local_data_paths(config_dict):
bind_mount_args += ["-v", f"{local_path}:{local_path}"]
def _get_docker_args(docker_args, bind_mount_args, outdir):
if filesystem.is_local_path(outdir):
bind_mount_args += ["-v", f"{os.path.abspath(outdir)}:{DOCKER_OUTDIR}"]
docker_args += ["--rm", "--user", f"{os.getuid()}:{os.getgid()}"]
def _get_credentials_args(keyfile, docker_args, bind_mount_args):
if keyfile is not None:
bind_mount_args += ["-v", f"{os.path.abspath(keyfile)}:{DOCKER_KEYFILE}"]
docker_args += ["-e", f"GOOGLE_APPLICATION_CREDENTIALS={DOCKER_KEYFILE}"]