thds.mops 3.9.20251007062727__py3-none-any.whl → 3.9.20251008224110__py3-none-any.whl

thds/mops/k8s/apply_yaml.py

@@ -1,22 +1,45 @@
+import subprocess
 import tempfile
 
+import kubernetes
 from kubernetes import client, utils
+from packaging import version
 
 
 def format_yaml(yaml_template_str: str, **template_values: str) -> str:
     return yaml_template_str.format(**template_values)
 
 
+def kubectl_apply_file(yaml_path: str) -> None:
+    subprocess.run(["kubectl", "apply", "-f", yaml_path], check=True)
+
+
+def kubectl_apply(yaml_string: str) -> None:
+    with tempfile.NamedTemporaryFile("w", prefix="kubectl-yaml") as f:
+        f.write(yaml_string)
+        f.flush()
+        kubectl_apply_file(f.name)
+
+
+def apply_yaml(yaml_path: str) -> None:
+    if version.parse(kubernetes.__version__) < version.parse("32.0.0"):
+        kubectl_apply_file(yaml_path)  # best effort
+        return
+
+    # NOTE: Prior to 32.0.0, this function doesn't actually server-side apply.
+    # https://github.com/kubernetes-client/python/pull/2252
+    # Hence the check above to use kubectl for older versions.
+    utils.create_from_yaml(client.ApiClient(), yaml_path)
+
+
 def create_yaml_template(yaml_str: str, **template_values: str) -> None:
     """Format a YAML template with the given keyword arguments, then apply it to the Kubernetes cluster.
 
     You must already have set up your SDK config.
 
-    NOTE: This function doesn't actually apply, and can't until the next release of the K8S SDK:
-    https://github.com/kubernetes-client/python/pull/2252
     """
     formatted_yaml = format_yaml(yaml_str, **template_values)
     with tempfile.NamedTemporaryFile("w", prefix="kubectl-yaml") as f:
         f.write(formatted_yaml)
         f.flush()
-        utils.create_from_yaml(client.ApiClient(), f.name)
+        apply_yaml(f.name)

thds/mops/pure/core/source.py

@@ -33,6 +33,7 @@ have a Source object returned to it while it performs low-level deserialization.
 """
 
 import io
+import json
 import sys
 import typing as ty
 from functools import partial
@@ -71,19 +72,44 @@ def _hashref_uri(hash: hashing.Hash, type: ty.Literal["local", "remote"]) -> str
     return to_uri(local_hashref)
 
 
-def _read_hashref(hashref_uri: str) -> str:
+class _HashrefMeta(ty.NamedTuple):
+    size: int
+
+    @classmethod
+    def empty(cls) -> "_HashrefMeta":
+        return cls(size=0)
+
+    def serialize(self) -> str:
+        serialized = json.dumps(self._asdict())
+        return serialized
+
+    @classmethod
+    def deserialize(cls, serialized: ty.Union[str, ty.Sequence[str]]) -> "_HashrefMeta":
+        s = serialized if isinstance(serialized, str) else serialized[0]
+        try:
+            return cls(**json.loads(s))
+        except json.JSONDecodeError:
+            logger.warning("Failed to deserialize hashref metadata '%s'", serialized)
+            return cls.empty()
+
+
+def _read_hashref(hashref_uri: str) -> ty.Tuple[str, _HashrefMeta]:
     """Return URI represented by this hashref. Performs IO."""
     uri_bytes = io.BytesIO()
     lookup_blob_store(hashref_uri).readbytesinto(hashref_uri, uri_bytes)
-    uri = uri_bytes.getvalue().decode()
+    content = uri_bytes.getvalue().decode()
+    uri, *rest = content.split("\n")
     assert uri, f"Hashref from {hashref_uri} is empty"
-    return uri
+    if not rest:
+        return uri, _HashrefMeta.empty()
+    return uri, _HashrefMeta.deserialize(rest)
 
 
-def _write_hashref(hashref_uri: str, uri: str) -> None:
+def _write_hashref(hashref_uri: str, uri: str, size: int) -> None:
     """Write URI to this hashref. Performs IO."""
     assert uri, f"Should never encode hashref ({hashref_uri}) pointing to empty URI"
-    lookup_blob_store(hashref_uri).putbytes(hashref_uri, uri.encode(), type_hint="text/plain")
+    content = "\n".join([uri, _HashrefMeta(size=size).serialize()])
+    lookup_blob_store(hashref_uri).putbytes(hashref_uri, content.encode(), type_hint="text/plain")
 
 
 def source_from_hashref(hash: hashing.Hash) -> Source:
@@ -92,7 +118,9 @@ def source_from_hashref(hash: hashing.Hash) -> Source:
     local_file_hashref_uri = _hashref_uri(hash, "local")
     remote_hashref_uri = _hashref_uri(hash, "remote")
 
-    def remote_uri(allow_blob_not_found: bool = True) -> str:
+    def remote_uri_and_meta(
+        allow_blob_not_found: bool = True,
+    ) -> ty.Tuple[str, _HashrefMeta]:
         try:
             return _read_hashref(remote_hashref_uri)
         except Exception as e:
@@ -102,7 +130,7 @@ def source_from_hashref(hash: hashing.Hash) -> Source:
                 # 'remote' blob not found is sometimes fine, but anything else is weird
                 # and we should raise.
                 raise
-            return ""
+            return "", _HashrefMeta.empty()
 
     try:
         # we might be on the same machine where this was originally invoked.
@@ -110,7 +138,9 @@ def source_from_hashref(hash: hashing.Hash) -> Source:
         # Then, there's no need to bother grabbing the remote_uri
         # - but for debugging's sake, it's quite nice to actually
         # have the full remote URI as well even if we're ultimately going to use the local copy.
-        return source.from_file(_read_hashref(local_file_hashref_uri), hash=hash, uri=remote_uri())
+        local_uri, _ = _read_hashref(local_file_hashref_uri)
+        remote_uri, _ = remote_uri_and_meta()
+        return source.from_file(local_uri, hash=hash, uri=remote_uri)
     except FileNotFoundError:
         # we are not on the same machine as the local ref. assume we need the remote URI.
         pass
@@ -120,14 +150,17 @@ def source_from_hashref(hash: hashing.Hash) -> Source:
             raise
 
     # no local file, so we assume there must be a remote URI.
-    return source.from_uri(remote_uri(False), hash=hash)
+    remote_uri, meta = remote_uri_and_meta(False)
+    return source.from_uri(remote_uri, hash=hash, size=meta.size)
 
 
-def _upload_and_create_remote_hashref(local_path: Path, remote_uri: str, hash: hashing.Hash) -> None:
+def _upload_and_create_remote_hashref(
+    local_path: Path, remote_uri: str, hash: hashing.Hash, size: int
+) -> None:
     # exists only to provide a local (non-serializable) closure around local_path and remote_uri.
     lookup_blob_store(remote_uri).putfile(local_path, remote_uri)
     # make sure we never overwrite a hashref until it's actually going to be valid.
-    _write_hashref(_hashref_uri(hash, "remote"), remote_uri)
+    _write_hashref(_hashref_uri(hash, "remote"), remote_uri, size)
 
 
 def _auto_remote_uri(hash: hashing.Hash) -> str:
@@ -155,7 +188,7 @@ def prepare_source_argument(source_: Source) -> ty.Union[str, hashing.Hash]:
         deferred_work.add(
             __name__ + "-localhashref",
             source_.hash,
-            partial(_write_hashref, _hashref_uri(source_.hash, "local"), str(local_path)),
+            partial(_write_hashref, _hashref_uri(source_.hash, "local"), str(local_path), source_.size),
         )
         # then also register pending upload - if the URI is a local file, we need to determine a
         # remote URI for this thing automagically; otherwise, use whatever was already
@@ -164,7 +197,9 @@ def prepare_source_argument(source_: Source) -> ty.Union[str, hashing.Hash]:
         deferred_work.add(
             __name__ + "-remotehashref",
             source_.hash,
-            partial(_upload_and_create_remote_hashref, local_path, remote_uri, source_.hash),
+            partial(
+                _upload_and_create_remote_hashref, local_path, remote_uri, source_.hash, source_.size
+            ),
         )
     else:
         # prepare to (later, if necessary) create a remote hashref, because this Source
@@ -172,7 +207,7 @@ def prepare_source_argument(source_: Source) -> ty.Union[str, hashing.Hash]:
         deferred_work.add(
             __name__,
             source_.hash,
-            partial(_write_hashref, _hashref_uri(source_.hash, "remote"), source_.uri),
+            partial(_write_hashref, _hashref_uri(source_.hash, "remote"), source_.uri, source_.size),
         )
 
     return hashing.Hash(algo=sys.intern(source_.hash.algo), bytes=source_.hash.bytes)
@@ -192,15 +227,17 @@ def prepare_source_argument(source_: Source) -> ty.Union[str, hashing.Hash]:
 # just that mops must detect Sources in the return value and must force an upload on them.
 # In essence, this creates a bifurcated code path for Sources during serialization; if
 # we're "on the way out", we avoid uploading until it is clear that the data will be used
-# in a remote environment. Whereas "on the way back", we must always upload, and nothing
-# can or should be deferred; upload should happen at the time of serialization.
+# in a remote environment. Whereas "on the way back", we must always upload -- there, we
+# defer uploads until everything is serialized, then we perform all deferred uploads in
+# parallel, prior to writing the serialized result.
 #
 # Nevertheless, a local caller should still be able to short-circuit the _download_ by
 # using a locally-created File, if on the same machine where the local file was created.
 
 
 class SourceResult(ty.NamedTuple):
-    """Contains the fully-specified local URI and remote URI, plus (probably) a Hash.
+    """Contains the fully-specified local URI and remote URI, plus (probably) a Hash
+    and a size.
 
     Everything is defined right here. No need for any kind of dynamic lookup, and
     optimization buys us nothing, since memoization only operates on arguments.
@@ -210,6 +247,10 @@ class SourceResult(ty.NamedTuple):
     hash: ty.Optional[hashing.Hash]
     file_uri: str
 
+    size: int = 0
+    # instances of older versions of this namedtuple will be missing this field.
+    # we supply a default for backward-compatibility.
+
 
 class DuplicateSourceBasenameError(ValueError):
     """This is not a catchable error - it will be raised inside the mops result-wrapping
@@ -247,7 +288,7 @@ def prepare_source_result(source_: Source, existing_uris: ty.Collection[str] = t
         else:
             file_uri = ""
             logger.debug("Creating a SourceResult for a URI that is presumed to already be uploaded.")
-        return SourceResult(source_.uri, source_.hash, file_uri)
+        return SourceResult(source_.uri, source_.hash, file_uri, source_.size)
 
     # by definition, if this is a file URI, it now needs to be uploaded, because we could
     # be transferring back to an orchestrator on a different machine, but also because a
@@ -277,15 +318,17 @@ def prepare_source_result(source_: Source, existing_uris: ty.Collection[str] = t
         partial(_put_file_to_blob_store, local_path, remote_uri),
     )
     # upload must _always_ happen on remotely-returned Sources, as detailed above.
-    return SourceResult(remote_uri, source_.hash, source_.uri)
+    return SourceResult(remote_uri, source_.hash, source_.uri, source_.size)
 
 
-def source_from_source_result(remote_uri: str, hash: ty.Optional[hashing.Hash], file_uri: str) -> Source:
+def source_from_source_result(
+    remote_uri: str, hash: ty.Optional[hashing.Hash], file_uri: str, size: int
+) -> Source:
     """Call when deserializing a remote function return value on the orchestrator side, to
     replace all SourceResults with the intended Source object.
     """
     if not file_uri:
-        return source.from_uri(remote_uri, hash=hash)
+        return source.from_uri(remote_uri, hash=hash, size=size)
 
     local_path = source.path_from_uri(file_uri)
 
@@ -305,7 +348,7 @@ def source_from_source_result(remote_uri: str, hash: ty.Optional[hashing.Hash],
             logger.warning(
                 f"Unable to reuse destination local path {local_path} when constructing Source {remote_uri}: {e}"
             )
-    return source.from_uri(remote_uri, hash=hash)
+    return source.from_uri(remote_uri, hash=hash, size=size)
 
 
 def create_source_at_uri(filename: StrOrPath, destination_uri: str) -> Source:

thds/mops/pure/pickling/pickles.py

@@ -14,7 +14,7 @@ from pathlib import Path
 from thds.core import hashing, log, source
 
 from ..core.script_support import add_main_module_function, get_main_module_function
-from ..core.source import source_from_hashref, source_from_source_result
+from ..core.source import SourceResult, source_from_hashref, source_from_source_result
 from ..core.uris import get_bytes, lookup_blob_store
 
 logger = log.getLogger(__name__)
@@ -131,17 +131,15 @@ class UnpickleSourceHashrefArgument(ty.NamedTuple):
         return source_from_hashref(self.hash)
 
 
-class UnpickleSourceResult(ty.NamedTuple):
+class UnpickleSourceResult(SourceResult):
     """Stability for this is not critical, as it will only ever exist in the result
     payload, which does not participate in memoization.
     """
 
-    remote_uri: str
-    hash: ty.Optional[hashing.Hash]
-    file_uri: str
-
     def __call__(self) -> source.Source:
-        return source_from_source_result(*self)
+        return source_from_source_result(
+            remote_uri=self.remote_uri, hash=self.hash, file_uri=self.file_uri, size=self.size
+        )
 
 
 class UnpickleFunctionWithLogicKey(ty.NamedTuple):

{thds_mops-3.9.20251007062727.dist-info → thds_mops-3.9.20251008224110.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: thds.mops
-Version: 3.9.20251007062727
+Version: 3.9.20251008224110
 Summary: ML Ops tools for Trilliant Health
 Author-email: Trilliant Health <info@trillianthealth.com>
 Project-URL: Repository, https://github.com/TrilliantHealth/ds-monorepo

{thds_mops-3.9.20251007062727.dist-info → thds_mops-3.9.20251008224110.dist-info}/RECORD

@@ -18,7 +18,7 @@ thds/mops/impure/runner.py,sha256=UI1NZWMZ_5TQHfFKLnoiSm2zDR3zCunTKFmJoybkyCo,28
 thds/mops/k8s/__init__.py,sha256=zl4GVcCFRvPscyo6gvv5Lx0OKB7d3QjtVFjYurnxMuE,764
 thds/mops/k8s/_launch.py,sha256=hgPty47CdwryPHKMmEnoxSsSvcSpXhHYSVYnLC2QJb0,10956
 thds/mops/k8s/_shared.py,sha256=MR-s6ijWUHZGjxK_fsOpHuRDB6kuofjo5xiIb7ul2VM,86
-thds/mops/k8s/apply_yaml.py,sha256=hVW6dIVbNdzHdbGlc2VAPGkdByv_rH2oPybyIm7tKIM,820
+thds/mops/k8s/apply_yaml.py,sha256=zMiQ_k1KcdlIg6diqp6FJD_74I_01YD3cZWA9lqu1fo,1462
 thds/mops/k8s/auth.py,sha256=0zs4TQgkD6VPrhDD43xt7JGwP6uWf3ctySGLcPKN7iw,1691
 thds/mops/k8s/batching.py,sha256=Djt17ffxWyTq4Q7XcAKQdCe9JIIfPahHwm0wqgFqevI,8368
 thds/mops/k8s/config.py,sha256=_znocX5BW8kfG_Cbq6f3apx5FqSihD7Tmic-SBkVjMQ,2992
@@ -58,7 +58,7 @@ thds/mops/pure/core/pipeline_id_mask.py,sha256=AVAy06TdNAmivxGec1gahBYvkJCn7yn-g
 thds/mops/pure/core/script_support.py,sha256=4VCBL5AfGSHcZWpOxMw6nnAbQyk1B-979G_OjvUg9B0,953
 thds/mops/pure/core/serialize_big_objs.py,sha256=YcOS1ccs82ZWO7nTbeumErMzYVe4hgXCTsfvMggYmd8,2332
 thds/mops/pure/core/serialize_paths.py,sha256=JoVXFGSA68QbL4oY8tQbp9MoizTCKj_nPRCuA3i03i8,6122
-thds/mops/pure/core/source.py,sha256=R36ajrCU1JdWF-8iD8YqAiP-q39ypZqf2DeBsqC9lYo,15105
+thds/mops/pure/core/source.py,sha256=i6SRgOFfkdaidx6uEq6poGFUU1zQWiqaXPIG-l5zY7Q,16580
 thds/mops/pure/core/types.py,sha256=_3gDwztDKV4Xeyw2jvyMRJAjmR6gRsmfYmsRCcZMUwI,5436
 thds/mops/pure/core/uris.py,sha256=qO9_f-ro7kax6haNOPTPe81-_aUSRFELeeZH4PMTTU4,2694
 thds/mops/pure/core/use_runner.py,sha256=m1Mu1XDr3xRf_u_VSiHfTG4TH6fnSg0IqwmtbLKG_oc,2103
@@ -88,7 +88,7 @@ thds/mops/pure/pickling/__init__.py,sha256=WNdG8PdJCk-kYaXkvvPa--hjYGoUlBXG3w2X8
 thds/mops/pure/pickling/_pickle.py,sha256=YB8xbqDiwdk8ccnVZ2_4kQn98V2JSrFqw2E3J-jEHlA,8081
 thds/mops/pure/pickling/memoize_only.py,sha256=oI5CMy6IEJc46Gb_BGWNUuAe3fysS7HxRSTajN0WssI,837
 thds/mops/pure/pickling/mprunner.py,sha256=VWYS_PXLgYJetK69CCZ0-b1109-QBHWssC0MskHww94,8831
-thds/mops/pure/pickling/pickles.py,sha256=CSlnjLssE0Ad8YzqyaKqWCSNyW5LiMFKiXO6hWAZmvU,5097
+thds/mops/pure/pickling/pickles.py,sha256=KYkPexi5mGWjrv9uZxt4iWuBUPyYlME2FQIwRiPlPqc,5134
 thds/mops/pure/pickling/remote.py,sha256=7JXZRGnLI5y5dqElIDrhIlaRv6Q_zQ_78aqNhO7O4KY,8478
 thds/mops/pure/pickling/sha256_b64.py,sha256=HL0cPixHPZYuZDVDBscxsnI-3a2amWEfw-LseOX-PyY,2916
 thds/mops/pure/runner/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -109,8 +109,8 @@ thds/mops/pure/tools/summarize/cli.py,sha256=7kDtn24ok8oBO3jFjlMmOK3jnZYpMoE_5Y8
 thds/mops/pure/tools/summarize/run_summary.py,sha256=w45qiQr7elrHDiK9Hgs85gtU3gwLuXa447ih1Y23BBY,5776
 thds/mops/testing/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 thds/mops/testing/deferred_imports.py,sha256=f0ezCgQAtzTqW1yAOb0OWgsB9ZrlztLB894LtpWDaVw,3780
-thds_mops-3.9.20251007062727.dist-info/METADATA,sha256=XPhDxj2piPOxhXs_SV5dtp3fjsev-jY539-mM9fEu6U,2225
-thds_mops-3.9.20251007062727.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-thds_mops-3.9.20251007062727.dist-info/entry_points.txt,sha256=qKvCAaB80syXfxVR3xx6x9J0YJdaQWkIbVSw-NwFgMw,322
-thds_mops-3.9.20251007062727.dist-info/top_level.txt,sha256=LTZaE5SkWJwv9bwOlMbIhiS-JWQEEIcjVYnJrt-CriY,5
-thds_mops-3.9.20251007062727.dist-info/RECORD,,
+thds_mops-3.9.20251008224110.dist-info/METADATA,sha256=XSLkz2hWjPoo5hbsr8buF_n883pZ43NaHu5lkRjmngI,2225
+thds_mops-3.9.20251008224110.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+thds_mops-3.9.20251008224110.dist-info/entry_points.txt,sha256=qKvCAaB80syXfxVR3xx6x9J0YJdaQWkIbVSw-NwFgMw,322
+thds_mops-3.9.20251008224110.dist-info/top_level.txt,sha256=LTZaE5SkWJwv9bwOlMbIhiS-JWQEEIcjVYnJrt-CriY,5
+thds_mops-3.9.20251008224110.dist-info/RECORD,,