diff --git a/CHANGES.md b/CHANGES.md index b3923d5..65d129b 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -6,7 +6,7 @@ `s:softwareVersion` (#90) * Allow creation of EOAP-only and xcube-server-only images, omitting unnecessary dependencies (#56) -* Add a utility function to read annotations from notebook code (#91) +* Add support for notebook-assisted stage-in (#23, #91, #95) ## Changes in 0.1.2 diff --git a/docs/notebook.md b/docs/notebook.md index a3a73bf..3e44682 100644 --- a/docs/notebook.md +++ b/docs/notebook.md @@ -35,7 +35,7 @@ name `xcengine_config`. Available configuration settings are: - `workflow_id`: a string identifier for the workflow in your Application Package. The runner or Application Package platform can use this - identifier to refer to you Application Package. By default, the name + identifier to refer to your Application Package. By default, the name of the notebook (without the `.ipynb` suffix) is used. - `environment_file`: the name of a YAML file defining a [conda environment](https://docs.conda.io/projects/conda/en/latest/user-guide/tasks/manage-environments.html) @@ -51,17 +51,45 @@ name `xcengine_config`. Available configuration settings are: Some of these configuration settings can also be set on the command line. +## Dataset input + +As well as the usual methods of dataset input, xcengine provides support for +the Application Package ‘stage-in’ process described in the [OGC Best +Practice document](https://docs.ogc.org/bp/20-089r1.html), in which an +Application Package Platform provides the Application Package with a +[STAC catalogue](https://stacspec.org/) of one or more input datasets. + +xcengine currently provides basic support for stage-in: in a generated +Application Package, the xcengine support code provides the notebook code with +the path to the STAC stage-in catalogue. The notebook code can then read this +catalogue (e.g. using [PySTAC](https://pystac.readthedocs.io/)) to find the +staged-in datasets. + +An input variable for a STAC stage-in catalogue can be defined in the +parameters cell (see above) as a string variable. The variable name can be +freely chosen, but the variable declaration must be annotated with the string +`"EOInput"` to distinguish it from an ordinary string parameter, like this: + +```python +dataset_inputs: "EOInput" = "/some/default/path" +``` + +When the converted notebook is run as an Application Package, the variable +`dataset_inputs` will be set to a string specifying a filesystem path +containing a STAC catalogue called `catalog.json`, which the notebook code +can use to find staged-in datasets. + ## Dataset output ### Selecting datasets for output -No additional code or configuration is needed for datasets to be written from -Application Packages or served when the container image is run in xcube -Server/Viewer mode. xcengine will automatically output or serve any instance -of `xarray.DataSet` which is in scope when the notebook's code has finished -executing. If you're created some datasets which you *don't* wish to be -written, you can use the Python -[`del` statement](https://docs.python.org/3/reference/simple_stmts.html#the-del-statement) +No additional code or configuration is needed for datasets to be written +(‘staged out’) from Application Packages or served when the container image is +run in xcube Server/Viewer mode. xcengine will automatically output or serve +any instance of `xarray.DataSet` which is in scope when the notebook's code +has finished executing. If you're created some datasets which you *don't* wish +to be written, you can use the Python [`del` +statement](https://docs.python.org/3/reference/simple_stmts.html#the-del-statement) to delete them at the end of the notebook to remove them, e.g. ```python @@ -77,3 +105,11 @@ like this: ```python my_dataset.attrs["xcengine_output_format"] = "netcdf" ``` + +## Determining whether your code is running in an xcengine container + +In an xcengine-derived container, the environment variable `XCENGINE_VERSION` +is always set to the version of xcengine that created the container image. If +your notebook code needs to determine whether it's running inside an xcengine +container, you can check whether this variable is set (e.g. using +`os.environ`). diff --git a/test/test_core.py b/test/test_core.py index 09153c1..030a296 100644 --- a/test/test_core.py +++ b/test/test_core.py @@ -444,7 +444,8 @@ def test_image_builder_build_dir( build_env_path = build_dir / "environment.yml" assert build_env_path.is_file() output_env = yaml.safe_load(build_env_path.read_text()) - assert {"name", "channels", "dependencies"} <= set(output_env) + assert {"name", "channels", "dependencies", "variables"} <= set(output_env) + assert type(output_env["variables"]["XCENGINE_VERSION"]) is str if env_type != "none": assert output_env["name"] == env_def["name"] assert output_env["channels"] == env_def["channels"] diff --git a/test/test_parameters.py b/test/test_parameters.py index 4275a15..8f2618a 100644 --- a/test/test_parameters.py +++ b/test/test_parameters.py @@ -19,6 +19,7 @@ def expected_vars(): "some_float": (float, 3.14159), "some_string": (str, "foo"), "some_bool": (bool, False), + "some_directory": ("Directory", "/some/path") } @@ -37,6 +38,9 @@ def params_yaml(): some_bool: type: bool default: false +some_directory: + type: Directory + default: "/some/path" """ @@ -130,6 +134,16 @@ def test_parameters_get_commandline_inputs(notebook_parameters): "doc": "some_bool", "inputBinding": {"prefix": "--some-bool"}, }, + "some_directory": { + "type": "Directory", + "default": { + "class": "Directory", + "location": "/some/path" + }, + "label": "some_directory", + "doc": "some_directory", + "inputBinding": {"prefix": "--some-directory"}, + }, } @@ -139,6 +153,7 @@ def test_parameters_get_cwl_step_inputs(notebook_parameters): "some_float": "some_float", "some_string": "some_string", "some_bool": "some_bool", + "some_directory": "some_directory", } @@ -148,23 +163,33 @@ def test_parameters_from_code(expected_vars): some_float = 3.14159 some_string = "foo" some_bool = False +some_directory: "EOInput" = "/some/path" """) assert parameters.params == expected_vars assert parameters.config == {} -def test_parameters_from_code_with_xce_config(expected_vars): - xce_config = dict(foo=1, bar="hi!", baz={}) +@pytest.mark.parametrize("config", [ + (dict(foo=1, bar="hi!", baz={}), True), + ("Not a dict", False), + ({42: "Wrong key type"}, False)]) +def test_parameters_from_code_with_xce_config(expected_vars, config): + xce_config, valid = config code = f""" some_int = 42 some_float = 3.14159 some_string = "foo" some_bool = False +some_directory: "EOInput" = "/some/path" {NotebookParameters.config_var_name} = {xce_config!r} """ - parameters = xcengine.parameters.NotebookParameters.from_code(code) - assert parameters.params == expected_vars - assert parameters.config == xce_config + if valid: + parameters = xcengine.parameters.NotebookParameters.from_code(code) + assert parameters.params == expected_vars + assert parameters.config == xce_config + else: + with pytest.raises(TypeError): + xcengine.parameters.NotebookParameters.from_code(code) def test_parameters_from_code_with_setup(expected_vars): @@ -175,11 +200,13 @@ def test_parameters_from_code_with_setup(expected_vars): some_float = 3.14159 some_string = some_uppercase_string.lower() some_bool = not not_some_bool +some_directory: "EOInput" = "/" + "/".join(some_path_components) """, setup_code=""" half_of_some_int = 21 some_uppercase_string = "FOO" not_some_bool = True +some_path_components = ["some", "path"] """, ).params == expected_vars @@ -212,6 +239,15 @@ def test_parameters_get_workflow_inputs(notebook_parameters): "label": "some_bool", "doc": "some_bool", }, + "some_directory": { + "type": "Directory", + "default": { + "class": "Directory", + "location": "/some/path", + }, + "label": "some_directory", + "doc": "some_directory", + } } @@ -221,9 +257,20 @@ def test_parameters_to_yaml(notebook_parameters): "some_float": {"type": "float", "default": 3.14159}, "some_string": {"type": "str", "default": "foo"}, "some_bool": {"type": "bool", "default": False}, + "some_directory": {"type": "Directory", "default": "/some/path"}, } +def test_parameters_to_yaml_unhandled_type(): + with pytest.raises(TypeError): + # Create empty parameters and modify them afterwards to avoid + # __init__ catching the mistake. + np = xcengine.parameters.NotebookParameters({}) + # Disable the inspection, since this is wrong on purpose. + # noinspection bad-assignment + np.params = {"foo": (42, 42)} + np.to_yaml() + def test_parameters_from_yaml(expected_vars, params_yaml): assert NotebookParameters.from_yaml(params_yaml).params == expected_vars @@ -241,6 +288,15 @@ def test_parameters_from_yaml_with_dataset(): } +def test_parameters_from_yaml_unknown_type(): + with pytest.raises(ValueError): + NotebookParameters.from_yaml(""" +some_input: + type: unsupported + default: null + """) + + def test_parameters_from_file(tmp_path, expected_vars, params_yaml): path = tmp_path / "params.yaml" path.write_text(params_yaml) @@ -264,12 +320,15 @@ def test_parameters_read_cli_arguments(notebook_parameters): "--some-float", "2.71828", "--some-bool", + "--some-directory", + "/a/different/path" ] ) == { "some_int": 23, "some_float": 2.71828, "some_string": "bar", "some_bool": True, + "some_directory": "/a/different/path" } assert notebook_parameters.read_params_from_cli([]) == {} diff --git a/xcengine/core.py b/xcengine/core.py index 32ed087..8fb373c 100644 --- a/xcengine/core.py +++ b/xcengine/core.py @@ -28,6 +28,7 @@ import nbformat import yaml +import xcengine from xcengine import util from xcengine.parameters import NotebookParameters @@ -322,6 +323,9 @@ def ensure_present(pkg: str): for package in packages: ensure_present(package) + + env_vars = conda_env.setdefault("variables", {}) + env_vars["XCENGINE_VERSION"] = xcengine.__version__ return conda_env def _build_image(self) -> docker.models.images.Image: @@ -346,7 +350,7 @@ def write_dockerfile(destination: pathlib.Path) -> None: destination.parent.mkdir(parents=True, exist_ok=True) with open(destination, "w") as fh: fh.write(textwrap.dedent("""\ - FROM mambaorg/micromamba:1.5.10-noble-cuda-12.6.0 + FROM mambaorg/micromamba:2.9-cuda13.2.1-ubuntu24.04 COPY Dockerfile Dockerfile COPY environment.yml environment.yml RUN micromamba install -y -n base -f environment.yml && \\ diff --git a/xcengine/parameters.py b/xcengine/parameters.py index c88f3dc..ad99565 100644 --- a/xcengine/parameters.py +++ b/xcengine/parameters.py @@ -1,8 +1,9 @@ +import builtins import logging import os import pathlib import typing -from typing import Any, ClassVar +from typing import Any, ClassVar, cast import xarray as xr import yaml @@ -13,7 +14,7 @@ class NotebookParameters: - params: dict[str, tuple[type, Any]] + params: dict[str, tuple[type | str, Any]] cwl_params: dict[str, tuple[type | str, Any]] dataset_inputs: list[str] config_var_name: ClassVar[str] = "xcengine_config" @@ -21,7 +22,7 @@ class NotebookParameters: def __init__( self, - params: dict[str, tuple[type, Any]], + params: dict[str, tuple[type | str, Any]], config: dict[str, Any] | None = None, ): self.params = params @@ -46,16 +47,27 @@ def from_code( ) -> "NotebookParameters": variables = cls.extract_variables(code, setup_code) config = variables.pop(cls.config_var_name, (None, None)) - # TODO: throw an error here if config has wrong type - return cls(variables, config[1]) + if config[1] is not None: + if type(config[1]) is not dict: + raise TypeError("Configuration variable must be a dict") + if not all(type(k) is str for k in cast(dict, config[1]).keys()): + raise TypeError("Configuration dict keys must be strings") + return cls(variables, cast(dict[str, Any], config[1])) @classmethod def from_yaml(cls, yaml_content: str | typing.IO) -> "NotebookParameters": input_data = yaml.safe_load(yaml_content) + def convert_type(yaml_spec: str) -> type | str: + match yaml_spec: + case "int" | "float" | "bool" | "str" | "Dataset": + return eval(yaml_spec, globals(), {"Dataset": xr.Dataset}) + case "Directory": + return "Directory" + raise ValueError(f'Unknown type in YAML: "{yaml_spec}"') return cls( { k: ( - eval(v["type"], globals(), {"Dataset": xr.Dataset}), + convert_type(v["type"]), v["default"], ) for k, v in input_data.items() @@ -70,7 +82,7 @@ def from_yaml_file(cls, path: str | os.PathLike) -> "NotebookParameters": @classmethod def extract_variables( cls, code: str, setup_code: str | None = None - ) -> dict[str, tuple[type, Any]]: + ) -> dict[str, tuple[type | str, Any]]: if setup_code is None: locals_: dict[str, object] = {} old_locals = {} @@ -78,14 +90,19 @@ def extract_variables( exec(setup_code, globals(), locals_ := {}) old_locals = locals_.copy() exec(code, globals(), locals_) + annotations = cls.read_annotations(code) new_vars = locals_.keys() - old_locals.keys() new_var_dict = { - k: cls.make_param_tuple(k, locals_[k]) for k in new_vars + k: cls.make_param_tuple(k, locals_[k]) for k in new_vars if not k.startswith("__") } + for k in new_var_dict: + if k in annotations and annotations[k] == "'EOInput'": + old_var = new_var_dict[k] + new_var_dict[k] = ("Directory", old_var[1]) return dict(sorted(new_var_dict.items())) @classmethod - def make_param_tuple(cls, key: str, value: Any) -> tuple[type, Any]: + def make_param_tuple(cls, key: str, value: Any) -> tuple[type | str, Any]: return ( t := type(value), ( @@ -116,7 +133,8 @@ def get_cwl_workflow_input(self, var_name: str) -> dict[str, Any]: "label": var_name, "doc": var_name, "type": self.cwl_type(type_), - "default": default_, + "default": {"class": "Directory", "location": default_} + if type_ == "Directory" else default_, } def get_cwl_commandline_input(self, var_name: str) -> dict[str, Any]: @@ -125,9 +143,17 @@ def get_cwl_commandline_input(self, var_name: str) -> dict[str, Any]: } def to_yaml(self) -> str: + def dump_type(type_: type | str) -> str: + match type_: + case type(): + return type_.__name__ + case str(): + return type_ + case _: + raise TypeError(f"Unhandled type {type_} for YAML export") return yaml.safe_dump( { - name: {"type": type_.__name__, "default": default_} + name: {"type": dump_type(type_), "default": default_} for name, (type_, default_) in self.params.items() } ) @@ -158,9 +184,15 @@ def read_params_from_cli(self, args: list[str]) -> dict[str, Any]: for param_name, (type_, _) in self.params.items(): arg_name = "--" + param_name.replace("_", "-") if arg_name in args and type_ != xr.Dataset: - values[param_name] = type_ is bool or type_( - args[args.index(arg_name) + 1] - ) + match type_: + case builtins.bool: + values[param_name] = True + case "Directory": + values[param_name] = args[args.index(arg_name) + 1] + case _: + values[param_name] = type_( + args[args.index(arg_name) + 1] + ) if "product" in self.cwl_params and "--product" in args: self.read_datasets_from_product( args[args.index("--product") + 1], values @@ -223,7 +255,7 @@ def read_staged_in_dataset( return xr.open_dataset(stage_in_path / asset.href) @staticmethod - def cwl_type(type_: type) -> str: + def cwl_type(type_: type | str) -> str: try: # noinspection PyTypeChecker return { @@ -231,6 +263,7 @@ def cwl_type(type_: type) -> str: float: "double", str: "string", bool: "boolean", + "Directory": "Directory", }[type_] except KeyError: raise ValueError(f"Unhandled type {type_}")