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

cascade/benchmarks/util.py

@@ -12,7 +12,6 @@
 
 import logging
 import logging.config
-import multiprocessing
 import os
 import subprocess
 import sys
@@ -177,7 +176,7 @@ def run_locally(
             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(
+            p = platform.get_mp_ctx("executor-loc").Process(
                 target=launch_executor,
                 args=(
                     job,

cascade/executor/executor.py

@@ -18,7 +18,6 @@ the tasks themselves.
 import atexit
 import logging
 import os
-from multiprocessing import get_context
 from multiprocessing.process import BaseProcess
 from typing import Iterable
 
@@ -97,12 +96,12 @@ class Executor:
         # TODO make the shm server params configurable
         shm_port = portBase + 2
         shm_api.publish_client_port(shm_port)
-        ctx = get_context("fork")
+        ctx = platform.get_mp_ctx("executor-aux")
         if log_base:
             shm_logging = logging_config_filehandler(f"{log_base}.shm.txt")
         else:
             shm_logging = logging_config
-        logger.debug("about to fork into shm process")
+        logger.debug("about to start an shm process")
         self.shm_process = ctx.Process(
             target=shm_server,
             args=(
@@ -118,7 +117,7 @@ class Executor:
             dsr_logging = logging_config_filehandler(f"{log_base}.dsr.txt")
         else:
             dsr_logging = logging_config
-        logger.debug("about to fork into data server")
+        logger.debug("about to start a data server process")
         self.data_server = ctx.Process(
             target=start_data_server,
             args=(
@@ -190,7 +189,7 @@ 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 fork would be better but causes issues on macos+torch with XPC_ERROR_CONNECTION_INVALID
-        ctx = get_context("forkserver")
+        ctx = platform.get_mp_ctx("worker")
         for worker in workers:
             runnerContext = RunnerContext(
                 workerId=worker,

cascade/executor/platform.py

@@ -8,9 +8,11 @@
 
 """Macos-vs-Linux specific code"""
 
+import multiprocessing as mp
 import os
 import socket
 import sys
+import typing
 
 
 def get_bindabble_self():
@@ -34,3 +36,28 @@ def gpu_init(worker_num: int):
         )
     else:
         pass  # no macos specific gpu init due to unified mem model
+
+
+MpSituation = typing.Literal[
+    "worker", "executor-loc", "executor-aux", "gateway", "other"
+]
+_MpSituation = typing.get_args(MpSituation)
+
+
+def get_mp_ctx(situation: MpSituation) -> mp.context.BaseContext:
+    """Generally, forking is safe everywhere as we try to be careful not to
+    initialize non-safe objects prior to forking. However, combination of
+    mac + mps + anemoi + pickled callables causes xpc_error_connection_invalid,
+    thus we stick to spawn on darwin platforms. Setting
+    OBJC_DISABLE_INITIALIZE_FORK_SAFETY=YES helps some but not fully.
+
+    We distinguish in which situation is this method called, as fine graining
+    may be (eventually) possible
+    """
+
+    if situation not in _MpSituation:
+        raise TypeError(f"{situation=} is not in {_MpSituation}")
+    if sys.platform == "darwin":
+        return mp.get_context("spawn")
+    else:
+        return mp.get_context("fork")

cascade/executor/runner/packages.py

@@ -42,7 +42,15 @@ class PackagesEnv(AbstractContextManager):
         logger.debug(
             f"installing {len(packages)} packages: {','.join(packages[:3])}{',...' if len(packages) > 3 else ''}"
         )
-        install_command = ["uv", "pip", "install", "--prefix", self.td.name]
+        install_command = [
+            "uv",
+            "pip",
+            "install",
+            "--prefix",
+            self.td.name,
+            "--prerelease",
+            "allow",
+        ]
         if os.environ.get("VENV_OFFLINE", "") == "YES":
             install_command += ["--offline"]
         if cache_dir := os.environ.get("VENV_CACHE", ""):

earthkit/workflows/_version.py

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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: earthkit-workflows
-Version: 0.4.4
+Version: 0.4.6
 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.4.dist-info → earthkit_workflows-0.4.6.dist-info}/RECORD

@@ -10,7 +10,7 @@ cascade/benchmarks/job1.py,sha256=MOcZZYgf36MzHCjtby0lQyenM1ODUlagG8wtt2CbpnI,46
 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/benchmarks/util.py,sha256=ivnd4olB10_F2BCCJan4gzYiyZXMn3gA2iNbacbu464,9812
 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
@@ -21,14 +21,14 @@ 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=3I9QnyX-YvJvGnMSM4kWfBJDgi_uUCv0M4ncXf4z85o,13659
+cascade/executor/executor.py,sha256=htHobVN0eAxofii2g29yN--qCkaQgA05MVKCgCW7Dhg,13644
 cascade/executor/msg.py,sha256=7HI0rKeCRaV1ONR4HWEa64nHbu-p6-QdBwJNitmst48,4340
-cascade/executor/platform.py,sha256=6uLdcH8mibvIQfF1nTSvtfyym4r6dLThNSF1JZ-6mLM,1393
+cascade/executor/platform.py,sha256=mRUauodvRle9rAbtFr5n9toKzIgt_pecNlhOjon4dvY,2348
 cascade/executor/serde.py,sha256=z6klTOZqW_BVGrbIRNz4FN0_XTfRiKBRQuvgsQIuyAo,2827
 cascade/executor/runner/__init__.py,sha256=30BM80ZyA7w3IrGiKKLSFuhRehbR2Mm99OJ8q5PJ63c,1547
 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/packages.py,sha256=zOMCKRX34j6vSIK741KV7kTFRgy3bZz_xeLrldslaL4,2633
 cascade/executor/runner/runner.py,sha256=zqpkvxdWLbwyUFaUbZmSj0KQEBNRpmF8gwVotiaamhc,4870
 cascade/gateway/__init__.py,sha256=1EzMKdLFXEucj0YWOlyVqLx4suOntitwM03T_rRubIk,829
 cascade/gateway/__main__.py,sha256=F_wft7ja5ckM0SqeXsy_u2j-Ch6OTlpbTTlYtDkvGMI,917
@@ -59,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=TdWpE_3Kcp1U42ekVTDBVAxSA3TKm_FWE5S-AZDssXw,72
+earthkit/workflows/_version.py,sha256=9B6r-UnSSKU-ri_i5lTQQPMIYjzVbLs7Yazh1HuJ9G8,72
 earthkit/workflows/decorators.py,sha256=DM4QAtQ2glUUcDecwPkXcdlu4dio7MvgpcdmU5LYvD8,937
 earthkit/workflows/fluent.py,sha256=IN_sqwr7W8wbwP7wTOklgnjVe34IUCmv1ku-DWVTCJc,30179
 earthkit/workflows/mark.py,sha256=PdsXmRfhw1SyyJ74mzFPsLRqMCdlYv556fFX4bqlh9Y,1319
@@ -89,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.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,,
+earthkit_workflows-0.4.6.dist-info/licenses/LICENSE,sha256=73MJ7twXMKnWwmzmrMiFwUeY7c6JTvxphVggeUq9Sq4,11381
+earthkit_workflows-0.4.6.dist-info/METADATA,sha256=5fpGN2qME_zCzZhyiPrixuMxzQjLS0PRwbWjppj-Zww,1571
+earthkit_workflows-0.4.6.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+earthkit_workflows-0.4.6.dist-info/top_level.txt,sha256=oNrH3Km3hK5kDkTOiM-8G8OQglvZcy-gUKy7rlooWXs,17
+earthkit_workflows-0.4.6.dist-info/RECORD,,