ob-metaflow 2.13.10.1__py2.py3-none-any.whl → 2.14.0.2rc0__py2.py3-none-any.whl

metaflow/plugins/argo/argo_workflows.py

@@ -1966,6 +1966,11 @@ class ArgoWorkflows(object):
                 resources["disk"],
             )
 
+            extended_resources = resources.get("extended_resources", {})
+
+            qos_requests = {**qos_requests, **extended_resources}
+            qos_limits = {**qos_limits, **extended_resources}
+
             # Create a ContainerTemplate for this node. Ideally, we would have
             # liked to inline this ContainerTemplate and avoid scanning the workflow
             # twice, but due to issues with variable substitution, we will have to
@@ -2022,6 +2027,7 @@ class ArgoWorkflows(object):
                     shared_memory=shared_memory,
                     port=port,
                     qos=resources["qos"],
+                    extended_resources=extended_resources,
                 )
 
                 for k, v in env.items():

metaflow/plugins/kubernetes/kubernetes.py

@@ -194,6 +194,7 @@ class Kubernetes(object):
         port=None,
         num_parallel=None,
         qos=None,
+        extended_resources=None,
     ):
         name = "js-%s" % str(uuid4())[:6]
         jobset = (
@@ -227,6 +228,7 @@ class Kubernetes(object):
                 port=port,
                 num_parallel=num_parallel,
                 qos=qos,
+                extended_resources=extended_resources,
             )
             .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,
+        extended_resources=None,
     ):
         if env is None:
             env = {}
@@ -546,6 +549,7 @@ class Kubernetes(object):
                 shared_memory=shared_memory,
                 port=port,
                 qos=qos,
+                extended_resources=extended_resources,
             )
             .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(
+    "--extended-resources",
+    default=None,
+    type=JSONTypeClass(),
+    multiple=False,
+)
 @click.pass_context
 def step(
     ctx,
@@ -176,6 +182,7 @@ def step(
     qos=None,
     labels=None,
     annotations=None,
+    extended_resources=None,
     **kwargs
 ):
     def echo(msg, stream="stderr", job_id=None, **kwargs):
@@ -319,6 +326,7 @@ def step(
                 qos=qos,
                 labels=labels,
                 annotations=annotations,
+                extended_resources=extended_resources,
             )
     except Exception:
         traceback.print_exc(chain=False)

metaflow/plugins/kubernetes/kubernetes_decorator.py

@@ -123,6 +123,9 @@ 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
+    extended_resources: Dict[str, str], optional, default None
+        Extended resources to be requested for the pod.
+        https://kubernetes.io/docs/tasks/administer-cluster/extended-resource-node/
     """
 
     name = "kubernetes"
@@ -153,6 +156,7 @@ class KubernetesDecorator(StepDecorator):
         "executable": None,
         "hostname_resolution_timeout": 10 * 60,
         "qos": KUBERNETES_QOS,
+        "extended_resources": {},
     }
     package_url = None
     package_sha = None
@@ -488,6 +492,7 @@ class KubernetesDecorator(StepDecorator):
                     "persistent_volume_claims",
                     "labels",
                     "annotations",
+                    "extended_resources",
                 ]:
                     cli_args.command_options[k] = json.dumps(v)
                 else:

metaflow/plugins/kubernetes/kubernetes_job.py

@@ -89,10 +89,11 @@ class KubernetesJob(object):
 
         additional_obp_configs = {
             "OBP_PERIMETER": initial_configs["OBP_PERIMETER"],
-            "OBP_INTEGRATIONS_URL": initial_configs[
-                "OBP_INTEGRATIONS_URL"
-            ],
+            "OBP_INTEGRATIONS_URL": initial_configs["OBP_INTEGRATIONS_URL"],
         }
+        extended_resources = self._kwargs.get("extended_resources", {})
+        qos_requests = {**qos_requests, **extended_resources}
+        qos_limits = {**qos_limits, **extended_resources}
 
         return client.V1JobSpec(
             # Retries are handled by Metaflow when it is responsible for
@@ -118,6 +119,7 @@ class KubernetesJob(object):
                     containers=[
                         client.V1Container(
                             command=self._kwargs["command"],
+                            security_context=client.V1SecurityContext(privileged=True),
                             termination_message_policy="FallbackToLogsOnError",
                             ports=(
                                 []

metaflow/plugins/kubernetes/kubernetes_jobsets.py

@@ -560,6 +560,11 @@ class JobSetSpec(object):
             self._kwargs["memory"],
             self._kwargs["disk"],
         )
+
+        extended_resources = self._kwargs.get("extended_resources", {})
+        qos_requests = {**qos_requests, **extended_resources}
+        qos_limits = {**qos_limits, **extended_resources}
+
         return dict(
             name=self.name,
             template=client.api_client.ApiClient().sanitize_for_serialization(
@@ -611,6 +616,9 @@ class JobSetSpec(object):
                                                 )
                                             ]
                                         ),
+                                        security_context=client.V1SecurityContext(
+                                            privileged=True
+                                        ),
                                         env=[
                                             client.V1EnvVar(name=k, value=str(v))
                                             for k, v in self._kwargs.get(

metaflow/plugins/metadata_providers/service.py

@@ -1,23 +1,18 @@
 import os
 import random
+import time
 
 import requests
-import time
 
 from metaflow.exception import (
     MetaflowException,
-    MetaflowTaggingError,
     MetaflowInternalError,
-)
-from metaflow.metaflow_config import (
-    SERVICE_RETRY_COUNT,
-    SERVICE_HEADERS,
-    SERVICE_URL,
+    MetaflowTaggingError,
 )
 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 metaflow.util import version_parse
 
 
@@ -39,6 +34,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
 
@@ -59,7 +71,9 @@ class ServiceMetadataProvider(MetadataProvider):
     def compute_info(cls, val):
         v = val.rstrip("/")
         try:
-            resp = requests.get(os.path.join(v, "ping"), headers=SERVICE_HEADERS.copy())
+            resp = cls._session.get(
+                os.path.join(v, "ping"), headers=SERVICE_HEADERS.copy()
+            )
             resp.raise_for_status()
         except:  # noqa E722
             raise ValueError("Metaflow service [%s] unreachable." % v)
@@ -412,27 +426,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 +489,6 @@ class ServiceMetadataProvider(MetadataProvider):
                         resp.text,
                     )
             time.sleep(2**i)
-
         if resp:
             raise ServiceException(
                 "Metadata request (%s) failed (code %s): %s"
@@ -499,9 +512,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"):

metaflow/runner/click_api.py

@@ -527,12 +527,6 @@ def extract_all_params(cmd_obj: Union[click.Command, click.Group]):
             )
             arg_parameters[each_param.name] = each_param
         elif isinstance(each_param, click.Option):
-            if each_param.hidden:
-                # Skip hidden options because users should not be setting those.
-                # These are typically internal only options (used by the Runner in part
-                # for example to pass state files or configs to pass local-config-file).
-                continue
-
             opt_params_sigs[each_param.name], annotations[each_param.name] = (
                 get_inspect_param_obj(each_param, inspect.Parameter.KEYWORD_ONLY)
             )

metaflow/version.py

@@ -1 +1 @@
-metaflow_version = "2.13.10.1"
+metaflow_version = "2.14.0.2rc0"

{ob_metaflow-2.13.10.1.dist-info → ob_metaflow-2.14.0.2rc0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: ob-metaflow
-Version: 2.13.10.1
+Version: 2.14.0.2rc0
 Summary: Metaflow: More Data Science, Less Engineering
 Author: Netflix, Outerbounds & the Metaflow Community
 Author-email: help@outerbounds.co
@@ -12,7 +12,7 @@ Requires-Dist: boto3
 Requires-Dist: pylint
 Requires-Dist: kubernetes
 Provides-Extra: stubs
-Requires-Dist: metaflow-stubs==2.13.10.1; extra == "stubs"
+Requires-Dist: metaflow-stubs==2.14.0.2rc0; extra == "stubs"
 Dynamic: author
 Dynamic: author-email
 Dynamic: description

{ob_metaflow-2.13.10.1.dist-info → ob_metaflow-2.14.0.2rc0.dist-info}/RECORD

@@ -36,7 +36,7 @@ metaflow/tuple_util.py,sha256=_G5YIEhuugwJ_f6rrZoelMFak3DqAR2tt_5CapS1XTY,830
 metaflow/unbounded_foreach.py,sha256=p184WMbrMJ3xKYHwewj27ZhRUsSj_kw1jlye5gA9xJk,387
 metaflow/util.py,sha256=mJBkV5tShIyCsLDeM1zygQGeciQVMrVPm_qI8Oi33G0,14656
 metaflow/vendor.py,sha256=FchtA9tH22JM-eEtJ2c9FpUdMn8sSb1VHuQS56EcdZk,5139
-metaflow/version.py,sha256=nEqMd26PFOk16IDyKeu-GgwaKxcoub6bcMATTlr6Rmg,31
+metaflow/version.py,sha256=oiICSOwc_b5wi5EFZxih7ogG7-vIVMe9j_VgVtch29E,33
 metaflow/_vendor/__init__.py,sha256=y_CiwUD3l4eAKvTVDZeqgVujMy31cAM1qjAB-HfI-9s,353
 metaflow/_vendor/typing_extensions.py,sha256=0nUs5p1A_UrZigrAVBoOEM6TxU37zzPDUtiij1ZwpNc,110417
 metaflow/_vendor/zipp.py,sha256=ajztOH-9I7KA_4wqDYygtHa6xUBVZgFpmZ8FE74HHHI,8425
@@ -181,7 +181,7 @@ metaflow/plugins/airflow/sensors/s3_sensor.py,sha256=iDReG-7FKnumrtQg-HY6cCUAAqN
 metaflow/plugins/argo/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 metaflow/plugins/argo/argo_client.py,sha256=PS_cYGnPw9h4X7TP_plObDH3clMw4reOsBLkkGPTd0Y,16282
 metaflow/plugins/argo/argo_events.py,sha256=_C1KWztVqgi3zuH57pInaE9OzABc2NnncC-zdwOMZ-w,5909
-metaflow/plugins/argo/argo_workflows.py,sha256=zx5aVyLLwnBzjhu7d9-DotQaw6WgHEoWmZ0KqRNdJCo,183449
+metaflow/plugins/argo/argo_workflows.py,sha256=yWyDlFRQX6A7pwI08Z1MGqY-gbFTnvSp5WJwg66ZdlM,183711
 metaflow/plugins/argo/argo_workflows_cli.py,sha256=kzZF2XRr5If9t135wMXXDTpIRupnWCK7mtDpvZ6K8P0,38424
 metaflow/plugins/argo/argo_workflows_decorator.py,sha256=ogCSBmwsC2C3eusydrgjuAJd4qK18f1sI4jJwA4Fd-o,7800
 metaflow/plugins/argo/argo_workflows_deployer.py,sha256=6kHxEnYXJwzNCM9swI8-0AckxtPWqwhZLerYkX8fxUM,4444
@@ -288,17 +288,17 @@ metaflow/plugins/gcp/gs_utils.py,sha256=ZmIGFse1qYyvAVrwga23PQUzF6dXEDLLsZ2F-YRm
 metaflow/plugins/gcp/includefile_support.py,sha256=OQO0IVWv4ObboL0VqEZwcDOyj9ORLdur66JToxQ84vU,3887
 metaflow/plugins/kubernetes/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 metaflow/plugins/kubernetes/kube_utils.py,sha256=jdFMGbEmIow-oli26v31W9CmbZXigx06b3D_xIobpk0,4140
-metaflow/plugins/kubernetes/kubernetes.py,sha256=4WHVs421w5JfFSRAdpiZ6X9w7xEK8UPYGNcc9e0JOFc,30420
-metaflow/plugins/kubernetes/kubernetes_cli.py,sha256=o_o0BDEJFpTuga7txRmkvZH8OIuTb5kI4UaG6xbzf84,13929
+metaflow/plugins/kubernetes/kubernetes.py,sha256=JngFWhz9l6sYORMX9-6mdrq0e_FWN5BP9QHUVW75wzI,30596
+metaflow/plugins/kubernetes/kubernetes_cli.py,sha256=fGCzz9_MT6KTKWAlpjzjFtYCSy3WctqJXhNIfTg8z5Y,14122
 metaflow/plugins/kubernetes/kubernetes_client.py,sha256=tuvXP-QKpdeSmzVolB2R_TaacOr5DIb0j642eKcjsiM,6491
-metaflow/plugins/kubernetes/kubernetes_decorator.py,sha256=Spz1vkIurf9Z9mENZEdXk2tLS3v2jqVJIF5wyYyqbJ4,31684
-metaflow/plugins/kubernetes/kubernetes_job.py,sha256=0PjcVgJkmSqjBNEkibk8y0xxBrgu99D8ar9RsXnzCeY,32833
-metaflow/plugins/kubernetes/kubernetes_jobsets.py,sha256=9kU43eE5IvIa7y-POzBdxnJOazWsedKhwQ51Tu1HN_A,42471
+metaflow/plugins/kubernetes/kubernetes_decorator.py,sha256=cy7nr7HlJS1se3WszOLEdE2i4VfpSOFxkiGcp__MX1M,31963
+metaflow/plugins/kubernetes/kubernetes_job.py,sha256=eU_c8YAY9IPygBo0MIqYA6xGvmXEBo_0AmV38S811TM,33083
+metaflow/plugins/kubernetes/kubernetes_jobsets.py,sha256=0ebS4H5Lnehq0-CXxqh30ZunukuYJ5gD_9mSTaRcXx4,42851
 metaflow/plugins/kubernetes/spot_metadata_cli.py,sha256=an0nWCxgflmqIPBCBrlb4m3DereDFFJBLt-KKhqcHc8,1670
 metaflow/plugins/kubernetes/spot_monitor_sidecar.py,sha256=zrWU-smQwPnL6MBHmzTxWyEA00R6iKKQbhhy50xFwQ8,3832
 metaflow/plugins/metadata_providers/__init__.py,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
 metaflow/plugins/metadata_providers/local.py,sha256=9UAxe9caN6kU1lkSlIoJbRGgTqsMa62cBTnyMwhqiaA,22446
-metaflow/plugins/metadata_providers/service.py,sha256=NKZfFMamx6upP6aFRJfXlfYIhySgFNzz6kbp1yPD7LA,20222
+metaflow/plugins/metadata_providers/service.py,sha256=XhXarANcgkQ4D5q8qhywvd8BjPN274DLGLVx1KtxWF0,20725
 metaflow/plugins/pypi/__init__.py,sha256=0YFZpXvX7HCkyBFglatual7XGifdA1RwC3U4kcizyak,1037
 metaflow/plugins/pypi/bootstrap.py,sha256=oRfJkAp99R338nYX2wq4FrV0Ax5h4QqqifNmAXrM3CY,14892
 metaflow/plugins/pypi/conda_decorator.py,sha256=piFcE4uGmWhhbGlxMK0GHd7BGEyqy6r9BFy8Mjoi80Q,15937
@@ -312,7 +312,7 @@ metaflow/plugins/secrets/__init__.py,sha256=mhJaN2eMS_ZZVewAMR2E-JdP5i0t3v9e6Dcw
 metaflow/plugins/secrets/inline_secrets_provider.py,sha256=EChmoBGA1i7qM3jtYwPpLZDBybXLergiDlN63E0u3x8,294
 metaflow/plugins/secrets/secrets_decorator.py,sha256=s-sFzPWOjahhpr5fMj-ZEaHkDYAPTO0isYXGvaUwlG8,11273
 metaflow/runner/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-metaflow/runner/click_api.py,sha256=dKTrom8j7yvBQfFNl4JH9WJcyBBHy39ROtVpilUdqkw,23429
+metaflow/runner/click_api.py,sha256=truFw2mbtnEKWVhfY5haIxJEyyoB6-cipoY5-ZWXFXo,23111
 metaflow/runner/deployer.py,sha256=Yas_SZCss3kfJw3hLC8_IyzgiytUFGoEGHz-l-rBBKk,8980
 metaflow/runner/deployer_impl.py,sha256=Kab9rLoA3EiBJDtTTulhPCeKzqiljW366nx2Tm0LYy0,6143
 metaflow/runner/metaflow_runner.py,sha256=L302ew_BPBPs-NnW8n92dqqbqmHwrwGL5D6kTZvl5vY,16074
@@ -360,9 +360,9 @@ metaflow/user_configs/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3h
 metaflow/user_configs/config_decorators.py,sha256=qCKVAvd0NKgaCxQ2OThes5-DYHXq6A1HqURubYNeFdw,20481
 metaflow/user_configs/config_options.py,sha256=m6jccSpzI4qUJ7vyYkYBIf8G3V0Caunxg_k7zg4Zlqg,21067
 metaflow/user_configs/config_parameters.py,sha256=oeJGVKu1ao_YQX6Lg6P2FEv5k5-_F4sARLlVpTW9ezM,15502
-ob_metaflow-2.13.10.1.dist-info/LICENSE,sha256=nl_Lt5v9VvJ-5lWJDT4ddKAG-VZ-2IaLmbzpgYDz2hU,11343
-ob_metaflow-2.13.10.1.dist-info/METADATA,sha256=DsnOlL1LI7oDUS-wnH7RHGzoiHmFHveNkCon-aOcFUA,5318
-ob_metaflow-2.13.10.1.dist-info/WHEEL,sha256=9Hm2OB-j1QcCUq9Jguht7ayGIIZBRTdOXD1qg9cCgPM,109
-ob_metaflow-2.13.10.1.dist-info/entry_points.txt,sha256=IKwTN1T3I5eJL3uo_vnkyxVffcgnRdFbKwlghZfn27k,57
-ob_metaflow-2.13.10.1.dist-info/top_level.txt,sha256=v1pDHoWaSaKeuc5fKTRSfsXCKSdW1zvNVmvA-i0if3o,9
-ob_metaflow-2.13.10.1.dist-info/RECORD,,
+ob_metaflow-2.14.0.2rc0.dist-info/LICENSE,sha256=nl_Lt5v9VvJ-5lWJDT4ddKAG-VZ-2IaLmbzpgYDz2hU,11343
+ob_metaflow-2.14.0.2rc0.dist-info/METADATA,sha256=EGri-n8c5GqNu91Vf_PHyq4pz4ZCwHqhimXpMJJYKQQ,5322
+ob_metaflow-2.14.0.2rc0.dist-info/WHEEL,sha256=9Hm2OB-j1QcCUq9Jguht7ayGIIZBRTdOXD1qg9cCgPM,109
+ob_metaflow-2.14.0.2rc0.dist-info/entry_points.txt,sha256=IKwTN1T3I5eJL3uo_vnkyxVffcgnRdFbKwlghZfn27k,57
+ob_metaflow-2.14.0.2rc0.dist-info/top_level.txt,sha256=v1pDHoWaSaKeuc5fKTRSfsXCKSdW1zvNVmvA-i0if3o,9
+ob_metaflow-2.14.0.2rc0.dist-info/RECORD,,