ob-metaflow 2.15.7.2__py2.py3-none-any.whl → 2.15.10.1__py2.py3-none-any.whl

metaflow/plugins/datatools/s3/s3op.py

@@ -1,5 +1,6 @@
 from __future__ import print_function
 
+import errno
 import json
 import time
 import math
@@ -15,7 +16,10 @@ from tempfile import NamedTemporaryFile
 from multiprocessing import Process, Queue
 from itertools import starmap, chain, islice
 
+from boto3.exceptions import RetriesExceededError, S3UploadFailedError
 from boto3.s3.transfer import TransferConfig
+from botocore.config import Config
+from botocore.exceptions import ClientError, SSLError
 
 try:
     # python2
@@ -46,13 +50,21 @@ from metaflow.plugins.datatools.s3.s3util import (
 import metaflow.tracing as tracing
 from metaflow.metaflow_config import (
     S3_WORKER_COUNT,
+    S3_CLIENT_RETRY_CONFIG,
 )
 
 DOWNLOAD_FILE_THRESHOLD = 2 * TransferConfig().multipart_threshold
 DOWNLOAD_MAX_CHUNK = 2 * 1024 * 1024 * 1024 - 1
 
+DEFAULT_S3_CLIENT_PARAMS = {"config": Config(retries=S3_CLIENT_RETRY_CONFIG)}
 RANGE_MATCH = re.compile(r"bytes (?P<start>[0-9]+)-(?P<end>[0-9]+)/(?P<total>[0-9]+)")
 
+# from botocore ClientError MSG_TEMPLATE:
+# https://github.com/boto/botocore/blob/68ca78f3097906c9231840a49931ef4382c41eea/botocore/exceptions.py#L521
+BOTOCORE_MSG_TEMPLATE_MATCH = re.compile(
+    r"An error occurred \((\w+)\) when calling the (\w+) operation.*: (.+)"
+)
+
 S3Config = namedtuple("S3Config", "role session_vars client_params")
 
 
@@ -97,6 +109,7 @@ ERROR_VERIFY_FAILED = 9
 ERROR_LOCAL_FILE_NOT_FOUND = 10
 ERROR_INVALID_RANGE = 11
 ERROR_TRANSIENT = 12
+ERROR_OUT_OF_DISK_SPACE = 13
 
 
 def format_result_line(idx, prefix, url="", local=""):
@@ -147,6 +160,7 @@ def normalize_client_error(err):
             "LimitExceededException",
             "RequestThrottled",
             "EC2ThrottledException",
+            "InternalError",
         ):
             return 503
     return error_code
@@ -221,54 +235,68 @@ def worker(result_file_name, queue, mode, s3config):
                 elif mode == "download":
                     tmp = NamedTemporaryFile(dir=".", mode="wb", delete=False)
                     try:
-                        if url.range:
-                            resp = s3.get_object(
-                                Bucket=url.bucket, Key=url.path, Range=url.range
-                            )
-                            range_result = resp["ContentRange"]
-                            range_result_match = RANGE_MATCH.match(range_result)
-                            if range_result_match is None:
-                                raise RuntimeError(
-                                    "Wrong format for ContentRange: %s"
-                                    % str(range_result)
+                        try:
+                            if url.range:
+                                resp = s3.get_object(
+                                    Bucket=url.bucket, Key=url.path, Range=url.range
                                 )
-                            range_result = {
-                                x: int(range_result_match.group(x))
-                                for x in ["total", "start", "end"]
-                            }
-                        else:
-                            resp = s3.get_object(Bucket=url.bucket, Key=url.path)
-                            range_result = None
-                        sz = resp["ContentLength"]
-                        if range_result is None:
-                            range_result = {"total": sz, "start": 0, "end": sz - 1}
-                        if not url.range and sz > DOWNLOAD_FILE_THRESHOLD:
-                            # In this case, it is more efficient to use download_file as it
-                            # will download multiple parts in parallel (it does it after
-                            # multipart_threshold)
-                            s3.download_file(url.bucket, url.path, tmp.name)
-                        else:
-                            read_in_chunks(tmp, resp["Body"], sz, DOWNLOAD_MAX_CHUNK)
-                        tmp.close()
-                        os.rename(tmp.name, url.local)
-                    except client_error as err:
-                        tmp.close()
-                        os.unlink(tmp.name)
-                        error_code = normalize_client_error(err)
-                        if error_code == 404:
-                            result_file.write("%d %d\n" % (idx, -ERROR_URL_NOT_FOUND))
+                                range_result = resp["ContentRange"]
+                                range_result_match = RANGE_MATCH.match(range_result)
+                                if range_result_match is None:
+                                    raise RuntimeError(
+                                        "Wrong format for ContentRange: %s"
+                                        % str(range_result)
+                                    )
+                                range_result = {
+                                    x: int(range_result_match.group(x))
+                                    for x in ["total", "start", "end"]
+                                }
+                            else:
+                                resp = s3.get_object(Bucket=url.bucket, Key=url.path)
+                                range_result = None
+                            sz = resp["ContentLength"]
+                            if range_result is None:
+                                range_result = {"total": sz, "start": 0, "end": sz - 1}
+                            if not url.range and sz > DOWNLOAD_FILE_THRESHOLD:
+                                # In this case, it is more efficient to use download_file as it
+                                # will download multiple parts in parallel (it does it after
+                                # multipart_threshold)
+                                s3.download_file(url.bucket, url.path, tmp.name)
+                            else:
+                                read_in_chunks(
+                                    tmp, resp["Body"], sz, DOWNLOAD_MAX_CHUNK
+                                )
+                            tmp.close()
+                            os.rename(tmp.name, url.local)
+                        except client_error as err:
+                            tmp.close()
+                            os.unlink(tmp.name)
+                            handle_client_error(err, idx, result_file)
                             continue
-                        elif error_code == 403:
-                            result_file.write(
-                                "%d %d\n" % (idx, -ERROR_URL_ACCESS_DENIED)
-                            )
+                        except RetriesExceededError as e:
+                            tmp.close()
+                            os.unlink(tmp.name)
+                            err = convert_to_client_error(e)
+                            handle_client_error(err, idx, result_file)
                             continue
-                        elif error_code == 503:
-                            result_file.write("%d %d\n" % (idx, -ERROR_TRANSIENT))
+                        except OSError as e:
+                            tmp.close()
+                            os.unlink(tmp.name)
+                            if e.errno == errno.ENOSPC:
+                                result_file.write(
+                                    "%d %d\n" % (idx, -ERROR_OUT_OF_DISK_SPACE)
+                                )
+                            else:
+                                result_file.write("%d %d\n" % (idx, -ERROR_TRANSIENT))
+                            result_file.flush()
                             continue
-                        else:
-                            raise
-                        # TODO specific error message for out of disk space
+                    except (SSLError, Exception) as e:
+                        tmp.close()
+                        os.unlink(tmp.name)
+                        # assume anything else is transient
+                        result_file.write("%d %d\n" % (idx, -ERROR_TRANSIENT))
+                        result_file.flush()
+                        continue
                     # If we need the metadata, get it and write it out
                     if pre_op_info:
                         with open("%s_meta" % url.local, mode="w") as f:
@@ -316,28 +344,67 @@ def worker(result_file_name, queue, mode, s3config):
                             if url.encryption is not None:
                                 extra["ServerSideEncryption"] = url.encryption
                         try:
-                            s3.upload_file(
-                                url.local, url.bucket, url.path, ExtraArgs=extra
-                            )
-                            # We indicate that the file was uploaded
-                            result_file.write("%d %d\n" % (idx, 0))
-                        except client_error as err:
-                            error_code = normalize_client_error(err)
-                            if error_code == 403:
-                                result_file.write(
-                                    "%d %d\n" % (idx, -ERROR_URL_ACCESS_DENIED)
+                            try:
+                                s3.upload_file(
+                                    url.local, url.bucket, url.path, ExtraArgs=extra
                                 )
+                                # We indicate that the file was uploaded
+                                result_file.write("%d %d\n" % (idx, 0))
+                            except client_error as err:
+                                # Shouldn't get here, but just in case.
+                                # Internally, botocore catches ClientError and returns a S3UploadFailedError.
+                                # See https://github.com/boto/boto3/blob/develop/boto3/s3/transfer.py#L377
+                                handle_client_error(err, idx, result_file)
                                 continue
-                            elif error_code == 503:
-                                result_file.write("%d %d\n" % (idx, -ERROR_TRANSIENT))
+                            except S3UploadFailedError as e:
+                                err = convert_to_client_error(e)
+                                handle_client_error(err, idx, result_file)
                                 continue
-                            else:
-                                raise
+                        except (SSLError, Exception) as e:
+                            # assume anything else is transient
+                            result_file.write("%d %d\n" % (idx, -ERROR_TRANSIENT))
+                            result_file.flush()
+                            continue
         except:
             traceback.print_exc()
+            result_file.flush()
             sys.exit(ERROR_WORKER_EXCEPTION)
 
 
+def convert_to_client_error(e):
+    match = BOTOCORE_MSG_TEMPLATE_MATCH.search(str(e))
+    if not match:
+        raise e
+    error_code = match.group(1)
+    operation_name = match.group(2)
+    error_message = match.group(3)
+    response = {
+        "Error": {
+            "Code": error_code,
+            "Message": error_message,
+        }
+    }
+    return ClientError(response, operation_name)
+
+
+def handle_client_error(err, idx, result_file):
+    error_code = normalize_client_error(err)
+    if error_code == 404:
+        result_file.write("%d %d\n" % (idx, -ERROR_URL_NOT_FOUND))
+        result_file.flush()
+    elif error_code == 403:
+        result_file.write("%d %d\n" % (idx, -ERROR_URL_ACCESS_DENIED))
+        result_file.flush()
+    elif error_code == 503:
+        result_file.write("%d %d\n" % (idx, -ERROR_TRANSIENT))
+        result_file.flush()
+    else:
+        # optimistically assume it is a transient error
+        result_file.write("%d %d\n" % (idx, -ERROR_TRANSIENT))
+        result_file.flush()
+    # TODO specific error message for out of disk space
+
+
 def start_workers(mode, urls, num_workers, inject_failure, s3config):
     # We start the minimum of len(urls) or num_workers to avoid starting
     # workers that will definitely do nothing
@@ -381,6 +448,22 @@ def start_workers(mode, urls, num_workers, inject_failure, s3config):
                 if proc.exitcode is not None:
                     if proc.exitcode != 0:
                         msg = "Worker process failed (exit code %d)" % proc.exitcode
+
+                        # IMPORTANT: if this process has put items on a queue, then it will not terminate
+                        # until all buffered items have been flushed to the pipe, causing a deadlock.
+                        # `cancel_join_thread()` allows it to exit without flushing the queue.
+                        # Without this line, the parent process would hang indefinitely when a subprocess
+                        # did not exit cleanly in the case of unhandled exceptions.
+                        #
+                        # The error situation is:
+                        # 1. this process puts stuff in queue
+                        # 2. subprocess dies so doesn't consume its end-of-queue marker (the None)
+                        # 3. other subprocesses consume all useful bits AND their end-of-queue marker
+                        # 4. one marker is left and not consumed
+                        # 5. this process cannot shut down until the queue is empty.
+                        # 6. it will never be empty because all subprocesses (workers) have died.
+                        queue.cancel_join_thread()
+
                         exit(msg, proc.exitcode)
                     # Read the output file if all went well
                     with open(out_path, "r") as out_file:
@@ -573,6 +656,8 @@ def exit(exit_code, url):
         msg = "Local file not found: %s" % url
     elif exit_code == ERROR_TRANSIENT:
         msg = "Transient error for url: %s" % url
+    elif exit_code == ERROR_OUT_OF_DISK_SPACE:
+        msg = "Out of disk space when downloading URL: %s" % url
     else:
         msg = "Unknown error"
     print("s3op failed:\n%s" % msg, file=sys.stderr)
@@ -745,7 +830,7 @@ def lst(
     s3config = S3Config(
         s3role,
         json.loads(s3sessionvars) if s3sessionvars else None,
-        json.loads(s3clientparams) if s3clientparams else None,
+        json.loads(s3clientparams) if s3clientparams else DEFAULT_S3_CLIENT_PARAMS,
     )
 
     urllist = []
@@ -878,7 +963,7 @@ def put(
     s3config = S3Config(
         s3role,
         json.loads(s3sessionvars) if s3sessionvars else None,
-        json.loads(s3clientparams) if s3clientparams else None,
+        json.loads(s3clientparams) if s3clientparams else DEFAULT_S3_CLIENT_PARAMS,
     )
 
     urls = list(starmap(_make_url, _files()))
@@ -1025,7 +1110,7 @@ def get(
     s3config = S3Config(
         s3role,
         json.loads(s3sessionvars) if s3sessionvars else None,
-        json.loads(s3clientparams) if s3clientparams else None,
+        json.loads(s3clientparams) if s3clientparams else DEFAULT_S3_CLIENT_PARAMS,
     )
 
     # Construct a list of URL (prefix) objects
@@ -1103,6 +1188,8 @@ def get(
             )
             if verify:
                 verify_info.append((url, sz))
+        elif sz == -ERROR_OUT_OF_DISK_SPACE:
+            exit(ERROR_OUT_OF_DISK_SPACE, url)
         elif sz == -ERROR_URL_ACCESS_DENIED:
             denied_url = url
             break
@@ -1172,7 +1259,7 @@ def info(
     s3config = S3Config(
         s3role,
         json.loads(s3sessionvars) if s3sessionvars else None,
-        json.loads(s3clientparams) if s3clientparams else None,
+        json.loads(s3clientparams) if s3clientparams else DEFAULT_S3_CLIENT_PARAMS,
     )
 
     # Construct a list of URL (prefix) objects

metaflow/plugins/kubernetes/kubernetes.py

@@ -194,6 +194,7 @@ class Kubernetes(object):
         port=None,
         num_parallel=None,
         qos=None,
+        security_context=None,
     ):
         name = "js-%s" % str(uuid4())[:6]
         jobset = (
@@ -227,6 +228,7 @@ class Kubernetes(object):
                 port=port,
                 num_parallel=num_parallel,
                 qos=qos,
+                security_context=security_context,
             )
             .environment_variable("METAFLOW_CODE_SHA", code_package_sha)
             .environment_variable("METAFLOW_CODE_URL", code_package_url)
@@ -504,6 +506,7 @@ class Kubernetes(object):
         name_pattern=None,
         qos=None,
         annotations=None,
+        security_context=None,
     ):
         if env is None:
             env = {}
@@ -546,6 +549,7 @@ class Kubernetes(object):
                 shared_memory=shared_memory,
                 port=port,
                 qos=qos,
+                security_context=security_context,
             )
             .environment_variable("METAFLOW_CODE_SHA", code_package_sha)
             .environment_variable("METAFLOW_CODE_URL", code_package_url)

metaflow/plugins/kubernetes/kubernetes_cli.py

@@ -145,6 +145,12 @@ def kubernetes():
     type=JSONTypeClass(),
     multiple=False,
 )
+@click.option(
+    "--security-context",
+    default=None,
+    type=JSONTypeClass(),
+    multiple=False,
+)
 @click.pass_context
 def step(
     ctx,
@@ -176,6 +182,7 @@ def step(
     qos=None,
     labels=None,
     annotations=None,
+    security_context=None,
     **kwargs
 ):
     def echo(msg, stream="stderr", job_id=None, **kwargs):
@@ -319,6 +326,7 @@ def step(
                 qos=qos,
                 labels=labels,
                 annotations=annotations,
+                security_context=security_context,
             )
     except Exception:
         traceback.print_exc(chain=False)

metaflow/plugins/kubernetes/kubernetes_decorator.py

@@ -124,6 +124,14 @@ class KubernetesDecorator(StepDecorator):
         Only applicable when @parallel is used.
     qos: str, default: Burstable
         Quality of Service class to assign to the pod. Supported values are: Guaranteed, Burstable, BestEffort
+
+    security_context: Dict[str, Any], optional, default None
+        Container security context. Applies to the task container. Allows the following keys:
+        - privileged: bool, optional, default None
+        - allow_privilege_escalation: bool, optional, default None
+        - run_as_user: int, optional, default None
+        - run_as_group: int, optional, default None
+        - run_as_non_root: bool, optional, default None
     """
 
     name = "kubernetes"
@@ -154,6 +162,7 @@ class KubernetesDecorator(StepDecorator):
         "executable": None,
         "hostname_resolution_timeout": 10 * 60,
         "qos": KUBERNETES_QOS,
+        "security_context": None,
     }
     package_url = None
     package_sha = None
@@ -489,6 +498,7 @@ class KubernetesDecorator(StepDecorator):
                     "persistent_volume_claims",
                     "labels",
                     "annotations",
+                    "security_context",
                 ]:
                     cli_args.command_options[k] = json.dumps(v)
                 else:

metaflow/plugins/kubernetes/kubernetes_job.py

@@ -94,6 +94,13 @@ class KubernetesJob(object):
             ],
         }
 
+        security_context = self._kwargs.get("security_context", {})
+        _security_context = {}
+        if security_context is not None and len(security_context) > 0:
+            _security_context = {
+                "security_context": client.V1SecurityContext(**security_context)
+            }
+
         return client.V1JobSpec(
             # Retries are handled by Metaflow when it is responsible for
             # executing the flow. The responsibility is moved to Kubernetes
@@ -224,6 +231,7 @@ class KubernetesJob(object):
                                 if self._kwargs["persistent_volume_claims"] is not None
                                 else []
                             ),
+                            **_security_context,
                         )
                     ],
                     node_selector=self._kwargs.get("node_selector"),

metaflow/plugins/kubernetes/kubernetes_jobsets.py

@@ -562,6 +562,12 @@ class JobSetSpec(object):
             self._kwargs["memory"],
             self._kwargs["disk"],
         )
+        security_context = self._kwargs.get("security_context", {})
+        _security_context = {}
+        if security_context is not None and len(security_context) > 0:
+            _security_context = {
+                "security_context": client.V1SecurityContext(**security_context)
+            }
         return dict(
             name=self.name,
             template=client.api_client.ApiClient().sanitize_for_serialization(
@@ -708,6 +714,7 @@ class JobSetSpec(object):
                                             is not None
                                             else []
                                         ),
+                                        **_security_context,
                                     )
                                 ],
                                 node_selector=self._kwargs.get("node_selector"),

metaflow/plugins/pypi/conda_decorator.py

@@ -227,7 +227,8 @@ class CondaStepDecorator(StepDecorator):
         self.interpreter = (
             self.environment.interpreter(self.step)
             if not any(
-                decorator.name in ["batch", "kubernetes", "nvidia", "snowpark", "slurm"]
+                decorator.name
+                in ["batch", "kubernetes", "nvidia", "snowpark", "slurm", "nvct"]
                 for decorator in next(
                     step for step in self.flow if step.name == self.step
                 ).decorators

metaflow/plugins/pypi/conda_environment.py

@@ -326,6 +326,7 @@ class CondaEnvironment(MetaflowEnvironment):
                 "nvidia",
                 "snowpark",
                 "slurm",
+                "nvct",
             ]:
                 target_platform = getattr(decorator, "target_platform", "linux-64")
                 break

metaflow/plugins/uv/bootstrap.py

@@ -0,0 +1,100 @@
+import os
+import subprocess
+import sys
+import time
+
+from metaflow.util import which
+from metaflow.metaflow_config import get_pinned_conda_libs
+from urllib.request import Request, urlopen
+from urllib.error import URLError
+
+# TODO: support version/platform/architecture selection.
+UV_URL = "https://github.com/astral-sh/uv/releases/download/0.6.11/uv-x86_64-unknown-linux-gnu.tar.gz"
+
+if __name__ == "__main__":
+
+    def run_cmd(cmd, stdin_str=None):
+        result = subprocess.run(
+            cmd,
+            shell=True,
+            input=stdin_str,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            text=True,
+        )
+        if result.returncode != 0:
+            print(f"Bootstrap failed while executing: {cmd}")
+            print("Stdout:", result.stdout)
+            print("Stderr:", result.stderr)
+            sys.exit(1)
+
+    def install_uv():
+        import tarfile
+
+        uv_install_path = os.path.join(os.getcwd(), "uv_install")
+        if which("uv"):
+            return
+
+        print("Installing uv...")
+
+        # Prepare directory once
+        os.makedirs(uv_install_path, exist_ok=True)
+
+        # Download and decompress in one go
+        headers = {
+            "Accept-Encoding": "gzip, deflate, br",
+            "Connection": "keep-alive",
+            "User-Agent": "python-urllib",
+        }
+
+        def _tar_filter(member: tarfile.TarInfo, path):
+            if os.path.basename(member.name) != "uv":
+                return None  # skip
+            member.path = os.path.basename(member.path)
+            return member
+
+        max_retries = 3
+        for attempt in range(max_retries):
+            try:
+                req = Request(UV_URL, headers=headers)
+                with urlopen(req) as response:
+                    with tarfile.open(fileobj=response, mode="r:gz") as tar:
+                        tar.extractall(uv_install_path, filter=_tar_filter)
+                break
+            except (URLError, IOError) as e:
+                if attempt == max_retries - 1:
+                    raise Exception(
+                        f"Failed to download UV after {max_retries} attempts: {e}"
+                    )
+                time.sleep(2**attempt)
+
+        # Update PATH only once at the end
+        os.environ["PATH"] += os.pathsep + uv_install_path
+
+    def get_dependencies(datastore_type):
+        # return required dependencies for Metaflow that must be added to the UV environment.
+        pinned = get_pinned_conda_libs(None, datastore_type)
+
+        # return only dependency names instead of pinned versions
+        return pinned.keys()
+
+    def sync_uv_project(datastore_type):
+        print("Syncing uv project...")
+        dependencies = " ".join(get_dependencies(datastore_type))
+        cmd = f"""set -e;
+            uv sync --frozen --no-install-package metaflow;
+            uv pip install {dependencies} --strict
+            """
+        run_cmd(cmd)
+
+    if len(sys.argv) != 2:
+        print("Usage: bootstrap.py <datastore_type>")
+        sys.exit(1)
+
+    try:
+        datastore_type = sys.argv[1]
+        install_uv()
+        sync_uv_project(datastore_type)
+    except Exception as e:
+        print(f"Error: {str(e)}", file=sys.stderr)
+        sys.exit(1)

metaflow/plugins/uv/uv_environment.py

@@ -0,0 +1,70 @@
+import os
+
+from metaflow.exception import MetaflowException
+from metaflow.metaflow_environment import MetaflowEnvironment
+
+
+class UVException(MetaflowException):
+    headline = "uv error"
+
+
+class UVEnvironment(MetaflowEnvironment):
+    TYPE = "uv"
+
+    def __init__(self, flow):
+        self.flow = flow
+
+    def validate_environment(self, logger, datastore_type):
+        self.datastore_type = datastore_type
+        self.logger = logger
+
+    def init_environment(self, echo, only_steps=None):
+        self.logger("Bootstrapping uv...")
+
+    def executable(self, step_name, default=None):
+        return "uv run python"
+
+    def add_to_package(self):
+        # NOTE: We treat uv.lock and pyproject.toml as regular project assets and ship these along user code as part of the code package
+        # These are the minimal required files to reproduce the UV environment on the remote platform.
+        def _find(filename):
+            current_dir = os.getcwd()
+            while True:
+                file_path = os.path.join(current_dir, filename)
+                if os.path.isfile(file_path):
+                    return file_path
+                parent_dir = os.path.dirname(current_dir)
+                if parent_dir == current_dir:  # Reached root
+                    raise UVException(
+                        f"Could not find {filename} in current directory or any parent directory"
+                    )
+                current_dir = parent_dir
+
+        pyproject_path = _find("pyproject.toml")
+        uv_lock_path = _find("uv.lock")
+        files = [
+            (uv_lock_path, "uv.lock"),
+            (pyproject_path, "pyproject.toml"),
+        ]
+        return files
+
+    def pylint_config(self):
+        config = super().pylint_config()
+        # Disable (import-error) in pylint
+        config.append("--disable=F0401")
+        return config
+
+    def bootstrap_commands(self, step_name, datastore_type):
+        return [
+            "echo 'Bootstrapping uv project...'",
+            "flush_mflogs",
+            # We have to prevent the tracing module from loading, as the bootstrapping process
+            # uses the internal S3 client which would fail to import tracing due to the required
+            # dependencies being bundled into the conda environment, which is yet to be
+            # initialized at this point.
+            'DISABLE_TRACING=True python -m metaflow.plugins.uv.bootstrap "%s"'
+            % datastore_type,
+            "echo 'uv project bootstrapped.'",
+            "flush_mflogs",
+            "export PATH=$PATH:$(pwd)/uv_install",
+        ]