diff --git a/CHANGES.md b/CHANGES.md index d77b339..0684e45 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -9,7 +9,8 @@ * Tweak CWL format (#24) * Use micromamba entry point in Docker image (#26) * Allow setting of CWL workflow ID (#29) -* Add in-notebook configuration interface (#30) +* Support in-notebook configuration of workflow ID, environment file, + and container image tag (#30, #33) * Support writing of stage-out STAC by notebook (#32) ## Changes in 0.1.0 diff --git a/test/data/paramtest.ipynb b/test/data/paramtest.ipynb index f5828b3..beeffc8 100644 --- a/test/data/paramtest.ipynb +++ b/test/data/paramtest.ipynb @@ -152,7 +152,11 @@ "source": [ "parameter_1 = my_constant * 2\n", "parameter_2 = \"default value\"\n", - "xcengine_config = dict(workflow_id=\"my-workflow\")" + "xcengine_config = dict(\n", + " workflow_id=\"my-workflow\",\n", + " environment_file=\"my-environment.yml\",\n", + " container_image_tag=\"my-tag\",\n", + ")" ] } ], diff --git a/test/test_core.py b/test/test_core.py index e2fd296..4bbc731 100644 --- a/test/test_core.py +++ b/test/test_core.py @@ -42,7 +42,9 @@ def test_image_builder_init(init_mock, tmp_path, tag): assert abs( datetime.datetime.now(datetime.UTC) - pytz.utc.localize( - datetime.datetime.strptime(ib.tag, ImageBuilder.tag_format) + datetime.datetime.strptime( + ib.tag, nb_path.stem + ":" + ImageBuilder.tag_format + ) ) ) < datetime.timedelta(seconds=10) else: @@ -198,3 +200,19 @@ def test_script_creator_cwl(tmp_path, nb_name): assert workflow["id"] == ( nb_path.stem if nb_name == "noparamtest" else "my-workflow" ) + + +def test_script_creator_notebook_config(): + nb_path = pathlib.Path(__file__).parent / "data" / "paramtest.ipynb" + script_creator = ScriptCreator(nb_path) + config = script_creator.nb_params.config + assert config["environment_file"] == "my-environment.yml" + assert config["container_image_tag"] == "my-tag" + + +def test_image_builder_notebook_config(tmp_path): + nb_path = pathlib.Path(__file__).parent / "data" / "paramtest.ipynb" + image_builder = ImageBuilder(nb_path, None, tmp_path, None) + config = image_builder.script_creator.nb_params.config + assert config["environment_file"] == "my-environment.yml" + assert config["container_image_tag"] == "my-tag" diff --git a/xcengine/core.py b/xcengine/core.py index fab6773..edca016 100755 --- a/xcengine/core.py +++ b/xcengine/core.py @@ -16,7 +16,7 @@ import uuid import datetime from collections.abc import Mapping, Generator, Iterable -from typing import Any +from typing import Any, ClassVar import docker from docker.errors import BuildError @@ -178,26 +178,49 @@ class ImageBuilder: runs a container initialized from that image. """ - tag_format = "xcengine:%Y.%m.%d.%H.%M.%S" + tag_format: ClassVar[str] = "%Y.%m.%d.%H.%M.%S" def __init__( self, notebook: pathlib.Path, - environment: pathlib.Path, + environment: pathlib.Path | None, build_dir: pathlib.Path, tag: str | None, ): self.notebook = notebook - self.environment = environment self.build_dir = build_dir + self.script_creator = ScriptCreator(self.notebook) + nb_config = self.script_creator.nb_params.config if tag is None: - self.tag = datetime.datetime.now(datetime.UTC).strftime( - self.tag_format - ) - LOGGER.info(f"No tag specified; using {self.tag}") + LOGGER.info("No tag specified; looking for one in the notebook.") + nb_tag = nb_config.get("container_image_tag", None) + if nb_tag is not None: + LOGGER.info(f"Using tag {nb_tag} from notebook.") + self.tag = nb_tag + else: + timestamp = datetime.datetime.now(datetime.UTC).strftime( + self.tag_format + ) + self.tag = f"{notebook.stem}:{timestamp}" + LOGGER.info(f"No tag in notebook; using {self.tag}") else: self.tag = tag - self.script_creator = ScriptCreator(self.notebook) + + if environment is None: + LOGGER.info( + "No environment file specified; " + "looking for one in the notebook." + ) + nb_env = nb_config.get("environment_file", None) + if nb_env is not None: + LOGGER.info( + f"Using environment file {nb_tag} as specified in notebook." + ) + self.environment = notebook.parent / nb_env + else: + LOGGER.info(f"No environment specified in notebook.") + else: + self.environment = environment def build(self) -> Image: self.script_creator.convert_notebook_to_script(self.build_dir) diff --git a/xcengine/parameters.py b/xcengine/parameters.py index 4bcd4dc..936df7f 100644 --- a/xcengine/parameters.py +++ b/xcengine/parameters.py @@ -7,7 +7,7 @@ import pystac import xarray as xr import yaml -from typing_extensions import ClassVar +from typing import ClassVar LOGGER = logging.getLogger(__name__) logging.basicConfig(level=logging.INFO)