ob-metaflow 2.11.13.1__py2.py3-none-any.whl → 2.19.7.1rc0__py2.py3-none-any.whl

metaflow/plugins/{metadata → metadata_providers}/local.py

@@ -2,20 +2,30 @@ import collections
 import glob
 import json
 import os
+import re
 import random
 import tempfile
 import time
 from collections import namedtuple
+from typing import List
 
 from metaflow.exception import MetaflowInternalError, MetaflowTaggingError
-from metaflow.metadata.metadata import ObjectOrder
+from metaflow.metadata_provider.metadata import ObjectOrder
 from metaflow.metaflow_config import DATASTORE_LOCAL_DIR
-from metaflow.metadata import MetadataProvider
+from metaflow.metadata_provider import MetadataProvider
 from metaflow.tagging_util import MAX_USER_TAG_SET_SIZE, validate_tags
 
 
 class LocalMetadataProvider(MetadataProvider):
     TYPE = "local"
+    DATASTORE_DIR = DATASTORE_LOCAL_DIR  # ".metaflow"
+
+    @classmethod
+    def _get_storage_class(cls):
+        # This method is meant to be overridden
+        from metaflow.plugins.datastores.local_storage import LocalStorage
+
+        return LocalStorage
 
     def __init__(self, environment, flow, event_logger, monitor):
         super(LocalMetadataProvider, self).__init__(
@@ -24,30 +34,28 @@ class LocalMetadataProvider(MetadataProvider):
 
     @classmethod
     def compute_info(cls, val):
-        from metaflow.plugins.datastores.local_storage import LocalStorage
+        storage_class = cls._get_storage_class()
 
-        v = os.path.realpath(os.path.join(val, DATASTORE_LOCAL_DIR))
+        v = os.path.realpath(os.path.join(val, cls.DATASTORE_DIR))
         if os.path.isdir(v):
-            LocalStorage.datastore_root = v
+            storage_class.datastore_root = v
             return val
         raise ValueError(
-            "Could not find directory %s in directory %s" % (DATASTORE_LOCAL_DIR, val)
+            "Could not find directory %s in directory %s" % (cls.DATASTORE_DIR, val)
         )
 
     @classmethod
     def default_info(cls):
-        from metaflow.plugins.datastores.local_storage import LocalStorage
+        storage_class = cls._get_storage_class()
 
         def print_clean(line, **kwargs):
             print(line)
 
-        v = LocalStorage.get_datastore_root_from_config(
+        v = storage_class.get_datastore_root_from_config(
             print_clean, create_on_absent=False
         )
         if v is None:
-            return (
-                "<No %s directory found in current working tree>" % DATASTORE_LOCAL_DIR
-            )
+            return "<No %s directory found in current working tree>" % cls.DATASTORE_DIR
         return os.path.dirname(v)
 
     def version(self):
@@ -100,7 +108,7 @@ class LocalMetadataProvider(MetadataProvider):
     def register_data_artifacts(
         self, run_id, step_name, task_id, attempt_id, artifacts
     ):
-        meta_dir = self._create_and_get_metadir(
+        meta_dir = self.__class__._create_and_get_metadir(
             self._flow_name, run_id, step_name, task_id
         )
         artlist = self._artifacts_to_json(
@@ -110,7 +118,7 @@ class LocalMetadataProvider(MetadataProvider):
         self._save_meta(meta_dir, artdict)
 
     def register_metadata(self, run_id, step_name, task_id, metadata):
-        meta_dir = self._create_and_get_metadir(
+        meta_dir = self.__class__._create_and_get_metadir(
             self._flow_name, run_id, step_name, task_id
         )
         metalist = self._metadata_to_json(run_id, step_name, task_id, metadata)
@@ -130,9 +138,7 @@ class LocalMetadataProvider(MetadataProvider):
 
         def _optimistically_mutate():
             # get existing tags
-            run = LocalMetadataProvider.get_object(
-                "run", "self", {}, None, flow_id, run_id
-            )
+            run = cls.get_object("run", "self", {}, None, flow_id, run_id)
             if not run:
                 raise MetaflowTaggingError(
                     msg="Run not found (%s, %s)" % (flow_id, run_id)
@@ -165,15 +171,13 @@ class LocalMetadataProvider(MetadataProvider):
             validate_tags(next_user_tags_set, existing_tags=existing_user_tag_set)
 
             # write new tag set to file system
-            LocalMetadataProvider._persist_tags_for_run(
+            cls._persist_tags_for_run(
                 flow_id, run_id, next_user_tags_set, existing_system_tag_set
             )
 
             # read tags back from file system to see if our optimism is misplaced
             # I.e. did a concurrent mutate overwrite our change
-            run = LocalMetadataProvider.get_object(
-                "run", "self", {}, None, flow_id, run_id
-            )
+            run = cls.get_object("run", "self", {}, None, flow_id, run_id)
             if not run:
                 raise MetaflowTaggingError(
                     msg="Run not found for read-back check (%s, %s)" % (flow_id, run_id)
@@ -202,6 +206,70 @@ class LocalMetadataProvider(MetadataProvider):
             "Tagging failed due to too many conflicting updates from other processes"
         )
 
+    @classmethod
+    def filter_tasks_by_metadata(
+        cls,
+        flow_name: str,
+        run_id: str,
+        step_name: str,
+        field_name: str,
+        pattern: str,
+    ) -> List[str]:
+        """
+        Filter tasks by metadata field and pattern, returning task pathspecs that match criteria.
+
+        Parameters
+        ----------
+        flow_name : str
+            Identifier for the flow
+        run_id : str
+            Identifier for the run
+        step_name : str
+            Name of the step to query tasks from
+        field_name : str
+            Name of metadata field to query
+        pattern : str
+            Pattern to match in metadata field value
+
+        Returns
+        -------
+        List[str]
+            List of task pathspecs that match the query criteria
+        """
+        tasks = cls.get_object("step", "task", {}, None, flow_name, run_id, step_name)
+        if not tasks:
+            return []
+
+        regex = re.compile(pattern)
+        matching_task_pathspecs = []
+
+        for task in tasks:
+            task_id = task.get("task_id")
+            if not task_id:
+                continue
+
+            if pattern == ".*":
+                # If the pattern is ".*", we can match all tasks without reading metadata
+                matching_task_pathspecs.append(
+                    f"{flow_name}/{run_id}/{step_name}/{task_id}"
+                )
+                continue
+
+            metadata = cls.get_object(
+                "task", "metadata", {}, None, flow_name, run_id, step_name, task_id
+            )
+
+            if any(
+                meta.get("field_name") == field_name
+                and regex.match(meta.get("value", ""))
+                for meta in metadata
+            ):
+                matching_task_pathspecs.append(
+                    f"{flow_name}/{run_id}/{step_name}/{task_id}"
+                )
+
+        return matching_task_pathspecs
+
     @classmethod
     def _get_object_internal(
         cls, obj_type, obj_order, sub_type, sub_order, filters, attempt, *args
@@ -213,8 +281,6 @@ class LocalMetadataProvider(MetadataProvider):
         if obj_type not in ("root", "flow", "run", "step", "task", "artifact"):
             raise MetaflowInternalError(msg="Unexpected object type %s" % obj_type)
 
-        from metaflow.plugins.datastores.local_storage import LocalStorage
-
         if obj_type == "artifact":
             # Artifacts are actually part of the tasks in the filesystem
             # E.g. we get here for (obj_type, sub_type) == (artifact, self)
@@ -241,13 +307,13 @@ class LocalMetadataProvider(MetadataProvider):
 
         # Special handling of self, artifact, and metadata
         if sub_type == "self":
-            meta_path = LocalMetadataProvider._get_metadir(*args[:obj_order])
+            meta_path = cls._get_metadir(*args[:obj_order])
             if meta_path is None:
                 return None
             self_file = os.path.join(meta_path, "_self.json")
             if os.path.isfile(self_file):
                 obj = MetadataProvider._apply_filter(
-                    [LocalMetadataProvider._read_json_file(self_file)], filters
+                    [cls._read_json_file(self_file)], filters
                 )[0]
                 # For non-descendants of a run, we are done
 
@@ -258,7 +324,7 @@ class LocalMetadataProvider(MetadataProvider):
                     raise MetaflowInternalError(
                         msg="Unexpected object type %s" % obj_type
                     )
-                run = LocalMetadataProvider.get_object(
+                run = cls.get_object(
                     "run", "self", {}, None, *args[:RUN_ORDER]  # *[flow_id, run_id]
                 )
                 if not run:
@@ -275,7 +341,7 @@ class LocalMetadataProvider(MetadataProvider):
             if obj_type not in ("root", "flow", "run", "step", "task"):
                 raise MetaflowInternalError(msg="Unexpected object type %s" % obj_type)
 
-            meta_path = LocalMetadataProvider._get_metadir(*args[:obj_order])
+            meta_path = cls._get_metadir(*args[:obj_order])
             result = []
             if meta_path is None:
                 return result
@@ -286,9 +352,7 @@ class LocalMetadataProvider(MetadataProvider):
                 attempts_done = sorted(glob.iglob(attempt_done_files))
                 if attempts_done:
                     successful_attempt = int(
-                        LocalMetadataProvider._read_json_file(attempts_done[-1])[
-                            "value"
-                        ]
+                        cls._read_json_file(attempts_done[-1])["value"]
                     )
             if successful_attempt is not None:
                 which_artifact = "*"
@@ -299,10 +363,10 @@ class LocalMetadataProvider(MetadataProvider):
                     "%d_artifact_%s.json" % (successful_attempt, which_artifact),
                 )
                 for obj in glob.iglob(artifact_files):
-                    result.append(LocalMetadataProvider._read_json_file(obj))
+                    result.append(cls._read_json_file(obj))
 
             # We are getting artifacts. We should overlay with ancestral run's tags
-            run = LocalMetadataProvider.get_object(
+            run = cls.get_object(
                 "run", "self", {}, None, *args[:RUN_ORDER]  # *[flow_id, run_id]
             )
             if not run:
@@ -322,12 +386,12 @@ class LocalMetadataProvider(MetadataProvider):
             if obj_type not in ("root", "flow", "run", "step", "task"):
                 raise MetaflowInternalError(msg="Unexpected object type %s" % obj_type)
             result = []
-            meta_path = LocalMetadataProvider._get_metadir(*args[:obj_order])
+            meta_path = cls._get_metadir(*args[:obj_order])
             if meta_path is None:
                 return result
             files = os.path.join(meta_path, "sysmeta_*")
             for obj in glob.iglob(files):
-                result.append(LocalMetadataProvider._read_json_file(obj))
+                result.append(cls._read_json_file(obj))
             return result
 
         # For the other types, we locate all the objects we need to find and return them
@@ -335,14 +399,13 @@ class LocalMetadataProvider(MetadataProvider):
             raise MetaflowInternalError(msg="Unexpected object type %s" % obj_type)
         if sub_type not in ("flow", "run", "step", "task"):
             raise MetaflowInternalError(msg="unexpected sub type %s" % sub_type)
-        obj_path = LocalMetadataProvider._make_path(
-            *args[:obj_order], create_on_absent=False
-        )
+        obj_path = cls._make_path(*args[:obj_order], create_on_absent=False)
         result = []
         if obj_path is None:
             return result
         skip_dirs = "*/" * (sub_order - obj_order)
-        all_meta = os.path.join(obj_path, skip_dirs, LocalStorage.METADATA_DIR)
+        storage_class = cls._get_storage_class()
+        all_meta = os.path.join(obj_path, skip_dirs, storage_class.METADATA_DIR)
         SelfInfo = collections.namedtuple("SelfInfo", ["filepath", "run_id"])
         self_infos = []
         for meta_path in glob.iglob(all_meta):
@@ -352,9 +415,7 @@ class LocalMetadataProvider(MetadataProvider):
             run_id = None
             # flow and run do not need info from ancestral run
             if sub_type in ("step", "task"):
-                run_id = LocalMetadataProvider._deduce_run_id_from_meta_dir(
-                    meta_path, sub_type
-                )
+                run_id = cls._deduce_run_id_from_meta_dir(meta_path, sub_type)
                 # obj_type IS run, or more granular than run, let's do sanity check vs args
                 if obj_order >= RUN_ORDER:
                     if run_id != args[RUN_ORDER - 1]:
@@ -364,10 +425,10 @@ class LocalMetadataProvider(MetadataProvider):
             self_infos.append(SelfInfo(filepath=self_file, run_id=run_id))
 
         for self_info in self_infos:
-            obj = LocalMetadataProvider._read_json_file(self_info.filepath)
+            obj = cls._read_json_file(self_info.filepath)
             if self_info.run_id:
                 flow_id_from_args = args[0]
-                run = LocalMetadataProvider.get_object(
+                run = cls.get_object(
                     "run",
                     "self",
                     {},
@@ -386,8 +447,8 @@ class LocalMetadataProvider(MetadataProvider):
 
         return MetadataProvider._apply_filter(result, filters)
 
-    @staticmethod
-    def _deduce_run_id_from_meta_dir(meta_dir_path, sub_type):
+    @classmethod
+    def _deduce_run_id_from_meta_dir(cls, meta_dir_path, sub_type):
         curr_order = ObjectOrder.type_to_order(sub_type)
         levels_to_ascend = curr_order - ObjectOrder.type_to_order("run")
         if levels_to_ascend < 0:
@@ -402,8 +463,8 @@ class LocalMetadataProvider(MetadataProvider):
             )
         return run_id
 
-    @staticmethod
-    def _makedirs(path):
+    @classmethod
+    def _makedirs(cls, path):
         # this is for python2 compatibility.
         # Python3 has os.makedirs(exist_ok=True).
         try:
@@ -415,17 +476,15 @@ class LocalMetadataProvider(MetadataProvider):
             else:
                 raise
 
-    @staticmethod
-    def _persist_tags_for_run(flow_id, run_id, tags, system_tags):
-        subpath = LocalMetadataProvider._create_and_get_metadir(
-            flow_name=flow_id, run_id=run_id
-        )
+    @classmethod
+    def _persist_tags_for_run(cls, flow_id, run_id, tags, system_tags):
+        subpath = cls._create_and_get_metadir(flow_name=flow_id, run_id=run_id)
         selfname = os.path.join(subpath, "_self.json")
         if not os.path.isfile(selfname):
             raise MetaflowInternalError(
                 msg="Could not verify Run existence on disk - missing %s" % selfname
             )
-        LocalMetadataProvider._save_meta(
+        cls._save_meta(
             subpath,
             {
                 "_self": MetadataProvider._run_to_json_static(
@@ -442,11 +501,11 @@ class LocalMetadataProvider(MetadataProvider):
             tags = set()
         if sys_tags is None:
             sys_tags = set()
-        subpath = self._create_and_get_metadir(
+        subpath = self.__class__._create_and_get_metadir(
             self._flow_name, run_id, step_name, task_id
         )
         selfname = os.path.join(subpath, "_self.json")
-        self._makedirs(subpath)
+        self.__class__._makedirs(subpath)
         if os.path.isfile(selfname):
             # There is a race here, but we are not aiming to make this as solid as
             # the metadata service. This is used primarily for concurrent resumes,
@@ -483,26 +542,31 @@ class LocalMetadataProvider(MetadataProvider):
         self._register_system_metadata(run_id, step_name, task_id, attempt)
         return to_return
 
-    @staticmethod
+    @classmethod
     def _make_path(
-        flow_name=None, run_id=None, step_name=None, task_id=None, create_on_absent=True
+        cls,
+        flow_name=None,
+        run_id=None,
+        step_name=None,
+        task_id=None,
+        create_on_absent=True,
     ):
 
-        from metaflow.plugins.datastores.local_storage import LocalStorage
+        storage_class = cls._get_storage_class()
 
-        if LocalStorage.datastore_root is None:
+        if storage_class.datastore_root is None:
 
             def print_clean(line, **kwargs):
                 print(line)
 
-            LocalStorage.datastore_root = LocalStorage.get_datastore_root_from_config(
+            storage_class.datastore_root = storage_class.get_datastore_root_from_config(
                 print_clean, create_on_absent=create_on_absent
             )
-        if LocalStorage.datastore_root is None:
+        if storage_class.datastore_root is None:
             return None
 
         if flow_name is None:
-            return LocalStorage.datastore_root
+            return storage_class.datastore_root
         components = []
         if flow_name:
             components.append(flow_name)
@@ -512,37 +576,35 @@ class LocalMetadataProvider(MetadataProvider):
                     components.append(step_name)
                     if task_id:
                         components.append(task_id)
-        return LocalStorage().full_uri(LocalStorage.path_join(*components))
+        return storage_class().full_uri(storage_class.path_join(*components))
 
-    @staticmethod
+    @classmethod
     def _create_and_get_metadir(
-        flow_name=None, run_id=None, step_name=None, task_id=None
+        cls, flow_name=None, run_id=None, step_name=None, task_id=None
     ):
-        from metaflow.plugins.datastores.local_storage import LocalStorage
+        storage_class = cls._get_storage_class()
 
-        root_path = LocalMetadataProvider._make_path(
-            flow_name, run_id, step_name, task_id
-        )
-        subpath = os.path.join(root_path, LocalStorage.METADATA_DIR)
-        LocalMetadataProvider._makedirs(subpath)
+        root_path = cls._make_path(flow_name, run_id, step_name, task_id)
+        subpath = os.path.join(root_path, storage_class.METADATA_DIR)
+        cls._makedirs(subpath)
         return subpath
 
-    @staticmethod
-    def _get_metadir(flow_name=None, run_id=None, step_name=None, task_id=None):
-        from metaflow.plugins.datastores.local_storage import LocalStorage
+    @classmethod
+    def _get_metadir(cls, flow_name=None, run_id=None, step_name=None, task_id=None):
+        storage_class = cls._get_storage_class()
 
-        root_path = LocalMetadataProvider._make_path(
+        root_path = cls._make_path(
             flow_name, run_id, step_name, task_id, create_on_absent=False
         )
         if root_path is None:
             return None
-        subpath = os.path.join(root_path, LocalStorage.METADATA_DIR)
+        subpath = os.path.join(root_path, storage_class.METADATA_DIR)
         if os.path.isdir(subpath):
             return subpath
         return None
 
-    @staticmethod
-    def _dump_json_to_file(filepath, data, allow_overwrite=False):
+    @classmethod
+    def _dump_json_to_file(cls, filepath, data, allow_overwrite=False):
         if os.path.isfile(filepath) and not allow_overwrite:
             return
         try:
@@ -556,15 +618,13 @@ class LocalMetadataProvider(MetadataProvider):
             if f and os.path.isfile(f.name):
                 os.remove(f.name)
 
-    @staticmethod
-    def _read_json_file(filepath):
+    @classmethod
+    def _read_json_file(cls, filepath):
         with open(filepath, "r") as f:
             return json.load(f)
 
-    @staticmethod
-    def _save_meta(root_dir, metadict, allow_overwrite=False):
+    @classmethod
+    def _save_meta(cls, root_dir, metadict, allow_overwrite=False):
         for name, datum in metadict.items():
             filename = os.path.join(root_dir, "%s.json" % name)
-            LocalMetadataProvider._dump_json_to_file(
-                filename, datum, allow_overwrite=allow_overwrite
-            )
+            cls._dump_json_to_file(filename, datum, allow_overwrite=allow_overwrite)

metaflow/plugins/{metadata → metadata_providers}/service.py

@@ -1,23 +1,20 @@
 import os
 import random
+import time
 
 import requests
-import time
 
+from typing import List
 from metaflow.exception import (
     MetaflowException,
-    MetaflowTaggingError,
     MetaflowInternalError,
+    MetaflowTaggingError,
 )
-from metaflow.metaflow_config import (
-    SERVICE_RETRY_COUNT,
-    SERVICE_HEADERS,
-    SERVICE_URL,
-)
-from metaflow.metadata import MetadataProvider
-from metaflow.metadata.heartbeat import HB_URL_KEY
+from metaflow.metadata_provider import MetadataProvider
+from metaflow.metadata_provider.heartbeat import HB_URL_KEY
+from metaflow.metaflow_config import SERVICE_HEADERS, SERVICE_RETRY_COUNT, SERVICE_URL
 from metaflow.sidecar import Message, MessageTypes, Sidecar
-
+from urllib.parse import urlencode
 from metaflow.util import version_parse
 
 
@@ -39,6 +36,23 @@ class ServiceException(MetaflowException):
 class ServiceMetadataProvider(MetadataProvider):
     TYPE = "service"
 
+    _session = requests.Session()
+    _session.mount(
+        "http://",
+        requests.adapters.HTTPAdapter(
+            pool_connections=20,
+            pool_maxsize=20,
+            max_retries=0,  # Handle retries explicitly
+            pool_block=False,
+        ),
+    )
+    _session.mount(
+        "https://",
+        requests.adapters.HTTPAdapter(
+            pool_connections=20, pool_maxsize=20, max_retries=0, pool_block=False
+        ),
+    )
+
     _supports_attempt_gets = None
     _supports_tag_mutation = None
 
@@ -58,12 +72,18 @@ class ServiceMetadataProvider(MetadataProvider):
     @classmethod
     def compute_info(cls, val):
         v = val.rstrip("/")
-        try:
-            resp = requests.get(os.path.join(v, "ping"), headers=SERVICE_HEADERS.copy())
-            resp.raise_for_status()
-        except:  # noqa E722
-            raise ValueError("Metaflow service [%s] unreachable." % v)
-        return v
+        for i in range(SERVICE_RETRY_COUNT):
+            try:
+                resp = cls._session.get(
+                    os.path.join(v, "ping"), headers=SERVICE_HEADERS.copy()
+                )
+                resp.raise_for_status()
+            except:  # noqa E722
+                time.sleep(2 ** (i - 1))
+            else:
+                return v
+
+        raise ValueError("Metaflow service [%s] unreachable." % v)
 
     @classmethod
     def default_info(cls):
@@ -304,6 +324,64 @@ class ServiceMetadataProvider(MetadataProvider):
             self._register_system_metadata(run_id, step_name, task["task_id"], attempt)
         return task["task_id"], did_create
 
+    @classmethod
+    def filter_tasks_by_metadata(
+        cls,
+        flow_name: str,
+        run_id: str,
+        step_name: str,
+        field_name: str,
+        pattern: str,
+    ) -> List[str]:
+        """
+        Filter tasks by metadata field and pattern, returning task pathspecs that match criteria.
+
+        Parameters
+        ----------
+        flow_name : str
+            Flow name, that the run belongs to.
+        run_id: str
+            Run id, together with flow_id, that identifies the specific Run whose tasks to query
+        step_name: str
+            Step name to query tasks from
+        field_name: str
+            Metadata field name to query
+        pattern: str
+            Pattern to match in metadata field value
+
+        Returns
+        -------
+        List[str]
+            List of task pathspecs that satisfy the query
+        """
+        query_params = {}
+
+        if pattern == ".*":
+            # we do not need to filter tasks at all if pattern allows 'any'
+            query_params = {}
+        else:
+            if field_name:
+                query_params["metadata_field_name"] = field_name
+            if pattern:
+                query_params["pattern"] = pattern
+
+        url = ServiceMetadataProvider._obj_path(flow_name, run_id, step_name)
+        url = f"{url}/filtered_tasks?{urlencode(query_params)}"
+
+        try:
+            resp, _ = cls._request(None, url, "GET")
+        except Exception as e:
+            if e.http_code == 404:
+                # filter_tasks_by_metadata endpoint does not exist in the version of metadata service
+                # deployed currently. Raise a more informative error message.
+                raise MetaflowInternalError(
+                    "The version of metadata service deployed currently does not support filtering tasks by metadata. "
+                    "Upgrade Metadata service to version 2.5.0 or greater to use this feature."
+                ) from e
+            # Other unknown exception
+            raise e
+        return resp
+
     @staticmethod
     def _obj_path(
         flow_name,
@@ -412,27 +490,27 @@ class ServiceMetadataProvider(MetadataProvider):
                 if method == "GET":
                     if monitor:
                         with monitor.measure("metaflow.service_metadata.get"):
-                            resp = requests.get(url, headers=SERVICE_HEADERS.copy())
+                            resp = cls._session.get(url, headers=SERVICE_HEADERS.copy())
                     else:
-                        resp = requests.get(url, headers=SERVICE_HEADERS.copy())
+                        resp = cls._session.get(url, headers=SERVICE_HEADERS.copy())
                 elif method == "POST":
                     if monitor:
                         with monitor.measure("metaflow.service_metadata.post"):
-                            resp = requests.post(
+                            resp = cls._session.post(
                                 url, headers=SERVICE_HEADERS.copy(), json=data
                             )
                     else:
-                        resp = requests.post(
+                        resp = cls._session.post(
                             url, headers=SERVICE_HEADERS.copy(), json=data
                         )
                 elif method == "PATCH":
                     if monitor:
                         with monitor.measure("metaflow.service_metadata.patch"):
-                            resp = requests.patch(
+                            resp = cls._session.patch(
                                 url, headers=SERVICE_HEADERS.copy(), json=data
                             )
                     else:
-                        resp = requests.patch(
+                        resp = cls._session.patch(
                             url, headers=SERVICE_HEADERS.copy(), json=data
                         )
                 else:
@@ -475,7 +553,6 @@ class ServiceMetadataProvider(MetadataProvider):
                         resp.text,
                     )
             time.sleep(2**i)
-
         if resp:
             raise ServiceException(
                 "Metadata request (%s) failed (code %s): %s"
@@ -499,9 +576,9 @@ class ServiceMetadataProvider(MetadataProvider):
             try:
                 if monitor:
                     with monitor.measure("metaflow.service_metadata.get"):
-                        resp = requests.get(url, headers=SERVICE_HEADERS.copy())
+                        resp = cls._session.get(url, headers=SERVICE_HEADERS.copy())
                 else:
-                    resp = requests.get(url, headers=SERVICE_HEADERS.copy())
+                    resp = cls._session.get(url, headers=SERVICE_HEADERS.copy())
             except:
                 if monitor:
                     with monitor.count("metaflow.service_metadata.failed_request"):
@@ -514,7 +591,7 @@ class ServiceMetadataProvider(MetadataProvider):
             else:
                 if resp.status_code < 300:
                     return resp.headers.get("METADATA_SERVICE_VERSION", None)
-                elif resp.status_code != 503:
+                elif resp.status_code not in (503, 500):
                     raise ServiceException(
                         "Metadata request (%s) failed"
                         " (code %s): %s" % (url, resp.status_code, resp.text),