ifed-server-abi 2.0.2__py3-none-any.whl

ifed_server_abi/__init__.py

@@ -0,0 +1,297 @@
+"""IFED runtime ABI — the executor-side contract a compiled server plan is loaded through.
+
+It never imports ``ifed``: the SDK travels inside the plan bundle, this package does not.
+``torch`` must be importable here but is deliberately undeclared, so a torch bump alone
+never forces an ABI release.
+"""
+
+import json
+import os
+import sys
+
+__all__ = [
+  "CALLBACKS_NAME",
+  "MAX_BUNDLE_BYTES",
+  "NAMES",
+  "BY_REFERENCE",
+  "ServerAbi",
+  "abi_tag",
+  "channel",
+  "check_phases",
+  "python_tag",
+  "torch_tag",
+  "close_channel",
+  "open_channel",
+  "run_phase",
+  "save",
+]
+
+CALLBACKS_NAME = "callbacks.pkl"
+
+NAMES = ("prepare", "aggregate", "finalize")
+
+# the bundle carries code, not data, so anything this big is not a plan
+MAX_BUNDLE_BYTES = 64 * 1024 * 1024
+
+# Top-level module names this ABI guarantees on the reader, so they may be named rather
+# than embedded; everything else, the SDK included, travels inside the bundle. Part of the
+# contract: changing it is a major release of this package. Import names, not distribution
+# names, because
+# that is what cloudpickle keys on and because packages_distributions() does not map
+# editable installs. torch and numpy cannot be pickled by value at all, and cloudpickle
+# pickling itself would recurse.
+BY_REFERENCE = frozenset({
+  "cloudpickle", "ifed_server_abi", "numpy", "torch",
+})
+
+_CHANNEL = None
+
+
+def open_channel() -> dict:
+  global _CHANNEL
+  _CHANNEL = {}
+  return _CHANNEL
+
+
+def close_channel():
+  global _CHANNEL
+  closed, _CHANNEL = _CHANNEL, None
+  return closed
+
+
+def channel():
+  """The scratch dict of the phase in flight, or None. Keys are namespaced by their owner."""
+  return _CHANNEL
+
+
+class ServerAbi:
+  """A compiled server plan: three adapters that move task state between paths."""
+
+  def __init__(self, bundle):
+    self.bundle_path = os.fspath(bundle)
+    self.plan_dir = os.path.realpath(os.path.dirname(os.path.abspath(self.bundle_path)))
+    self._phases = _load(self.bundle_path)
+
+  @property
+  def phases(self) -> dict:
+    """The phase adapters by name. A copy, so the cached bundle stays immutable."""
+    return dict(self._phases)
+
+  def prepare(self, server_state, agent_state, inter_state, sink=None) -> None:
+    self._run("prepare", {
+      "server_state_path": os.fspath(server_state),
+      "agent_state_path": os.fspath(agent_state),
+      "inter_state_path": os.fspath(inter_state),
+    }, sink)
+
+  def aggregate(self, agent_states, aggr_state, sink=None) -> None:
+    self._run("aggregate", {
+      "agent_state_paths": [os.fspath(p) for p in agent_states],
+      "aggr_state_path": os.fspath(aggr_state),
+    }, sink)
+
+  def finalize(self, aggr_state, inter_state, server_state, sink=None) -> None:
+    self._run("finalize", {
+      "aggr_state_path": os.fspath(aggr_state),
+      "inter_state_path": os.fspath(inter_state),
+      "server_state_path": os.fspath(server_state),
+    }, sink)
+
+  def _run(self, phase, paths, sink) -> None:
+    scratch = open_channel()
+    try:
+      # plan_dir is how the adapter finds its own artifacts without the ABI naming them
+      self._phases[phase](self.plan_dir, paths, scratch, sink)
+    finally:
+      # a leaked channel would cross-contaminate the next phase of a pooled worker
+      close_channel()
+
+
+def run_phase(phase, plan_dir, context_json) -> None:
+  """The entry point the embedded-CPython bridge calls, once per phase.
+
+  ``plan_dir`` is the unpacked server plan; the context JSON names only state paths.
+  """
+  if phase not in NAMES:
+    raise ValueError(f"Unknown phase: {phase}")
+  context = json.loads(context_json)
+  abi = ServerAbi(os.path.join(plan_dir, CALLBACKS_NAME))
+  sink = context.get("metrics_sink_path")
+  if phase == "prepare":
+    abi.prepare(
+      context["server_state_path"],
+      context["agent_state_path"],
+      context["inter_state_path"],
+      sink,
+    )
+  elif phase == "aggregate":
+    abi.aggregate(context["agent_state_paths"], context["aggr_state_path"], sink)
+  else:
+    abi.finalize(
+      context["aggr_state_path"],
+      context["inter_state_path"],
+      context["server_state_path"],
+      sink,
+    )
+
+
+def abi_tag() -> str:
+  """This package's version. Compatibility is keyed on major, so a patch or minor
+  release never invalidates a bundle already compiled against it."""
+  from importlib.metadata import version
+
+  return version("ifed-server-abi")
+
+
+def _abi_major(tag) -> str:
+  return str(tag).split(".")[0]
+
+
+def python_tag() -> str:
+  """The interpreter minor. Bundles carry code objects, which are minor-specific."""
+  return f"{sys.version_info.major}.{sys.version_info.minor}"
+
+
+def torch_tag() -> str:
+  """The torch minor that compiled the plan's TorchScript. Patch level does not matter."""
+  import torch
+
+  return ".".join(torch.__version__.split(".")[:2])
+
+
+def save(phases: dict, path) -> None:
+  """Write a plan bundle.
+
+  Everything outside BY_REFERENCE travels inside the bundle: the writer's own
+  installation cannot answer what is installed on the reading side.
+
+  Writing lives here so the envelope and the version marker come from the same place
+  that validates them on read.
+  """
+  import cloudpickle
+
+  check_phases(phases, f"refusing to write '{path}'")
+  embedded = _embedded_modules()
+  for module in embedded:
+    cloudpickle.register_pickle_by_value(module)
+  try:
+    payload = cloudpickle.dumps(
+      {
+        "abi": abi_tag(),
+        "python": python_tag(),
+        "torch": torch_tag(),
+        "phases": dict(phases),
+      },
+      protocol=4,
+    )
+  finally:
+    for module in embedded:
+      cloudpickle.unregister_pickle_by_value(module)
+  parent = os.path.dirname(os.path.abspath(path))
+  if parent:
+    os.makedirs(parent, exist_ok=True)
+  with open(path, "wb") as f:
+    f.write(payload)
+
+
+def _embedded_modules() -> list:
+  """Registration is lazy — a module only costs bytes if the pickled graph reaches it."""
+  import sys
+  import types
+
+  out, seen = [], set()
+  for module in list(sys.modules.values()):
+    if not isinstance(module, types.ModuleType) or id(module) in seen:
+      continue
+    name = getattr(module, "__name__", "") or ""
+    # cloudpickle registers by name and rejects a name it cannot find in sys.modules
+    if sys.modules.get(name) is not module:
+      continue
+    top = name.split(".")[0]
+    if top in BY_REFERENCE or top in sys.stdlib_module_names:
+      continue
+    # dunder roots are synthetic namespaces, not importable source: __main__ (already by
+    # value to cloudpickle) and __torch__ (TorchScript's compilation unit)
+    if top.startswith("__"):
+      continue
+    src = getattr(module, "__file__", None)
+    # _codegen fabricates __file__, and the compiler chdirs into the plan dir
+    if not src or not os.path.isfile(os.path.realpath(src)):
+      continue
+    seen.add(id(module))
+    out.append(module)
+  return out
+
+
+_bundle_cache = {}
+
+
+def _load(path) -> dict:
+  import cloudpickle
+
+  if not os.path.exists(path):
+    raise ValueError(
+      f"No server plan at '{path}'. The compiler emits it next to the state "
+      f"codec; recompile the plan to produce it."
+    )
+  size = os.path.getsize(path)
+  if size > MAX_BUNDLE_BYTES:
+    raise ValueError(
+      f"server plan at '{path}' is too large: {size} bytes (cap: {MAX_BUNDLE_BYTES})"
+    )
+  # every phase loads the plan, and unpickling executes code
+  key = (path, os.path.getmtime(path))
+  phases = _bundle_cache.get(key)
+  if phases is None:
+    with open(path, "rb") as f:
+      # unpickling runs plan-owner code, the same trust the adapters themselves carry
+      bundle = cloudpickle.loads(f.read())
+    phases = _check(bundle, path)
+    _bundle_cache[key] = phases
+  return phases
+
+
+def _check(bundle, path) -> dict:
+  context = f"server plan at '{path}' is unusable"
+  if not isinstance(bundle, dict):
+    raise ValueError(f"{context}: expected a dict, got {type(bundle).__name__}.")
+  version = bundle.get("abi")
+  if version is None or _abi_major(version) != _abi_major(abi_tag()):
+    raise ValueError(
+      f"{context}: it declares ABI {version!r}, this runtime speaks {abi_tag()}. "
+      f"Recompile the plan with an ifed client whose ifed-server-abi major matches, "
+      f"then re-submit."
+    )
+  built_with = bundle.get("python")
+  if built_with != python_tag():
+    raise ValueError(
+      f"{context}: it was compiled for Python {built_with}, this runtime is "
+      f"{python_tag()}. A plan carries version-specific bytecode, so it must be "
+      f"recompiled on the interpreter that will load it."
+    )
+  built_against = bundle.get("torch")
+  if built_against != torch_tag():
+    raise ValueError(
+      f"{context}: its TorchScript was compiled against torch {built_against}, this "
+      f"runtime has {torch_tag()}. Recompile the plan on a matching torch."
+    )
+  phases = bundle.get("phases")
+  check_phases(phases, context)
+  return phases
+
+
+def check_phases(phases, context: str) -> None:
+  """The shape both ends agree on: exactly the three names, each callable."""
+  if not isinstance(phases, dict):
+    raise ValueError(
+      f"{context}: expected a dict of server callbacks, got {type(phases).__name__}."
+    )
+  if sorted(phases) != sorted(NAMES):
+    raise ValueError(
+      f"{context}: expected exactly {sorted(NAMES)}, got {sorted(phases)}."
+    )
+  for name in NAMES:
+    if not callable(phases[name]):
+      raise ValueError(
+        f"{context}: '{name}' is {type(phases[name]).__name__}, not a callable."
+      )

ifed_server_abi-2.0.2.dist-info/METADATA

@@ -0,0 +1,20 @@
+Metadata-Version: 2.3
+Name: ifed-server-abi
+Version: 2.0.2
+Summary: IFED runtime ABI: the contract the executor loads a compiled server plan through
+Project-URL: Source, https://github.com/JetBrains-Research/ifed
+Project-URL: Issues, https://github.com/JetBrains-Research/ifed/issues
+License: Apache-2.0
+Requires-Python: >=3.10
+Requires-Dist: cloudpickle==3.1.2
+Provides-Extra: test
+Requires-Dist: pytest>=7.0; extra == 'test'
+Description-Content-Type: text/markdown
+
+# IFED server ABI
+
+This package defines the Python callback contract used by the IFED executor.
+Install it through `ifed` unless you are developing the IFED runtime itself.
+
+See the [IFED source repository](https://github.com/JetBrains-Research/ifed) for
+documentation and development instructions.

ifed_server_abi-2.0.2.dist-info/RECORD

@@ -0,0 +1,5 @@
+ifed_server_abi/__init__.py,sha256=paeWT-XRICciaHiQEgL5YkI9WqMft6kfBU-yKd92Ag0,9423
+ifed_server_abi-2.0.2.dist-info/METADATA,sha256=-7W1Bte70rja7uKkzyAIn_0DBFDpYI_LbWEKi9Jq02c,767
+ifed_server_abi-2.0.2.dist-info/WHEEL,sha256=C2FUgwZgiLbznR-k0b_5k3Ai_1aASOXDss3lzCUsUug,87
+ifed_server_abi-2.0.2.dist-info/licenses/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ifed_server_abi-2.0.2.dist-info/RECORD,,

ifed_server_abi-2.0.2.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.26.3
+Root-Is-Purelib: true
+Tag: py3-none-any

ifed_server_abi-2.0.2.dist-info/licenses/LICENSE

@@ -0,0 +1,202 @@
+
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.