ob-metaflow 2.14.0.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/version.py

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

{ob_metaflow-2.14.0.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.14.0.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.14.0.1; extra == "stubs"
+Requires-Dist: metaflow-stubs==2.14.0.2rc0; extra == "stubs"
 Dynamic: author
 Dynamic: author-email
 Dynamic: description

{ob_metaflow-2.14.0.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=OaCJjK-uyp1C0aFiz2JTkqP14Wyxki5dmByGYREkaeI,30
+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,12 +288,12 @@ 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
@@ -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.14.0.1.dist-info/LICENSE,sha256=nl_Lt5v9VvJ-5lWJDT4ddKAG-VZ-2IaLmbzpgYDz2hU,11343
-ob_metaflow-2.14.0.1.dist-info/METADATA,sha256=-RYFWzqmKGJhhK1Ruy5UyjceMVZLPc6GEcokwJqoO48,5316
-ob_metaflow-2.14.0.1.dist-info/WHEEL,sha256=9Hm2OB-j1QcCUq9Jguht7ayGIIZBRTdOXD1qg9cCgPM,109
-ob_metaflow-2.14.0.1.dist-info/entry_points.txt,sha256=IKwTN1T3I5eJL3uo_vnkyxVffcgnRdFbKwlghZfn27k,57
-ob_metaflow-2.14.0.1.dist-info/top_level.txt,sha256=v1pDHoWaSaKeuc5fKTRSfsXCKSdW1zvNVmvA-i0if3o,9
-ob_metaflow-2.14.0.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,,