ob-metaflow 2.11.13.1__py2.py3-none-any.whl → 2.19.7.1rc0__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,9 @@ 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.exceptions import ClientError, SSLError
 
 try:
     # python2
@@ -47,12 +50,19 @@ import metaflow.tracing as tracing
 from metaflow.metaflow_config import (
     S3_WORKER_COUNT,
 )
+from metaflow.exception import MetaflowException
 
 DOWNLOAD_FILE_THRESHOLD = 2 * TransferConfig().multipart_threshold
 DOWNLOAD_MAX_CHUNK = 2 * 1024 * 1024 * 1024 - 1
 
 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 +107,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=""):
@@ -119,9 +130,9 @@ def normalize_client_error(err):
     try:
         return int(error_code)
     except ValueError:
-        if error_code in ("AccessDenied", "AllAccessDisabled"):
+        if error_code in ("AccessDenied", "AllAccessDisabled", "InvalidAccessKeyId"):
             return 403
-        if error_code == "NoSuchKey":
+        if error_code in ("NoSuchKey", "NoSuchBucket"):
             return 404
         if error_code == "InvalidRange":
             return 416
@@ -147,6 +158,7 @@ def normalize_client_error(err):
             "LimitExceededException",
             "RequestThrottled",
             "EC2ThrottledException",
+            "InternalError",
         ):
             return 503
     return error_code
@@ -155,7 +167,7 @@ def normalize_client_error(err):
 # S3 worker pool
 
 
-@tracing.cli_entrypoint("s3op/worker")
+@tracing.cli("s3op/worker")
 def worker(result_file_name, queue, mode, s3config):
     # Interpret mode, it can either be a single op or something like
     # info_download or info_upload which implies:
@@ -221,54 +233,77 @@ 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 %s\n" % (idx, -ERROR_TRANSIENT, "OSError")
+                                )
+                            result_file.flush()
                             continue
-                        else:
-                            raise
-                        # TODO specific error message for out of disk space
+                    except MetaflowException:
+                        # Re-raise Metaflow exceptions (including TimeoutException)
+                        tmp.close()
+                        os.unlink(tmp.name)
+                        raise
+                    except (SSLError, Exception) as e:
+                        tmp.close()
+                        os.unlink(tmp.name)
+                        # assume anything else is transient
+                        result_file.write(
+                            "%d %d %s\n" % (idx, -ERROR_TRANSIENT, type(e).__name__)
+                        )
+                        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 +351,77 @@ 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 MetaflowException:
+                            # Re-raise Metaflow exceptions (including TimeoutException)
+                            raise
+                        except (SSLError, Exception) as e:
+                            # assume anything else is transient
+                            result_file.write(
+                                "%d %d %s\n" % (idx, -ERROR_TRANSIENT, type(e).__name__)
+                            )
+                            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):
+    # Handle all MetaflowExceptions as fatal
+    if isinstance(err, MetaflowException):
+        raise err
+
+    error_code = normalize_client_error(err)
+    original_error_code = err.response["Error"]["Code"]
+
+    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 %s\n" % (idx, -ERROR_TRANSIENT, original_error_code))
+        result_file.flush()
+    else:
+        # optimistically assume it is a transient error
+        result_file.write("%d %d %s\n" % (idx, -ERROR_TRANSIENT, original_error_code))
+        result_file.flush()
+
+
 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
@@ -347,6 +431,7 @@ def start_workers(mode, urls, num_workers, inject_failure, s3config):
     random.seed()
 
     sz_results = []
+    transient_error_type = None
     # 1. push sources and destinations to the queue
     # We only push if we don't inject a failure; otherwise, we already set the sz_results
     # appropriately with the result of the injected failure.
@@ -381,17 +466,39 @@ 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:
                         for line in out_file:
-                            line_split = line.split(" ")
-                            sz_results[int(line_split[0])] = int(line_split[1])
+                            line_split = line.split(" ", 2)
+                            idx = int(line_split[0])
+                            size = int(line_split[1])
+                            sz_results[idx] = size
+
+                            # For transient errors, store the transient error type (should be the same for all)
+                            if size == -ERROR_TRANSIENT and len(line_split) > 2:
+                                transient_error_type = line_split[2].strip()
                 else:
                     # Put this process back in the processes to check
                     new_procs[proc] = out_path
             procs = new_procs
-    return sz_results
+    return sz_results, transient_error_type
 
 
 def process_urls(mode, urls, verbose, inject_failure, num_workers, s3config):
@@ -400,7 +507,9 @@ def process_urls(mode, urls, verbose, inject_failure, num_workers, s3config):
         print("%sing %d files.." % (mode.capitalize(), len(urls)), file=sys.stderr)
 
     start = time.time()
-    sz_results = start_workers(mode, urls, num_workers, inject_failure, s3config)
+    sz_results, transient_error_type = start_workers(
+        mode, urls, num_workers, inject_failure, s3config
+    )
     end = time.time()
 
     if verbose:
@@ -417,7 +526,7 @@ def process_urls(mode, urls, verbose, inject_failure, num_workers, s3config):
             ),
             file=sys.stderr,
         )
-    return sz_results
+    return sz_results, transient_error_type
 
 
 # Utility functions
@@ -502,11 +611,12 @@ class S3Ops(object):
                 # - the trailing slash is significant in S3
                 if "Contents" in page:
                     for key in page.get("Contents", []):
-                        url = url_base + key["Key"]
+                        key_path = key["Key"].lstrip("/")
+                        url = url_base + key_path
                         urlobj = S3Url(
                             url=url,
                             bucket=prefix_url.bucket,
-                            path=key["Key"],
+                            path=key_path,
                             local=generate_local_path(url),
                             prefix=prefix_url.url,
                         )
@@ -573,6 +683,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)
@@ -612,9 +724,21 @@ def generate_local_path(url, range="whole", suffix=None):
     quoted = url_quote(url)
     fname = quoted.split(b"/")[-1].replace(b".", b"_").replace(b"-", b"_")
     sha = sha1(quoted).hexdigest()
+
+    # Truncate fname to ensure the final filename doesn't exceed filesystem limits.
+    # Most filesystems have a 255 character limit. The structure is:
+    # <40-char-sha>-<fname>-<range>[-<suffix>]
+    # We need to leave room for: sha (40) + hyphens (2-3) + range (~10) + suffix (~10)
+    # This leaves roughly 190 characters for fname. We use 150 to be safe.
+    fname_decoded = fname.decode("utf-8")
+    max_fname_len = 150
+    if len(fname_decoded) > max_fname_len:
+        # Truncate and add an ellipsis to indicate truncation
+        fname_decoded = fname_decoded[:max_fname_len] + "..."
+
     if suffix:
-        return "-".join((sha, fname.decode("utf-8"), range, suffix))
-    return "-".join((sha, fname.decode("utf-8"), range))
+        return "-".join((sha, fname_decoded, range, suffix))
+    return "-".join((sha, fname_decoded, range))
 
 
 def parallel_op(op, lst, num_workers):
@@ -722,8 +846,8 @@ def cli():
     pass
 
 
-@tracing.cli_entrypoint("s3op/list")
 @cli.command("list", help="List S3 objects")
+@tracing.cli("s3op/list")
 @click.option(
     "--recursive/--no-recursive",
     default=False,
@@ -751,11 +875,16 @@ def lst(
     urllist = []
     to_iterate, _ = _populate_prefixes(prefixes, inputs)
     for _, prefix, url, _ in to_iterate:
-        src = urlparse(url)
+        src = urlparse(url, allow_fragments=False)
+        # We always consider the path being passed in to be a directory path so
+        # we add a trailing slash to the path if it doesn't already have one.
+        path_with_slash = src.path.lstrip("/")
+        if not path_with_slash.endswith("/"):
+            path_with_slash += "/"
         url = S3Url(
             url=url,
             bucket=src.netloc,
-            path=src.path.lstrip("/"),
+            path=path_with_slash,
             local=None,
             prefix=prefix,
         )
@@ -782,8 +911,8 @@ def lst(
             print(format_result_line(idx, url.prefix, url.url, str(size)))
 
 
-@tracing.cli_entrypoint("s3op/put")
 @cli.command(help="Upload files to S3")
+@tracing.cli("s3op/put")
 @click.option(
     "--file",
     "files",
@@ -857,7 +986,7 @@ def put(
                 yield input_line_idx, local, url, content_type, metadata, encryption
 
     def _make_url(idx, local, user_url, content_type, metadata, encryption):
-        src = urlparse(user_url)
+        src = urlparse(user_url, allow_fragments=False)
         url = S3Url(
             url=user_url,
             bucket=src.netloc,
@@ -885,7 +1014,7 @@ def put(
     ul_op = "upload"
     if not overwrite:
         ul_op = "info_upload"
-    sz_results = process_urls(
+    sz_results, transient_error_type = process_urls(
         ul_op, urls, verbose, inject_failure, num_workers, s3config
     )
     retry_lines = []
@@ -903,19 +1032,17 @@ def put(
         elif listing and sz == 0:
             out_lines.append(format_result_line(url.idx, url.url) + "\n")
         elif sz == -ERROR_TRANSIENT:
-            retry_lines.append(
-                json.dumps(
-                    {
-                        "idx": url.idx,
-                        "url": url.url,
-                        "local": url.local,
-                        "content_type": url.content_type,
-                        "metadata": url.metadata,
-                        "encryption": url.encryption,
-                    }
-                )
-                + "\n"
-            )
+            retry_data = {
+                "idx": url.idx,
+                "url": url.url,
+                "local": url.local,
+                "content_type": url.content_type,
+                "metadata": url.metadata,
+                "encryption": url.encryption,
+            }
+            if transient_error_type:
+                retry_data["transient_error_type"] = transient_error_type
+            retry_lines.append(json.dumps(retry_data) + "\n")
             # Output something to get a total count the first time around
             if not is_transient_retry:
                 out_lines.append("%d %s\n" % (url.idx, TRANSIENT_RETRY_LINE_CONTENT))
@@ -953,22 +1080,21 @@ def _populate_prefixes(prefixes, inputs):
             for idx, l in enumerate(f, start=len(prefixes)):
                 s = l.split(b" ")
                 if len(s) == 1:
+                    # User input format: <url>
                     url = url_unquote(s[0].strip())
                     prefixes.append((idx, url, url, None))
                 elif len(s) == 2:
+                    # User input format: <url> <range>
                     url = url_unquote(s[0].strip())
                     prefixes.append((idx, url, url, url_unquote(s[1].strip())))
-                else:
+                elif len(s) in (4, 5):
+                    # Retry format: <idx> <prefix> <url> <range> [<transient_error_type>]
+                    # The transient_error_type (5th field) is optional and only used for logging.
+                    # Lines with other field counts (e.g., 3) are silently ignored as invalid.
                     is_transient_retry = True
-                    if len(s) == 3:
-                        prefix = url = url_unquote(s[1].strip())
-                        range_info = url_unquote(s[2].strip())
-                    else:
-                        # Special case when we have both prefix and URL -- this is
-                        # used in recursive gets for example
-                        prefix = url_unquote(s[1].strip())
-                        url = url_unquote(s[2].strip())
-                        range_info = url_unquote(s[3].strip())
+                    prefix = url_unquote(s[1].strip())
+                    url = url_unquote(s[2].strip())
+                    range_info = url_unquote(s[3].strip())
                     if range_info == "<norange>":
                         range_info = None
                     prefixes.append(
@@ -977,8 +1103,8 @@ def _populate_prefixes(prefixes, inputs):
     return prefixes, is_transient_retry
 
 
-@tracing.cli_entrypoint("s3op/get")
 @cli.command(help="Download files from S3")
+@tracing.cli("s3op/get")
 @click.option(
     "--recursive/--no-recursive",
     default=False,
@@ -1032,7 +1158,7 @@ def get(
     urllist = []
     to_iterate, is_transient_retry = _populate_prefixes(prefixes, inputs)
     for idx, prefix, url, r in to_iterate:
-        src = urlparse(url)
+        src = urlparse(url, allow_fragments=False)
         url = S3Url(
             url=url,
             bucket=src.netloc,
@@ -1079,7 +1205,7 @@ def get(
 
     # exclude the non-existent files from loading
     to_load = [url for url, size in urls if size is not None]
-    sz_results = process_urls(
+    sz_results, transient_error_type = process_urls(
         dl_op, to_load, verbose, inject_failure, num_workers, s3config
     )
     # We check if there is any access denied
@@ -1103,6 +1229,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
@@ -1113,19 +1241,19 @@ def get(
                 break
             out_lines.append(format_result_line(url.idx, url.url) + "\n")
         elif sz == -ERROR_TRANSIENT:
-            retry_lines.append(
-                " ".join(
-                    [
-                        str(url.idx),
-                        url_quote(url.prefix).decode(encoding="utf-8"),
-                        url_quote(url.url).decode(encoding="utf-8"),
-                        url_quote(url.range).decode(encoding="utf-8")
-                        if url.range
-                        else "<norange>",
-                    ]
-                )
-                + "\n"
-            )
+            retry_line_parts = [
+                str(url.idx),
+                url_quote(url.prefix).decode(encoding="utf-8"),
+                url_quote(url.url).decode(encoding="utf-8"),
+                (
+                    url_quote(url.range).decode(encoding="utf-8")
+                    if url.range
+                    else "<norange>"
+                ),
+            ]
+            if transient_error_type:
+                retry_line_parts.append(transient_error_type)
+            retry_lines.append(" ".join(retry_line_parts) + "\n")
             # First time around, we output something to indicate the total length
             if not is_transient_retry:
                 out_lines.append("%d %s\n" % (url.idx, TRANSIENT_RETRY_LINE_CONTENT))
@@ -1177,7 +1305,7 @@ def info(
     urllist = []
     to_iterate, is_transient_retry = _populate_prefixes(prefixes, inputs)
     for idx, prefix, url, _ in to_iterate:
-        src = urlparse(url)
+        src = urlparse(url, allow_fragments=False)
         url = S3Url(
             url=url,
             bucket=src.netloc,
@@ -1191,7 +1319,7 @@ def info(
             exit(ERROR_INVALID_URL, url)
         urllist.append(url)
 
-    sz_results = process_urls(
+    sz_results, transient_error_type = process_urls(
         "info", urllist, verbose, inject_failure, num_workers, s3config
     )
 
@@ -1204,10 +1332,15 @@ def info(
                 format_result_line(url.idx, url.prefix, url.url, url.local) + "\n"
             )
         else:
-            retry_lines.append(
-                "%d %s <norange>\n"
-                % (url.idx, url_quote(url.url).decode(encoding="utf-8"))
-            )
+            retry_line_parts = [
+                str(url.idx),
+                url_quote(url.prefix).decode(encoding="utf-8"),
+                url_quote(url.url).decode(encoding="utf-8"),
+                "<norange>",
+            ]
+            if transient_error_type:
+                retry_line_parts.append(transient_error_type)
+            retry_lines.append(" ".join(retry_line_parts) + "\n")
             if not is_transient_retry:
                 out_lines.append("%d %s\n" % (url.idx, TRANSIENT_RETRY_LINE_CONTENT))
 

metaflow/plugins/env_escape/__init__.py

@@ -124,9 +124,9 @@ def load():
     cur_path = os.path.dirname(__file__)
     sys.path = [p for p in old_paths if p != cur_path]
     # Handle special case where we launch a shell (including with a command)
-    # and we are in the CWD (searched if '' is the first element of sys.path)
-    if cur_path == os.getcwd() and sys.path[0] == '':
-        sys.path = sys.path[1:]
+    # and we are in the CWD (searched if '' is present in sys.path)
+    if cur_path == os.getcwd() and '' in sys.path:
+        sys.path.remove("")
 
     # Remove the module (this file) to reload it properly. Do *NOT* update sys.modules but
     # modify directly since it may be referenced elsewhere