earthkit-workflows 0.5.0__tar.gz → 0.6.0__tar.gz

earthkit_workflows-0.6.0/.github/workflows/macos-test.yml

@@ -0,0 +1,35 @@
+name: CI # TODO rename the file in a standalone PR
+
+on:
+  # Trigger the workflow manually
+  workflow_dispatch: ~
+
+  push:
+    branches:
+      - 'main'
+      - 'develop'
+
+  pull_request: ~
+
+jobs:
+  ci:
+    strategy:
+      fail-fast: true
+      matrix:
+        arch_type: ["macos-ARM64", "linux-x86"]
+        python_version: ["3.10", "3.11", "3.12", "3.13"]
+    runs-on: "${{ fromJSON('{\"linux-x86\": [\"self-hosted\", \"Linux\", \"platform-builder-Rocky-8.6\"], \"macos-ARM64\": [\"self-hosted\", \"macOS\", \"ARM64\"]}')[matrix.arch_type] }}"
+    timeout-minutes: 20
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: ${{ inputs.python-version }}
+      - uses: astral-sh/setup-uv@v6
+        with:
+          version: 0.7.19
+      - uses: extractions/setup-just@v3
+      - run: |
+          uv sync --python "${{ matrix.python_version }}"
+          just fmt
+          just val

earthkit_workflows-0.6.0/.pre-commit-config.yaml

@@ -0,0 +1,29 @@
+repos:
+- repo: https://github.com/pre-commit/pre-commit-hooks
+  rev: v6.0.0
+  hooks:
+  - id: check-yaml # Check YAML files for syntax errors only
+    args: [--unsafe, --allow-multiple-documents]
+#  - id: no-commit-to-branch # NOTE prevents commit to main/master, but since we run prek on that branch it makes no sense
+  - id: check-added-large-files
+    exclude: |
+      (?x)(
+        .*uv.lock|
+        .*pylock.toml
+      )
+  - id: check-merge-conflict # Check for files that contain merge conflict
+- repo: https://github.com/astral-sh/ruff-pre-commit
+  rev: v0.12.2
+  hooks:
+  - id: ruff # better black/flake/isort
+    files: ^src/
+    args:
+    - --select
+    - I # isorting
+    - --fix
+    - --exit-non-zero-on-fix
+  - id: ruff-format
+    files: ^backend/
+ci:
+  autoupdate_schedule: monthly
+  autoupdate_commit_msg: "chore(deps): pre-commit.ci autoupdate"

earthkit_workflows-0.6.0/.python_version

@@ -0,0 +1 @@
+3.11.9

earthkit_workflows-0.6.0/PKG-INFO

@@ -0,0 +1,132 @@
+Metadata-Version: 2.4
+Name: earthkit-workflows
+Version: 0.6.0
+Summary: Earthkit Workflows is a Python library for declaring earthkit task DAGs, as well as scheduling and executing them on heterogeneous computing systems.
+Author-email: "European Centre for Medium-Range Weather Forecasts (ECMWF)" <software.support@ecmwf.int>
+License-Expression: Apache-2.0
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: earthkit-data
+Requires-Dist: cloudpickle
+Requires-Dist: numpy
+Requires-Dist: xarray
+Requires-Dist: networkx
+Requires-Dist: array-api-compat
+Requires-Dist: sortedcontainers
+Requires-Dist: pyvis
+Requires-Dist: dill
+Requires-Dist: pyrsistent
+Requires-Dist: pydantic
+Requires-Dist: pyzmq
+Requires-Dist: fire
+Requires-Dist: orjson
+Dynamic: license-file
+
+<p align="center">
+  <picture>
+    <source srcset="https://github.com/ecmwf/logos/raw/refs/heads/main/logos/earthkit/earthkit-workflows-dark.svg" media="(prefers-color-scheme: dark)">
+    <img src="https://github.com/ecmwf/logos/raw/refs/heads/main/logos/earthkit/earthkit-workflows-light.svg" height="120">
+  </picture>
+</p>
+
+<p align="center">
+  <a href="https://github.com/ecmwf/codex/raw/refs/heads/main/ESEE">
+    <img src="https://github.com/ecmwf/codex/raw/refs/heads/main/ESEE/production_chain_badge.svg" alt="ECMWF Software EnginE">
+  </a>
+  <a href="https://github.com/ecmwf/codex/raw/refs/heads/main/Project Maturity">
+    <img src="https://github.com/ecmwf/codex/raw/refs/heads/main/Project Maturity/emerging_badge.svg" alt="Maturity Level">
+  </a>
+  <a href="https://opensource.org/licenses/apache-2-0">
+    <img src="https://img.shields.io/badge/Licence-Apache 2.0-blue.svg" alt="Licence">
+  </a>
+  <a href="https://github.com/ecmwf/earthkit-workflows/tags">
+    <img src="https://img.shields.io/github/v/tag/ecmwf/earthkit-workflows?color=purple&label=Release" alt="Latest Release">
+  </a>
+</p>
+
+<p align="center">
+  <a href="#installation">Installation</a>
+  •
+  <a href="#quick-start">Quick Start</a>
+  •
+  <a href="#documentation">Documentation</a>
+</p>
+
+> \[!IMPORTANT\]
+> This software is **Emerging** and subject to ECMWF's guidelines on [Software Maturity](https://github.com/ecmwf/codex/raw/refs/heads/main/Project%20Maturity).
+
+**earthkit-workflows** is a Python library for declaring earthkit task as DAGs.
+It contains an internal `cascade` engine for scheduling and executing task graphs almost optimally across heterogeneous platforms with complex network technologies and topologies.
+It effectively performs task-based parallelism across CPUs, GPUs, distributed systems (HPC), and any combination thereof.
+It is designed for a no-IO approach, where expensive storage of intermediate data is minimised whilst maximising all available transport technologies between different hardware.
+
+Cascade is designed to work on well-profiled task graphs, where:
+* the task graph is a static DAG,
+* the DAG nodes are defined by tasks with well-known execution times,
+* the DAG edges are defined by data dependencies with well-known data sizes,
+* the characteristics of the hardware (processors, network connections) are known.
+
+earthkit-workflows allows for declaring such task graphs using a neat fluent API, and interoperates pleasantly with the rest of the [earthkit](https://github.com/ecmwf/earthkit) ecosystem.
+
+## Installation
+
+Install via `pip` with:
+
+```
+$ pip install 'earthkit-workflows[all]'
+```
+
+For development, you can use `pip install -e .` though there is currently an issue with earthkit masking. Additionally you may want to install pre-commit hooks via
+```
+$ pip install pre-commit
+$ pre-commit install
+```
+
+## Quick Start
+
+*Note*: this section is moderately outdated.
+
+We support two regimes for cascade executions -- local mode (ideal for developing and debugging small graphs) and distributed mode (assumed for slurm & HPC).
+
+To launch in local mode, in your python repl / jupyno:
+```
+import cascade.benchmarks.job1 as j1
+import cascade.benchmarks.distributed as di
+import cloudpickle
+
+spec = di.ZmqClusterSpec.local(j1.get_prob())
+print(spec.controller.outputs)
+# prints out:
+# {DatasetId(task='mean:dc9d90 ...
+# defaults to all "sinks", but can be overridden
+
+rv = di.launch_from_specs(spec, None)
+
+for key, value in rv.outputs.items():
+    deser = cloudpickle.loads(value)
+    print(f"output {key} is of type {type(deser)}")
+```
+
+For distributed mode, launch
+```
+./scripts/launch_slurm.sh ./localConfigs/<your_config.sh>
+```
+Inside the `<your_config.sh>`, you define size of the cluster, logging directory output, which job to run... Pay special attention to definitions of your `venv` and `LD_LIBRARY_PATH` etc -- this is not autotamed.
+
+Both of these examples hardcode particular job, `"job1"`, which is a benchmarking thing.
+Most likely, you want to define your own -- for the local mode, just pass `cascade.Graph` instance to the call; in the dist mode, you need to provide that instance in the `cascade.benchmarks.__main__` modules instead (ideally by extending the `get_job` function).
+
+There is also `python -m cascade.benchmarks local <..>` -- you may use that as an alternative path to local mode, for your own e2e tests.
+
+## Documentation
+
+Not yet available.
+
+## Contributions and Support
+Due to the maturity and status of the project, there is no support provided -- unless the usage of this project happens within some higher-status initiative that ECMWF participates at.
+External contributions and created issues will be looked at, but are not guaranteed to be accepted or responded to.
+In general, follow ECMWF's guidelines for [external contributions](https://github.com/ecmwf/codex/tree/main/External%20Contributions).
+
+## License
+See [license](./LICENSE).

{earthkit_workflows-0.5.0 → earthkit_workflows-0.6.0}/justfile

@@ -10,12 +10,9 @@
 set dotenv-path := ".env"
 
 val:
-	mypy src/cascade --ignore-missing-imports
-	mypy tests --ignore-missing-imports
-	pytest -n 8 tests
-
+    uv run ty check src/cascade
+    uv run ty check tests/cascade
+    # TODO eventually broaden type coverage to ekw as well
+    uv run pytest -n8 tests
 fmt:
-    # TODO replace with pre-commit
-    isort --profile black .
-    black .
-    flake8 .
+    uv run prek --all-files

{earthkit_workflows-0.5.0 → earthkit_workflows-0.6.0}/pyproject.toml

@@ -7,7 +7,7 @@
 # nor does it submit to any jurisdiction.
 
 [build-system]
-requires = ["setuptools>=65", "setuptools_scm[toml]>=6.2"]
+requires = ["setuptools>=80", "setuptools_scm[toml]>=6.2", "packaging>=25"]
 build-backend = "setuptools.build_meta"
 
 [project]
@@ -31,7 +31,7 @@ dependencies = [
     "dill",
     "pyrsistent",
     "pydantic",
-    "zmq",
+    "pyzmq",
     "fire",
     "orjson",
 ]
@@ -39,21 +39,13 @@ dependencies = [
 dynamic = ["version"]
 readme = "README.md"
 
-[project.optional-dependencies]
-tests = ["pytest", "pytest-xdist", "earthkit-data"]
-lint = ["black", "isort", "flake8"]
-gpu = ["jax[cpu]", "jax[cuda11_pip]", "cupy-cuda11x", "numba"]
-examples = ["cftime", "bokeh"]
-earthkit = ["earthkit-data"]
-all = ["cascade[tests,kubernetes,lint,gpu,examples,earthkit]"]
+[dependency-groups]
+dev = ["pytest", "pytest-xdist>=3.8", "prek", "ty==0.0.2", "build", "bokeh"]
 
 [tool.setuptools]
 include-package-data = true
 zip-safe = false
 
-[tool.setuptools.dynamic]
-readme = {file = ["readme.md"], content-type = "text/markdown"}
-
 [tool.setuptools_scm]
 write_to = "src/earthkit/workflows/_version.py"
 write_to_template = '''# Do not change! Do not track in version control!
@@ -65,17 +57,12 @@ local_scheme = "no-local-version"
 include = ["earthkit.workflows*", "cascade*"]
 where = ["src"]
 
-[tool.isort]
-profile = "black"
-
-[tool.mypy]
-plugins = "pydantic.mypy"
-
 [tool.pytest.ini_options]
 log_cli = true
 log_cli_level = "DEBUG"
 testpaths = ["tests/"]
 addopts = "-n8"
 
-[tool.ruff.lint]
-ignore = ["E722", "E731", "E741"]
+[tool.ruff]
+line-length = 140
+lint.ignore = [ "E722", "E731", "E741" ]

{earthkit_workflows-0.5.0 → earthkit_workflows-0.6.0}/src/cascade/benchmarks/anemoi.py

@@ -10,7 +10,7 @@ from earthkit.workflows import Cascade
 
 
 def get_graph(lead_time, ensemble_members, CKPT=None, date="2024-12-02T00:00"):
-    import anemoicascade as ac
+    import anemoicascade as ac  # ty: ignore[unresolved-import]
 
     CKPT = (
         CKPT

{earthkit_workflows-0.5.0 → earthkit_workflows-0.6.0}/src/cascade/benchmarks/dask.py

@@ -5,9 +5,9 @@ from cascade.low.core import JobInstance
 from cascade.low.dask import graph2job
 
 
-def get_job(job: str) -> JobInstance:
+def get_job(job_name: str) -> JobInstance:
 
-    if job == "add":
+    if job_name == "add":
 
         def add(x, y):
             result = x + y
@@ -21,7 +21,7 @@ def get_job(job: str) -> JobInstance:
             dataset for task in job.tasks for dataset in job.outputs_of(task)
         ]
         return job
-    elif job == "groupby":
+    elif job_name == "groupby":
         df = dd.DataFrame.from_dict({"x": [0, 0, 1, 1], "y": [1, 2, 3, 4]})
         df = df.groupby("x").sum()
         job = graph2job(df.__dask_graph__())
@@ -30,4 +30,4 @@ def get_job(job: str) -> JobInstance:
         ]
         return job
     else:
-        raise NotImplementedError(job)
+        raise NotImplementedError(job_name)

{earthkit_workflows-0.5.0 → earthkit_workflows-0.6.0}/src/cascade/benchmarks/dist.py

@@ -26,7 +26,7 @@ def dist_func_torch(a: int) -> int:
     import datetime as dt
 
     import numpy as np
-    import torch.distributed as dist
+    import torch.distributed as dist  # ty: ignore[unresolved-import]
 
     world_size = int(os.environ["CASCADE_GANG_WORLD_SIZE"])
     rank = int(os.environ["CASCADE_GANG_RANK"])
@@ -61,8 +61,8 @@ def dist_func_jax(a: int) -> int:
     os.environ["JAX_NUM_CPU_DEVICES"] = "1"
     os.environ["JAX_PLATFORM_NAME"] = "cpu"
     os.environ["JAX_PLATFORMS"] = "cpu"
-    import jax
-    import jax.numpy as jp
+    import jax  # ty: ignore[unresolved-import]
+    import jax.numpy as jp  # ty: ignore[unresolved-import]
 
     jax.config.update("jax_platforms", "cpu")
     jax.config.update("jax_platform_name", "cpu")

{earthkit_workflows-0.5.0 → earthkit_workflows-0.6.0}/src/cascade/benchmarks/job1.py

@@ -16,10 +16,9 @@ Controlled by env var params: JOB1_{DATA_ROOT, GRID, ...}, see below
 import os
 
 import earthkit.data
-
 from earthkit.workflows.fluent import Payload
-from earthkit.workflows.plugins.pproc.fluent import from_source
-from earthkit.workflows.plugins.pproc.utils.window import Range
+from earthkit.workflows.plugins.pproc.fluent import from_source  # ty: ignore
+from earthkit.workflows.plugins.pproc.utils.window import Range  # ty: ignore
 
 # *** PARAMS ***
 
@@ -137,7 +136,7 @@ def download_inputs():
             }
             data = earthkit.data.from_source("mars", **ekp)
             with open(f"{data_root}/data_{number}_{step}.grib", "wb") as f:
-                data.write(f)
+                data.write(f) # ty: ignore
 
 
 def download_climatology():
@@ -157,7 +156,7 @@ def download_climatology():
         }
         data = earthkit.data.from_source("mars", **ekp)
         with open(f"{data_root}/data_clim_{step}.grib", "wb") as f:
-            data.write(f)
+            data.write(f) # ty: ignore
 
 
 if __name__ == "__main__":

{earthkit_workflows-0.5.0 → earthkit_workflows-0.6.0}/src/cascade/benchmarks/matmul.py

@@ -1,9 +1,9 @@
 import os
 from typing import Any
 
-import jax
-import jax.numpy as jp
-import jax.random as jr
+import jax  # ty: ignore[unresolved-import]
+import jax.numpy as jp  # ty: ignore[unresolved-import]
+import jax.random as jr  # ty: ignore[unresolved-import]
 
 from cascade.low.builders import JobBuilder, TaskBuilder
 from cascade.low.core import JobInstance
@@ -65,7 +65,7 @@ def execute_locally():
 
     from multiprocessing.shared_memory import SharedMemory
 
-    mem = SharedMemory("benchmark_tmp", create=True, size=m0.nbytes)
+    mem = SharedMemory("benchmark_tmp", create=True, size=m0.nbytes); assert mem.buf is not None
     mem.buf[:] = m0.tobytes()
 
 

{earthkit_workflows-0.5.0 → earthkit_workflows-0.6.0}/src/cascade/benchmarks/tests.py

@@ -32,7 +32,7 @@ from cascade.executor.runner.memory import Memory, ds2shmid
 from cascade.executor.runner.packages import PackagesEnv
 from cascade.executor.runner.runner import ExecutionContext, run
 from cascade.low.builders import TaskBuilder
-from cascade.low.core import DatasetId
+from cascade.low.core import DatasetId, WorkerId
 from cascade.shm.server import entrypoint as shm_server
 
 logger = logging.getLogger(__name__)
@@ -75,7 +75,7 @@ def simple_runner(callback: BackboneAddress, executionContext: ExecutionContext)
         raise ValueError(f"expected 1 task, gotten {len(tasks)}")
     taskId = tasks[0]
     taskInstance = executionContext.tasks[taskId]
-    with Memory(callback, "testWorker") as memory, PackagesEnv() as pckg:
+    with Memory(callback, WorkerId(host="testHost", worker="testWorker")) as memory, PackagesEnv() as pckg:
         # for key, value in taskSequence.extra_env.items():
         #    os.environ[key] = value
 
@@ -142,7 +142,7 @@ def run_test(
         while perf_counter_ns() < end:
             mess = listener.recv_messages()
             if mess == [
-                DatasetPublished(origin="testWorker", ds=output, transmit_idx=None)
+                DatasetPublished(origin=WorkerId(host="testHost", worker="testWorker"), ds=output, transmit_idx=None)
             ]:
                 break
             elif not mess:

{earthkit_workflows-0.5.0 → earthkit_workflows-0.6.0}/src/cascade/benchmarks/util.py

@@ -29,7 +29,7 @@ from cascade.executor.comms import callback
 from cascade.executor.config import logging_config, logging_config_filehandler
 from cascade.executor.executor import Executor
 from cascade.executor.msg import BackboneAddress, ExecutorShutdown
-from cascade.low.core import DatasetId, JobInstance
+from cascade.low.core import DatasetId, JobInstance, JobInstanceRich
 from cascade.low.func import msum
 from cascade.scheduler.precompute import precompute
 from earthkit.workflows.graph import Graph, deduplicate_nodes
@@ -37,15 +37,16 @@ from earthkit.workflows.graph import Graph, deduplicate_nodes
 logger = logging.getLogger("cascade.benchmarks")
 
 
-def get_job(benchmark: str | None, instance_path: str | None) -> JobInstance:
+def get_job(benchmark: str | None, instance_path: str | None) -> JobInstanceRich:
     # NOTE because of os.environ, we don't import all... ideally we'd have some file-based init/config mech instead
     if benchmark is not None and instance_path is not None:
         raise TypeError("specified both benchmark name and job instance")
     elif instance_path is not None:
         with open(instance_path, "rb") as f:
             d = orjson.loads(f.read())
-            return JobInstance(**d)
+            return JobInstanceRich(**d)
     elif benchmark is not None:
+        instance: JobInstance
         if benchmark.startswith("j1"):
             import cascade.benchmarks.job1 as job1
 
@@ -58,25 +59,26 @@ def get_job(benchmark: str | None, instance_path: str | None) -> JobInstance:
                 msum((v for k, v in graphs.items() if k.startswith(prefix)), Graph)
             )
             graphs["j1.all"] = union("j1.")
-            return cascade.low.into.graph2job(graphs[benchmark])
+            instance = cascade.low.into.graph2job(graphs[benchmark])
         elif benchmark.startswith("generators"):
             import cascade.benchmarks.generators as generators
 
-            return generators.get_job()
+            instance = generators.get_job()
         elif benchmark.startswith("matmul"):
             import cascade.benchmarks.matmul as matmul
 
-            return matmul.get_job()
+            instance = matmul.get_job()
         elif benchmark.startswith("dist"):
             import cascade.benchmarks.dist as dist
 
-            return dist.get_job()
+            instance = dist.get_job()
         elif benchmark.startswith("dask"):
             import cascade.benchmarks.dask as dask
 
-            return dask.get_job(benchmark[len("dask.") :])
+            instance = dask.get_job(benchmark[len("dask.") :])
         else:
             raise NotImplementedError(benchmark)
+        return JobInstanceRich(jobInstance=instance, checkpointSpec=None)
     else:
         raise TypeError("specified neither benchmark name nor job instance")
 
@@ -116,7 +118,7 @@ def get_gpu_count(host_idx: int, worker_count: int) -> int:
 
 
 def launch_executor(
-    job_instance: JobInstance,
+    job: JobInstanceRich,
     controller_address: BackboneAddress,
     workers_per_host: int,
     portBase: int,
@@ -136,7 +138,7 @@ def launch_executor(
         logger.info(f"will set {gpu_count} gpus on host {i}")
         os.environ["CASCADE_GPU_COUNT"] = str(gpu_count)
         executor = Executor(
-            job_instance,
+            job.jobInstance,
             controller_address,
             workers_per_host,
             f"h{i}",
@@ -154,7 +156,7 @@ def launch_executor(
 
 
 def run_locally(
-    job: JobInstance,
+    job: JobInstanceRich,
     hosts: int,
     workers: int,
     portBase: int = 12345,
@@ -195,7 +197,7 @@ def run_locally(
             ps.append(p)
 
         # compute preschedule
-        preschedule = precompute(job)
+        preschedule = precompute(job.jobInstance)
 
         # check processes started healthy
         for i, p in enumerate(ps):
@@ -240,9 +242,9 @@ def main_local(
     port_base: int = 12345,
     log_base: str | None = None,
 ) -> None:
-    jobInstance = get_job(job, instance)
+    jobInstanceRich = get_job(job, instance)
     run_locally(
-        jobInstance,
+        jobInstanceRich,
         hosts,
         workers_per_host,
         report_address=report_address,
@@ -266,17 +268,17 @@ def main_dist(
     """
     launch = perf_counter_ns()
 
-    jobInstance = get_job(job, instance)
+    jobInstanceRich = get_job(job, instance)
 
     if idx == 0:
         logging.config.dictConfig(logging_config)
         tp = ThreadPoolExecutor(max_workers=1)
-        preschedule_fut = tp.submit(precompute, jobInstance)
+        preschedule_fut = tp.submit(precompute, jobInstanceRich.jobInstance)
         b = Bridge(controller_url, hosts)
         preschedule = preschedule_fut.result()
         tp.shutdown()
         start = perf_counter_ns()
-        run(jobInstance, b, preschedule, report_address=report_address)
+        run(jobInstanceRich, b, preschedule, report_address=report_address)
         end = perf_counter_ns()
         print(
             f"compute took {(end-start)/1e9:.3f}s, including startup {(end-launch)/1e9:.3f}s"
@@ -284,12 +286,13 @@ def main_dist(
     else:
         gpu_count = get_gpu_count(0, workers_per_host)
         launch_executor(
-            jobInstance,
+            jobInstanceRich,
             controller_url,
             workers_per_host,
             12345,
             idx,
             shm_vol_gb,
             gpu_count,
-            f"tcp://{platform.get_bindabble_self()}",
+            log_base = None, # TODO handle log collection for dist scenario
+            url_base = f"tcp://{platform.get_bindabble_self()}",
         )

{earthkit_workflows-0.5.0 → earthkit_workflows-0.6.0}/src/cascade/controller/act.py

@@ -10,6 +10,7 @@
 
 import logging
 
+import cascade.executor.checkpoints as checkpoints
 from cascade.controller.core import State
 from cascade.executor.bridge import Bridge
 from cascade.executor.msg import TaskSequence
@@ -76,6 +77,12 @@ def flush_queues(bridge: Bridge, state: State, context: JobExecutionContext):
     for dataset, host in state.drain_fetching_queue():
         bridge.fetch(dataset, host)
 
+    for dataset, host in state.drain_persist_queue():
+        if context.checkpoint_spec is None:
+            raise TypeError(f"unexpected persist need when checkpoint storage not configured")
+        persist_params = checkpoints.serialize_persist_params(context.checkpoint_spec)
+        bridge.persist(dataset, host, context.checkpoint_spec.storage_type, persist_params)
+
     for ds in state.drain_purging_queue():
         for host in context.purge_dataset(ds):
             logger.debug(f"issuing purge of {ds=} to {host=}")