earthkit-workflows 0.4.3__py3-none-any.whl → 0.4.4__py3-none-any.whl

cascade/benchmarks/__main__.py

@@ -19,285 +19,9 @@ Make sure you correctly configure:
  - your venv (cascade, fiab, pproc-cascade, compatible version of earthkit-data, ...)
 """
 
-# TODO rework, simplify
-
-import logging
-import logging.config
-import multiprocessing
-import os
-import subprocess
-import sys
-from concurrent.futures import ThreadPoolExecutor
-from time import perf_counter_ns
-
 import fire
-import orjson
-
-import cascade.executor.platform as platform
-import cascade.low.into
-from cascade.controller.impl import run
-from cascade.executor.bridge import Bridge
-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 JobInstance
-from cascade.low.func import msum
-from cascade.scheduler.precompute import precompute
-from earthkit.workflows.graph import Graph, deduplicate_nodes
-
-logger = logging.getLogger("cascade.benchmarks")
-
-
-def get_job(benchmark: str | None, instance_path: str | None) -> JobInstance:
-    # 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)
-    elif benchmark is not None:
-        if benchmark.startswith("j1"):
-            import cascade.benchmarks.job1 as job1
-
-            graphs = {
-                "j1.prob": job1.get_prob(),
-                "j1.ensms": job1.get_ensms(),
-                "j1.efi": job1.get_efi(),
-            }
-            union = lambda prefix: deduplicate_nodes(
-                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])
-        elif benchmark.startswith("generators"):
-            import cascade.benchmarks.generators as generators
-
-            return generators.get_job()
-        elif benchmark.startswith("matmul"):
-            import cascade.benchmarks.matmul as matmul
-
-            return matmul.get_job()
-        elif benchmark.startswith("dist"):
-            import cascade.benchmarks.dist as dist
-
-            return dist.get_job()
-        else:
-            raise NotImplementedError(benchmark)
-    else:
-        raise TypeError("specified neither benchmark name nor job instance")
-
-
-def get_cuda_count() -> int:
-    try:
-        if "CUDA_VISIBLE_DEVICES" in os.environ:
-            # TODO we dont want to just count, we want to actually use literally these ids
-            # NOTE this is particularly useful for "" value -- careful when refactoring
-            visible = os.environ["CUDA_VISIBLE_DEVICES"]
-            visible_count = sum(1 for e in visible if e == ",") + (1 if visible else 0)
-            return visible_count
-        gpus = sum(
-            1
-            for l in subprocess.run(
-                ["nvidia-smi", "--list-gpus"], check=True, capture_output=True
-            )
-            .stdout.decode("ascii")
-            .split("\n")
-            if "GPU" in l
-        )
-    except:
-        logger.exception("unable to determine available gpus")
-        gpus = 0
-    return gpus
-
-
-def get_gpu_count(host_idx: int, worker_count: int) -> int:
-    if sys.platform == "darwin":
-        # we should inspect some gpu capabilities details to prevent overcommit
-        return worker_count
-    else:
-        if host_idx == 0:
-            return get_cuda_count()
-        else:
-            return 0
-
-
-def launch_executor(
-    job_instance: JobInstance,
-    controller_address: BackboneAddress,
-    workers_per_host: int,
-    portBase: int,
-    i: int,
-    shm_vol_gb: int | None,
-    gpu_count: int,
-    log_base: str | None,
-    url_base: str,
-):
-    if log_base is not None:
-        log_base = f"{log_base}.host{i}"
-        log_path = f"{log_base}.txt"
-        logging.config.dictConfig(logging_config_filehandler(log_path))
-    else:
-        logging.config.dictConfig(logging_config)
-    try:
-        logger.info(f"will set {gpu_count} gpus on host {i}")
-        os.environ["CASCADE_GPU_COUNT"] = str(gpu_count)
-        executor = Executor(
-            job_instance,
-            controller_address,
-            workers_per_host,
-            f"h{i}",
-            portBase,
-            shm_vol_gb,
-            log_base,
-            url_base,
-        )
-        executor.register()
-        executor.recv_loop()
-    except Exception:
-        # NOTE we log this to get the stacktrace into the logfile
-        logger.exception("executor failure")
-        raise
-
-
-def run_locally(
-    job: JobInstance,
-    hosts: int,
-    workers: int,
-    portBase: int = 12345,
-    log_base: str | None = None,
-    report_address: str | None = None,
-):
-    if log_base is not None:
-        log_path = f"{log_base}.controller.txt"
-        logging.config.dictConfig(logging_config_filehandler(log_path))
-    else:
-        logging.config.dictConfig(logging_config)
-    logger.debug(f"local run starting with {hosts=} and {workers=} on {portBase=}")
-    launch = perf_counter_ns()
-    c = f"tcp://localhost:{portBase}"
-    m = f"tcp://localhost:{portBase+1}"
-    ps = []
-    try:
-        # executors forking
-        for i, executor in enumerate(range(hosts)):
-            gpu_count = get_gpu_count(i, workers)
-            # NOTE forkserver/spawn seem to forget venv, we need fork
-            logger.debug(f"forking into executor on host {i}")
-            p = multiprocessing.get_context("fork").Process(
-                target=launch_executor,
-                args=(
-                    job,
-                    c,
-                    workers,
-                    portBase + 1 + i * 10,
-                    i,
-                    None,
-                    gpu_count,
-                    log_base,
-                    "tcp://localhost",
-                ),
-            )
-            p.start()
-            ps.append(p)
-
-        # compute preschedule
-        preschedule = precompute(job)
-
-        # check processes started healthy
-        for i, p in enumerate(ps):
-            if not p.is_alive():
-                # TODO ideally we would somehow connect this with the Register message
-                # consumption in the Controller -- but there we don't assume that
-                # executors are on the same physical host
-                raise ValueError(f"executor {i} failed to live due to {p.exitcode}")
-
-        # start bridge itself
-        logger.debug("starting bridge")
-        b = Bridge(c, hosts)
-        start = perf_counter_ns()
-        run(job, 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"
-        )
-    except Exception:
-        # NOTE we log this to get the stacktrace into the logfile
-        logger.exception("controller failure, proceed with executor shutdown")
-        for p in ps:
-            if p.is_alive():
-                callback(m, ExecutorShutdown())
-                import time
-
-                time.sleep(1)
-                p.kill()
-        raise
-
-
-def main_local(
-    workers_per_host: int,
-    hosts: int = 1,
-    report_address: str | None = None,
-    job: str | None = None,
-    instance: str | None = None,
-    port_base: int = 12345,
-    log_base: str | None = None,
-) -> None:
-    jobInstance = get_job(job, instance)
-    run_locally(
-        jobInstance,
-        hosts,
-        workers_per_host,
-        report_address=report_address,
-        portBase=port_base,
-        log_base=log_base,
-    )
-
-
-def main_dist(
-    idx: int,
-    controller_url: str,
-    hosts: int = 3,
-    workers_per_host: int = 10,
-    shm_vol_gb: int = 64,
-    job: str | None = None,
-    instance: str | None = None,
-    report_address: str | None = None,
-) -> None:
-    """Entrypoint for *both* controller and worker -- they are on different hosts! Distinguished by idx: 0 for
-    controller, 1+ for worker. Assumed to come from slurm procid.
-    """
-    launch = perf_counter_ns()
-
-    jobInstance = get_job(job, instance)
-
-    if idx == 0:
-        logging.config.dictConfig(logging_config)
-        tp = ThreadPoolExecutor(max_workers=1)
-        preschedule_fut = tp.submit(precompute, jobInstance)
-        b = Bridge(controller_url, hosts)
-        preschedule = preschedule_fut.result()
-        tp.shutdown()
-        start = perf_counter_ns()
-        run(jobInstance, 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"
-        )
-    else:
-        gpu_count = get_gpu_count(0, workers_per_host)
-        launch_executor(
-            jobInstance,
-            controller_url,
-            workers_per_host,
-            12345,
-            idx,
-            shm_vol_gb,
-            gpu_count,
-            f"tcp://{platform.get_bindabble_self()}",
-        )
 
+from cascade.benchmarks.util import main_dist, main_local
 
 if __name__ == "__main__":
     fire.Fire({"local": main_local, "dist": main_dist})

cascade/benchmarks/dask.py

@@ -0,0 +1,33 @@
+import dask.dataframe as dd
+from dask._task_spec import convert_legacy_graph
+
+from cascade.low.core import JobInstance
+from cascade.low.dask import graph2job
+
+
+def get_job(job: str) -> JobInstance:
+
+    if job == "add":
+
+        def add(x, y):
+            result = x + y
+            print(f"da {result=}")
+            return result
+
+        dl = {"a": 1, "b": 2, "c": (add, "a", "b")}
+        dn = convert_legacy_graph(dl)
+        job = graph2job(dn)
+        job.ext_outputs = [
+            dataset for task in job.tasks for dataset in job.outputs_of(task)
+        ]
+        return job
+    elif job == "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__())
+        job.ext_outputs = [
+            dataset for task in job.tasks for dataset in job.outputs_of(task)
+        ]
+        return job
+    else:
+        raise NotImplementedError(job)

cascade/benchmarks/util.py

@@ -0,0 +1,294 @@
+# (C) Copyright 2025- ECMWF.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+
+"""Contains utility methods for benchmark definitions and cluster starting"""
+
+# TODO rework, simplify, split into benchmark.util and cluster.setup or smth
+
+import logging
+import logging.config
+import multiprocessing
+import os
+import subprocess
+import sys
+from concurrent.futures import ThreadPoolExecutor
+from time import perf_counter_ns
+
+import orjson
+
+import cascade.executor.platform as platform
+import cascade.low.into
+from cascade.controller.impl import run
+from cascade.executor.bridge import Bridge
+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 JobInstance
+from cascade.low.func import msum
+from cascade.scheduler.precompute import precompute
+from earthkit.workflows.graph import Graph, deduplicate_nodes
+
+logger = logging.getLogger("cascade.benchmarks")
+
+
+def get_job(benchmark: str | None, instance_path: str | None) -> JobInstance:
+    # 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)
+    elif benchmark is not None:
+        if benchmark.startswith("j1"):
+            import cascade.benchmarks.job1 as job1
+
+            graphs = {
+                "j1.prob": job1.get_prob(),
+                "j1.ensms": job1.get_ensms(),
+                "j1.efi": job1.get_efi(),
+            }
+            union = lambda prefix: deduplicate_nodes(
+                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])
+        elif benchmark.startswith("generators"):
+            import cascade.benchmarks.generators as generators
+
+            return generators.get_job()
+        elif benchmark.startswith("matmul"):
+            import cascade.benchmarks.matmul as matmul
+
+            return matmul.get_job()
+        elif benchmark.startswith("dist"):
+            import cascade.benchmarks.dist as dist
+
+            return dist.get_job()
+        elif benchmark.startswith("dask"):
+            import cascade.benchmarks.dask as dask
+
+            return dask.get_job(benchmark[len("dask.") :])
+        else:
+            raise NotImplementedError(benchmark)
+    else:
+        raise TypeError("specified neither benchmark name nor job instance")
+
+
+def get_cuda_count() -> int:
+    try:
+        if "CUDA_VISIBLE_DEVICES" in os.environ:
+            # TODO we dont want to just count, we want to actually use literally these ids
+            # NOTE this is particularly useful for "" value -- careful when refactoring
+            visible = os.environ["CUDA_VISIBLE_DEVICES"]
+            visible_count = sum(1 for e in visible if e == ",") + (1 if visible else 0)
+            return visible_count
+        gpus = sum(
+            1
+            for l in subprocess.run(
+                ["nvidia-smi", "--list-gpus"], check=True, capture_output=True
+            )
+            .stdout.decode("ascii")
+            .split("\n")
+            if "GPU" in l
+        )
+    except:
+        logger.exception("unable to determine available gpus")
+        gpus = 0
+    return gpus
+
+
+def get_gpu_count(host_idx: int, worker_count: int) -> int:
+    if sys.platform == "darwin":
+        # we should inspect some gpu capabilities details to prevent overcommit
+        return worker_count
+    else:
+        if host_idx == 0:
+            return get_cuda_count()
+        else:
+            return 0
+
+
+def launch_executor(
+    job_instance: JobInstance,
+    controller_address: BackboneAddress,
+    workers_per_host: int,
+    portBase: int,
+    i: int,
+    shm_vol_gb: int | None,
+    gpu_count: int,
+    log_base: str | None,
+    url_base: str,
+):
+    if log_base is not None:
+        log_base = f"{log_base}.host{i}"
+        log_path = f"{log_base}.txt"
+        logging.config.dictConfig(logging_config_filehandler(log_path))
+    else:
+        logging.config.dictConfig(logging_config)
+    try:
+        logger.info(f"will set {gpu_count} gpus on host {i}")
+        os.environ["CASCADE_GPU_COUNT"] = str(gpu_count)
+        executor = Executor(
+            job_instance,
+            controller_address,
+            workers_per_host,
+            f"h{i}",
+            portBase,
+            shm_vol_gb,
+            log_base,
+            url_base,
+        )
+        executor.register()
+        executor.recv_loop()
+    except Exception:
+        # NOTE we log this to get the stacktrace into the logfile
+        logger.exception("executor failure")
+        raise
+
+
+def run_locally(
+    job: JobInstance,
+    hosts: int,
+    workers: int,
+    portBase: int = 12345,
+    log_base: str | None = None,
+    report_address: str | None = None,
+) -> None:
+    if log_base is not None:
+        log_path = f"{log_base}.controller.txt"
+        logging.config.dictConfig(logging_config_filehandler(log_path))
+    else:
+        logging.config.dictConfig(logging_config)
+    logger.debug(f"local run starting with {hosts=} and {workers=} on {portBase=}")
+    launch = perf_counter_ns()
+    c = f"tcp://localhost:{portBase}"
+    m = f"tcp://localhost:{portBase+1}"
+    ps = []
+    try:
+        # executors forking
+        for i, executor in enumerate(range(hosts)):
+            gpu_count = get_gpu_count(i, workers)
+            # NOTE forkserver/spawn seem to forget venv, we need fork
+            logger.debug(f"forking into executor on host {i}")
+            p = multiprocessing.get_context("fork").Process(
+                target=launch_executor,
+                args=(
+                    job,
+                    c,
+                    workers,
+                    portBase + 1 + i * 10,
+                    i,
+                    None,
+                    gpu_count,
+                    log_base,
+                    "tcp://localhost",
+                ),
+            )
+            p.start()
+            ps.append(p)
+
+        # compute preschedule
+        preschedule = precompute(job)
+
+        # check processes started healthy
+        for i, p in enumerate(ps):
+            if not p.is_alive():
+                # TODO ideally we would somehow connect this with the Register message
+                # consumption in the Controller -- but there we don't assume that
+                # executors are on the same physical host
+                raise ValueError(f"executor {i} failed to live due to {p.exitcode}")
+
+        # start bridge itself
+        logger.debug("starting bridge")
+        b = Bridge(c, hosts)
+        start = perf_counter_ns()
+        result = run(job, 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"
+        )
+        if os.environ.get("CASCADE_DEBUG_PRINT"):
+            for key, value in result.outputs.items():
+                print(f"{key} => {value}")
+    except Exception:
+        # NOTE we log this to get the stacktrace into the logfile
+        logger.exception("controller failure, proceed with executor shutdown")
+        for p in ps:
+            if p.is_alive():
+                callback(m, ExecutorShutdown())
+                import time
+
+                time.sleep(1)
+                p.kill()
+        raise
+
+
+def main_local(
+    workers_per_host: int,
+    hosts: int = 1,
+    report_address: str | None = None,
+    job: str | None = None,
+    instance: str | None = None,
+    port_base: int = 12345,
+    log_base: str | None = None,
+) -> None:
+    jobInstance = get_job(job, instance)
+    run_locally(
+        jobInstance,
+        hosts,
+        workers_per_host,
+        report_address=report_address,
+        portBase=port_base,
+        log_base=log_base,
+    )
+
+
+def main_dist(
+    idx: int,
+    controller_url: str,
+    hosts: int = 3,
+    workers_per_host: int = 10,
+    shm_vol_gb: int = 64,
+    job: str | None = None,
+    instance: str | None = None,
+    report_address: str | None = None,
+) -> None:
+    """Entrypoint for *both* controller and worker -- they are on different hosts! Distinguished by idx: 0 for
+    controller, 1+ for worker. Assumed to come from slurm procid.
+    """
+    launch = perf_counter_ns()
+
+    jobInstance = get_job(job, instance)
+
+    if idx == 0:
+        logging.config.dictConfig(logging_config)
+        tp = ThreadPoolExecutor(max_workers=1)
+        preschedule_fut = tp.submit(precompute, jobInstance)
+        b = Bridge(controller_url, hosts)
+        preschedule = preschedule_fut.result()
+        tp.shutdown()
+        start = perf_counter_ns()
+        run(jobInstance, 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"
+        )
+    else:
+        gpu_count = get_gpu_count(0, workers_per_host)
+        launch_executor(
+            jobInstance,
+            controller_url,
+            workers_per_host,
+            12345,
+            idx,
+            shm_vol_gb,
+            gpu_count,
+            f"tcp://{platform.get_bindabble_self()}",
+        )

cascade/executor/executor.py

@@ -22,6 +22,8 @@ from multiprocessing import get_context
 from multiprocessing.process import BaseProcess
 from typing import Iterable
 
+import cloudpickle
+
 import cascade.executor.platform as platform
 import cascade.shm.api as shm_api
 import cascade.shm.client as shm_client
@@ -187,9 +189,8 @@ class Executor:
 
     def start_workers(self, workers: Iterable[WorkerId]) -> None:
         # TODO this method assumes no other message will arrive to mlistener! Thus cannot be used for workers now
-        # NOTE a bit risky, but forkserver is too slow. Has unhealthy side effects, like preserving imports caused
-        # by JobInstance deser
-        ctx = get_context("fork")
+        # NOTE fork would be better but causes issues on macos+torch with XPC_ERROR_CONNECTION_INVALID
+        ctx = get_context("forkserver")
         for worker in workers:
             runnerContext = RunnerContext(
                 workerId=worker,
@@ -198,7 +199,11 @@ class Executor:
                 callback=self.mlistener.address,
                 log_base=self.log_base,
             )
-            p = ctx.Process(target=entrypoint, kwargs={"runnerContext": runnerContext})
+            # NOTE we need to cloudpickle because runnerContext contains some lambdas
+            p = ctx.Process(
+                target=entrypoint,
+                kwargs={"runnerContext": cloudpickle.dumps(runnerContext)},
+            )
             p.start()
             self.workers[worker] = p
             logger.debug(f"started process {p.pid} for worker {worker}")

cascade/executor/platform.py

@@ -8,6 +8,7 @@
 
 """Macos-vs-Linux specific code"""
 
+import os
 import socket
 import sys
 
@@ -22,3 +23,14 @@ def get_bindabble_self():
     else:
         # NOTE not sure if fqdn or hostname is better -- all we need is for it to be resolvable within cluster
         return socket.gethostname()  # socket.getfqdn()
+
+
+def gpu_init(worker_num: int):
+    if sys.platform != "darwin":
+        # TODO there is implicit coupling with executor.executor and benchmarks.main -- make it cleaner!
+        gpus = int(os.environ.get("CASCADE_GPU_COUNT", "0"))
+        os.environ["CUDA_VISIBLE_DEVICES"] = (
+            str(worker_num) if worker_num < gpus else ""
+        )
+    else:
+        pass  # no macos specific gpu init due to unified mem model

cascade/executor/runner/entrypoint.py

@@ -11,11 +11,13 @@
 import logging
 import logging.config
 import os
-import sys
 from dataclasses import dataclass
+from typing import Any
 
+import cloudpickle
 import zmq
 
+import cascade.executor.platform as platform
 import cascade.executor.serde as serde
 from cascade.executor.comms import callback
 from cascade.executor.config import logging_config, logging_config_filehandler
@@ -118,7 +120,9 @@ def execute_sequence(
         )
 
 
-def entrypoint(runnerContext: RunnerContext):
+def entrypoint(runnerContext: Any):
+    """runnerContext is a cloudpickled instance of RunnerContext -- needed for forkserver mp context due to defautdicts"""
+    runnerContext = cloudpickle.loads(runnerContext)
     if runnerContext.log_base:
         log_path = f"{runnerContext.log_base}.{runnerContext.workerId.worker}"
         logging.config.dictConfig(logging_config_filehandler(log_path))
@@ -134,16 +138,8 @@ def entrypoint(runnerContext: RunnerContext):
     ):
         label("worker", repr(runnerContext.workerId))
         worker_num = runnerContext.workerId.worker_num()
-        gpus = int(os.environ.get("CASCADE_GPU_COUNT", "0"))
-        if sys.platform != "darwin":
-            os.environ["CUDA_VISIBLE_DEVICES"] = (
-                str(worker_num) if worker_num < gpus else ""
-            )
-            # NOTE check any(task.definition.needs_gpu) anywhere?
-            # TODO configure OMP_NUM_THREADS, blas, mkl, etc -- not clear how tho
-        else:
-            if gpus != 1:
-                logger.warning("unexpected absence of gpu on darwin")
+        platform.gpu_init(worker_num)
+        # TODO configure OMP_NUM_THREADS, blas, mkl, etc -- not clear how tho
 
         for serdeTypeEnc, (serdeSer, serdeDes) in runnerContext.job.serdes.items():
             serde.SerdeRegistry.register(type_dec(serdeTypeEnc), serdeSer, serdeDes)

cascade/gateway/__main__.py

@@ -14,13 +14,15 @@ from cascade.executor.config import logging_config, logging_config_filehandler
 from cascade.gateway.server import serve
 
 
-def main(url: str, log_base: str | None = None) -> None:
+def main(
+    url: str, log_base: str | None = None, troika_config: str | None = None
+) -> None:
     if log_base:
         log_path = f"{log_base}/gateway.txt"
         logging.config.dictConfig(logging_config_filehandler(log_path))
     else:
         logging.config.dictConfig(logging_config)
-    serve(url, log_base)
+    serve(url, log_base, troika_config)
 
 
 if __name__ == "__main__":

cascade/gateway/api.py

@@ -19,6 +19,18 @@ from cascade.low.core import DatasetId, JobInstance
 CascadeGatewayAPI = BaseModel
 
 
+@dataclass
+class TroikaSpec:
+    """Requires the gateway to have been started with --troika_config pointing
+    to some config.yml troika file. The connection must work (passwordlessly),
+    and must allow for script being copied. The remote host must have a venv
+    already in place, and must be able to resolve gateway's fqdn
+    """
+
+    venv: str  # remote host path to venv -- *do* include the bin/activate
+    conn: str  # which connection from config.yml to pick
+
+
 @dataclass
 class JobSpec:
     # job benchmark + envvars -- set to None/{} if using custom jobs instead
@@ -33,6 +45,7 @@ class JobSpec:
     workers_per_host: int
     hosts: int
     use_slurm: bool
+    troika: TroikaSpec | None = None
 
 
 class SubmitJobRequest(CascadeGatewayAPI):

cascade/gateway/router.py

@@ -8,9 +8,11 @@
 
 """Represents information about submitted jobs. The main business logic of `cascade.gateway`"""
 
+import base64
 import itertools
 import logging
 import os
+import stat
 import subprocess
 import uuid
 from dataclasses import dataclass
@@ -22,7 +24,7 @@ import zmq
 import cascade.executor.platform as platform
 from cascade.controller.report import JobId, JobProgress, JobProgressStarted
 from cascade.executor.comms import get_context
-from cascade.gateway.api import JobSpec
+from cascade.gateway.api import JobSpec, TroikaSpec
 from cascade.low.core import DatasetId
 from cascade.low.func import next_uuid
 
@@ -43,6 +45,65 @@ class Job:
 local_job_port = 12345
 
 
+def _spawn_troika_singlehost(
+    job_spec: JobSpec, addr: str, job_id: JobId, troika: TroikaSpec, troika_config: str
+) -> subprocess.Popen:
+    script = "#!/bin/bash\n"
+    script += f"source {troika.venv}\n"
+    for k, v in job_spec.envvars.items():
+        script += f"export {k}={v}\n"
+    if job_spec.benchmark_name is not None:
+        if job_spec.job_instance is not None:
+            raise TypeError("specified both benchmark name and job instance")
+        script += "python -m cascade.benchmarks local"
+        script += f" --job {job_spec.benchmark_name}"
+    else:
+        if job_spec.job_instance is None:
+            raise TypeError("specified neither benchmark name nor job instance")
+        job_desc_raw = orjson.dumps(job_spec.job_instance.dict())
+        job_desc_enc = base64.b64encode(job_desc_raw).decode("ascii")
+        script += f'JOB_ENC="{job_desc_enc}"'
+        job_json_path = f"/tmp/cascJob.{job_id}.json"
+        script += f'echo "$JOB_ENC" | base64 --decode > {job_json_path}'
+        script += "python -m cascade.benchmarks local"
+        script += f" --instance {job_json_path}"
+
+    script += (
+        f" --workers_per_host {job_spec.workers_per_host} --hosts {job_spec.hosts}"
+    )
+    script += f" --report_address {addr},{job_id}"
+    # NOTE technically not needed to be globally unique, but we cant rely on troika environment isolation...
+    global local_job_port
+    script += f" --port_base {local_job_port}"
+    local_job_port += 1 + job_spec.hosts * job_spec.workers_per_host * 10
+    script += "\n"
+    script_path = f"/tmp/troikascade.{job_id}.sh"
+    with open(script_path, "w") as f:
+        f.write(script)
+    os.chmod(
+        script_path,
+        stat.S_IRUSR
+        | stat.S_IRGRP
+        | stat.S_IROTH
+        | stat.S_IWUSR
+        | stat.S_IXUSR
+        | stat.S_IXGRP
+        | stat.S_IXOTH,
+    )
+    return subprocess.Popen(
+        [
+            "troika",
+            "-c",
+            troika_config,
+            "submit",
+            "-o",
+            f"/tmp/output.{job_id}.txt",
+            troika.conn,
+            script_path,
+        ]
+    )
+
+
 def _spawn_local(
     job_spec: JobSpec, addr: str, job_id: JobId, log_base: str | None
 ) -> subprocess.Popen:
@@ -112,9 +173,26 @@ def _spawn_slurm(job_spec: JobSpec, addr: str, job_id: JobId) -> subprocess.Pope
 
 
 def _spawn_subprocess(
-    job_spec: JobSpec, addr: str, job_id: JobId, log_base: str | None
+    job_spec: JobSpec,
+    addr: str,
+    job_id: JobId,
+    log_base: str | None,
+    troika_config: str | None,
 ) -> subprocess.Popen:
-    if job_spec.use_slurm:
+    if job_spec.troika is not None:
+        if log_base is not None:
+            raise ValueError(f"unexpected {log_base=}")
+        if troika_config is None:
+            raise ValueError("cant spawn troika job without troika config")
+        if not job_spec.use_slurm:
+            return _spawn_troika_singlehost(
+                job_spec, addr, job_id, job_spec.troika, troika_config
+            )
+        else:
+            # TODO create a slurm script like in spawn_slurm, but dont refer to any other file
+            raise NotImplementedError
+
+    elif job_spec.use_slurm:
         if log_base is not None:
             raise ValueError(f"unexpected {log_base=}")
         return _spawn_slurm(job_spec, addr, job_id)
@@ -123,11 +201,14 @@ def _spawn_subprocess(
 
 
 class JobRouter:
-    def __init__(self, poller: zmq.Poller, log_base: str | None):
+    def __init__(
+        self, poller: zmq.Poller, log_base: str | None, troika_config: str | None
+    ):
         self.poller = poller
         self.jobs: dict[str, Job] = {}
         self.procs: dict[str, subprocess.Popen] = {}
         self.log_base = log_base
+        self.troika_config = troika_config
 
     def spawn_job(self, job_spec: JobSpec) -> JobId:
         job_id = next_uuid(self.jobs.keys(), lambda: str(uuid.uuid4()))
@@ -139,7 +220,7 @@ class JobRouter:
         self.poller.register(socket, flags=zmq.POLLIN)
         self.jobs[job_id] = Job(socket, JobProgressStarted, -1, {})
         self.procs[job_id] = _spawn_subprocess(
-            job_spec, full_addr, job_id, self.log_base
+            job_spec, full_addr, job_id, self.log_base, self.troika_config
         )
         return job_id
 

cascade/gateway/server.py

@@ -79,14 +79,16 @@ def handle_controller(socket: zmq.Socket, jobs: JobRouter) -> None:
         jobs.put_result(report.job_id, dataset_id, result)
 
 
-def serve(url: str, log_base: str | None = None) -> None:
+def serve(
+    url: str, log_base: str | None = None, troika_config: str | None = None
+) -> None:
     ctx = get_context()
     poller = zmq.Poller()
 
     fe = ctx.socket(zmq.REP)
     fe.bind(url)
     poller.register(fe, flags=zmq.POLLIN)
-    jobs = JobRouter(poller, log_base)
+    jobs = JobRouter(poller, log_base, troika_config)
 
     logger.debug("entering recv loop")
     is_break = False

cascade/low/dask.py

@@ -0,0 +1,99 @@
+# (C) Copyright 2025- ECMWF.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+
+"""Experimental module to convert dask graphs into cascade jobs. May not preserve
+all semantics.
+
+We don't explicitly support legacy dask graph -- you need to invoke
+`dask._task_spec.convert_legacy_graph` yourself.
+For higher level dask objects, extract the graph via `__dask_graph__()` (defined
+eg on dd.DataFrame or dask.delayed objects).
+"""
+
+import logging
+from typing import Any
+
+from dask._task_spec import Alias, DataNode, Task, TaskRef
+
+from cascade.low.builders import TaskBuilder
+from cascade.low.core import DatasetId, JobInstance, Task2TaskEdge, TaskInstance
+from earthkit.workflows.graph import Node
+
+logger = logging.getLogger(__name__)
+
+
+def daskKeyRepr(key: str | int | float | tuple) -> str:
+    return repr(key)
+
+
+def task2task(key: str, task: Task) -> tuple[TaskInstance, list[Task2TaskEdge]]:
+    instance = TaskBuilder.from_callable(task.func)
+    edges: list[Task2TaskEdge] = []
+
+    for i, v in enumerate(task.args):
+        if isinstance(v, Alias | TaskRef):
+            edge = Task2TaskEdge(
+                source=DatasetId(task=daskKeyRepr(v.key), output=Node.DEFAULT_OUTPUT),
+                sink_task=key,
+                sink_input_ps=str(i),
+                sink_input_kw=None,
+            )
+            edges.append(edge)
+        elif isinstance(v, Task):
+            # TODO
+            raise NotImplementedError
+        else:
+            instance.static_input_ps[f"{i}"] = v
+    for k, v in task.kwargs.items():
+        if isinstance(v, Alias | TaskRef):
+            edge = Task2TaskEdge(
+                source=DatasetId(task=daskKeyRepr(v.key), output=Node.DEFAULT_OUTPUT),
+                sink_task=key,
+                sink_input_kw=k,
+                sink_input_ps=None,
+            )
+            edges.append(edge)
+        elif isinstance(v, Task):
+            # TODO
+            raise NotImplementedError
+        else:
+            instance.static_input_kw[k] = v
+
+    return instance, edges
+
+
+def graph2job(dask: dict) -> JobInstance:
+    task_nodes = {}
+    edges = []
+
+    for node, value in dask.items():
+        key = daskKeyRepr(node)
+        if isinstance(value, DataNode):
+
+            def provider() -> Any:
+                return value.value
+
+            task_nodes[key] = TaskBuilder.from_callable(provider)
+        elif isinstance(value, Task):
+            node, _edges = task2task(key, value)
+            task_nodes[key] = node
+            edges.extend(_edges)
+        elif isinstance(value, list | tuple | set):
+            # TODO implement, consult further:
+            # https://docs.dask.org/en/stable/spec.html
+            # https://docs.dask.org/en/stable/custom-graphs.html
+            # https://github.com/dask/dask/blob/main/dask/_task_spec.py#L829
+            logger.warning("encountered nested container => confused ostrich")
+            continue
+        else:
+            raise NotImplementedError
+
+    return JobInstance(
+        tasks=task_nodes,
+        edges=edges,
+    )

cascade/low/into.py

@@ -6,7 +6,7 @@
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
 
-"""Lowering of the cascade.graph structures into cascade.low representation"""
+"""Lowering of the earthkit.workflows.graph structures into cascade.low representation"""
 
 import logging
 from typing import Any, Callable, cast

earthkit/workflows/_version.py

@@ -1,2 +1,2 @@
 # Do not change! Do not track in version control!
-__version__ = "0.4.3"
+__version__ = "0.4.4"

{earthkit_workflows-0.4.3.dist-info → earthkit_workflows-0.4.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: earthkit-workflows
-Version: 0.4.3
+Version: 0.4.4
 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

{earthkit_workflows-0.4.3.dist-info → earthkit_workflows-0.4.4.dist-info}/RECORD

@@ -1,14 +1,16 @@
 cascade/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 cascade/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 cascade/benchmarks/__init__.py,sha256=Gu8kEApmJ2zsIhT2zpm1-6n84-OwWnz-0vO8UHYtBzo,528
-cascade/benchmarks/__main__.py,sha256=UnIGFykZNw_WsuBbgWIxSEUZB_U4_Yra0VrmyMTEFQ8,9914
+cascade/benchmarks/__main__.py,sha256=z3Ib0NlIgMrn2zjrJhqqnJkjCIb4xKDSpO5vF9j-Onc,966
 cascade/benchmarks/anemoi.py,sha256=qtAI03HdtAmcksCgjIEZyNyUNzMp370KF4lAh5g4cOk,1077
+cascade/benchmarks/dask.py,sha256=U0B0jpLIeIs4Zl0SX_opMypXQXOIS5ER6mGdtPCgqkQ,953
 cascade/benchmarks/dist.py,sha256=ngXJJzegnMUVwDFPvGMG6997lamB-aSEHi74oBbayrE,4116
 cascade/benchmarks/generators.py,sha256=NK4fFisWsZdMkA2Auzrn-P7G5D9AKpo2JVnqXE44YT8,2169
 cascade/benchmarks/job1.py,sha256=MOcZZYgf36MzHCjtby0lQyenM1ODUlagG8wtt2CbpnI,4640
 cascade/benchmarks/matmul.py,sha256=5STuvPY6Q37E2pKRCde9dQjL5M6tx7tkES9cBLZ6eK4,1972
 cascade/benchmarks/plotting.py,sha256=vSz9HHbqZwMXHpBUS-In6xsXGgK7QIoQTTiYfSwYwZs,4428
 cascade/benchmarks/reporting.py,sha256=MejaM-eekbMYLAnuBxGv_t4dR1ODJs4Rpc0fiZSGjyw,5410
+cascade/benchmarks/util.py,sha256=obgRxtRcz023lvrtnI8vzDtgVqlqlRrWCDy_lwuML30,9835
 cascade/controller/__init__.py,sha256=p4C2p3S_0nUGamP9Mi6cSa5bvpiWbI6sVWtGhFnNqjw,1278
 cascade/controller/act.py,sha256=WHIsk4H-Bbyl_DABX2VWhyKy_cNnp12x1nilatPCL8I,2981
 cascade/controller/core.py,sha256=NqvZ5g5GNphwOpzdXbCI0_fxIzzmO97_n2xZKswK72Q,3589
@@ -19,27 +21,28 @@ cascade/executor/bridge.py,sha256=WDE-GM2Bv7nUk1-nV-otMGuaRYw1-Vmd7PWploXBp6Y,82
 cascade/executor/comms.py,sha256=-9qrKwva6WXkHRQtzSnLFy5gB3bOWuxYJP5fL6Uavw8,8736
 cascade/executor/config.py,sha256=8azy_sXdvDGO0zTNqA0pdtkXsyihM4FQ4U1W_3Dhua0,1571
 cascade/executor/data_server.py,sha256=xLIbLkWn8PnJl4lMP8ADHa2S0EgPwr0-bH7_Sib_Y70,13701
-cascade/executor/executor.py,sha256=ufl20jq1lObNONHiYkXDFFCNfv4xKVvFsC_gzhTLdKQ,13524
+cascade/executor/executor.py,sha256=3I9QnyX-YvJvGnMSM4kWfBJDgi_uUCv0M4ncXf4z85o,13659
 cascade/executor/msg.py,sha256=7HI0rKeCRaV1ONR4HWEa64nHbu-p6-QdBwJNitmst48,4340
-cascade/executor/platform.py,sha256=Ym-J6v0Wx56r2z_1ZfK5cX1Zs8eMuu7l1SsNtVvT1Go,959
+cascade/executor/platform.py,sha256=6uLdcH8mibvIQfF1nTSvtfyym4r6dLThNSF1JZ-6mLM,1393
 cascade/executor/serde.py,sha256=z6klTOZqW_BVGrbIRNz4FN0_XTfRiKBRQuvgsQIuyAo,2827
 cascade/executor/runner/__init__.py,sha256=30BM80ZyA7w3IrGiKKLSFuhRehbR2Mm99OJ8q5PJ63c,1547
-cascade/executor/runner/entrypoint.py,sha256=32i2U4fmEvQnsV1MTl0Xf8mK_1nbk1BEVJqIidd6MRM,8042
+cascade/executor/runner/entrypoint.py,sha256=WyxOFGAYDQD_fXsM4H9_6xBrnAmQrCTUnljfcW6-BoM,7918
 cascade/executor/runner/memory.py,sha256=jkAV9T7-imciVcGvkV7OhRfosEpOQJU1OME7z-4ztAs,6371
 cascade/executor/runner/packages.py,sha256=OZjEOvKy8LQ2uguGZU1L7TVYz1415JOUGySRfU_D_sc,2513
 cascade/executor/runner/runner.py,sha256=zqpkvxdWLbwyUFaUbZmSj0KQEBNRpmF8gwVotiaamhc,4870
 cascade/gateway/__init__.py,sha256=1EzMKdLFXEucj0YWOlyVqLx4suOntitwM03T_rRubIk,829
-cascade/gateway/__main__.py,sha256=x6-DQin6ICvalHT9YcghGyVMoykEATOdN5ON9IeHPYA,862
-cascade/gateway/api.py,sha256=-7HTUhK9idszVCwiVwyHMcNx7n6qRcyPWsLx2e19n3A,2511
+cascade/gateway/__main__.py,sha256=F_wft7ja5ckM0SqeXsy_u2j-Ch6OTlpbTTlYtDkvGMI,917
+cascade/gateway/api.py,sha256=-Vuo9fDqFNFIofcHZ79UB1rTWnQR3D9Pna2CjqdyHaE,3021
 cascade/gateway/client.py,sha256=1p4Tvrf-BH0LQHOES5rY1z3JNIfmXcqWG2kYl4rpcE0,4061
-cascade/gateway/router.py,sha256=81kS_ZtWaDV-WLcxaYgsYv-1WHy8aQxYGLmVL5CzOto,7549
-cascade/gateway/server.py,sha256=tsOyKtVFs5EZmWrjKdi9JwWxK0DG207oSa9OQ-4zN3M,3772
+cascade/gateway/router.py,sha256=RcDniyPOZnu6_HuMUrQjZ4P-PoUbUVezvYXG_ryBLUg,10399
+cascade/gateway/server.py,sha256=vb3z0TfoMvSHqczhmYgzeXGVcw2M9yGpyW0t6d57Oag,3827
 cascade/low/__init__.py,sha256=5cw2taOGITK_gFbICftzK2YLdEAnLUY5OzblFzdHss4,769
 cascade/low/builders.py,sha256=_u5X8G_EF00hFt8Anv9AXo6yPf1O8MHDmqs2kKmREl0,7073
 cascade/low/core.py,sha256=_3x4ka_pmCgZbfwFeyhq8S4M6wmh0s24VRCLhk5yQFM,6444
+cascade/low/dask.py,sha256=xToT_vyfkgUUxSFN7dS7qLttxzuBbBZfDylPzGg7sPg,3319
 cascade/low/execution_context.py,sha256=cdDJLYhreo4T7t4qXgFBosncubZpTrm0hELo7q4miqo,6640
 cascade/low/func.py,sha256=ihL5n3cK-IJnATgP4Dub2m-Mp_jHMxJzCA1v4uMEsi8,5211
-cascade/low/into.py,sha256=QvjrcBuHfu7qpEkeB0EJu1EAaRxOEZskUnyjkRJ_9gA,3391
+cascade/low/into.py,sha256=lDOpO4gX-154BgLJWonVQZiGRbUqv-GhYy8qWBqJ1QQ,3402
 cascade/low/tracing.py,sha256=qvGVKB1huwcYoyvMYN-2wQ92pLQTErocTjpIjWv9glA,4511
 cascade/low/views.py,sha256=UwafO2EQHre17GjG8hdzO8b6qBRtTRtDlhOc1pTf8Io,1822
 cascade/scheduler/__init__.py,sha256=VT2qQ0gOQWHC4-T0FcCs59w8WZ94j2nUn7tiGm5XepA,1148
@@ -56,7 +59,7 @@ cascade/shm/disk.py,sha256=Fdl_pKOseaXroRp01OwqWVsdI-sSmiFizIFCdxBuMWM,2653
 cascade/shm/func.py,sha256=ZWikgnSLCmbSoW2LDRJwtjxdwTxkR00OUHAsIRQ-ChE,638
 cascade/shm/server.py,sha256=LnnNX0F6QJt5V_JLfmC3ZMHGNL5WpLY44wpB_pYDr7Y,5042
 earthkit/workflows/__init__.py,sha256=-p4anEn0YQbYWM2tbXb0Vc3wq4-m6kFhcNEgAVu5Jis,1948
-earthkit/workflows/_version.py,sha256=42UO3Zhb1-CO_KjQGniDtgDR8xL_a28l7aAWQcA9o9Y,72
+earthkit/workflows/_version.py,sha256=TdWpE_3Kcp1U42ekVTDBVAxSA3TKm_FWE5S-AZDssXw,72
 earthkit/workflows/decorators.py,sha256=DM4QAtQ2glUUcDecwPkXcdlu4dio7MvgpcdmU5LYvD8,937
 earthkit/workflows/fluent.py,sha256=IN_sqwr7W8wbwP7wTOklgnjVe34IUCmv1ku-DWVTCJc,30179
 earthkit/workflows/mark.py,sha256=PdsXmRfhw1SyyJ74mzFPsLRqMCdlYv556fFX4bqlh9Y,1319
@@ -86,8 +89,8 @@ earthkit/workflows/graph/split.py,sha256=t-Sji5eZb01QO1szqmDNTodDDALqdo-0R0x1ESs
 earthkit/workflows/graph/transform.py,sha256=BZ8n7ePUnuGgoHkMqZC3SLzifu4oq6q6t6vka0khFtg,3842
 earthkit/workflows/graph/visit.py,sha256=MP-aFSqOl7aqJY2i7QTgY4epqb6yM7_lK3ofvOqfahw,1755
 earthkit/workflows/plugins/__init__.py,sha256=nhMAC0eMLxoJamjqB5Ns0OWy0OuxEJ_YvaDFGEQITls,129
-earthkit_workflows-0.4.3.dist-info/licenses/LICENSE,sha256=73MJ7twXMKnWwmzmrMiFwUeY7c6JTvxphVggeUq9Sq4,11381
-earthkit_workflows-0.4.3.dist-info/METADATA,sha256=whowFjotkctfsy1v7mH-CsoCIucWMZpFcf8bqajhmQk,1571
-earthkit_workflows-0.4.3.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-earthkit_workflows-0.4.3.dist-info/top_level.txt,sha256=oNrH3Km3hK5kDkTOiM-8G8OQglvZcy-gUKy7rlooWXs,17
-earthkit_workflows-0.4.3.dist-info/RECORD,,
+earthkit_workflows-0.4.4.dist-info/licenses/LICENSE,sha256=73MJ7twXMKnWwmzmrMiFwUeY7c6JTvxphVggeUq9Sq4,11381
+earthkit_workflows-0.4.4.dist-info/METADATA,sha256=DUjbGRAmD_GJrn0XwYMLFUMpCg8cLlNII4lz7WLWbms,1571
+earthkit_workflows-0.4.4.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+earthkit_workflows-0.4.4.dist-info/top_level.txt,sha256=oNrH3Km3hK5kDkTOiM-8G8OQglvZcy-gUKy7rlooWXs,17
+earthkit_workflows-0.4.4.dist-info/RECORD,,