earthkit-workflows 0.3.6__py3-none-any.whl → 0.4.1__py3-none-any.whl

cascade/benchmarks/__main__.py

@@ -26,7 +26,9 @@ import logging.config
 import multiprocessing
 import os
 import subprocess
+import sys
 from concurrent.futures import ThreadPoolExecutor
+from socket import getfqdn
 from time import perf_counter_ns
 
 import fire
@@ -41,7 +43,7 @@ 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.graph import precompute
+from cascade.scheduler.precompute import precompute
 from earthkit.workflows.graph import Graph, deduplicate_nodes
 
 logger = logging.getLogger("cascade.benchmarks")
@@ -73,14 +75,28 @@ def get_job(benchmark: str | None, instance_path: str | None) -> JobInstance:
             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_gpu_count() -> int:
+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(
@@ -91,12 +107,22 @@ def get_gpu_count() -> int:
             if "GPU" in l
         )
     except:
-        # TODO support macos
         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,
@@ -106,6 +132,7 @@ def launch_executor(
     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}"
@@ -123,6 +150,7 @@ def launch_executor(
         portBase,
         shm_vol_gb,
         log_base,
+        url_base,
     )
     executor.register()
     executor.recv_loop()
@@ -147,14 +175,21 @@ def run_locally(
     m = f"tcp://localhost:{portBase+1}"
     ps = []
     for i, executor in enumerate(range(hosts)):
-        if i == 0:
-            gpu_count = get_gpu_count()
-        else:
-            gpu_count = 0
+        gpu_count = get_gpu_count(i, workers)
         # NOTE forkserver/spawn seem to forget venv, we need fork
         p = multiprocessing.get_context("fork").Process(
             target=launch_executor,
-            args=(job, c, workers, portBase + 1 + i * 10, i, None, gpu_count, log_base),
+            args=(
+                job,
+                c,
+                workers,
+                portBase + 1 + i * 10,
+                i,
+                None,
+                gpu_count,
+                log_base,
+                "tcp://localhost",
+            ),
         )
         p.start()
         ps.append(p)
@@ -228,7 +263,7 @@ def main_dist(
             f"compute took {(end-start)/1e9:.3f}s, including startup {(end-launch)/1e9:.3f}s"
         )
     else:
-        gpu_count = get_gpu_count()
+        gpu_count = get_gpu_count(0, workers_per_host)
         launch_executor(
             jobInstance,
             controller_url,
@@ -237,6 +272,7 @@ def main_dist(
             idx,
             shm_vol_gb,
             gpu_count,
+            f"tcp://{getfqdn()}",
         )
 
 

cascade/benchmarks/dist.py

@@ -0,0 +1,123 @@
+"""Demonstrates gang scheduling capabilities, ie, multiple nodes capable of mutual communication.
+
+The job is a source -> (dist group) -> sink, where:
+    source just returns an int,
+    dist group is L nodes to be scheduled as a single gang
+        rank=0 node broadcasts a buffer containing the node's input
+        each node returns its input multiplied by broadcasted buffer
+    sink returns the sum of all inputs
+
+There are multiple implementations of that:
+    torch
+    jax (actually does a mesh-shard global sum instead of broadcast -- the point is to showcase dist init)
+"""
+
+import os
+
+from cascade.low.builders import JobBuilder, TaskBuilder
+from cascade.low.core import JobInstance, SchedulingConstraint
+
+
+def source_func() -> int:
+    return 42
+
+
+def dist_func_torch(a: int) -> int:
+    import datetime as dt
+
+    import numpy as np
+    import torch.distributed as dist
+
+    world_size = int(os.environ["CASCADE_GANG_WORLD_SIZE"])
+    rank = int(os.environ["CASCADE_GANG_RANK"])
+    coordinator = os.environ["CASCADE_GANG_COORDINATOR"]
+    print(f"starting with envvars: {rank=}/{world_size=}, {coordinator=}")
+    dist.init_process_group(
+        backend="gloo",
+        init_method=coordinator,
+        timeout=dt.timedelta(minutes=1),
+        world_size=world_size,
+        rank=rank,
+    )
+    group_ranks = np.arange(world_size, dtype=int)
+    group = dist.new_group(group_ranks)
+
+    if rank == 0:
+        buf = [a]
+        dist.broadcast_object_list(buf, src=0, group=group)
+        print("broadcast ok")
+    else:
+        buf = np.array([0], dtype=np.uint64)
+        dist.broadcast_object_list(buf, src=0, group=group)
+        print(f"broadcast recevied {buf}")
+
+    return a * buf[0]
+
+
+def dist_func_jax(a: int) -> int:
+    world_size = int(os.environ["CASCADE_GANG_WORLD_SIZE"])
+    rank = int(os.environ["CASCADE_GANG_RANK"])
+    coordinator = os.environ["CASCADE_GANG_COORDINATOR"]
+    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
+
+    jax.config.update("jax_platforms", "cpu")
+    jax.config.update("jax_platform_name", "cpu")
+    # NOTE neither of the above seems to actually help with an init error message :(
+    print(f"starting with envvars: {rank=}/{world_size=}, {coordinator=}")
+    if coordinator.startswith("tcp://"):
+        coordinator = coordinator[len("tcp://") :]
+    jax.distributed.initialize(coordinator, num_processes=world_size, process_id=rank)
+    assert jax.device_count() == world_size
+
+    mesh = jax.make_mesh((world_size,), ("i",))
+    global_data = jp.arange(world_size)
+    sharding = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec("i"))
+    global_array = jax.device_put(global_data, sharding)
+    result = jp.sum(global_array)
+    print(f"worker {rank}# got result {result=}")
+    return a + result
+
+
+def build_dist_func(impl: str):
+    if impl == "torch":
+        return dist_func_torch
+    elif impl == "jax":
+        return dist_func_jax
+    else:
+        raise NotImplementedError(impl)
+
+
+def sink_func(**kwargs) -> int:
+    c = 0
+    for _, v in kwargs.items():
+        c += v
+    print(f"sink accumulated {c}")
+    return c
+
+
+def get_job() -> JobInstance:
+    source_node = TaskBuilder.from_callable(source_func)
+    sink_node = TaskBuilder.from_callable(sink_func)
+    job = JobBuilder().with_node("source", source_node).with_node("sink", sink_node)
+    L = int(os.environ["DIST_L"])
+    IMPL = os.environ["DIST_IMPL"]
+    node = TaskBuilder.from_callable(build_dist_func(IMPL))
+
+    for i in range(L):
+        job = (
+            job.with_node(f"proc{i}", node)
+            .with_edge("source", f"proc{i}", "a")
+            .with_edge(f"proc{i}", "sink", f"v{i}")
+        )
+        job.nodes["sink"].definition.input_schema[
+            f"v{i}"
+        ] = "int"  # TODO put some allow_kw into TaskDefinition instead to allow this
+
+    job = job.build().get_or_raise()
+    job.ext_outputs = list(job.outputs_of("sink"))
+    job.constraints = [SchedulingConstraint(gang=[f"proc{i}" for i in range(L)])]
+    return job

cascade/benchmarks/job1.py

@@ -16,10 +16,10 @@ Controlled by env var params: JOB1_{DATA_ROOT, GRID, ...}, see below
 import os
 
 import earthkit.data
-from ppcascade.fluent import from_source
-from ppcascade.utils.window import Range
 
 from earthkit.workflows.fluent import Payload
+from earthkit.workflows.plugins.pproc.fluent import from_source
+from earthkit.workflows.plugins.pproc.utils.window import Range
 
 # *** PARAMS ***
 

cascade/benchmarks/matmul.py

@@ -0,0 +1,73 @@
+import os
+from typing import Any
+
+import jax
+import jax.numpy as jp
+import jax.random as jr
+
+from cascade.low.builders import JobBuilder, TaskBuilder
+from cascade.low.core import JobInstance
+
+
+def get_funcs():
+    K = int(os.environ["MATMUL_K"])
+    size = (2**K, 2**K)
+    E = int(os.environ["MATMUL_E"])
+
+    def source() -> Any:
+        k0 = jr.key(0)
+        m = jr.uniform(key=k0, shape=size)
+        return m
+
+    def powr(m: Any) -> Any:
+        print(f"powr device is {m.device}")
+        return m**E * jp.percentile(m, 0.7)
+
+    return source, powr
+
+
+def get_job() -> JobInstance:
+    L = int(os.environ["MATMUL_L"])
+    # D = os.environ["MATMUL_D"]
+    # it would be tempting to with jax.default_device(jax.devices(D)):
+    # alas, it doesn't work because we can't inject this at deser time
+
+    source, powr = get_funcs()
+    source_node = TaskBuilder.from_callable(source)
+    if os.environ.get("CUDA_VISIBLE_DEVICES", "") != "":
+        source_node.definition.needs_gpu = True
+    # currently no need to set True downstream since scheduler prefers no transfer
+
+    job = JobBuilder().with_node("source", source_node)
+    prv = "source"
+    for i in range(L):
+        cur = f"pow{i}"
+        node = TaskBuilder.from_callable(powr)
+        job = job.with_node(cur, node).with_edge(prv, cur, 0)
+        prv = cur
+
+    job = job.build().get_or_raise()
+    job.ext_outputs = list(job.outputs_of(cur))
+    return job
+
+
+def execute_locally():
+    L = int(os.environ["MATMUL_L"])
+
+    source, powr = get_funcs()
+
+    device = "gpu" if os.environ.get("CUDA_VISIBLE_DEVICES", "") != "" else "cpu"
+    print(f"device is {device}")
+    with jax.default_device(jax.devices(device)[0]):
+        m0 = source()
+        for _ in range(L):
+            m0 = powr(m0)
+
+    from multiprocessing.shared_memory import SharedMemory
+
+    mem = SharedMemory("benchmark_tmp", create=True, size=m0.nbytes)
+    mem.buf[:] = m0.tobytes()
+
+
+if __name__ == "__main__":
+    execute_locally()

cascade/controller/act.py

@@ -51,6 +51,7 @@ def act(bridge: Bridge, assignment: Assignment) -> None:
         worker=assignment.worker,
         tasks=assignment.tasks,
         publish=assignment.outputs,
+        extra_env=assignment.extra_env,
     )
 
     for task in assignment.tasks:

cascade/controller/impl.py

@@ -43,6 +43,11 @@ def run(
     reporter = Reporter(report_address)
 
     try:
+        total_gpus = sum(worker.gpu for worker in env.workers.values())
+        needs_gpus = any(task.definition.needs_gpu for task in job.tasks.values())
+        if needs_gpus and total_gpus == 0:
+            raise ValueError("environment contains no gpu yet job demands one")
+
         while (
             state.has_awaitable()
             or context.has_awaitable()

cascade/controller/notify.py

@@ -22,6 +22,7 @@ from cascade.low.core import DatasetId, HostId, WorkerId
 from cascade.low.execution_context import DatasetStatus, JobExecutionContext
 from cascade.low.func import assert_never
 from cascade.low.tracing import TaskLifecycle, TransmitLifecycle, mark
+from cascade.scheduler.api import gang_check_ready
 from cascade.scheduler.assign import set_worker2task_overhead
 from cascade.scheduler.core import Schedule
 
@@ -67,6 +68,7 @@ def consider_computable(
                     # NOTE this is a task newly made computable, so we need to calc
                     # `overhead` for all hosts/workers assigned to the component
                     set_worker2task_overhead(schedule, context, worker, child_task)
+                gang_check_ready(child_task, component.gang_preparation)
 
 
 # TODO refac less explicit mutation of context, use class methods

cascade/executor/bridge.py

@@ -46,7 +46,7 @@ class Bridge:
         self.transmit_idx_counter = 0
         self.sender = ReliableSender(self.mlistener.address, resend_grace_ms)
         registered = 0
-        self.environment = Environment(workers={})
+        self.environment = Environment(workers={}, host_url_base={})
         logger.debug("about to start receiving registrations")
         registration_grace = time.time_ns() + 3 * 60 * 1_000_000_000
         while registered < expected_executors:
@@ -69,6 +69,7 @@ class Bridge:
                     self.environment.workers[worker.worker_id] = Worker(
                         cpu=worker.cpu, gpu=worker.gpu, memory_mb=worker.memory_mb
                     )
+                self.environment.host_url_base[message.host] = message.url_base
                 registered += 1
                 self.heartbeat_checker[message.host] = GraceWatcher(
                     2 * executor_heartbeat_grace_ms

cascade/executor/config.py

@@ -27,6 +27,7 @@ logging_config = {
         "cascade.executor": {"level": "DEBUG"},
         "cascade.scheduler": {"level": "DEBUG"},
         "cascade.gateway": {"level": "DEBUG"},
+        "earthkit.workflows": {"level": "DEBUG"},
         "httpcore": {"level": "ERROR"},
         "httpx": {"level": "ERROR"},
         "": {"level": "WARNING", "handlers": ["default"]},

cascade/executor/executor.py

@@ -69,8 +69,9 @@ class Executor:
         workers: int,
         host: HostId,
         portBase: int,
-        shm_vol_gb: int | None = None,
-        log_base: str | None = None,
+        shm_vol_gb: int | None,
+        log_base: str | None,
+        url_base: str,
     ) -> None:
         self.job_instance = job_instance
         self.param_source = param_source(job_instance.edges)
@@ -138,6 +139,7 @@ class Executor:
                 )
                 for idx, worker_id in enumerate(self.workers.keys())
             ],
+            url_base=url_base,
         )
         logger.debug("constructed executor")
 

cascade/executor/msg.py

@@ -71,6 +71,7 @@ class TaskSequence:
     worker: WorkerId  # worker for running those tasks
     tasks: list[TaskId]  # to be executed in the given order
     publish: set[DatasetId]  # set of outputs to be published
+    extra_env: list[tuple[str, str]]  # extra env var to set
 
 
 @dataclass(frozen=True)
@@ -147,6 +148,7 @@ class ExecutorRegistration:
     host: HostId
     maddress: BackboneAddress
     daddress: BackboneAddress
+    url_base: str  # used for eg dist comms init
     workers: list[Worker]
 
 

cascade/executor/runner/entrypoint.py

@@ -11,6 +11,7 @@
 import logging
 import logging.config
 import os
+import sys
 from dataclasses import dataclass
 
 import zmq
@@ -67,6 +68,25 @@ class RunnerContext:
         )
 
 
+class Config:
+    """Some parameters to drive behaviour. Currently not exposed externally -- no clear argument
+    that they should be. As is, just a means of code experimentation.
+    """
+
+    # flushing approach -- when we finish a computation of task sequence, there is a question what
+    # to do with the output. We could either publish & drop, or publish and retain in memory. The
+    # former is is slower -- if the next task sequence needs this output, it requires a fetch & deser
+    # from cashme. But the latter is more risky -- we effectively have the same dataset twice in
+    # system memory. The `posttask_flush` below goes the former way, the `pretask_flush` is a careful
+    # way of latter -- we drop the output from memory only if the *next* task sequence does not need
+    # it, ie, we retain a cache of age 1. We could ultimately have controller decide about this, or
+    # decide dynamically based on memory pressure -- but neither is easy.
+    posttask_flush = False  # after task is done, drop all outputs from memory
+    pretask_flush = (
+        True  # when we receive a task, we drop those in memory that wont be needed
+    )
+
+
 def worker_address(workerId: WorkerId) -> BackboneAddress:
     return f"ipc:///tmp/{repr(workerId)}.socket"
 
@@ -79,11 +99,17 @@ def execute_sequence(
 ) -> None:
     taskId: TaskId | None = None
     try:
+        for key, value in taskSequence.extra_env.items():
+            os.environ[key] = value
         executionContext = runnerContext.project(taskSequence)
         for taskId in taskSequence.tasks:
             pckg.extend(executionContext.tasks[taskId].definition.environment)
             run(taskId, executionContext, memory)
-        memory.flush()
+        if Config.posttask_flush:
+            memory.flush()
+        for key in taskSequence.extra_env.keys():
+            # NOTE we should in principle restore the previous value, but we dont expect collisions
+            del os.environ[key]
     except Exception as e:
         logger.exception("runner failure, about to report")
         callback(
@@ -107,10 +133,17 @@ def entrypoint(runnerContext: RunnerContext):
         PackagesEnv() as pckg,
     ):
         label("worker", repr(runnerContext.workerId))
-        gpu_id = str(runnerContext.workerId.worker_num())
-        os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(gpu_id)
-        # NOTE check any(task.definition.needs_gpu) anywhere?
-        # TODO configure OMP_NUM_THREADS, blas, mkl, etc -- not clear how tho
+        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")
 
         for serdeTypeEnc, (serdeSer, serdeDes) in runnerContext.job.serdes.items():
             serde.SerdeRegistry.register(type_dec(serdeTypeEnc), serdeSer, serdeDes)
@@ -151,6 +184,9 @@ def entrypoint(runnerContext: RunnerContext):
                     for key, _ in runnerContext.job.tasks[task].definition.output_schema
                 }
                 missing_ds = required - availab_ds
+                if Config.pretask_flush:
+                    extraneous_ds = availab_ds - required
+                    memory.flush(extraneous_ds)
                 if missing_ds:
                     waiting_ts = mDes
                     for ds in availab_ds.intersection(required):

cascade/executor/runner/memory.py

@@ -51,7 +51,6 @@ class Memory(AbstractContextManager):
             else:
                 outputValue = "ok"
 
-        # TODO how do we purge from here over time?
         self.local[outputId] = outputValue
 
         if isPublish:
@@ -68,6 +67,18 @@ class Memory(AbstractContextManager):
                 self.callback,
                 DatasetPublished(ds=outputId, origin=self.worker, transmit_idx=None),
             )
+        else:
+            # NOTE even if its not actually published, we send the message to allow for
+            # marking the task itself as completed -- its odd, but arguably better than
+            # introducing a TaskCompleted message. TODO we should fine-grain host-wide
+            # and worker-only publishes at the `controller.notify` level, to not cause
+            # incorrect shm.purge calls at worklow end, which log an annoying key error
+            logger.debug(f"fake publish of {outputId} for the sake of task completion")
+            shmid = ds2shmid(outputId)
+            callback(
+                self.callback,
+                DatasetPublished(ds=outputId, origin=self.worker, transmit_idx=None),
+            )
 
     def provide(self, inputId: DatasetId, annotation: str) -> Any:
         if inputId not in self.local:
@@ -85,18 +96,24 @@ class Memory(AbstractContextManager):
 
     def pop(self, ds: DatasetId) -> None:
         if ds in self.local:
+            logger.debug(f"popping local {ds}")
             val = self.local.pop(ds)  # noqa: F841
             del val
         if ds in self.bufs:
+            logger.debug(f"popping buf {ds}")
             buf = self.bufs.pop(ds)
             buf.close()
 
-    def flush(self) -> None:
-        # NOTE poor man's memory management -- just drop those locals that weren't published. Called
+    def flush(self, datasets: set[DatasetId] = set()) -> None:
+        # NOTE poor man's memory management -- just drop those locals that didn't come from cashme. Called
         # after every taskSequence. In principle, we could purge some locals earlier, and ideally scheduler
         # would invoke some targeted purges to also remove some published ones earlier (eg, they are still
         # needed somewhere but not here)
-        purgeable = [inputId for inputId in self.local if inputId not in self.bufs]
+        purgeable = [
+            inputId
+            for inputId in self.local
+            if inputId not in self.bufs and (not datasets or inputId in datasets)
+        ]
         logger.debug(f"will flush {len(purgeable)} datasets")
         for inputId in purgeable:
             self.local.pop(inputId)
@@ -115,6 +132,8 @@ class Memory(AbstractContextManager):
                     free, total = torch.cuda.mem_get_info()
                     logger.debug(f"cuda mem avail post cache empty: {free/total:.2%}")
                     if free / total < 0.8:
+                        # NOTE this ofc makes low sense if there is any other application (like browser or ollama)
+                        # that the user may be running
                         logger.warning("cuda mem avail low despite cache empty!")
                         logger.debug(torch.cuda.memory_summary())
         except ImportError:

cascade/low/core.py

@@ -106,15 +106,26 @@ def type_enc(t: Type) -> str:
     return b64encode(cloudpickle.dumps(t)).decode("ascii")
 
 
+class SchedulingConstraint(BaseModel):
+    gang: list[TaskId] = Field(
+        description="this set of TaskIds must be started at the same time, with ranks and address list as envvar",
+    )
+
+
 class JobInstance(BaseModel):
     tasks: dict[TaskId, TaskInstance]
     edges: list[Task2TaskEdge]
     serdes: dict[str, tuple[str, str]] = Field(
-        {},
+        default_factory=lambda: {},
         description="for each Type with custom serde, add entry here. The string is fully qualified name of the ser/des functions",
     )
     ext_outputs: list[DatasetId] = Field(
-        [], description="ids to externally materialize"
+        default_factory=lambda: [],
+        description="ids to externally materialize",
+    )
+    constraints: list[SchedulingConstraint] = Field(
+        default_factory=lambda: [],
+        description="constraints for the scheduler such as gangs",
     )
 
     def outputs_of(self, task_id: TaskId) -> set[DatasetId]:
@@ -157,6 +168,7 @@ class Worker(BaseModel):
 
 class Environment(BaseModel):
     workers: dict[WorkerId, Worker]
+    host_url_base: dict[HostId, str]
 
 
 class TaskExecutionRecord(BaseModel):

cascade/low/execution_context.py

@@ -108,6 +108,12 @@ class JobExecutionContext:
             self.idle_workers.add(worker)
 
     def dataset_preparing(self, dataset: DatasetId, worker: WorkerId) -> None:
+        # NOTE Currently this is invoked during `build_assignment`, as we need
+        # some state tranisition to allow fusing opportunities as well as
+        # preventing double transmits. This may not be the best idea, eg for long
+        # fusing chains -- instead, we may execute this transition at the time
+        # it actually happens, granularize the preparing state into
+        # (will_appear, is_appearing), etc
         # NOTE Currently, these `if`s are necessary because we issue transmit
         # command when host *has* DS but worker does *not*. This ends up no-op,
         # but we totally dont want host state to reset -- it wouldnt recover